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