d229a4b · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

implement StreamD style tapAsync for fusion

d229a4bHarendra Kumar7 years ago

src/Streamly/Internal/Data/SVar.hs+14

compare with master

unifiedsplit

@@ -65,12 +65,14 @@ module Streamly.Internal.Data.SVar

65 , sendToProducer

66 , sendYield

67 , sendStop

68+ , sendStopToProducer

68 , enqueueLIFO

69 , enqueueLIFO

69 , enqueueFIFO

70 , enqueueFIFO

70 , enqueueAhead

71 , enqueueAhead

71 , reEnqueueAhead

72 , reEnqueueAhead

72 , pushWorkerPar

73 , pushWorkerPar

73 , handleChildException

74 , handleChildException

75+ , handleFoldException

75 , queueEmptyAhead

77 , queueEmptyAhead

76 , dequeueAhead

78 , dequeueAhead

@@ -1101,6 +1103,12 @@ sendToProducer sv msg = do

1101 sendWithDoorBell (outputQueueFromConsumer sv)

1103 sendWithDoorBell (outputQueueFromConsumer sv)

1102 (outputDoorBellFromConsumer sv) msg

1104 (outputDoorBellFromConsumer sv) msg

1103

1105

1106+-- {-# NOINLINE sendStopToProducer #-}

1107+sendStopToProducer :: MonadIO m => SVar t m a -> m ()

1108+sendStopToProducer sv = liftIO $ do

1109+ tid <- myThreadId

1110+ void $ sendToProducer sv (ChildStop tid Nothing)

1111+

1104workerCollectLatency :: WorkerInfo -> IO (Maybe (Count, NanoSecond64))

1112workerCollectLatency :: WorkerInfo -> IO (Maybe (Count, NanoSecond64))

1105workerCollectLatency winfo = do

1113workerCollectLatency winfo = do

1106 (cnt0, t0) <- readIORef (workerLatencyStart winfo)

1114 (cnt0, t0) <- readIORef (workerLatencyStart winfo)

@@ -1528,6 +1536,12 @@ handleChildException sv e = do

1528 tid <- myThreadId

1536 tid <- myThreadId

1529 void $ send sv (ChildStop tid (Just e))

1537 void $ send sv (ChildStop tid (Just e))

1530

1538

1539+{-# NOINLINE handleFoldException #-}

1540+handleFoldException :: SVar t m a -> SomeException -> IO ()

1541+handleFoldException sv e = do

1542+ tid <- myThreadId

1543+ void $ sendToProducer sv (ChildStop tid (Just e))

1544+

1531{-# NOINLINE recordMaxWorkers #-}

1545{-# NOINLINE recordMaxWorkers #-}

1532recordMaxWorkers :: MonadIO m => SVar t m a -> m ()

1546recordMaxWorkers :: MonadIO m => SVar t m a -> m ()

1533recordMaxWorkers sv = liftIO $ do

1547recordMaxWorkers sv = liftIO $ do

src/Streamly/Internal/Data/Stream/Parallel.hs+5−40

@@ -35,7 +35,6 @@ module Streamly.Internal.Data.Stream.Parallel

35where

37import Control.Concurrent (myThreadId, takeMVar)

38−import Control.Exception (SomeException(..))

39import Control.Monad (when)

38import Control.Monad (when)

40import Control.Monad.Base (MonadBase(..), liftBaseDefault)

39import Control.Monad.Base (MonadBase(..), liftBaseDefault)

41import Control.Monad.Catch (MonadThrow, throwM)

40import Control.Monad.Catch (MonadThrow, throwM)

@@ -55,7 +54,7 @@ import Prelude hiding (map)

55import qualified Data.Set as Set

54import qualified Data.Set as Set

57import Streamly.Internal.Data.Stream.SVar

56import Streamly.Internal.Data.Stream.SVar

58− (fromSVar, fromProducer, fromConsumer)

57+ (fromSVar, fromProducer, fromConsumer, pushToFold)

59import Streamly.Internal.Data.Stream.StreamK

58import Streamly.Internal.Data.Stream.StreamK

60 (IsStream(..), Stream, mkStream, foldStream, foldStreamShared, adapt)

59 (IsStream(..), Stream, mkStream, foldStream, foldStreamShared, adapt)

@@ -277,30 +276,6 @@ foldParallel f m = do

277-- Clone and distribute a stream in parallel

276-- Clone and distribute a stream in parallel

278------------------------------------------------------------------------------

277------------------------------------------------------------------------------

279

278

280−-- push values to a fold worker via an SVar. Returns whether the fold is done.

281−{-# INLINE pushToFold #-}

282−pushToFold :: MonadAsync m => SVar Stream m a -> a -> m Bool

283−pushToFold sv a = do

284− -- Check for exceptions before decrement so that we do not

285− -- block forever if the child already exited with an exception.

286− --

287− -- We avoid a race between the consumer fold sending an event and we

288− -- blocking on decrementBufferLimit by waking up the producer thread in

289− -- sendToProducer before any event is sent by the fold to the producer

290− -- stream.

291− let qref = outputQueueFromConsumer sv

292− done <- do

293− (_, n) <- liftIO $ readIORef qref

294− if (n > 0)

295− then fromConsumer sv

296− else return False

297− if done

298− then return True

299− else liftIO $ do

300− decrementBufferLimit sv

301− void $ send sv (ChildYield a)

302− return False

303−

304-- Tap a stream and send the elements to the specified SVar in addition to

279-- Tap a stream and send the elements to the specified SVar in addition to

305-- yielding them again.

280-- yielding them again.

306--

281--

@@ -347,18 +322,6 @@ teeToSVar svr m = mkStream $ \st yld sng stp -> do

347-- Note: If we terminate due to an exception, we do not actively terminate the

322-- Note: If we terminate due to an exception, we do not actively terminate the

348-- fold. It gets cleaned up by the GC.

323-- fold. It gets cleaned up by the GC.

349

324

350−{-# NOINLINE handleFoldException #-}

351−handleFoldException :: SVar t m a -> SomeException -> IO ()

352−handleFoldException sv e = do

353− tid <- myThreadId

354− void $ sendToProducer sv (ChildStop tid (Just e))

355−

356−{-# NOINLINE sendStopToConsumer #-}

357−sendStopToConsumer :: MonadIO m => SVar t m a -> m ()

358−sendStopToConsumer sv = liftIO $ do

359− tid <- myThreadId

360− void $ sendToProducer sv (ChildStop tid Nothing)

361−

362-- | Create an SVar with a fold consumer that will fold any elements sent to it

325-- | Create an SVar with a fold consumer that will fold any elements sent to it

363-- using the supplied fold function.

326-- using the supplied fold function.

364{-# INLINE newFoldSVar #-}

327{-# INLINE newFoldSVar #-}

@@ -371,11 +334,13 @@ newFoldSVar stt f = do

371 -- Add the producer thread-id to the SVar.

334 -- Add the producer thread-id to the SVar.

372 liftIO myThreadId >>= modifyThread sv

335 liftIO myThreadId >>= modifyThread sv

373

336

337+ -- XXX move this in fromProducer itself.

338+ --

374 -- A wrapper to send a Stop event back to the producer when the fold

339 -- A wrapper to send a Stop event back to the producer when the fold

375 -- receives a stop from the producer.

340 -- receives a stop from the producer.

376 let pull m = mkStream $ \st yld _ stp -> do

341 let pull m = mkStream $ \st yld _ stp -> do

377− let stop = sendStopToConsumer sv >> stp

342+ let stop = sendStopToProducer sv >> stp

378− single a = yld a (K.nilM (sendStopToConsumer sv))

343+ single a = yld a (K.nilM (sendStopToProducer sv))

379 yieldk a r = yld a (pull r)

344 yieldk a r = yld a (pull r)

380 in foldStreamShared st yieldk single stop m

345 in foldStreamShared st yieldk single stop m

381

346

src/Streamly/Internal/Data/Stream/SVar.hs+26−1

@@ -17,11 +17,12 @@ module Streamly.Internal.Data.Stream.SVar

17 , fromProducer

18 , fromConsumer

19 , toSVar

20+ , pushToFold

20 )

21 )

21where

22where

23import Control.Exception (fromException)

24import Control.Exception (fromException)

24−import Control.Monad (when)

25+import Control.Monad (when, void)

25import Control.Monad.Catch (throwM)

26import Control.Monad.Catch (throwM)

26import Control.Monad.IO.Class (MonadIO(liftIO))

27import Control.Monad.IO.Class (MonadIO(liftIO))

27import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

28import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

@@ -177,3 +178,27 @@ fromConsumer sv = do

177 Nothing -> return True

178 Nothing -> return True

178 Just ex -> throwM ex

179 Just ex -> throwM ex

179 ChildYield _ -> error "Bug: fromConsumer: invalid ChildYield event"

180 ChildYield _ -> error "Bug: fromConsumer: invalid ChildYield event"

181+

182+-- push values to a fold worker via an SVar. Returns whether the fold is done.

183+{-# INLINE pushToFold #-}

184+pushToFold :: MonadAsync m => SVar Stream m a -> a -> m Bool

185+pushToFold sv a = do

186+ -- Check for exceptions before decrement so that we do not

187+ -- block forever if the child already exited with an exception.

188+ --

189+ -- We avoid a race between the consumer fold sending an event and we

190+ -- blocking on decrementBufferLimit by waking up the producer thread in

191+ -- sendToProducer before any event is sent by the fold to the producer

192+ -- stream.

193+ let qref = outputQueueFromConsumer sv

194+ done <- do

195+ (_, n) <- liftIO $ readIORef qref

196+ if (n > 0)

197+ then fromConsumer sv

198+ else return False

199+ if done

200+ then return True

201+ else liftIO $ do

202+ decrementBufferLimit sv

203+ void $ send sv (ChildYield a)

204+ return False

src/Streamly/Internal/Data/Stream/StreamD.hs+112

@@ -125,6 +125,7 @@ module Streamly.Internal.Data.Stream.StreamD

125

126 -- ** Specialized Folds

127 , tap

128+ , tapAsync

128 , drain

129 , drain

129 , null

130 , null

130 , head

131 , head

@@ -299,6 +300,7 @@ module Streamly.Internal.Data.Stream.StreamD

299 )

300 )

300where

301where

301

302

303+import Control.Concurrent (myThreadId, takeMVar)

302import Control.Exception (Exception, SomeException, fromException)

304import Control.Exception (Exception, SomeException, fromException)

303import Control.Monad (void, when)

305import Control.Monad (void, when)

304import Control.Monad.Catch (MonadCatch, throwM)

306import Control.Monad.Catch (MonadCatch, throwM)

@@ -336,6 +338,7 @@ import Streamly.Internal.Data.Unfold.Types (Unfold(..))

336

338

337import Streamly.Internal.Data.Stream.StreamD.Type

339import Streamly.Internal.Data.Stream.StreamD.Type

338import Streamly.Internal.Data.SVar

340import Streamly.Internal.Data.SVar

341+import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)

339

342

340import qualified Streamly.Internal.Data.Pipe.Types as Pipe

343import qualified Streamly.Internal.Data.Pipe.Types as Pipe

341import qualified Streamly.Internal.Memory.Array.Types as A

344import qualified Streamly.Internal.Memory.Array.Types as A

@@ -732,6 +735,44 @@ fromSVar svar = Stream step FromSVarInit

732 liftIO $ printSVar sv "SVar Done"

735 liftIO $ printSVar sv "SVar Done"

733 return Stop

736 return Stop

734

737

738+-------------------------------------------------------------------------------

739+-- Process events received by a fold consumer from a stream producer

740+-------------------------------------------------------------------------------

741+

742+{-# INLINE_NORMAL fromProducer #-}

743+fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a

744+fromProducer svar = Stream step (FromSVarRead svar)

745+ where

746+

747+ {-# INLINE_LATE step #-}

748+ step _ (FromSVarRead sv) = do

749+ list <- readOutputQ sv

750+ -- Reversing the output is important to guarantee that we process the

751+ -- outputs in the same order as they were generated by the constituent

752+ -- streams.

753+ return $ Skip $ FromSVarLoop sv (Prelude.reverse list)

754+

755+ step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv

756+ step _ (FromSVarLoop sv (ev : es)) = do

757+ case ev of

758+ ChildYield a -> return $ Yield a (FromSVarLoop sv es)

759+ ChildStop tid e -> do

760+ accountThread sv tid

761+ case e of

762+ Nothing -> do

763+ sendStopToProducer sv

764+ return $ Skip (FromSVarDone sv)

765+ Just _ -> error "Bug: fromProducer: received exception"

766+

767+ step _ (FromSVarDone sv) = do

768+ when (svarInspectMode sv) $ do

769+ t <- liftIO $ getTime Monotonic

770+ liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)

771+ liftIO $ printSVar sv "SVar Done"

772+ return Stop

773+

774+ step _ FromSVarInit = undefined

775+

735-------------------------------------------------------------------------------

776-------------------------------------------------------------------------------

736-- Hoisting the inner monad

777-- Hoisting the inner monad

737-------------------------------------------------------------------------------

778-------------------------------------------------------------------------------

@@ -3732,3 +3773,74 @@ foldParallel f m = do

3732 sv <- newParallelVar StopNone defState

3773 sv <- newParallelVar StopNone defState

3733 toSVarParallel defState sv (toStreamD m)

3774 toSVarParallel defState sv (toStreamD m)

3734 f $ fromStreamD $ fromSVar sv

3775 f $ fromStreamD $ fromSVar sv

3776+

3777+-------------------------------------------------------------------------------

3778+-- Concurrent tap

3779+-------------------------------------------------------------------------------

3780+

3781+-- | Create an SVar with a fold consumer that will fold any elements sent to it

3782+-- using the supplied fold function.

3783+{-# INLINE newFoldSVar #-}

3784+newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)

3785+newFoldSVar stt f = do

3786+ -- Buffer size for the SVar is derived from the current state

3787+ sv <- newParallelVar StopAny (adaptState stt)

3788+ -- Add the producer thread-id to the SVar.

3789+ liftIO myThreadId >>= modifyThread sv

3790+ void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)

3791+ return sv

3792+

3793+ where

3794+

3795+ {-# NOINLINE work #-}

3796+ work sv = void $ runFold f $ fromProducer sv

3797+

3798+data TapState sv st = TapInit | Tapping sv st | TapDone st

3799+

3800+{-# INLINE_NORMAL tapAsync #-}

3801+tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a

3802+tapAsync f (Stream step1 state1) = Stream step TapInit

3803+ where

3804+

3805+ drainFold svr = do

3806+ -- In general, a Stop event would come equipped with the result

3807+ -- of the fold. It is not used here but it would be useful in

3808+ -- applicative and distribute.

3809+ done <- fromConsumer svr

3810+ when (not done) $ do

3811+ liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"

3812+ $ takeMVar (outputDoorBellFromConsumer svr)

3813+ drainFold svr

3814+

3815+ stopFold svr = do

3816+ liftIO $ sendStop svr Nothing

3817+ -- drain/wait until a stop event arrives from the fold.

3818+ drainFold svr

3819+

3820+ {-# INLINE_LATE step #-}

3821+ step gst TapInit = do

3822+ sv <- newFoldSVar gst f

3823+ return $ Skip (Tapping sv state1)

3824+

3825+ step gst (Tapping sv st) = do

3826+ r <- step1 gst st

3827+ case r of

3828+ Yield a s -> do

3829+ done <- pushToFold sv a

3830+ if done

3831+ then do

3832+ -- XXX we do not need to wait synchronously here

3833+ stopFold sv

3834+ return $ Yield a (TapDone s)

3835+ else return $ Yield a (Tapping sv s)

3836+ Skip s -> return $ Skip (Tapping sv s)

3837+ Stop -> do

3838+ stopFold sv

3839+ return $ Stop

3840+

3841+ step gst (TapDone st) = do

3842+ r <- step1 gst st

3843+ return $ case r of

3844+ Yield a s -> Yield a (TapDone s)

3845+ Skip s -> Skip (TapDone s)

3846+ Stop -> Stop

src/Streamly/Internal/Prelude.hs+34−1

@@ -341,7 +341,7 @@ module Streamly.Internal.Prelude

341 -- ** Distributing

342 , trace

343 , tap

344− , Par.tapAsync

344+ , tapAsync

345

346 -- * Windowed Classification

347

@@ -3567,9 +3567,41 @@ reassembleBy = undefined

3567-- Compare with 'trace'.

3568--

3569-- @since 0.7.0

3570+{-# INLINE tap #-}

3570tap :: (IsStream t, Monad m) => FL.Fold m a b -> t m a -> t m a

3571tap :: (IsStream t, Monad m) => FL.Fold m a b -> t m a -> t m a

3571tap f xs = D.fromStreamD $ D.tap f (D.toStreamD xs)

3572tap f xs = D.fromStreamD $ D.tap f (D.toStreamD xs)

3572

3573

3574+-- | Redirect a copy of the stream to a supplied fold and run it concurrently

3575+-- in an independent thread. The fold may buffer some elements. The buffer size

3576+-- is determined by the prevailing 'maxBuffer' setting.

3577+--

3578+-- @

3579+-- Stream m a -> m b

3580+-- |

3581+-- -----stream m a ---------------stream m a-----

3582+--

3583+-- @

3584+--

3585+-- @

3586+-- > S.drain $ S.tapAsync (S.mapM_ print) (S.enumerateFromTo 1 2)

3587+-- 1

3588+-- 2

3589+-- @

3590+--

3591+-- Exceptions from the concurrently running fold are propagated to the current

3592+-- computation. Note that, because of buffering in the fold, exceptions may be

3593+-- delayed and may not correspond to the current element being processed in the

3594+-- parent stream, but we guarantee that before the parent stream stops the tap

3595+-- finishes and all exceptions from it are drained.

3596+--

3597+--

3598+-- Compare with 'tap'.

3599+--

3600+-- /Internal/

3601+{-# INLINE tapAsync #-}

3602+tapAsync :: (IsStream t, MonadAsync m) => FL.Fold m a b -> t m a -> t m a

3603+tapAsync f xs = D.fromStreamD $ D.tapAsync f (D.toStreamD xs)

3604+

3573-- | Apply a monadic function to each element flowing through the stream and

3605-- | Apply a monadic function to each element flowing through the stream and

3574-- discard the results.

3606-- discard the results.

3575--

3607--

@@ -3582,6 +3614,7 @@ tap f xs = D.fromStreamD $ D.tap f (D.toStreamD xs)

3582-- Compare with 'tap'.

3614-- Compare with 'tap'.

3583--

3615--

3584-- @since 0.7.0

3616-- @since 0.7.0

3617+{-# INLINE trace #-}

3585trace :: (IsStream t, MonadAsync m) => (a -> m b) -> t m a -> t m a

3618trace :: (IsStream t, MonadAsync m) => (a -> m b) -> t m a -> t m a

3586trace f = mapM (\x -> void (f x) >> return x)

3619trace f = mapM (\x -> void (f x) >> return x)

3587

3620