on a296fac help sign in
609 lines 23 kB history raw
1{-# LANGUAGE UndecidableInstances #-}
3#include "inline.hs"
5-- |
6-- Module : Streamly.Internal.Data.Stream.Parallel
7-- Copyright : (c) 2017 Composewell Technologies
8--
9-- License : BSD3
10-- Maintainer : streamly@composewell.com
11-- Stability : experimental
12-- Portability : GHC
13--
14-- To run examples in this module:
15--
16-- >>> import qualified Streamly.Prelude as Stream
17-- >>> import Control.Concurrent (threadDelay)
18-- >>> :{
19-- delay n = do
20-- threadDelay (n * 1000000) -- sleep for n seconds
21-- putStrLn (show n ++ " sec") -- print "n sec"
22-- return n -- IO Int
23-- :}
24--
25module Streamly.Internal.Data.Stream.Parallel
26 (
27 -- * Parallel Stream Type
28 ParallelT
29 , Parallel
30 , fromParallel
32 -- * Merge Concurrently
33 , parallel
34 , parallelFst
35 , parallelMin
37 -- * Evaluate Concurrently
38 , mkParallel
39 , mkParallelD
40 , mkParallelK
42 -- * Tap Concurrently
43 , tapAsync
44 , tapAsyncF
45 , distributeAsync_
47 -- * Callbacks
48 , newCallbackStream
49 )
50where
52import Control.Concurrent (myThreadId, takeMVar)
53import Control.Monad (when)
54import Control.Monad.Base (MonadBase(..), liftBaseDefault)
55import Control.Monad.Catch (MonadThrow, throwM)
56-- import Control.Monad.Error.Class (MonadError(..))
57import Control.Monad.IO.Class (MonadIO(..))
58import Control.Monad.Reader.Class (MonadReader(..))
59import Control.Monad.State.Class (MonadState(..))
60import Control.Monad.Trans.Class (MonadTrans(lift))
61import Data.Functor (void)
62import Data.IORef (readIORef, writeIORef)
63import Data.Maybe (fromJust)
64#if __GLASGOW_HASKELL__ < 808
65import Data.Semigroup (Semigroup(..))
66#endif
67import Prelude hiding (map)
69import qualified Data.Set as Set
71import Streamly.Internal.Control.Concurrent (MonadAsync)
72import Streamly.Internal.Data.Fold.Type (Fold)
73import Streamly.Internal.Data.Stream.StreamD.Type (Step(..))
74import Streamly.Internal.Data.Stream.StreamK
75 (IsStream(..), Stream, mkStream, foldStream, foldStreamShared, adapt)
77import qualified Streamly.Internal.Data.Stream.StreamK as K (withLocal)
78import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
79import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
80import qualified Streamly.Internal.Data.Stream.SVar.Generate as SVar
81import qualified Streamly.Internal.Data.Stream.SVar.Eliminate as SVar
83import Streamly.Internal.Data.SVar
85#include "Instances.hs"
87--
88-- $setup
89-- >>> import qualified Streamly.Prelude as Stream
90-- >>> import Control.Concurrent (threadDelay)
91-- >>> :{
92-- delay n = do
93-- threadDelay (n * 1000000) -- sleep for n seconds
94-- putStrLn (show n ++ " sec") -- print "n sec"
95-- return n -- IO Int
96-- :}
98-------------------------------------------------------------------------------
99-- Parallel
100-------------------------------------------------------------------------------
102-------------------------------------------------------------------------------
103-- StreamK based worker routines
104-------------------------------------------------------------------------------
106{-# NOINLINE runOne #-}
107runOne
108 :: MonadIO m
109 => State Stream m a -> Stream m a -> Maybe WorkerInfo -> m ()
110runOne st m0 winfo =
111 case getYieldLimit st of
112 Nothing -> go m0
113 Just _ -> runOneLimited st m0 winfo
115 where
117 go m = do
118 liftIO $ decrementBufferLimit sv
119 foldStreamShared st yieldk single stop m
121 sv = fromJust $ streamVar st
123 stop = liftIO $ do
124 incrementBufferLimit sv
125 sendStop sv winfo
126 sendit a = liftIO $ void $ send sv (ChildYield a)
127 single a = sendit a >> (liftIO $ sendStop sv winfo)
128 yieldk a r = sendit a >> go r
130runOneLimited
131 :: MonadIO m
132 => State Stream m a -> Stream m a -> Maybe WorkerInfo -> m ()
133runOneLimited st m0 winfo = go m0
135 where
137 go m = do
138 yieldLimitOk <- liftIO $ decrementYieldLimit sv
139 if yieldLimitOk
140 then do
141 liftIO $ decrementBufferLimit sv
142 foldStreamShared st yieldk single stop m
143 else do
144 liftIO $ cleanupSVarFromWorker sv
145 liftIO $ sendStop sv winfo
147 sv = fromJust $ streamVar st
149 stop = liftIO $ do
150 incrementBufferLimit sv
151 incrementYieldLimit sv
152 sendStop sv winfo
153 sendit a = liftIO $ void $ send sv (ChildYield a)
154 single a = sendit a >> (liftIO $ sendStop sv winfo)
155 yieldk a r = sendit a >> go r
157-------------------------------------------------------------------------------
158-- Consing and appending a stream in parallel style
159-------------------------------------------------------------------------------
161-- Note that consing and appending requires StreamK as it would not scale well
162-- with StreamD unless we are only consing a very small number of streams or
163-- elements in a stream. StreamK allows us to manipulate control flow in a way
164-- which StreamD cannot allow. StreamK can make a jump without having to
165-- remember the past state.
167{-# NOINLINE forkSVarPar #-}
168forkSVarPar :: (IsStream t, MonadAsync m)
169 => SVarStopStyle -> t m a -> t m a -> t m a
170forkSVarPar ss m r = mkStream $ \st yld sng stp -> do
171 sv <- newParallelVar ss st
172 pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream m)
173 case ss of
174 StopBy -> liftIO $ do
175 set <- readIORef (workerThreads sv)
176 writeIORef (svarStopBy sv) $ Set.elemAt 0 set
177 _ -> return ()
178 pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream r)
179 foldStream st yld sng stp (SVar.fromSVar sv)
181{-# INLINE joinStreamVarPar #-}
182joinStreamVarPar :: (IsStream t, MonadAsync m)
183 => SVarStyle -> SVarStopStyle -> t m a -> t m a -> t m a
184joinStreamVarPar style ss m1 m2 = mkStream $ \st yld sng stp ->
185 case streamVar st of
186 Just sv | svarStyle sv == style && svarStopStyle sv == ss -> do
187 -- Here, WE ARE IN THE WORKER/PRODUCER THREAD, we know that because
188 -- the SVar exists. We are running under runOne and the output we
189 -- produce ultimately will be sent to the SVar by runOne.
190 --
191 -- If we came here the worker/runOne is evaluating a `parallel`
192 -- combinator. In this case, we always fork a new worker for the
193 -- first component (m1) in the parallel composition and continue to
194 -- evaluate the second component (m2) in the current worker thread.
195 --
196 -- When m1 is serially composed, the worker would evaluate it
197 -- without any further forks and the resulting output is sent to
198 -- the SVar and the evaluation terminates. If m1 is a `parallel`
199 -- composition of two streams the worker would again recurses here.
200 --
201 -- Similarly, when m2 is serially composed it gets evaluated here
202 -- and the resulting output is sent to the SVar by the runOne
203 -- wrapper. When m2 is composed with `parallel` it will again
204 -- recurse here and so on until it finally terminates.
205 --
206 -- When we create a right associated expression using `parallel`,
207 -- then m1 would always terminate without further forks or
208 -- recursion into this routine, therefore, the worker returns
209 -- immediately after evaluating it. And m2 would continue to
210 -- fork/recurse, therefore, the current thread always recurses and
211 -- forks new workers one after the other. This is a tail recursive
212 -- style execution, m2, the recursive expression always executed at
213 -- the tail.
214 --
215 -- When the expression is left associated, the worker spawned would
216 -- get the forking/recursing responsibility and then again the
217 -- worker spawned by that worker would fork, thus creating layer
218 -- over layer of workers and a chain of threads leading to a very
219 -- inefficient execution.
220 pushWorkerPar sv (runOne st $ toStream m1)
221 foldStreamShared st yld sng stp m2
222 _ ->
223 -- Here WE ARE IN THE CONSUMER THREAD, we create a new SVar, fork
224 -- worker threads to execute m1 and m2 and this thread starts
225 -- pulling the stream from the SVar.
226 foldStreamShared st yld sng stp (forkSVarPar ss m1 m2)
228-------------------------------------------------------------------------------
229-- User facing APIs
230-------------------------------------------------------------------------------
232-- | XXX we can implement it more efficienty by directly implementing instead
233-- of combining streams using parallel.
234{-# INLINE consMParallel #-}
235{-# SPECIALIZE consMParallel :: IO a -> ParallelT IO a -> ParallelT IO a #-}
236consMParallel :: MonadAsync m => m a -> ParallelT m a -> ParallelT m a
237consMParallel m r = fromStream $ K.fromEffect m `parallel` (toStream r)
239infixr 6 `parallel`
241-- | Like 'Streamly.Prelude.async' except that the execution is much more
242-- strict. There is no limit on the number of threads. While
243-- 'Streamly.Prelude.async' may not schedule a stream if there is no demand
244-- from the consumer, 'parallel' always evaluates both the streams immediately.
245-- The only limit that applies to 'parallel' is 'Streamly.Prelude.maxBuffer'.
246-- Evaluation may block if the output buffer becomes full.
247--
248-- >>> import Streamly.Prelude (parallel)
249-- >>> stream = Stream.fromEffect (delay 2) `parallel` Stream.fromEffect (delay 1)
250-- >>> Stream.toList stream -- IO [Int]
251-- 1 sec
252-- 2 sec
253-- [1,2]
254--
255-- 'parallel' guarantees that all the streams are scheduled for execution
256-- immediately, therefore, we could use things like starting timers inside the
257-- streams and relying on the fact that all timers were started at the same
258-- time.
259--
260-- Unlike 'async' this operation cannot be used to fold an infinite lazy
261-- container of streams, because it schedules all the streams strictly
262-- concurrently.
263--
264-- /Since: 0.2.0 ("Streamly")/
265--
266-- @since 0.8.0
267{-# INLINE parallel #-}
268parallel :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
269parallel = joinStreamVarPar ParallelVar StopNone
271-- This is a co-parallel like combinator for streams, where first stream is the
272-- main stream and the rest are just supporting it, when the first ends
273-- everything ends.
274--
275-- | Like `parallel` but stops the output as soon as the first stream stops.
276--
277-- /Pre-release/
278{-# INLINE parallelFst #-}
279parallelFst :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
280parallelFst = joinStreamVarPar ParallelVar StopBy
282-- This is a race like combinator for streams.
283--
284-- | Like `parallel` but stops the output as soon as any of the two streams
285-- stops.
286--
287-- /Pre-release/
288{-# INLINE parallelMin #-}
289parallelMin :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
290parallelMin = joinStreamVarPar ParallelVar StopAny
292------------------------------------------------------------------------------
293-- Convert a stream to parallel
294------------------------------------------------------------------------------
296-- | Like 'mkParallel' but uses StreamK internally.
297--
298-- /Pre-release/
299--
300mkParallelK :: (IsStream t, MonadAsync m) => t m a -> t m a
301mkParallelK m = mkStream $ \st yld sng stp -> do
302 sv <- newParallelVar StopNone (adaptState st)
303 -- pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream m)
304 SVar.toSVarParallel st sv $ D.toStreamD m
305 foldStream st yld sng stp $ SVar.fromSVar sv
307-------------------------------------------------------------------------------
308-- Concurrent application and fold
309-------------------------------------------------------------------------------
311-- | Same as 'mkParallel' but for StreamD stream.
312--
313{-# INLINE_NORMAL mkParallelD #-}
314mkParallelD :: MonadAsync m => D.Stream m a -> D.Stream m a
315mkParallelD m = D.Stream step Nothing
316 where
318 step gst Nothing = do
319 sv <- newParallelVar StopNone gst
320 SVar.toSVarParallel gst sv m
321 -- XXX use unfold instead?
322 return $ Skip $ Just $ SVar.fromSVarD sv
324 step gst (Just (D.UnStream step1 st)) = do
325 r <- step1 gst st
326 return $ case r of
327 Yield a s -> Yield a (Just $ D.Stream step1 s)
328 Skip s -> Skip (Just $ D.Stream step1 s)
329 Stop -> Stop
331-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
332-- parallel style SVar for evaluation. Currently, parallel style cannot use
333-- rate control whereas Async style can use rate control. In async style SVar
334-- the worker thread terminates when the buffer is full whereas in Parallel
335-- style it blocks.
336--
337-- | Make the stream producer and consumer run concurrently by introducing a
338-- buffer between them. The producer thread evaluates the input stream until
339-- the buffer fills, it blocks if the buffer is full until there is space in
340-- the buffer. The consumer consumes the stream lazily from the buffer.
341--
342-- @mkParallel = D.fromStreamD . mkParallelD . D.toStreamD@
343--
344-- /Pre-release/
345--
346{-# INLINE_NORMAL mkParallel #-}
347mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
348mkParallel = D.fromStreamD . mkParallelD . D.toStreamD
350-------------------------------------------------------------------------------
351-- Concurrent tap
352-------------------------------------------------------------------------------
354-- NOTE: In regular pull style streams, the consumer stream is pulling elements
355-- from the SVar and we have several workers producing elements and pushing to
356-- SVar. In case of folds, we, the parent stream driving the fold, are the
357-- stream producing worker, we start an SVar and start pushing to the SVar, the
358-- fold on the other side of the SVar is the consumer stream.
359--
360-- In the pull stream case exceptions are propagated from the producing workers
361-- to the consumer stream, the exceptions are propagated on the same channel as
362-- the produced stream elements. However, in case of push style folds the
363-- current stream itself is the worker and the fold is the consumer, in this
364-- case we have to propagate the exceptions from the consumer to the producer.
365-- This is reverse of the pull case and we need a reverse direction channel
366-- to propagate the exception.
367--
368-- | Redirect a copy of the stream to a supplied fold and run it concurrently
369-- in an independent thread. The fold may buffer some elements. The buffer size
370-- is determined by the prevailing 'Streamly.Prelude.maxBuffer' setting.
371--
372-- @
373-- Stream m a -> m b
374-- |
375-- -----stream m a ---------------stream m a-----
376--
377-- @
378--
379-- @
380-- > S.drain $ S.tapAsync (S.mapM_ print) (S.enumerateFromTo 1 2)
381-- 1
382-- 2
383-- @
384--
385-- Exceptions from the concurrently running fold are propagated to the current
386-- computation. Note that, because of buffering in the fold, exceptions may be
387-- delayed and may not correspond to the current element being processed in the
388-- parent stream, but we guarantee that before the parent stream stops the tap
389-- finishes and all exceptions from it are drained.
390--
391--
392-- Compare with 'tap'.
393--
394-- /Pre-release/
395{-# INLINE tapAsync #-}
396tapAsync :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> t m a
397tapAsync f m = mkStream $ \st yld sng stp -> do
398 sv <- SVar.newFoldSVar st f
399 foldStreamShared st yld sng stp (SVar.teeToSVar sv m)
401data TapState fs st a = TapInit | Tapping !fs st | TapDone st
403-- | Like 'tapAsync' but uses a 'Fold' instead of a fold function.
404--
405{-# INLINE_NORMAL tapAsyncF #-}
406tapAsyncF :: MonadAsync m => Fold m a b -> D.Stream m a -> D.Stream m a
407tapAsyncF f (D.Stream step1 state1) = D.Stream step TapInit
408 where
410 drainFold svr = do
411 -- In general, a Stop event would come equipped with the result
412 -- of the fold. It is not used here but it would be useful in
413 -- applicative and distribute.
414 done <- SVar.fromConsumer svr
415 when (not done) $ do
416 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
417 $ takeMVar (outputDoorBellFromConsumer svr)
418 drainFold svr
420 stopFold svr = do
421 liftIO $ sendStop svr Nothing
422 -- drain/wait until a stop event arrives from the fold.
423 drainFold svr
425 {-# INLINE_LATE step #-}
426 step gst TapInit = do
427 sv <- SVar.newFoldSVarF gst f
428 return $ Skip (Tapping sv state1)
430 step gst (Tapping sv st) = do
431 r <- step1 gst st
432 case r of
433 Yield a s -> do
434 done <- SVar.pushToFold sv a
435 if done
436 then do
437 -- XXX we do not need to wait synchronously here
438 stopFold sv
439 return $ Yield a (TapDone s)
440 else return $ Yield a (Tapping sv s)
441 Skip s -> return $ Skip (Tapping sv s)
442 Stop -> do
443 stopFold sv
444 return Stop
446 step gst (TapDone st) = do
447 r <- step1 gst st
448 return $ case r of
449 Yield a s -> Yield a (TapDone s)
450 Skip s -> Skip (TapDone s)
451 Stop -> Stop
453-- | Concurrently distribute a stream to a collection of fold functions,
454-- discarding the outputs of the folds.
455--
456-- @
457-- > Stream.drain $ Stream.distributeAsync_ [Stream.mapM_ print, Stream.mapM_ print] (Stream.enumerateFromTo 1 2)
458-- 1
459-- 2
460-- 1
461-- 2
462--
463-- @
464--
465-- @
466-- distributeAsync_ = flip (foldr tapAsync)
467-- @
468--
469-- /Pre-release/
470--
471{-# INLINE distributeAsync_ #-}
472distributeAsync_ :: (Foldable f, IsStream t, MonadAsync m)
473 => f (t m a -> m b) -> t m a -> t m a
474distributeAsync_ = flip (foldr tapAsync)
476------------------------------------------------------------------------------
477-- ParallelT
478------------------------------------------------------------------------------
480-- | For 'ParallelT' streams:
481--
482-- @
483-- (<>) = 'Streamly.Prelude.parallel'
484-- (>>=) = flip . 'Streamly.Prelude.concatMapWith' 'Streamly.Prelude.parallel'
485-- @
486--
487-- See 'Streamly.Prelude.AsyncT', 'ParallelT' is similar except that all
488-- iterations are strictly concurrent while in 'AsyncT' it depends on the
489-- consumer demand and available threads. See 'parallel' for more details.
490--
491-- /Since: 0.1.0 ("Streamly")/
492--
493-- /Since: 0.7.0 (maxBuffer applies to ParallelT streams)/
494--
495-- @since 0.8.0
496newtype ParallelT m a = ParallelT {getParallelT :: Stream m a}
497 deriving (MonadTrans)
499-- | A parallely composing IO stream of elements of type @a@.
500-- See 'ParallelT' documentation for more details.
501--
502-- /Since: 0.2.0 ("Streamly")/
503--
504-- @since 0.8.0
505type Parallel = ParallelT IO
507-- | Fix the type of a polymorphic stream as 'ParallelT'.
508--
509-- /Since: 0.1.0 ("Streamly")/
510--
511-- @since 0.8.0
512fromParallel :: IsStream t => ParallelT m a -> t m a
513fromParallel = adapt
515instance IsStream ParallelT where
516 toStream = getParallelT
517 fromStream = ParallelT
519 {-# INLINE consM #-}
520 {-# SPECIALIZE consM :: IO a -> ParallelT IO a -> ParallelT IO a #-}
521 consM = consMParallel
523 {-# INLINE (|:) #-}
524 {-# SPECIALIZE (|:) :: IO a -> ParallelT IO a -> ParallelT IO a #-}
525 (|:) = consM
527------------------------------------------------------------------------------
528-- Semigroup
529------------------------------------------------------------------------------
531{-# INLINE mappendParallel #-}
532{-# SPECIALIZE mappendParallel :: ParallelT IO a -> ParallelT IO a -> ParallelT IO a #-}
533mappendParallel :: MonadAsync m => ParallelT m a -> ParallelT m a -> ParallelT m a
534mappendParallel m1 m2 = fromStream $ parallel (toStream m1) (toStream m2)
536instance MonadAsync m => Semigroup (ParallelT m a) where
537 (<>) = mappendParallel
539------------------------------------------------------------------------------
540-- Monoid
541------------------------------------------------------------------------------
543instance MonadAsync m => Monoid (ParallelT m a) where
544 mempty = K.nil
545 mappend = (<>)
547------------------------------------------------------------------------------
548-- Applicative
549------------------------------------------------------------------------------
551{-# INLINE apParallel #-}
552{-# SPECIALIZE apParallel :: ParallelT IO (a -> b) -> ParallelT IO a -> ParallelT IO b #-}
553apParallel :: MonadAsync m => ParallelT m (a -> b) -> ParallelT m a -> ParallelT m b
554apParallel (ParallelT m1) (ParallelT m2) =
555 let f x1 = K.concatMapBy parallel (pure . x1) m2
556 in ParallelT $ K.concatMapBy parallel f m1
558instance (Monad m, MonadAsync m) => Applicative (ParallelT m) where
559 {-# INLINE pure #-}
560 pure = ParallelT . K.fromPure
561 {-# INLINE (<*>) #-}
562 (<*>) = apParallel
564------------------------------------------------------------------------------
565-- Monad
566------------------------------------------------------------------------------
568{-# INLINE bindParallel #-}
569{-# SPECIALIZE bindParallel :: ParallelT IO a -> (a -> ParallelT IO b) -> ParallelT IO b #-}
570bindParallel :: MonadAsync m => ParallelT m a -> (a -> ParallelT m b) -> ParallelT m b
571bindParallel m f = fromStream $ K.bindWith parallel (K.adapt m) (\a -> K.adapt $ f a)
573instance MonadAsync m => Monad (ParallelT m) where
574 return = pure
575 (>>=) = bindParallel
577------------------------------------------------------------------------------
578-- Other instances
579------------------------------------------------------------------------------
581MONAD_COMMON_INSTANCES(ParallelT, MONADPARALLEL)
583-------------------------------------------------------------------------------
584-- From callback
585-------------------------------------------------------------------------------
587-- Note: we can use another API with two callbacks stop and yield if we want
588-- the callback to be able to indicate end of stream.
589--
590-- | Generates a callback and a stream pair. The callback returned is used to
591-- queue values to the stream. The stream is infinite, there is no way for the
592-- callback to indicate that it is done now.
593--
594-- /Pre-release/
595--
596{-# INLINE_NORMAL newCallbackStream #-}
597newCallbackStream :: (K.IsStream t, MonadAsync m) => m (a -> m (), t m a)
598newCallbackStream = do
599 sv <- newParallelVar StopNone defState
601 -- XXX Add our own thread-id to the SVar as we can not know the callback's
602 -- thread-id and the callback is not run in a managed worker. We need to
603 -- handle this better.
604 liftIO myThreadId >>= modifyThread sv
606 let callback a = liftIO $ void $ send sv (ChildYield a)
607 -- XXX we can return an SVar and then the consumer can unfold from the
608 -- SVar?
609 return (callback, D.fromStreamD (SVar.fromSVarD sv))