@@ -34,9 +34,9 @@ module Streamly.Internal.Data.Stream.Parallel
37−import Control.Concurrent (myThreadId)
37+import Control.Concurrent (myThreadId, takeMVar)
38−import Control.Exception (SomeException(..), throwIO)
38+import Control.Exception (SomeException(..))
39−import Control.Monad (ap)
39+import Control.Monad (ap, when)
40import Control.Monad.Base (MonadBase(..), liftBaseDefault)
40import Control.Monad.Base (MonadBase(..), liftBaseDefault)
41import Control.Monad.Catch (MonadThrow, throwM)
41import Control.Monad.Catch (MonadThrow, throwM)
42-- import Control.Monad.Error.Class (MonadError(..))
42-- import Control.Monad.Error.Class (MonadError(..))
@@ -54,12 +54,13 @@ import Prelude hiding (map)
55import qualified Data.Set as Set
55import qualified Data.Set as Set
57−import Streamly.Internal.Data.Stream.SVar (fromSVar, fromStreamVar)
57+import Streamly.Internal.Data.Stream.SVar
58+ (fromSVar, fromProducer, fromConsumer)
58import Streamly.Internal.Data.Stream.Serial (map)
59import Streamly.Internal.Data.Stream.Serial (map)
59−import Streamly.Internal.Data.Stream.StreamK (IsStream(..), Stream, mkStream, foldStream,
60+import Streamly.Internal.Data.Stream.StreamK
60− foldStreamShared, adapt)
61+ (IsStream(..), Stream, mkStream, foldStream, foldStreamShared, adapt)
62−import Streamly.Internal.Data.SVar hiding (handleChildException)
63+import Streamly.Internal.Data.SVar
64import qualified Streamly.Internal.Data.Stream.StreamK as K
65import qualified Streamly.Internal.Data.Stream.StreamK as K
65import qualified Streamly.Internal.Data.Stream.StreamD as D
66import qualified Streamly.Internal.Data.Stream.StreamD as D
@@ -274,49 +275,130 @@ foldParallel f m = do
274 f $ fromSVar sv
275 f $ fromSVar sv
276------------------------------------------------------------------------------
277------------------------------------------------------------------------------
277−-- Parallel tee
278+-- Clone and distribute a stream in parallel
278------------------------------------------------------------------------------
279------------------------------------------------------------------------------
281+-- push values to a fold worker via an SVar. Returns whether the fold is done.
282+{-# INLINE pushToFold #-}
283+pushToFold :: MonadAsync m => SVar Stream m a -> a -> m Bool
285+ -- Check for exceptions before decrement so that we do not
286+ -- block forever if the child already exited with an exception.
288+ -- We avoid a race between the consumer fold sending an event and we
289+ -- blocking on decrementBufferLimit by waking up the producer thread in
290+ -- sendToProducer before any event is sent by the fold to the producer
292+ let qref = outputQueueFromConsumer sv
294+ (_, n) <- liftIO $ readIORef qref
296+ then fromConsumer sv
301+ decrementBufferLimit sv
302+ void $ send sv (ChildYield a)
305+-- Tap a stream and send the elements to the specified SVar in addition to
306+-- yielding them again.
308+-- XXX this could be written in StreamD style for better efficiency with fusion.
280{-# INLINE teeToSVar #-}
310{-# INLINE teeToSVar #-}
281teeToSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> t m a
311teeToSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> t m a
282teeToSVar svr m = mkStream $ \st yld sng stp -> do
312teeToSVar svr m = mkStream $ \st yld sng stp -> do
283− foldStreamShared st yld sng stp (go svr m)
313+ foldStreamShared st yld sng stp (go False m)
287− go sv m0 = mkStream $ \st yld sng stp -> do
317+ go False m0 = mkStream $ \st yld _ stp -> do
288− let stop = do
318+ let drain = do
289− liftIO $ do
319+ -- In general, a Stop event would come equipped with the result
290− incrementBufferLimit sv
320+ -- of the fold. It is not used here but it would be useful in
291− sendStop sv Nothing
321+ -- applicative and distribute.
292− -- XXX wait for Stop response event on exception channel,
322+ done <- fromConsumer svr
293− -- drain the exception channel.
323+ when (not done) $ do
294− stp
324+ liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
295− sendit a = liftIO $ do
325+ $ takeMVar (outputDoorBellFromConsumer svr)
296− -- XXX check for exceptions before decrement so that we do not
326+ drain
297− -- block forever if the child already exited with an exception.
327+
298− decrementBufferLimit sv
328+ stopFold = do
299− void $ send sv (ChildYield a)
329+ liftIO $ sendStop svr Nothing
300− single a = sendit a >> (liftIO $ sendStop sv Nothing) >> sng a
330+ -- drain/wait until a stop event arrives from the fold.
333+ stop = stopFold >> stp
335+ done <- pushToFold svr a
336+ yld a (go done (K.nilM stopFold))
301− yieldk a r = sendit a >> yld a (go sv r)
337+ yieldk a r = pushToFold svr a >>= \done -> yld a (go done r)
302 in foldStreamShared st yieldk single stop m0
338 in foldStreamShared st yieldk single stop m0
304-- In case of folds the roles of worker and parent on an SVar are reversed. The
342-- In case of folds the roles of worker and parent on an SVar are reversed. The
305-- parent stream pushes values to an SVar instead of pulling from it and a
343-- parent stream pushes values to an SVar instead of pulling from it and a
306-- worker thread running the fold pulls from the SVar and folds the stream. We
344-- worker thread running the fold pulls from the SVar and folds the stream. We
307-- keep a separate channel for pushing exceptions in the reverse direction i.e.
345-- keep a separate channel for pushing exceptions in the reverse direction i.e.
308-- from the fold to the parent stream.
346-- from the fold to the parent stream.
310−-- NOTE: If we use fromSVar here it will kill the main computation (the parent)
348+-- Note: If we terminate due to an exception, we do not actively terminate the
311−-- when the SVar goes away so we use fromStreamVar instead.
349+-- fold. It gets cleaned up by the GC.
351+{-# NOINLINE handleFoldException #-}
352+handleFoldException :: SVar t m a -> SomeException -> IO ()
353+handleFoldException sv e = do
355+ void $ sendToProducer sv (ChildStop tid (Just e))
357+{-# NOINLINE sendStopToConsumer #-}
358+sendStopToConsumer :: MonadIO m => SVar t m a -> m ()
359+sendStopToConsumer sv = liftIO $ do
361+ void $ sendToProducer sv (ChildStop tid Nothing)
363+-- | Create an SVar with a fold consumer that will fold any elements sent to it
364+-- using the supplied fold function.
365+{-# INLINE newFoldSVar #-}
366+newFoldSVar :: (IsStream t, MonadAsync m)
367+ => State Stream m a -> (t m a -> m b) -> m (SVar Stream m a)
368+newFoldSVar stt f = do
369+ -- Buffer size for the SVar is derived from the current state
370+ sv <- newParallelVar StopAny (adaptState stt)
313−{-# NOINLINE handleChildException #-}
372+ -- Add the producer thread-id to the SVar.
314−handleChildException :: SVar t m a -> SomeException -> IO ()
315−handleChildException _sv e = do
316− -- tid <- myThreadId
373+ liftIO myThreadId >>= modifyThread sv
317− -- void $ sendReverse sv (ChildStop tid (Just e))
375+ -- A wrapper to send a Stop event back to the producer when the fold
376+ -- receives a stop from the producer.
377+ let pull m = mkStream $ \st yld _ stp -> do
378+ let stop = sendStopToConsumer sv >> stp
379+ single a = yld a (K.nilM (sendStopToConsumer sv))
380+ yieldk a r = yld a (pull r)
381+ in foldStreamShared st yieldk single stop m
383+ void $ doFork (void $ f $ fromStream $ pull $ fromProducer sv)
385+ (handleFoldException sv)
388+-- NOTE: In regular pull style streams, the consumer stream is pulling elements
389+-- from the SVar and we have several workers producing elements and pushing to
390+-- SVar. In case of folds, we, the parent stream driving the fold, are the
391+-- stream producing worker, we start an SVar and start pushing to the SVar, the
392+-- fold on the other side of the SVar is the consumer stream.
394+-- In the pull stream case exceptions are propagated from the producing workers
395+-- to the consumer stream, the exceptions are propagated on the same channel as
396+-- the produced stream elements. However, in case of push style folds the
397+-- current stream itself is the worker and the fold is the consumer, in this
398+-- case we have to propagate the exceptions from the consumer to the producer.
399+-- This is reverse of the pull case and we need a reverse direction channel
400+-- to propagate the exception.
320-- | Redirect a copy of the stream to a supplied fold and run it concurrently
402-- | Redirect a copy of the stream to a supplied fold and run it concurrently
321-- in an independent thread. The fold may buffer some elements. The buffer size
403-- in an independent thread. The fold may buffer some elements. The buffer size
322-- is determined by the prevailing 'maxBuffer' setting.
404-- is determined by the prevailing 'maxBuffer' setting.
@@ -337,8 +419,8 @@ handleChildException _sv e = do
337-- Exceptions from the concurrently running fold are propagated to the current
419-- Exceptions from the concurrently running fold are propagated to the current
338-- computation. Note that, because of buffering in the fold, exceptions may be
420-- computation. Note that, because of buffering in the fold, exceptions may be
339-- delayed and may not correspond to the current element being processed in the
421-- delayed and may not correspond to the current element being processed in the
340−-- parent stream, but we guarantee that the tap finishes and all exceptions
422+-- parent stream, but we guarantee that before the parent stream stops the tap
341−-- from it are drained before the parent stream stops.
423+-- finishes and all exceptions from it are drained.
344-- Compare with 'tap'.
426-- Compare with 'tap'.
@@ -347,14 +429,7 @@ handleChildException _sv e = do
347{-# INLINE tapAsync #-}
429{-# INLINE tapAsync #-}
348tapAsync :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> t m a
430tapAsync :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> t m a
349tapAsync f m = mkStream $ \st yld sng stp -> do
431tapAsync f m = mkStream $ \st yld sng stp -> do
350− -- Buffer size for the SVar is derived from the current state
432+ sv <- newFoldSVar st f
351− sv <- newParallelVar StopNone (adaptState st)
352− -- XXX exception handling
353− -- XXX if we terminate due to an exception, do we need to actively
354− -- terminate the fold?
355− liftIO myThreadId >>= modifyThread sv
356− void $ doFork (void $ f $ fromStream $ fromStreamVar sv)
357− (svarMrun sv) (handleChildException sv)
358 foldStreamShared st yld sng stp (teeToSVar sv m)
433 foldStreamShared st yld sng stp (teeToSVar sv m)
360------------------------------------------------------------------------------
435------------------------------------------------------------------------------