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