stm-conduit 2.2 → 2.2.1
raw patch · 3 files changed
+139/−25 lines, 3 filesdep +cerealdep +cereal-conduitdep +directorydep ~resourcet
Dependencies added: cereal, cereal-conduit, directory
Dependency ranges changed: resourcet
Files
- Data/Conduit/Async.hs +114/−13
- stm-conduit.cabal +17/−12
- test/Test.hs +8/−0
Data/Conduit/Async.hs view
@@ -1,6 +1,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TupleSections #-} {-# LANGUAGE RankNTypes #-}+{-# LANGUAGE RecordWildCards #-} -- | * Introduction --@@ -9,20 +10,30 @@ -- the consumer is concurrently consuming. module Data.Conduit.Async ( buffer , ($$&)+ , bufferToFile , gatherFrom , drainTo ) where -import Control.Applicative-import Control.Concurrent.Async.Lifted-import Control.Concurrent.STM-import Control.Exception.Lifted-import Control.Monad.IO.Class-import Control.Monad.Loops-import Control.Monad.Trans.Class-import Control.Monad.Trans.Control-import Data.Conduit-import Data.Conduit.List as CL+import Control.Applicative+import Control.Concurrent.Async.Lifted+import Control.Concurrent.STM+import Control.Concurrent.STM.TBChan+import Control.Exception.Lifted+import Control.Monad hiding (forM_)+import Control.Monad.IO.Class+import Control.Monad.Loops+import Control.Monad.Trans.Class+import Control.Monad.Trans.Control+import Control.Monad.Trans.Resource+import Data.Conduit+import qualified Data.Conduit.Binary as CB+import qualified Data.Conduit.Cereal as C+import qualified Data.Conduit.List as CL+import Data.Foldable (forM_)+import Data.Serialize as Cereal+import System.Directory (removeFile)+import System.IO -- | Concurrently join the producer and consumer, using a bounded queue of the -- given size. The producer will block when the queue is full, if it is@@ -38,9 +49,9 @@ chan <- liftIO $ newTBQueueIO size control $ \runInIO -> withAsync (runInIO $ sender chan) $ \input' ->- withAsync (runInIO $ recv chan $$ output) $ \output' -> do- link2 input' output'- wait output'+ withAsync (runInIO $ recv chan $$ output) $ \output' -> do+ link2 input' output'+ wait output' where send chan = liftIO . atomically . writeTBQueue chan @@ -60,6 +71,96 @@ ($$&) :: (MonadIO m, MonadBaseControl IO m) => Producer m a -> Consumer a m b -> m b ($$&) = buffer 64++data BufferContext m a = BufferContext+ { chan :: TBChan a+ , restore :: TChan (Source m a)+ , slotsFree :: TVar (Maybe Int)+ , done :: TVar Bool+ }++-- | Like 'buffer', except that when the bounded queue is overflowed, the+-- excess is cached in a local file so that consumption from upstream may+-- continue. When the queue becomes exhausted by yielding, it is filled+-- from the cache until all elements have been yielded.+--+-- Note that the maximum amount of memory consumed is equal to (2 *+-- memorySize + 1), so take this into account when picking a chunking size.+bufferToFile :: (MonadBaseControl IO m, MonadIO m, MonadResource m, Serialize a)+ => Int -- ^ Size of the bounded queue in memory+ -> Maybe Int -- ^ Max elements to keep on disk at one time+ -> FilePath -- ^ Directory to write temp files to+ -> Producer m a+ -> Consumer a m b+ -> m b+bufferToFile memorySize fileMax tempDir input output = do+ context <- liftIO $ BufferContext+ <$> newTBChanIO memorySize+ <*> newTChanIO+ <*> newTVarIO fileMax+ <*> newTVarIO False+ control $ \runInIO ->+ withAsync (runInIO $ sender context) $ \input' ->+ withAsync (runInIO $ recv context $$ output) $ \output' -> do+ link2 input' output'+ wait output'+ where+ sender BufferContext {..} = do+ input $$ awaitForever $ \x -> join $ liftIO $ atomically $ do+ written <- tryWriteTBChan chan x+ if written+ then return $ return ()+ else do+ action <- persistChan+ writeTBChan chan x+ return action+ liftIO $ atomically $ writeTVar done True+ where+ persistChan = do+ -- Empty the pending chan and return an action that writes the+ -- overflow to a disk file.+ xs <- exhaust chan+ mslots <- readTVar slotsFree+ let len = length xs+ forM_ mslots $ \slots -> check (len < slots)++ filePath <- newEmptyTMVar+ writeTChan restore $ do+ (path, key) <- liftIO $ atomically $ takeTMVar filePath+ CB.sourceFile path $= do+ C.conduitGet Cereal.get+ liftIO $ atomically $+ modifyTVar slotsFree (fmap (+ len))+ release key++ case xs of+ [] -> return $ return ()+ _ -> do+ modifyTVar slotsFree (fmap (+ (-len)))+ return $ do+ (key, (path, h)) <- allocate+ (openTempFile tempDir "conduit.bin")+ (\(path, h) -> hClose h >> removeFile path)+ liftIO $ do+ CL.sourceList xs $= C.conduitPut put+ $$ CB.sinkHandle h+ hClose h+ atomically $ putTMVar filePath (path, key)++ recv BufferContext {..} = loop where+ loop = do+ (src, exit) <- liftIO $ atomically $ do+ maction <- tryReadTChan restore+ case maction of+ Just action -> return (action, False)+ Nothing -> do+ xs <- exhaust chan+ isDone <- readTVar done+ return (CL.sourceList xs, isDone)+ src+ unless exit loop++ exhaust chan = whileM (not <$> isEmptyTBChan chan) (readTBChan chan) -- | Gather output values asynchronously from an action in the base monad and -- then yield them downstream. This provides a means of working around the
stm-conduit.cabal view
@@ -1,5 +1,5 @@ Name: stm-conduit-Version: 2.2+Version: 2.2.1 Synopsis: Introduces conduits to channels, and promotes using conduits concurrently. Description: Provides two simple conduit wrappers around STM@@ -22,17 +22,20 @@ Data.Conduit.TQueue build-depends:- base == 4.*- , transformers >= 0.2 && <= 0.4- , stm == 2.4.*- , stm-chans >= 2.0 && < 3.1- , conduit == 1.0.*- , resourcet >= 0.3 && < 0.5- , async >= 2.0.1- , monad-control >= 0.3.2- , monad-loops >= 0.4.2- , lifted-base >= 0.2.1- , lifted-async >= 0.1+ base == 4.*+ , transformers >= 0.2 && <= 0.4+ , stm == 2.4.*+ , stm-chans >= 2.0 && < 3.1+ , cereal >= 0.4.0.1+ , cereal-conduit >= 0.7.2+ , conduit == 1.0.*+ , directory >= 1.2+ , resourcet >= 0.3 && < 0.5+ , async >= 2.0.1+ , monad-control >= 0.3.2+ , monad-loops >= 0.4.2+ , lifted-base >= 0.2.1+ , lifted-async >= 0.1 ghc-options: -Wall -fwarn-tabs -fwarn-unused-imports @@ -55,6 +58,8 @@ , conduit , transformers , stm-chans+ , resourcet+ , directory source-repository head type: git
test/Test.hs view
@@ -12,6 +12,7 @@ import qualified Control.Monad as Monad import Control.Monad.IO.Class (liftIO) import Control.Monad.Trans.Class (lift)+import Control.Monad.Trans.Resource (runResourceT) import Control.Concurrent import Control.Concurrent.STM import Control.Concurrent.STM.TMQueue@@ -20,6 +21,7 @@ import Data.Conduit.Async import Data.Conduit.TMChan import Data.Conduit.TQueue+import System.Directory main = defaultMain tests @@ -31,6 +33,7 @@ ], testGroup "Async functions" [ testCase "buffer" test_buffer+ , testCase "bufferToFile" test_bufferToFile , testCase "gatherFrom" test_gatherFrom , testCase "drainTo" test_drainTo ],@@ -82,6 +85,11 @@ test_buffer = do sum' <- buffer 128 (CL.sourceList [1..100]) (CL.fold (+) 0) assertEqual "sum computed using buffer" sum' 5050++test_bufferToFile = do+ tempDir <- getTemporaryDirectory+ sum' <- runResourceT $ bufferToFile 16 (Just 5) tempDir (CL.sourceList [1 :: Int .. 100]) (CL.fold (+) 0)+ assertEqual "sum computed using bufferToFile" sum' 5050 test_gatherFrom = do sum' <- gatherFrom 128 gen $$ CL.fold (+) 0