@@ -34,9 +34,9 @@ module Streamly.Internal.Data.Stream.Parallel
3434 )
3535where
3636
37−import Control.Concurrent (myThreadId)
38−import Control.Exception (SomeException(..), throwIO)
39−import Control.Monad (ap)
37+import Control.Concurrent (myThreadId, takeMVar)
38+import Control.Exception (SomeException(..))
39+import Control.Monad (ap, when)
4040import Control.Monad.Base (MonadBase(..), liftBaseDefault)
4141import Control.Monad.Catch (MonadThrow, throwM)
4242-- import Control.Monad.Error.Class (MonadError(..))
@@ -54,12 +54,13 @@ import Prelude hiding (map)
5454
5555import qualified Data.Set as Set
5656
57−import Streamly.Internal.Data.Stream.SVar (fromSVar, fromStreamVar)
57+import Streamly.Internal.Data.Stream.SVar
58+ (fromSVar, fromProducer, fromConsumer)
5859import Streamly.Internal.Data.Stream.Serial (map)
59−import Streamly.Internal.Data.Stream.StreamK (IsStream(..), Stream, mkStream, foldStream,
60− foldStreamShared, adapt)
60+import Streamly.Internal.Data.Stream.StreamK
61+ (IsStream(..), Stream, mkStream, foldStream, foldStreamShared, adapt)
6162
62−import Streamly.Internal.Data.SVar hiding (handleChildException)
63+import Streamly.Internal.Data.SVar
6364
6465import qualified Streamly.Internal.Data.Stream.StreamK as K
6566import qualified Streamly.Internal.Data.Stream.StreamD as D
@@ -274,49 +275,130 @@ foldParallel f m = do
274275 f $ fromSVar sv
275276
276277------------------------------------------------------------------------------
277−-- Parallel tee
278+-- Clone and distribute a stream in parallel
278279------------------------------------------------------------------------------
279280
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
284+pushToFold sv a = do
285+ -- Check for exceptions before decrement so that we do not
286+ -- block forever if the child already exited with an exception.
287+ --
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
291+ -- stream.
292+ let qref = outputQueueFromConsumer sv
293+ done <- do
294+ (_, n) <- liftIO $ readIORef qref
295+ if (n > 0)
296+ then fromConsumer sv
297+ else return False
298+ if done
299+ then return True
300+ else liftIO $ do
301+ decrementBufferLimit sv
302+ void $ send sv (ChildYield a)
303+ return False
304+
305+-- Tap a stream and send the elements to the specified SVar in addition to
306+-- yielding them again.
307+--
308+-- XXX this could be written in StreamD style for better efficiency with fusion.
309+--
280310{-# INLINE teeToSVar #-}
281311teeToSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> t m a
282312teeToSVar 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)
284314
285315 where
286316
287− go sv m0 = mkStream $ \st yld sng stp -> do
288− let stop = do
289− liftIO $ do
290− incrementBufferLimit sv
291− sendStop sv Nothing
292− -- XXX wait for Stop response event on exception channel,
293− -- drain the exception channel.
294− stp
295− sendit a = liftIO $ do
296− -- XXX check for exceptions before decrement so that we do not
297− -- block forever if the child already exited with an exception.
298− decrementBufferLimit sv
299− void $ send sv (ChildYield a)
300− single a = sendit a >> (liftIO $ sendStop sv Nothing) >> sng a
301− yieldk a r = sendit a >> yld a (go sv r)
317+ go False m0 = mkStream $ \st yld _ stp -> do
318+ let drain = do
319+ -- In general, a Stop event would come equipped with the result
320+ -- of the fold. It is not used here but it would be useful in
321+ -- applicative and distribute.
322+ done <- fromConsumer svr
323+ when (not done) $ do
324+ liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
325+ $ takeMVar (outputDoorBellFromConsumer svr)
326+ drain
327+
328+ stopFold = do
329+ liftIO $ sendStop svr Nothing
330+ -- drain/wait until a stop event arrives from the fold.
331+ drain
332+
333+ stop = stopFold >> stp
334+ single a = do
335+ done <- pushToFold svr a
336+ yld a (go done (K.nilM stopFold))
337+ yieldk a r = pushToFold svr a >>= \done -> yld a (go done r)
302338 in foldStreamShared st yieldk single stop m0
303339
340+ go True m0 = m0
341+
304342-- In case of folds the roles of worker and parent on an SVar are reversed. The
305343-- parent stream pushes values to an SVar instead of pulling from it and a
306344-- worker thread running the fold pulls from the SVar and folds the stream. We
307345-- keep a separate channel for pushing exceptions in the reverse direction i.e.
308346-- from the fold to the parent stream.
309347--
310−-- NOTE: If we use fromSVar here it will kill the main computation (the parent)
311−-- when the SVar goes away so we use fromStreamVar instead.
348+-- Note: If we terminate due to an exception, we do not actively terminate the
349+-- fold. It gets cleaned up by the GC.
350+
351+{-# NOINLINE handleFoldException #-}
352+handleFoldException :: SVar t m a -> SomeException -> IO ()
353+handleFoldException sv e = do
354+ tid <- myThreadId
355+ void $ sendToProducer sv (ChildStop tid (Just e))
356+
357+{-# NOINLINE sendStopToConsumer #-}
358+sendStopToConsumer :: MonadIO m => SVar t m a -> m ()
359+sendStopToConsumer sv = liftIO $ do
360+ tid <- myThreadId
361+ void $ sendToProducer sv (ChildStop tid Nothing)
362+
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)
312371
313−{-# NOINLINE handleChildException #-}
314−handleChildException :: SVar t m a -> SomeException -> IO ()
315−handleChildException _sv e = do
316− -- tid <- myThreadId
317− -- void $ sendReverse sv (ChildStop tid (Just e))
318− throwIO e
372+ -- Add the producer thread-id to the SVar.
373+ liftIO myThreadId >>= modifyThread sv
319374
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
382+
383+ void $ doFork (void $ f $ fromStream $ pull $ fromProducer sv)
384+ (svarMrun sv)
385+ (handleFoldException sv)
386+ return sv
387+
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.
393+--
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.
401+--
320402-- | Redirect a copy of the stream to a supplied fold and run it concurrently
321403-- in an independent thread. The fold may buffer some elements. The buffer size
322404-- is determined by the prevailing 'maxBuffer' setting.
@@ -337,8 +419,8 @@ handleChildException _sv e = do
337419-- Exceptions from the concurrently running fold are propagated to the current
338420-- computation. Note that, because of buffering in the fold, exceptions may be
339421-- 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
341−-- from it are drained before the parent stream stops.
422+-- parent stream, but we guarantee that before the parent stream stops the tap
423+-- finishes and all exceptions from it are drained.
342424--
343425--
344426-- Compare with 'tap'.
@@ -347,14 +429,7 @@ handleChildException _sv e = do
347429{-# INLINE tapAsync #-}
348430tapAsync :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> t m a
349431tapAsync f m = mkStream $ \st yld sng stp -> do
350− -- Buffer size for the SVar is derived from the current state
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)
432+ sv <- newFoldSVar st f
358433 foldStreamShared st yld sng stp (teeToSVar sv m)
359434
360435------------------------------------------------------------------------------