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 --
  14 -- WARNING: this is not a general solution, you probably do not want to use
  15 -- this in your project! It handles only the simplest case of a single reader
  16 -- and a single writer, and will fail horribly in more general cases.
  17 -- One thread must only call writeChan, while the other only calls
  18 -- getBSChanContents.
  19 --
  20 -- Note also that only the lengths of the bytestrings in the chan count
  21 -- towards the total, so it will accept unboundedly many `BS.empty`s
  22 -- (consuming unbounded memory!)
  23 --
  24 -- If you're reading this and know of a nice light library I could have used
  25 -- instead of hacking this together, please let me know!
  26 -- There's stm-sbchan, but it seems a bit heavy, and would add stm as a
  27 -- dependency.
  28 module BoundedBSChan
  29     ( newBSChan
  30     , writeBSChan
  31     , getBSChanContents
  32     ) where
  33
  34 import Control.Concurrent
  35 import Control.Monad
  36 import Data.Maybe (fromMaybe)
  37 import System.IO.Unsafe (unsafeInterleaveIO)
  38
  39 import qualified Data.ByteString as BS
  40
  41 data BoundedBSChan = BoundedBSChan
  42     Int -- ^bound
  43     (MVar Int) -- ^tally of bytes in queue
  44     (MVar Int) -- ^bytes removed from queue, not yet incorporated into tally
  45     (Chan BS.ByteString) -- ^underlying unbounded chan
  46
  47 newBSChan :: Int -> IO BoundedBSChan
  48 newBSChan maxSize = liftM3 (BoundedBSChan maxSize) (newMVar 0) (newMVar 0) newChan
  49
  50 writeBSChan :: BoundedBSChan -> BS.ByteString -> IO ()
  51 writeBSChan c@(BoundedBSChan maxSize wv rv ch) b = do
  52     let len = BS.length b
  53     done <- modifyMVar wv $ \w ->
  54         if w > 0 && w + len > maxSize
  55         then takeMVar rv >>= \r -> return (w - r, False)
  56         else writeChan ch b >> return (w + len, True)
  57     unless done $ writeBSChan c b
  58
  59 readBSChan :: BoundedBSChan -> IO BS.ByteString
  60 readBSChan (BoundedBSChan _ _ rv ch) = do
  61     b <- readChan ch
  62     r <- fromMaybe 0 <$> tryTakeMVar rv
  63     putMVar rv $ r + BS.length b
  64     return b
  65
  66 getBSChanContents :: BoundedBSChan -> IO [BS.ByteString]
  67 getBSChanContents c = unsafeInterleaveIO $
  68     liftM2 (:) (readBSChan c) (getBSChanContents c)