diff --git a/Data/Conduit/Async.hs b/Data/Conduit/Async.hs --- a/Data/Conduit/Async.hs +++ b/Data/Conduit/Async.hs @@ -1,189 +1,47 @@ +{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE PackageImports #-} {-# LANGUAGE TupleSections #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE RecordWildCards #-} -- | * Introduction -- --- Contains a combinator for concurrently joining a producer and a consumer, --- such that the producer may continue to produce (up to the queue size) as --- the consumer is concurrently consuming. -module Data.Conduit.Async ( buffer - , ($$&) - , bufferToFile +-- Contain combinators for concurrently joining conduits, such that +-- the producing side may continue to produce (up to the queue size) +-- as the consumer is concurrently consuming. +module Data.Conduit.Async ( module Data.Conduit.Async.Composition , gatherFrom , drainTo ) where -import Control.Applicative -import qualified Control.Concurrent.Async as A -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 --- producing faster than the consumers is taking from it. Likewise, if the --- consumer races ahead, it will block until more input is available. --- --- Exceptions are properly managed and propagated between the two sides, so --- the net effect should be equivalent to not using buffer at all, save for --- the concurrent interleaving of effects. -buffer :: (MonadBaseControl IO m, MonadIO m) - => Int -> Source m a -> Sink a m r -> m r -buffer size input output = do - chan <- liftIO $ newTBQueueIO size - control $ \runInIO -> - A.withAsync (runInIO $ sender chan) $ \input' -> - A.withAsync (runInIO $ recv chan $$ output) $ \output' -> do - A.link2 input' output' - A.wait output' - where - send chan = liftIO . atomically . writeTBQueue chan - - sender chan = do - input $$ CL.mapM_ (send chan . Just) - send chan Nothing - - recv chan = do - mx <- liftIO $ atomically $ readTBQueue chan - case mx of - Nothing -> return () - Just x -> yield x >> recv chan - --- | An operator form of 'buffer'. In general you should be able to replace --- any use of 'Data.Conduit.$$' with '$$&' and suddenly reap the benefit of --- concurrency, if your conduits were spending time waiting on each other. -($$&) :: (MonadIO m, MonadBaseControl IO m) - => Source m a -> Sink a m b -> m b -($$&) = buffer 64 - -infixr 0 $$& - -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 -> - A.withAsync (runInIO $ sender context) $ \input' -> - A.withAsync (runInIO $ recv context $$ output) $ \output' -> do - A.link2 input' output' - A.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) +import Control.Monad.IO.Class +import Control.Monad.Loops +import Control.Monad.Trans.Class +import Data.Conduit +import Data.Conduit.Async.Composition +import Data.Foldable +import UnliftIO -- | Gather output values asynchronously from an action in the base monad and -- then yield them downstream. This provides a means of working around the -- restriction that 'ConduitM' cannot be an instance of 'MonadBaseControl' -- in order to, for example, yield values from within a Haskell callback -- function called from a C library. -gatherFrom :: (MonadIO m, MonadBaseControl IO m) +gatherFrom :: (MonadIO m, MonadUnliftIO m) => Int -- ^ Size of the queue to create -> (TBQueue o -> m ()) -- ^ Action that generates output values - -> Producer m o + -> ConduitT () o m () gatherFrom size scatter = do - chan <- liftIO $ newTBQueueIO size + chan <- liftIO $ newTBQueueIO (fromIntegral size) worker <- lift $ async (scatter chan) - lift . restoreM =<< gather worker chan + gather worker chan where gather worker chan = do (xs, mres) <- liftIO $ atomically $ do xs <- whileM (not <$> isEmptyTBQueue chan) (readTBQueue chan) (xs,) <$> pollSTM worker - Prelude.mapM_ yield xs + traverse_ yield xs case mres of Just (Left e) -> liftIO $ throwIO (e :: SomeException) Just (Right r) -> return r @@ -191,14 +49,14 @@ -- | Drain input values into an asynchronous action in the base monad via a -- bounded 'TBQueue'. This is effectively the dual of 'gatherFrom'. -drainTo :: (MonadIO m, MonadBaseControl IO m) +drainTo :: (MonadIO m, MonadUnliftIO m) => Int -- ^ Size of the queue to create -> (TBQueue (Maybe i) -> m r) -- ^ Action to consume input values - -> Consumer i m r + -> ConduitT i Void m r drainTo size gather = do - chan <- liftIO $ newTBQueueIO size + chan <- liftIO $ newTBQueueIO (fromIntegral size) worker <- lift $ async (gather chan) - lift . restoreM =<< scatter worker chan + scatter worker chan where scatter worker chan = do mval <- await diff --git a/Data/Conduit/Async/Composition.hs b/Data/Conduit/Async/Composition.hs new file mode 100644 --- /dev/null +++ b/Data/Conduit/Async/Composition.hs @@ -0,0 +1,421 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PackageImports #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + +module Data.Conduit.Async.Composition ( CConduit + , CFConduit + , ($=&) + , (=$&) + , (=$=&) + , ($$&) + , buffer + , buffer' + , bufferToFile + , bufferToFile' + , runCConduit + ) where + +import Conduit +import Control.Concurrent.STM (orElse, check) +import Control.Monad hiding (forM_) +import Control.Monad.Loops +import Control.Monad.Trans.Resource +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 +import GHC.Exts (Constraint) +import System.Directory (removeFile) +import System.IO (openBinaryTempFile) +import UnliftIO + +-- | 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 +-- producing faster than the consumers is taking from it. Likewise, if the +-- consumer races ahead, it will block until more input is available. +-- +-- Exceptions are properly managed and propagated between the two sides, so +-- the net effect should be equivalent to not using buffer at all, save for +-- the concurrent interleaving of effects. +-- +-- The underlying monad must always be an instance of +-- 'MonadBaseControl IO'. If at least one of the two conduits is a +-- 'CFConduit', it must additionally be a in instance of +-- 'MonadResource'. +-- +-- This function is similar to '$$'; for one more like '=$=', see +-- 'buffer''. +-- +-- >>> buffer 1 (CL.sourceList [1,2,3]) CL.consume +-- [1,2,3] +buffer :: (CCatable c1 c2 c3, CRunnable c3, RunConstraints c3 m) + => Int -- ^ Size of the bounded queue in memory. + -> c1 () x m () + -> c2 x Void m r + -> m r +buffer i c1 c2 = runCConduit (buffer' i c1 c2) + +-- | An operator form of 'buffer'. In general you should be able to replace +-- any use of '$$' with '$$&' and suddenly reap the benefit of +-- concurrency, if your conduits were spending time waiting on each other. +-- +-- The underlying monad must always be an instance of +-- 'MonadBaseControl IO'. If at least one of the two conduits is a +-- 'CFConduit', it must additionally be a in instance of +-- 'MonadResource'. +-- +-- >>> CL.sourceList [1,2,3] $$& CL.consume +-- [1,2,3] +-- +-- It can be combined with '$=&' and '$='. This creates two threads; +-- the first thread produces the list and the second thread does the +-- map and the consume: +-- +-- >>> CL.sourceList [1,2,3] $$& mapC (*2) $= CL.consume +-- [2,4,6] +-- +-- This creates three threads. The three conduits all run in their +-- own threads: +-- +-- >>> CL.sourceList [1,2,3] $$& mapC (*2) $=& CL.consume +-- [2,4,6] +-- +-- >>> CL.sourceList [1,2,3] $$& (mapC (*2) $= mapC (+1)) $=& CL.consume +-- [3,5,7] +($$&) :: (CCatable c1 c2 c3, CRunnable c3, RunConstraints c3 m) => c1 () x m () -> c2 x Void m r -> m r +a $$& b = runCConduit (a =$=& b) +infixr 0 $$& + +-- | An operator form of 'buffer''. In general you should be able to replace +-- any use of '=$=' with '=$=&' and '$$' either with '$$&' or '=$=' +-- and 'runCConduit' and suddenly reap the benefit of concurrency, if +-- your conduits were spending time waiting on each other. +-- +-- >>> runCConduit $ CL.sourceList [1,2,3] =$=& CL.consume +-- [1,2,3] +(=$=&) :: (CCatable c1 c2 c3) => c1 i x m () -> c2 x o m r -> c3 i o m r +a =$=& b = buffer' 64 a b +infixr 2 =$=& + +-- | An alias for '=$=&' by analogy with '=$=' and '$='. +($=&) :: (CCatable c1 c2 c3) => c1 i x m () -> c2 x o m r -> c3 i o m r +($=&) = (=$=&) +infixl 1 $=& + +-- | An alias for '=$=&' by analogy with '=$=' and '=$'. +(=$&) :: (CCatable c1 c2 c3) => c1 i x m () -> c2 x o m r -> c3 i o m r +(=$&) = (=$=&) +infixr 2 =$& + +-- | Conduits are concatenable; this class describes how. +-- class CCatable (c1 :: * -> * -> (* -> *) -> * -> *) (c2 :: * -> * -> (* -> *) -> * -> *) (c3 :: * -> * -> (* -> *) -> * -> *) | c1 c2 -> c3 where +class CCatable c1 c2 (c3 :: * -> * -> (* -> *) -> * -> *) | c1 c2 -> c3 where + -- | 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 + -- producing faster than the consumers is taking from it. Likewise, if the + -- consumer races ahead, it will block until more input is available. + -- + -- Exceptions are properly managed and propagated between the two sides, so + -- the net effect should be equivalent to not using buffer at all, save for + -- the concurrent interleaving of effects. + -- + -- This function is similar to '=$='; for one more like '$$', see + -- 'buffer'. + -- + -- >>> runCConduit $ buffer' 1 (CL.sourceList [1,2,3]) CL.consume + -- [1,2,3] + buffer' :: Int -- ^ Size of the bounded queue in memory + -> c1 i x m () + -> c2 x o m r + -> c3 i o m r + +-- | 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. +-- +-- This function is similar to '$$'; for one more like '=$=', see +-- 'bufferToFile''. +-- +-- >>> runResourceT $ bufferToFile 1 Nothing "/tmp" (CL.sourceList [1,2,3]) CL.consume +-- [1,2,3] +bufferToFile :: (CFConduitLike c1, CFConduitLike c2, Serialize x, MonadUnliftIO m, MonadResource m, MonadThrow m) + => 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 + -> c1 () x m () + -> c2 x Void m r + -> m r +bufferToFile bufsz dsksz tmpDir c1 c2 = runCConduit (bufferToFile' bufsz dsksz tmpDir c1 c2) + +-- | 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. +-- +-- This function is similar to '=$='; for one more like '$$', see +-- 'bufferToFile'. +-- +-- >>> runResourceT $ runCConduit $ bufferToFile' 1 Nothing "/tmp" (CL.sourceList [1,2,3]) CL.consume +-- [1,2,3] +-- +-- It is frequently convenient to define local function to use this in operator form: +-- +-- >>> :{ +-- runResourceT $ do +-- let buf c = bufferToFile' 10 Nothing "/tmp" c -- eta-conversion to avoid monomorphism restriction +-- runCConduit $ CL.sourceList [0x30, 0x31, 0x32] `buf` mapC (toEnum :: Int -> Char) `buf` CL.consume +-- :} +-- "012" +bufferToFile' :: (CFConduitLike c1, CFConduitLike c2, Serialize x) + => 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 + -> c1 i x m () + -> c2 x o m r + -> CFConduit i o m r +bufferToFile' bufsz dsksz tmpDir c1 c2 = combine (asCFConduit c1) (asCFConduit c2) + where combine (FSingle a) b = FMultipleF bufsz dsksz tmpDir a b + combine (FMultiple i a as) b = FMultiple i a (bufferToFile' bufsz dsksz tmpDir as b) + combine (FMultipleF bufsz' dsksz' tmpDir' a as) b = FMultipleF bufsz' dsksz' tmpDir' a (bufferToFile' bufsz dsksz tmpDir as b) + +-- | Conduits are, once there's a producer on one end and a consumer +-- on the other, runnable. +class CRunnable c where + type RunConstraints c (m :: * -> *) :: Constraint + -- | Execute a conduit concurrently. This is the concurrent + -- equivalent of 'runConduit'. + -- + -- The underlying monad must always be an instance of + -- 'MonadUnliftIO'. If the conduits is a 'CFConduit', it must + -- additionally be a in instance of 'MonadResource'. + runCConduit :: (RunConstraints c m) => c () Void m r -> m r + +instance CCatable ConduitT ConduitT CConduit where + buffer' i a b = buffer' i (Single a) (Single b) + +instance CCatable ConduitT CConduit CConduit where + buffer' i a b = buffer' i (Single a) b + +instance CCatable ConduitT CFConduit CFConduit where + buffer' i a b = buffer' i (asCFConduit a) b + +instance CCatable CConduit ConduitT CConduit where + buffer' i a b = buffer' i a (Single b) + +instance CCatable CConduit CConduit CConduit where + buffer' i (Single a) b = Multiple i a b + buffer' i (Multiple i' a as) b = Multiple i' a (buffer' i as b) + +instance CCatable CConduit CFConduit CFConduit where + buffer' i a b = buffer' i (asCFConduit a) b + +instance CCatable CFConduit ConduitT CFConduit where + buffer' i a b = buffer' i a (asCFConduit b) + +instance CCatable CFConduit CConduit CFConduit where + buffer' i a b = buffer' i a (asCFConduit b) + +instance CCatable CFConduit CFConduit CFConduit where + buffer' i (FSingle a) b = FMultiple i a b + buffer' i (FMultiple i' a as) b = FMultiple i' a (buffer' i as b) + buffer' i (FMultipleF bufsz dsksz tmpDir a as) b = FMultipleF bufsz dsksz tmpDir a (buffer' i as b) + +instance CRunnable ConduitT where + type RunConstraints ConduitT m = (MonadUnliftIO m, Monad m) + runCConduit = runConduit + +instance CRunnable CConduit where + type RunConstraints CConduit m = (MonadUnliftIO m, MonadIO m) + runCConduit (Single c) = runConduit c + runCConduit (Multiple bufsz c cs) = do + chan <- liftIO $ newTBQueueIO (fromIntegral bufsz) + withAsync (sender chan c) $ \c' -> + stage chan c' cs + +instance CRunnable CFConduit where + type RunConstraints CFConduit m = (MonadThrow m, MonadUnliftIO m, MonadIO m, MonadResource m) + runCConduit (FSingle c) = runConduit c + runCConduit (FMultiple bufsz c cs) = do + chan <- liftIO $ newTBQueueIO (fromIntegral bufsz) + withAsync (sender chan c) $ \c' -> + fstage (receiver chan) c' cs + runCConduit (FMultipleF bufsz filemax tempDir c cs) = do + context <- liftIO $ BufferContext <$> newTBQueueIO (fromIntegral bufsz) + <*> newTQueueIO + <*> newTVarIO filemax + <*> newTVarIO False + <*> pure tempDir + withAsync (fsender context c) $ \c' -> + fstage (freceiver context) c' cs + +-- | A "concurrent conduit", in which the stages run in parallel with +-- a buffering queue between them. +data CConduit i o m r where + Single :: ConduitT i o m r -> CConduit i o m r + Multiple :: Int -> ConduitT i x m () -> CConduit x o m r -> CConduit i o m r + +-- Combines a producer with a queue, sending it everything the +-- producer produces. +sender :: (MonadIO m) => TBQueue (Maybe o) -> ConduitT () o m () -> m () +sender chan input = do + runConduit $ input .| mapM_C (send chan . Just) + send chan Nothing + +-- One "layer" of withAsync in a CConduit run. +stage :: (MonadUnliftIO m) => TBQueue (Maybe i) -> Async x -> CConduit i Void m r -> m r +stage chan prevAsync (Single c) = + -- The last layer; feed the output of "chan" into the conduit and + -- wait for the result. + withAsync (runConduit $ receiver chan .| c) $ \c' -> do + link2 prevAsync c' + wait c' +stage chan prevAsync (Multiple bufsz c cs) = do + -- not the last layer, so take the input from "chan", have this + -- layer's conduit process it, and send the conduit's output to the + -- next layer. + chan' <- liftIO $ newTBQueueIO (fromIntegral bufsz) + withAsync (sender chan' $ receiver chan .| c) $ \c' -> do + link2 prevAsync c' + stage chan' c' cs + +-- A Producer which produces the values of the given channel until +-- Nothing is received. This is the other half of "sender". +receiver :: (MonadIO m) => TBQueue (Maybe o) -> ConduitT () o m () +receiver chan = do + mx <- recv chan + case mx of + Nothing -> return () + Just x -> yield x >> receiver chan + +-- | A "concurrent conduit", in which the stages run in parallel with +-- a buffering queue and possibly a disk file between them. +data CFConduit i o m r where + FSingle :: ConduitT i o m r -> CFConduit i o m r + FMultiple :: Int -> ConduitT i x m () -> CFConduit x o m r -> CFConduit i o m r + FMultipleF :: (Serialize x) => Int -> Maybe Int -> FilePath -> ConduitT i x m () -> CFConduit x o m r -> CFConduit i o m r + +class CFConduitLike a where + asCFConduit :: a i o m r -> CFConduit i o m r + +instance CFConduitLike ConduitT where + asCFConduit = FSingle + +instance CFConduitLike CConduit where + asCFConduit (Single c) = FSingle c + asCFConduit (Multiple i c cs) = FMultiple i c (asCFConduit cs) + +instance CFConduitLike CFConduit where + asCFConduit = id + +data BufferContext m a = BufferContext { chan :: TBQueue a + , restore :: TQueue (ConduitT () a m ()) + , slotsFree :: TVar (Maybe Int) + , done :: TVar Bool + , tempDir :: FilePath + } + +-- The file-backed equivlent of "sender". This sends the values +-- generated by "input" to the "chan" in the BufferContext until it +-- gets full, then flushes it to disk via "persistChan". +fsender :: (MonadThrow m, MonadResource m, Serialize x) => BufferContext m x -> ConduitT () x m () -> m () +fsender bc@BufferContext{..} input = do + runConduit $ input .| (mapM_C $ \x -> join $ liftIO $ atomically $ do + (writeTBQueue chan x >> return (return ())) `orElse` do + action <- persistChan bc + writeTBQueue chan x + return action) + liftIO $ atomically $ writeTVar done True + +-- Connect a stage to another stage via either an in-memory queue or a +-- disk buffer. This is the file-backed equivalent of "stage". +fstage :: (MonadThrow m, MonadUnliftIO m, MonadResource m) => ConduitT () i m () -> Async x -> CFConduit i Void m r -> m r +fstage prevStage prevAsync (FSingle c) = + -- The final conduit in the chain; just accept everything from + -- the previous stage and wait for the result. + withAsync (runConduit $ prevStage .| c) $ \c' -> do + link2 prevAsync c' + wait c' +fstage prevStage prevAsync (FMultiple bufsz c cs) = do + -- This stage is connected to the next via a non-file-backed + -- channel, so it just uses "sender" and "reciever" in the same way + -- "stage" does. + chan' <- liftIO $ newTBQueueIO (fromIntegral bufsz) + withAsync (sender chan' $ prevStage .| c) $ \c' -> do + link2 prevAsync c' + fstage (receiver chan') c' cs +fstage prevStage prevAsync (FMultipleF bufsz dsksz tempDir c cs) = do + -- This potentially needs to write its output to a file, so it uses + -- "fsender" send and tells the next stage to use "freceiver" to read. + bc <- liftIO $ BufferContext <$> newTBQueueIO (fromIntegral bufsz) + <*> newTQueueIO + <*> newTVarIO dsksz + <*> newTVarIO False + <*> pure tempDir + withAsync (fsender bc $ prevStage .| c) $ \c' -> do + link2 prevAsync c' + fstage (freceiver bc) c' cs + +-- Receives from disk files or the in-memory queue if no spill-to-disk +-- has occurred. +freceiver :: (MonadIO m) => BufferContext m o -> ConduitT () o m () +freceiver BufferContext{..} = loop where + loop = do + (src, exit) <- liftIO $ atomically $ do + (readTQueue restore >>= (\action -> return (action, False))) `orElse` do + xs <- exhaust chan + isDone <- readTVar done + return (CL.sourceList xs, isDone) + src + unless exit loop + +-- The channel is full, so (return an action which will) spill it to disk, unless too +-- many items are there already. +persistChan :: (MonadThrow m, MonadResource m, Serialize o) => BufferContext m o -> STM (m ()) +persistChan BufferContext{..} = do + xs <- exhaust chan + mslots <- readTVar slotsFree + let len = length xs + forM_ mslots $ \slots -> check (len < slots) + filePath <- newEmptyTMVar + writeTQueue restore $ do + (path, key) <- liftIO $ atomically $ takeTMVar filePath + CB.sourceFile path .| do + C.conduitGet2 get + liftIO $ atomically $ modifyTVar slotsFree (fmap (+ len)) + release key + case xs of + [] -> return (return ()) + _ -> do + modifyTVar slotsFree (fmap (subtract len)) + return $ do + (key, (path, h)) <- allocate (openBinaryTempFile tempDir "conduit.bin") (\(path, h) -> hClose h `finally` removeFile path) + liftIO $ do + runConduit $ CL.sourceList xs .| C.conduitPut put .| CB.sinkHandle h + hClose h + atomically $ putTMVar filePath (path, key) + +exhaust :: TBQueue a -> STM [a] +exhaust chan = whileM (not <$> isEmptyTBQueue chan) (readTBQueue chan) + +recv :: (MonadIO m) => TBQueue a -> m a +recv c = liftIO . atomically $ readTBQueue c + +send :: (MonadIO m) => TBQueue a -> a -> m () +send c = liftIO . atomically . writeTBQueue c diff --git a/Data/Conduit/TMChan.hs b/Data/Conduit/TMChan.hs --- a/Data/Conduit/TMChan.hs +++ b/Data/Conduit/TMChan.hs @@ -1,4 +1,9 @@ -{-# LANGUAGE NoMonomorphismRestriction, FlexibleContexts, RankNTypes,KindSignatures #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE NoMonomorphismRestriction #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE KindSignatures #-} + -- | * Introduction -- -- Contains a simple source and sink for linking together conduits in @@ -13,8 +18,14 @@ -- data (pictures, in this case) will be streamed down the channel to whatever -- is on the other side. -- --- > _ <- forkIO . runResourceT $ loadTextures lotsOfPictures $$ sinkTBMChan chan +-- > _ <- forkIO . runResourceT $ do +-- > _ <- register $ atomically $ closeTBMChan chan +-- > loadTextures lotsOfPictures $$ sinkTBMChan chan -- +-- We register closing function explicitly, because starting with version +-- @1.3.0@ @conduits@ library no longer maintain resources, so this is the +-- only way to safely close channel in case of exceptions. +-- -- Finally, we connect something to the other end of the channel. In this -- case, we connect a sink which uploads the textures one by one to the -- graphics card. @@ -53,44 +64,48 @@ , mergeConduits ) where -import Control.Applicative import Control.Monad import Control.Monad.IO.Class ( liftIO, MonadIO ) import Control.Monad.Trans.Class import Control.Monad.Trans.Resource -import Control.Concurrent.STM +import Control.Concurrent (killThread, forkIOWithUnmask) +import Control.Concurrent.STM hiding (atomically) import Control.Concurrent.STM.TBMChan import Control.Concurrent.STM.TMChan import Data.Conduit +import Data.Foldable import qualified Data.Conduit.List as CL +import UnliftIO as Lifted +-- | Convert channel into the source. +-- +-- *N.B* Since version 4.0 this function does not close the +-- channel if downstream is closed. chanSource :: MonadIO m => chan -- ^ The channel. -> (chan -> STM (Maybe a)) -- ^ The 'read' function. - -> (chan -> STM ()) -- ^ The 'close' function. - -> Source m a -chanSource ch reader closer = + -> ConduitT z a m () +chanSource ch reader = loop where loop = do a <- liftSTM $ reader ch case a of - Just x -> yieldOr x close >> loop + Just x -> yield x >> loop Nothing -> return () - close = liftSTM $ closer ch {-# INLINE chanSource #-} +-- | Convert channel into the consumer. +-- +-- *N.B* chanSink :: MonadIO m => chan -- ^ The channel. -> (chan -> a -> STM ()) -- ^ The 'write' function. - -> (chan -> STM ()) -- ^ The 'close' function. - -> Sink a m () -chanSink ch writer closer = do - CL.mapM_ $ liftIO . atomically . writer ch - liftSTM $ closer ch + -> ConduitT a z m () +chanSink ch writer = CL.mapM_ $ liftIO . atomically . writer ch {-# INLINE chanSink #-} -- | A simple wrapper around a TBMChan. As data is pushed into the channel, the @@ -98,15 +113,15 @@ -- channel is closed, the source will close also. -- -- If the channel fills up, the pipeline will stall until values are read. -sourceTBMChan :: MonadIO m => TBMChan a -> Source m a -sourceTBMChan ch = chanSource ch readTBMChan closeTBMChan +sourceTBMChan :: MonadIO m => TBMChan a -> ConduitT () a m () +sourceTBMChan ch = chanSource ch readTBMChan {-# INLINE sourceTBMChan #-} -- | A simple wrapper around a TMChan. As data is pushed into the channel, the -- source will read it and pass it down the conduit pipeline. When the -- channel is closed, the source will close also. -sourceTMChan :: MonadIO m => TMChan a -> Source m a -sourceTMChan ch = chanSource ch readTMChan closeTMChan +sourceTMChan :: MonadIO m => TMChan a -> ConduitT () a m () +sourceTMChan ch = chanSource ch readTMChan {-# INLINE sourceTMChan #-} -- | A simple wrapper around a TBMChan. As data is pushed into the sink, it @@ -114,27 +129,26 @@ -- the sink will block until space frees up. sinkTBMChan :: MonadIO m => TBMChan a - -> Bool -- ^ Should the channel be closed when the sink is closed? - -> Sink a m () -sinkTBMChan ch close = chanSink ch writeTBMChan (when close . closeTBMChan) + -> ConduitT a z m () +sinkTBMChan ch = chanSink ch writeTBMChan {-# INLINE sinkTBMChan #-} -- | A simple wrapper around a TMChan. As data is pushed into this sink, it -- will magically begin to appear in the channel. sinkTMChan :: MonadIO m => TMChan a - -> Bool -- ^ Should the channel be closed when the sink is closed? - -> Sink a m () -sinkTMChan ch close = chanSink ch writeTMChan (when close . closeTMChan) + -> ConduitT a z m () +sinkTMChan ch = chanSink ch writeTMChan {-# INLINE sinkTMChan #-} infixl 5 >=< -- | Modifies a TVar, returning its new value. modifyTVar'' :: TVar a -> (a -> a) -> STM a -modifyTVar'' tv f = do x <- f <$> readTVar tv - writeTVar tv x - return x +modifyTVar'' tv f = do + !x <- f <$> readTVar tv + writeTVar tv x + return x liftSTM :: forall (m :: * -> *) a. MonadIO m => STM a -> m a liftSTM = liftIO . atomically @@ -144,17 +158,18 @@ -- -- The order of the new source's data is undefined, but it will be some -- combination of the two given sources. -(>=<) :: (MonadResource mi, MonadIO mo, MonadBaseControl IO mi) - => Source mi a - -> Source mi a - -> mi (Source mo a) +(>=<) :: (MonadResource mi, MonadIO mo, MonadUnliftIO mi) + => ConduitT () a mi () + -> ConduitT () a mi () + -> mo (ConduitT () a mi ()) sa >=< sb = mergeSources [ sa, sb ] 16 {-# INLINE (>=<) #-} decRefcount :: TVar Int -> TBMChan a -> STM () -decRefcount tv chan = do n <- modifyTVar'' tv (subtract 1) - when (n == 0) $ - closeTBMChan chan +decRefcount tv chan = do + n <- modifyTVar'' tv (subtract 1) + when (n == 0) $ + closeTBMChan chan -- | Merges a list of sources, putting them all into a bounded channel, and -- returns a source which can be pulled from to pull from all the given @@ -163,25 +178,44 @@ -- The order of the new source's data is undefined, but it will be some -- combination of the given sources. The monad of the resultant source -- (@mo@) is independent of the monads of the input sources (@mi@). -mergeSources :: (MonadResource mi, MonadIO mo, MonadBaseControl IO mi) - => [Source mi a] -- ^ The sources to merge. +-- +-- @since 3.0 +-- All spawned threads will be removed when source is closed or upon an +-- exit from 'ResourceT' region. This means that result can only be used +-- within a 'runResourceT' scope. +-- +-- @before 3.0 +-- Spawned threads are not guaranteed to be closed. This may happen if +-- Source was closed before all it's input were closed. +mergeSources :: (MonadResource mi, MonadIO mo, MonadUnliftIO mi) + => [ConduitT () a mi ()] -- ^ The sources to merge. -> Int -- ^ The bound of the intermediate channel. - -> mi (Source mo a) -mergeSources sx bound = do c <- liftSTM $ newTBMChan bound - refcount <- liftSTM . newTVar $ length sx - mapM_ (\s -> runResourceT $ resourceForkIO $ s $$ chanSink c writeTBMChan $ decRefcount refcount) (map (transPipe lift) sx) - return $ sourceTBMChan c + -> mo (ConduitT () a mi ()) +mergeSources sx bound = do + return $ do + (chkey, c) <- allocate (liftSTM $ newTBMChan bound) + (liftSTM . closeTBMChan) + refcount <- liftSTM . newTVar $ length sx + st <- lift $ askUnliftIO + regs <- forM sx $ \s -> + register . killThread =<< + (liftIO $ forkIOWithUnmask $ \unmask -> + (unmask $ unliftIO st $ + runConduit $ s .| chanSink c writeTBMChan) + `Lifted.finally` (liftSTM $ decRefcount refcount c)) + chanSource c readTBMChan + release chkey + traverse_ release regs -- | Combines two conduits with unbounded channels, creating a new conduit -- which pulls data from a mix of the two: whichever produces first. -- -- The order of the new conduit's output is undefined, but it will be some -- combination of the two given conduits. -(<=>) :: (MonadIO mi, MonadIO mo, MonadBaseControl IO mi) - => Show i - => Conduit i (ResourceT mi) i - -> Conduit i (ResourceT mi) i - -> ResourceT mi (Conduit i mo i) +(<=>) :: (MonadThrow mi, MonadIO mo, MonadUnliftIO mi) + => ConduitT i i (ResourceT mi) () + -> ConduitT i i (ResourceT mi) () + -> ResourceT mi (ConduitT i i mo ()) sa <=> sb = mergeConduits [ sa, sb ] 16 {-# INLINE (<=>) #-} @@ -192,22 +226,46 @@ -- The order of the new conduits's outputs is undefined, but it will be some -- combination of the given conduits. The monad of the resultant conduit -- (@mo@) is independent of the monads of the input conduits (@mi@). -mergeConduits :: (MonadIO mi, MonadIO mo, MonadBaseControl IO mi) - => [Conduit i (ResourceT mi) o] -- ^ The conduits to merge. +-- +-- @since 3.0 +-- Closes all worker processes when resulting conduit is closed or when execution +-- leaves ResourceT context. This means that conduit is only valid inside +-- 'runResouceT' scope. +-- +-- @before 3.0 +-- Spawned threads are not guaranteed to be closed, This may happen if threads +-- Conduit was closed before all threads have finished execution. +{-# DEPRECATED mergeConduits "This method will dissapear in the next version." #-} +mergeConduits :: (MonadIO mo, MonadUnliftIO mi) + => [ConduitT i o (ResourceT mi) ()] -- ^ The conduits to merge. -> Int -- ^ The bound for the channels. - -> ResourceT mi (Conduit i mo o) + -> ResourceT mi (ConduitT i o mo ()) mergeConduits conduits bound = do let len = length conduits refcount <- liftSTM $ newTVar len - iChannels <- replicateM len $ liftSTM $ newTBMChan bound - oChannel <- liftSTM $ newTBMChan bound - forM_ (zip iChannels conduits) - $ \(iChannel, conduit) -> resourceForkIO - $ sourceTBMChan iChannel - $$ conduit - =$ chanSink oChannel (writeTBMChans . (:[])) (decRefcount refcount) + (iregs, iChannels) <- fmap unzip + $ replicateM len $ allocate (liftSTM $ newTBMChan bound) + (liftSTM . closeTBMChan) + (oreg, oChannel) <- allocate (liftSTM $ newTBMChan bound) + (liftSTM . closeTBMChan) + regs <- forM (zip iChannels conduits) + $ \(iChannel, conduit) -> Lifted.mask_ $ + register . killThread <=< resourceForkIO + $ (runConduit $ sourceTBMChan iChannel + .| conduit + .| chanSink oChannel writeTBMChan) + `finally` + (liftIO $ atomically $ decRefcount refcount oChannel >> closeTBMChan iChannel) + treg <- register $ do + traverse_ release regs + traverse_ release iregs + release oreg return - $ toConsumer (chanSink iChannels writeTBMChans (mapM_ closeTBMChan)) - >> toProducer (sourceTBMChan oChannel) + $ do chanSink iChannels writeTBMChans + -- release internal channels effiniently closing input channels + traverse_ release iregs + chanSource oChannel readTBMChan + -- release internals everything + release treg where writeTBMChans channels a = forM_ channels $ \c -> writeTBMChan c a diff --git a/Data/Conduit/TQueue.hs b/Data/Conduit/TQueue.hs --- a/Data/Conduit/TQueue.hs +++ b/Data/Conduit/TQueue.hs @@ -14,7 +14,7 @@ -- -- -- Here is short description of data structures: --- +-- -- * TQueue - unbounded infinite queue -- -- * TBQueue - bounded infinite queue @@ -24,7 +24,7 @@ -- * TBMQueue - bounded finite (closable) queue -- -- Caveats --- +-- -- Infinite operations means that source doesn't know when stream is -- ended so one need to use other methods of finishing stream like -- sending an exception or finish conduit in downstream. @@ -62,35 +62,37 @@ -- | A simple wrapper around a "TQueue". As data is pushed into the queue, the -- source will read it and pass it down the conduit pipeline. -sourceTQueue :: MonadIO m => TQueue a -> Source m a +sourceTQueue :: MonadIO m => TQueue a -> ConduitT z a m () sourceTQueue q = forever $ liftSTM (readTQueue q) >>= yield -- | A simple wrapper around a "TQueue". As data is pushed into this sink, it -- will magically begin to appear in the queue. -sinkTQueue :: MonadIO m => TQueue a -> Sink a m () +sinkTQueue :: MonadIO m => TQueue a -> ConduitT a z m () sinkTQueue q = CL.mapM_ (liftSTM . writeTQueue q) -- | A simple wrapper around a "TBQueue". As data is pushed into the queue, the -- source will read it and pass it down the conduit pipeline. -sourceTBQueue :: MonadIO m => TBQueue a -> Source m a +sourceTBQueue :: MonadIO m => TBQueue a -> ConduitT z a m () sourceTBQueue q = forever $ liftSTM (readTBQueue q) >>= yield -- | A simple wrapper around a "TBQueue". As data is pushed into this sink, it --- will magically begin to appear in the queue. Boolean argument is used --- to specify if queue should be closed when the sink is closed. -sinkTBQueue :: MonadIO m => TBQueue a -> Sink a m () +-- will magically begin to appear in the queue. +sinkTBQueue :: MonadIO m => TBQueue a -> ConduitT a z m () sinkTBQueue q = CL.mapM_ (liftSTM . writeTBQueue q) -- | A convenience wrapper for creating a source and sink TBQueue of the given -- size at once, without exposing the underlying queue. -entangledPair :: MonadIO m => Int -> m (Source m a, Sink a m ()) +-- +-- Returns release key that can be used for premature close of the communication +-- channel, otherwise channel will be closed when the ResourceT scope will be closed. +entangledPair :: MonadIO m => Int -> m (ConduitT z a m (), ConduitT a l m ()) entangledPair size = liftM (liftM2 (,) sourceTBQueue sinkTBQueue) $ - liftIO $ atomically $ newTBQueue size + liftIO $ atomically $ newTBQueue (fromIntegral size) -- | A simple wrapper around a "TMQueue". As data is pushed into the queue, the -- source will read it and pass it down the conduit pipeline. When the -- queue is closed, the source will close also. -sourceTMQueue :: MonadIO m => TMQueue a -> Source m a +sourceTMQueue :: MonadIO m => TMQueue a -> ConduitT z a m () sourceTMQueue q = loop where @@ -98,23 +100,19 @@ mx <- liftSTM $ readTMQueue q case mx of Nothing -> return () - Just x -> yieldOr x close >> loop - close = liftSTM $ closeTMQueue q + Just x -> yield x >> loop -- | A simple wrapper around a "TMQueue". As data is pushed into this sink, it -- will magically begin to appear in the queue. sinkTMQueue :: MonadIO m => TMQueue a - -> Bool -- ^ Should the queue be closed when the sink is closed? - -> Sink a m () -sinkTMQueue q shouldClose = do - CL.mapM_ (liftSTM . writeTMQueue q) - when shouldClose (liftSTM $ closeTMQueue q) + -> ConduitT a z m () +sinkTMQueue q = CL.mapM_ (liftSTM . writeTMQueue q) -- | A simple wrapper around a "TBMQueue". As data is pushed into the queue, the -- source will read it and pass it down the conduit pipeline. When the -- queue is closed, the source will close also. -sourceTBMQueue :: MonadIO m => TBMQueue a -> Source m a +sourceTBMQueue :: MonadIO m => TBMQueue a -> ConduitT z a m () sourceTBMQueue q = loop where @@ -122,18 +120,14 @@ mx <- liftSTM $ readTBMQueue q case mx of Nothing -> return () - Just x -> yieldOr x close >> loop - close = liftSTM $ closeTBMQueue q + Just x -> yield x >> loop -- | A simple wrapper around a "TBMQueue". As data is pushed into this sink, it -- will magically begin to appear in the queue. sinkTBMQueue :: MonadIO m => TBMQueue a - -> Bool -- ^ Should the queue be closed when the sink is closed? - -> Sink a m () -sinkTBMQueue q shouldClose = do - CL.mapM_ (liftSTM . writeTBMQueue q) - when shouldClose (liftSTM $ closeTBMQueue q) + -> ConduitT a z m () +sinkTBMQueue q = CL.mapM_ (liftSTM . writeTBMQueue q) liftSTM :: forall (m :: * -> *) a. MonadIO m => STM a -> m a diff --git a/Data/Conduit/Utils.hs b/Data/Conduit/Utils.hs --- a/Data/Conduit/Utils.hs +++ b/Data/Conduit/Utils.hs @@ -1,8 +1,8 @@ {-# LANGUAGE FunctionalDependencies #-} {-# LANGUAGE KindSignatures #-} {-# LANGUAGE FlexibleInstances #-} -{-| -This module provide different utility functions that allow to use safe higher +{-| +This module provide different utility functions that allow to use safe higher level usage. Conduit pairs allow creation of an internal datastructure that acts as a bridge, @@ -25,8 +25,8 @@ > mkSource = sourceTBMQueue > mkSink = flip sinkTBMQueue True - * Use "pair" or "pairBounded" to create a bridge. Because bridge data structure - is hidden and not seen in parameters, we need proxy type to help compiler to + * Use "pair" or "pairBounded" to create a bridge. Because bridge data structure + is hidden and not seen in parameters, we need proxy type to help compiler to choose type, we use "Proxy2" for that. > pairTBMQueue = pairBounded (proxy2 :: Proxy2 TBMQueue a) @@ -40,15 +40,15 @@ conduit code. This package provides predefined pairs for all STM types that are used -in the package. +in the package. -} module Data.Conduit.Utils - ( + ( -- * Conduit pairs -- ** Low level functions pairBounded -- MonadIO m => m (Source m a, Sink m a ()) - , pair -- MonadIO m => Int -> m (Source m a, Sink m a ()) + , pair -- MonadIO m => Int -> m (Source m a, Sink m a ()) -- ** Classes , UnboundedStream(..) , BoundedStream(..) @@ -90,14 +90,14 @@ -- | Class for structures that can handle unbounded stream of values. -- Such streams break conduit assumptions that constant memory will be --- used, because if receiver is slower then sender than values will +-- used, because if receiver is slower then sender than values will -- be accumulated. class UnboundedStream i o | i -> o where mkUStream :: i a -> IO (o a) -- | Class for structures that can handle bounded stream of values i.e. --- there is exists 'Int' value that sets an upper limit on the number --- of values that can be handled by structure. Exact meaning of this +-- there is exists 'Int' value that sets an upper limit on the number +-- of values that can be handled by structure. Exact meaning of this -- limit may depend on the carrier type. class BoundedStream i o | i -> o where mkBStream :: i a -> Int -> IO (o a) @@ -105,14 +105,14 @@ -- | Class that describes how we can make conduit out of the carrier -- value. class MonadIO m => IsConduit m (x :: * -> *) where - mkSink :: x a -> Sink a m () - mkSource :: x a -> Source m a + mkSink :: x a -> ConduitT a Void m () + mkSource :: x a -> ConduitT () a m () -- | Create bounded conduit pair, see "BoundedStream" class description. pairBounded :: (MonadIO m, IsConduit m o, BoundedStream i o) => i a -- ^ Type description. -> Int -- ^ Conduit size. - -> m (Source m a, Sink a m ()) + -> m (ConduitT () a m (), ConduitT a Void m ()) pairBounded p s = do q <- liftIO $ mkBStream p s return (mkSource q, mkSink q) @@ -120,7 +120,7 @@ -- | Create unbounded pair, see "UnboundedStream" class description. pair :: (MonadIO m, IsConduit m o, UnboundedStream i o) => i a -- ^ Type description. - -> m (Source m a, Sink a m ()) + -> m (ConduitT () a m (), ConduitT a Void m ()) pair p = do q <- liftIO $ mkUStream p return (mkSource q, mkSink q) @@ -129,13 +129,13 @@ -- Instances ------------------------------------------------------------------------------- instance BoundedStream (Proxy2 TBQueue) TBQueue where - mkBStream _ i = atomically $ newTBQueue i + mkBStream _ i = atomically $ newTBQueue (fromIntegral i) instance BoundedStream (Proxy2 TBMQueue) TBMQueue where - mkBStream _ i = atomically $ newTBMQueue i + mkBStream _ i = atomically $ newTBMQueue (fromIntegral i) instance BoundedStream (Proxy2 TBMChan) TBMChan where - mkBStream _ i = atomically $ newTBMChan i + mkBStream _ i = atomically $ newTBMChan (fromIntegral i) instance UnboundedStream (Proxy2 TMQueue) TMQueue where mkUStream _ = atomically $ newTMQueue @@ -152,11 +152,11 @@ instance MonadIO m => IsConduit m TBMQueue where mkSource = sourceTBMQueue - mkSink = flip sinkTBMQueue True + mkSink = sinkTBMQueue instance MonadIO m => IsConduit m TMQueue where mkSource = sourceTMQueue - mkSink = flip sinkTMQueue True + mkSink = sinkTMQueue instance MonadIO m => IsConduit m TQueue where mkSource = sourceTQueue @@ -164,11 +164,11 @@ instance MonadIO m => IsConduit m TBMChan where mkSource = sourceTBMChan - mkSink = flip sinkTBMChan True + mkSink = sinkTBMChan instance MonadIO m => IsConduit m TMChan where mkSource = sourceTMChan - mkSink = flip sinkTMChan True + mkSink = sinkTMChan ------------------------------------------------------------------------------- -- Specialized functions @@ -179,12 +179,12 @@ -- is not closable then there is no way to notify receiver side that bridge -- is closed, so it's possible to use it only in infinite streams of when -- some other mechanism of notification is used. -pairTQueue, pairTMQueue, pairTMChan :: MonadIO m => m (Source m a, Sink a m ()) +pairTQueue, pairTMQueue, pairTMChan :: MonadIO m => m (ConduitT () a m (), ConduitT a Void m ()) pairTQueue = pair (proxy2 :: Proxy2 TQueue a) pairTMQueue = pair (proxy2 :: Proxy2 TMQueue a) pairTMChan = pair (proxy2 :: Proxy2 TMChan a) -pairTBQueue, pairTBMQueue, pairTBMChan :: MonadIO m => Int -> m (Source m a, Sink a m ()) +pairTBQueue, pairTBMQueue, pairTBMChan :: MonadIO m => Int -> m (ConduitT () a m (), ConduitT a Void m ()) pairTBQueue = pairBounded (proxy2 :: Proxy2 TBQueue a) pairTBMQueue = pairBounded (proxy2 :: Proxy2 TBMQueue a) pairTBMChan = pairBounded (proxy2 :: Proxy2 TBMChan a) diff --git a/stm-conduit.cabal b/stm-conduit.cabal --- a/stm-conduit.cabal +++ b/stm-conduit.cabal @@ -1,18 +1,19 @@ Name: stm-conduit -Version: 2.5.4 -Synopsis: Introduces conduits to channels, and promotes using - conduits concurrently. -Description: Provides two simple conduit wrappers around STM - channels - a source and a sink. -Homepage: https://github.com/wowus/stm-conduit +Version: 4.0.1 +Synopsis: Introduces conduits to channels, and promotes using conduits concurrently. +Description: Provides two simple conduit wrappers around STM channels - a source and a sink. + +Homepage: https://github.com/cgaebel/stm-conduit License: BSD3 License-file: LICENSE Author: Clark Gaebel -Maintainer: cgaebel@uwaterloo.ca +Maintainer: cg.wowus.cg@gmail.com Category: Concurrency, Conduit Build-type: Simple +tested-with: GHC == 8.0.1, GHC == 8.2.1 + Cabal-version: >=1.8 Library @@ -22,25 +23,35 @@ Data.Conduit.TQueue Data.Conduit.Utils + other-modules: + Data.Conduit.Async.Composition + build-depends: - base == 4.* - , transformers >= 0.2 && < 0.5 - , stm == 2.4.* - , stm-chans >= 2.0 && < 3.1 - , cereal >= 0.4.0.1 - , cereal-conduit >= 0.7.2 - , conduit >= 1.0 && < 1.3 - , conduit-extra >= 1.0 && < 1.2 - , directory >= 1.1 - , resourcet >= 0.3 && < 1.2 - , 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.6 + , stm >= 2.4 && < 2.6 + , stm-chans >= 2.0 && < 3.1 + , cereal >= 0.4.0.1 + , cereal-conduit >= 0.8 + , conduit >= 1.0 && < 1.4 + , conduit-extra >= 1.0 && < 1.4 + , directory >= 1.1 + , exceptions + , resourcet >= 0.3 && < 1.3 + , async >= 2.0.1 + , monad-loops >= 0.4.2 + , unliftio >= 0.2.0 && < 0.3.0 ghc-options: -Wall -fwarn-tabs -fwarn-unused-imports +test-suite stm-conduit-doctests + type: exitcode-stdio-1.0 + main-is: DocTests.hs + ghc-options: -threaded + hs-source-dirs: test/ + ./ + build-depends: base + , doctest test-suite stm-conduit-tests type: exitcode-stdio-1.0 @@ -63,7 +74,8 @@ , stm-chans , resourcet , directory + , unliftio source-repository head type: git - location: git://github.com/wowus/stm-conduit.git + location: git://github.com/cgaebel/stm-conduit.git diff --git a/test/DocTests.hs b/test/DocTests.hs new file mode 100644 --- /dev/null +++ b/test/DocTests.hs @@ -0,0 +1,14 @@ +module Main where + +import Test.DocTest + +main :: IO () +main = doctest [ +-- "-packageghc" +-- , "-isrc" +-- , "-idist/build/autogen/" +-- , "-optP-include" +-- , "-optPdist/build/autogen/cabal_macros.h" +-- , "-cpp" +-- , "Data/Conduit/Async/Composition.hs" + ] diff --git a/test/Test.hs b/test/Test.hs --- a/test/Test.hs +++ b/test/Test.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE ScopedTypeVariables, RankNTypes, FlexibleContexts, BangPatterns #-} + module Main ( main ) where import Data.List (sort) @@ -10,8 +12,8 @@ import Test.HUnit import qualified Control.Monad as Monad -import Control.Monad.Trans.Resource (runResourceT) -import Control.Concurrent (forkIO) +import Control.Monad.Trans.Resource (runResourceT, allocate, release) +import Control.Concurrent (forkIO, threadDelay) import Control.Concurrent.STM import Control.Concurrent.STM.TMQueue import Data.Conduit @@ -20,18 +22,23 @@ import Data.Conduit.TMChan import Data.Conduit.TQueue import System.Directory +import Conduit +import qualified UnliftIO as Lifted main = defaultMain tests tests = [ testGroup "Behaves to spec" [ - testCase "simpleList using TMChan" test_simpleList + testCase "simpleList using TMChan" test_simpleList , testCase "simpleList using TQueue" test_simpleQueue , testCase "simpleList using TMQueue" test_simpleMQueue ], testGroup "Async functions" [ testCase "buffer" test_buffer + , testCase "multiple buffer" test_multi_buffer , testCase "bufferToFile" test_bufferToFile + , testCase "multiple bufferToFile" test_multi_bufferToFile + , testCase "mixed buffer" test_mixed_buffer , testCase "gatherFrom" test_gatherFrom , testCase "drainTo" test_drainTo , testCase "mergeConduits" test_mergeConduits @@ -42,38 +49,65 @@ ] ] -test_simpleList = do chan <- atomically $ newTMChan - forkIO . runResourceT $ sourceList testList $$ sinkTMChan chan True - lst' <- runResourceT $ sourceTMChan chan $$ consume - assertEqual "for the numbers [1..10000]," testList lst' - closed <- atomically $ isClosedTMChan chan - assertBool "channel is closed after running" closed +test_simpleList :: IO () +test_simpleList = + -- We start global resource region because we want to protect + -- the channel. This is not strictly required under some cases + -- but is needed in order to have exception safety. We need + -- to cleanup properly in case of exception, but it depending + -- on the use case it may be implemented by the other means. + runResourceT $ do + (reg, chan) <- allocate (newTMChanIO) + (atomically . closeTMChan) + _ <- Lifted.async $ do + runConduit $ sourceList testList .| sinkTMChan chan + -- We have to release resource explicitly in order to + -- show that no data will be sent here ever. + release reg + + lst' <- runConduit $ sourceTMChan chan .| consume + liftIO $ do + assertEqual "for the numbers [1..10000]," testList lst' + closed <- atomically $ isClosedTMChan chan + assertBool "channel is closed after running" closed where testList = [1..10000] -test_simpleQueue = do q <- atomically $ newTQueue - forkIO . runResourceT $ sourceList testList $$ sinkTQueue q - lst' <- runResourceT $ sourceTQueue q $$ CL.take (length testList) - assertEqual "for the numbers [1..10000]," testList lst' +test_simpleQueue :: IO () +test_simpleQueue = runResourceT $ do + q <- liftIO $ atomically $ newTQueue + liftIO $ forkIO . runConduit $ sourceList testList .| sinkTQueue q + lst' <- runConduit $ sourceTQueue q .| CL.take (length testList) + liftIO $ assertEqual "for the numbers [1..10000]," testList lst' where testList = [1..10000] -test_simpleMQueue = do q <- atomically $ newTMQueue - forkIO . runResourceT $ sourceList testList $$ sinkTMQueue q True - lst' <- runResourceT $ sourceTMQueue q $$ consume - assertEqual "for the numbers [1..10000]," testList lst' - closed <- atomically $ isClosedTMQueue q - assertBool "channel is closed after running" closed - where - testList = [1..10000] +test_simpleMQueue :: IO () +test_simpleMQueue = runResourceT $ do + (reg, q) <- allocate (newTMQueueIO) (atomically . closeTMQueue) + _ <- Lifted.async $ do + runConduit $ sourceList testList .| sinkTMQueue q + release reg + + lst' <- runConduit $ sourceTMQueue q .| consume + liftIO $ do + assertEqual "for the numbers [1..10000]," testList lst' + closed <- atomically $ isClosedTMQueue q + assertBool "channel is closed after running" closed + where + testList = [1..10000] -test_multipleWriters = do ms <- runResourceT $ mergeSources [ sourceList ([1..10]::[Integer]) - , sourceList ([11..20]::[Integer]) - ] 3 - xs <- ms $$ consume - assertEqual "for the numbers [1..10] and [11..20]," [1..20] $ sort xs +test_multipleWriters :: IO () +test_multipleWriters = do + xs <- runResourceT $ do + ms <- mergeSources [ sourceList ([1..10]::[Integer]) + , sourceList ([11..20]::[Integer]) + ] 3 + runConduit $ ms .| consume + assertEqual "for the numbers [1..10] and [11..20]," [1..20] $ sort xs -test_asyncOperator = do sum' <- CL.sourceList [1..n] $$ CL.fold (+) 0 +test_asyncOperator :: IO () +test_asyncOperator = do sum' <- runConduit $ CL.sourceList [1..n] .| CL.fold (+) 0 assertEqual ("for the sum of 1 to " ++ show n) sum sum' sum'' <- CL.sourceList [1..n] $$& CL.fold (+) 0 assertEqual "for the sum computed with the $$ and the $$&" sum' sum'' @@ -82,17 +116,56 @@ n = 100 sum = (n * (n+1)) `div` 2 +test_buffer :: IO () test_buffer = do sum' <- buffer 128 (CL.sourceList [1..100]) (CL.fold (+) 0) assertEqual "sum computed using buffer" sum' (5050 :: Integer) +test_multi_buffer :: IO () +test_multi_buffer = do + sumDoubles <- CL.sourceList [1..100] $$& mapC (* 2) $=& CL.fold (+) 0 + assertEqual "sum of doubles computed using two buffers" sumDoubles (10100 :: Integer) + +-- When we're testing file-buffering, we have to make sure to consume +-- slowly enough to ensure the incoming data piles up enough to be +-- flushed to disk.. +slowDown :: (MonadIO m) => Int -> ConduitT x x m () +slowDown delay = awaitForever $ \x -> do + liftIO $ threadDelay delay + yield x + +aLot, aLittle :: Int +aLot = 10000 +aLittle = 5000 + +test_bufferToFile :: IO () test_bufferToFile = do tempDir <- getTemporaryDirectory - sum' <- runResourceT $ bufferToFile 16 (Just 5) tempDir (CL.sourceList [1 :: Int .. 100]) (CL.fold (+) 0) + sum' <- runResourceT $ bufferToFile 16 (Just 25) tempDir (CL.sourceList [1 :: Int .. 100]) (slowDown aLittle .| CL.fold (+) 0) assertEqual "sum computed using bufferToFile" sum' 5050 +test_multi_bufferToFile :: IO () +test_multi_bufferToFile = do + tempDir <- getTemporaryDirectory + sumDoubles <- let buf c = bufferToFile' 16 (Just 25) tempDir c -- "c" avoids monomorphism restriction + in runResourceT $ runCConduit $ CL.sourceList [1 :: Int .. 100] `buf` (slowDown aLittle .| mapC (* 2)) `buf` (slowDown aLot .| CL.fold (+) 0) + assertEqual "sum of doubles computed using bufferToFile" sumDoubles 10100 + +test_mixed_buffer :: IO () +test_mixed_buffer = do + tempDir <- getTemporaryDirectory + sumDoubles <- + let buf = bufferToFile' 16 (Just 25) tempDir + in runResourceT $ + CL.sourceList [1 :: Int .. 100] $$& mapC (* 2) `buf` (slowDown aLittle .| CL.fold (+) 0) + assertEqual "sum of doubles computed using mixed buffers" sumDoubles 10100 + sumTriples <- let buf = bufferToFile' 16 (Just 25) tempDir + in runResourceT $ CL.sourceList [1 :: Int .. 100] `buf` (slowDown aLittle .| mapC (* 3)) $$& CL.fold (+) 0 + assertEqual "sum of triples computed using mixed buffers" sumTriples 15150 + +test_gatherFrom :: IO () test_gatherFrom = do - sum' <- gatherFrom 128 gen $$ CL.fold (+) 0 + sum' <- runConduit $ gatherFrom 128 gen .| CL.fold (+) 0 assertEqual "sum computed using gatherFrom" sum' 5050 where gen queue = Monad.void $ Monad.foldM f queue [1..100] @@ -101,16 +174,18 @@ atomically $ writeTBQueue q x return q +test_drainTo :: IO () test_drainTo = do - sum' <- CL.sourceList [1..100] $$ drainTo 128 (go 0) + sum' <- runConduit $ CL.sourceList [1..100] .| drainTo 128 (go 0) assertEqual "sum computed using drainTo" sum' 5050 where - go acc queue = do + go !acc queue = do mres <- atomically $ readTBQueue queue case mres of Nothing -> return acc Just res -> go (acc + res) queue +test_mergeConduits :: IO () test_mergeConduits = do merged <- runResourceT $ mergeConduits [ CL.map (* 2) @@ -119,5 +194,5 @@ let input = [1..10] expected = Prelude.map (2 *) input ++ tail (Prelude.scanl (+) 0 input) - xs <- sourceList ([1..10] :: [Integer]) $$ merged =$ consume + xs <- runConduit $ sourceList ([1..10] :: [Integer]) .| merged .| consume assertEqual "merged results" (sort expected) (sort xs)