BoundedBSChan.hs

   1 -- This file is part of Diohsc
   2 -- Copyright (C) 2020 Martin Bays <mbays@sdf.org>
   3 --
   4 -- This program is free software: you can redistribute it and/or modify
   5 -- it under the terms of version 3 of the GNU General Public License as
   6 -- published by the Free Software Foundation, or any later version.
   7 --
   8 -- You should have received a copy of the GNU General Public License
   9 -- along with this program.  If not, see http://www.gnu.org/licenses/.
  10
  11 -- |Wrapper around `Chan ByteString` which bounds the total size of the
  12 -- bytestring in the queue.
  13 -- Handles ^C in interleaved IO by truncating.
  14 --
  15 -- WARNING: this is not a general solution, you probably do not want to use
  16 -- this in your project! It handles only the simplest case of a single reader
  17 -- and a single writer, and will fail horribly in more general cases.
  18 -- One thread must only call writeChan, while the other only calls
  19 -- getBSChanContents.
  20 --
  21 -- Note also that only the lengths of the bytestrings in the chan count
  22 -- towards the total, so it will accept unboundedly many `BS.empty`s
  23 -- (consuming unbounded memory!)
  24 --
  25 -- If you're reading this and know of a nice light library I could have used
  26 -- instead of hacking this together, please let me know!
  27 -- There's stm-sbchan, but it seems a bit heavy, and would add stm as a
  28 -- dependency.
  29 module BoundedBSChan
  30     ( newBSChan
  31     , writeBSChan
  32     , getBSChanContents
  33     ) where
  34
  35 import           Control.Concurrent
  36 import           Control.Monad
  37 import           Data.Maybe               (fromMaybe)
  38 import           System.Console.Haskeline (handleInterrupt)
  39 import           System.IO.Unsafe         (unsafeInterleaveIO)
  40
  41 import qualified Data.ByteString          as BS
  42
  43 data BoundedBSChan = BoundedBSChan
  44     Int -- ^bound
  45     (MVar Int) -- ^tally of bytes in queue
  46     (MVar Int) -- ^bytes removed from queue, not yet incorporated into tally
  47     (Chan BS.ByteString) -- ^underlying unbounded chan
  48
  49 newBSChan :: Int -> IO BoundedBSChan
  50 newBSChan maxSize = liftM3 (BoundedBSChan maxSize) (newMVar 0) (newMVar 0) newChan
  51
  52 writeBSChan :: BoundedBSChan -> BS.ByteString -> IO ()
  53 writeBSChan c@(BoundedBSChan maxSize wv rv ch) b = do
  54     let len = BS.length b
  55     done <- modifyMVar wv $ \w ->
  56         if w > 0 && w + len > maxSize
  57         then takeMVar rv >>= \r -> return (w - r, False)
  58         else writeChan ch b >> return (w + len, True)
  59     unless done $ writeBSChan c b
  60
  61 readBSChan :: BoundedBSChan -> IO BS.ByteString
  62 readBSChan (BoundedBSChan _ _ rv ch) = handleInterrupt (return BS.empty) $ do
  63     b <- readChan ch
  64     r <- fromMaybe 0 <$> tryTakeMVar rv
  65     putMVar rv $ r + BS.length b
  66     return b
  67
  68 getBSChanContents :: BoundedBSChan -> IO [BS.ByteString]
  69 getBSChanContents c = unsafeInterleaveIO $
  70     liftM2 (:) (readBSChan c) (getBSChanContents c)