diohsc-0.1.7: BoundedBSChan.hs

-- This file is part of Diohsc
-- Copyright (C) 2020 Martin Bays <mbays@sdf.org>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of version 3 of the GNU General Public License as
-- published by the Free Software Foundation, or any later version.
--
-- You should have received a copy of the GNU General Public License
-- along with this program.  If not, see http://www.gnu.org/licenses/.

-- |Wrapper around `Chan ByteString` which bounds the total size of the
-- bytestring in the queue.
-- Handles ^C in interleaved IO by truncating.
--
-- WARNING: this is not a general solution, you probably do not want to use
-- this in your project! It handles only the simplest case of a single reader
-- and a single writer, and will fail horribly in more general cases.
-- One thread must only call writeChan, while the other only calls
-- getBSChanContents.
--
-- Note also that only the lengths of the bytestrings in the chan count
-- towards the total, so it will accept unboundedly many `BS.empty`s
-- (consuming unbounded memory!)
--
-- If you're reading this and know of a nice light library I could have used
-- instead of hacking this together, please let me know!
-- There's stm-sbchan, but it seems a bit heavy, and would add stm as a
-- dependency.
module BoundedBSChan
    ( newBSChan
    , writeBSChan
    , getBSChanContents
    ) where

import           Control.Concurrent
import           Control.Monad
import           Data.Maybe               (fromMaybe)
import           System.Console.Haskeline (handleInterrupt)
import           System.IO.Unsafe         (unsafeInterleaveIO)

import qualified Data.ByteString          as BS

data BoundedBSChan = BoundedBSChan
    Int -- ^bound
    (MVar Int) -- ^tally of bytes in queue
    (MVar Int) -- ^bytes removed from queue, not yet incorporated into tally
    (Chan BS.ByteString) -- ^underlying unbounded chan

newBSChan :: Int -> IO BoundedBSChan
newBSChan maxSize = liftM3 (BoundedBSChan maxSize) (newMVar 0) (newMVar 0) newChan

writeBSChan :: BoundedBSChan -> BS.ByteString -> IO ()
writeBSChan c@(BoundedBSChan maxSize wv rv ch) b = do
    let len = BS.length b
    done <- modifyMVar wv $ \w ->
        if w > 0 && w + len > maxSize
        then takeMVar rv >>= \r -> return (w - r, False)
        else writeChan ch b >> return (w + len, True)
    unless done $ writeBSChan c b

readBSChan :: BoundedBSChan -> IO BS.ByteString
readBSChan (BoundedBSChan _ _ rv ch) = handleInterrupt (return BS.empty) $ do
    b <- readChan ch
    r <- fromMaybe 0 <$> tryTakeMVar rv
    putMVar rv $ r + BS.length b
    return b

getBSChanContents :: BoundedBSChan -> IO [BS.ByteString]
getBSChanContents c = unsafeInterleaveIO $
    liftM2 (:) (readBSChan c) (getBSChanContents c)