git.iko.soy/streamly
on e3e60ca help sign in
streamly/src/Streamly/Internal/Data/Stream/Parallel.hs
518 lines 19 kB history raw
1{-# LANGUAGE UndecidableInstances #-}
3-- |
4-- Module : Streamly.Internal.Data.Stream.Parallel
5-- Copyright : (c) 2017 Composewell Technologies
6--
7-- License : BSD3
8-- Maintainer : streamly@composewell.com
9-- Stability : experimental
10-- Portability : GHC
11--
12-- To run examples in this module:
13--
14-- >>> import qualified Streamly.Prelude as Stream
15-- >>> import Control.Concurrent (threadDelay)
16-- >>> :{
17-- delay n = do
18-- threadDelay (n * 1000000) -- sleep for n seconds
19-- putStrLn (show n ++ " sec") -- print "n sec"
20-- return n -- IO Int
21-- :}
22--
23module Streamly.Internal.Data.Stream.Parallel
24 (
25 -- * Parallel Stream Type
26 ParallelT(..)
27 , Parallel
28 , consM
30 -- * Merge Concurrently
31 , parallelK
32 , parallelFstK
33 , parallelMinK
35 -- * Evaluate Concurrently
36 , mkParallelD
37 , mkParallelK
39 -- * Tap Concurrently
40 , tapAsyncK
41 , tapAsyncF
43 -- * Callbacks
44 , newCallbackStream
45 )
46where
48import Control.Concurrent (myThreadId, takeMVar)
49import Control.Monad (when)
50import Control.Monad.Base (MonadBase(..), liftBaseDefault)
51import Control.Monad.Catch (MonadThrow, throwM)
52-- import Control.Monad.Error.Class (MonadError(..))
53import Control.Monad.IO.Class (MonadIO(..))
54import Control.Monad.Reader.Class (MonadReader(..))
55import Control.Monad.State.Class (MonadState(..))
56import Control.Monad.Trans.Class (MonadTrans(lift))
57import Data.Functor (void)
58import Data.IORef (readIORef, writeIORef)
59import Data.Maybe (fromJust)
60#if __GLASGOW_HASKELL__ < 808
61import Data.Semigroup (Semigroup(..))
62#endif
63import Prelude hiding (map)
65import qualified Data.Set as Set
67import Streamly.Internal.Control.Concurrent (MonadAsync)
68import Streamly.Internal.Data.Fold.Type (Fold)
69import Streamly.Internal.Data.Stream.StreamD.Type (Step(..))
70import Streamly.Internal.Data.Stream.StreamK.Type (Stream)
72import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
73import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
74import qualified Streamly.Internal.Data.Stream.SVar.Generate as SVar
75import qualified Streamly.Internal.Data.Stream.SVar.Eliminate as SVar
76import qualified Streamly.Internal.Data.Stream.Type as Stream
78import Streamly.Internal.Data.SVar
80#include "inline.hs"
81#include "Instances.hs"
83--
84-- $setup
85-- >>> import qualified Streamly.Prelude as Stream
86-- >>> import Control.Concurrent (threadDelay)
87-- >>> :{
88-- delay n = do
89-- threadDelay (n * 1000000) -- sleep for n seconds
90-- putStrLn (show n ++ " sec") -- print "n sec"
91-- return n -- IO Int
92-- :}
94-------------------------------------------------------------------------------
95-- Parallel
96-------------------------------------------------------------------------------
98-------------------------------------------------------------------------------
99-- StreamK based worker routines
100-------------------------------------------------------------------------------
102{-# NOINLINE runOne #-}
103runOne
104 :: MonadIO m
105 => State Stream m a -> Stream m a -> Maybe WorkerInfo -> m ()
106runOne st m0 winfo =
107 case getYieldLimit st of
108 Nothing -> go m0
109 Just _ -> runOneLimited st m0 winfo
111 where
113 go m = do
114 liftIO $ decrementBufferLimit sv
115 K.foldStreamShared st yieldk single stop m
117 sv = fromJust $ streamVar st
119 stop = liftIO $ do
120 incrementBufferLimit sv
121 sendStop sv winfo
122 sendit a = liftIO $ void $ send sv (ChildYield a)
123 single a = sendit a >> liftIO (sendStop sv winfo)
124 yieldk a r = sendit a >> go r
126runOneLimited
127 :: MonadIO m
128 => State Stream m a -> Stream m a -> Maybe WorkerInfo -> m ()
129runOneLimited st m0 winfo = go m0
131 where
133 go m = do
134 yieldLimitOk <- liftIO $ decrementYieldLimit sv
135 if yieldLimitOk
136 then do
137 liftIO $ decrementBufferLimit sv
138 K.foldStreamShared st yieldk single stop m
139 else do
140 liftIO $ cleanupSVarFromWorker sv
141 liftIO $ sendStop sv winfo
143 sv = fromJust $ streamVar st
145 stop = liftIO $ do
146 incrementBufferLimit sv
147 incrementYieldLimit sv
148 sendStop sv winfo
149 sendit a = liftIO $ void $ send sv (ChildYield a)
150 single a = sendit a >> liftIO (sendStop sv winfo)
151 yieldk a r = sendit a >> go r
153-------------------------------------------------------------------------------
154-- Consing and appending a stream in parallel style
155-------------------------------------------------------------------------------
157-- Note that consing and appending requires StreamK as it would not scale well
158-- with StreamD unless we are only consing a very small number of streams or
159-- elements in a stream. StreamK allows us to manipulate control flow in a way
160-- which StreamD cannot allow. StreamK can make a jump without having to
161-- remember the past state.
163{-# NOINLINE forkSVarPar #-}
164forkSVarPar :: MonadAsync m
165 => SVarStopStyle -> Stream m a -> Stream m a -> Stream m a
166forkSVarPar ss m r = K.mkStream $ \st yld sng stp -> do
167 sv <- newParallelVar ss st
168 pushWorkerPar sv (runOne st{streamVar = Just sv} m)
169 case ss of
170 StopBy -> liftIO $ do
171 set <- readIORef (workerThreads sv)
172 writeIORef (svarStopBy sv) $ Set.elemAt 0 set
173 _ -> return ()
174 pushWorkerPar sv (runOne st{streamVar = Just sv} r)
175 K.foldStream st yld sng stp $ Stream.toStreamK (SVar.fromSVar sv)
177{-# INLINE joinStreamVarPar #-}
178joinStreamVarPar :: MonadAsync m
179 => SVarStyle -> SVarStopStyle -> Stream m a -> Stream m a -> Stream m a
180joinStreamVarPar style ss m1 m2 = K.mkStream $ \st yld sng stp ->
181 case streamVar st of
182 Just sv | svarStyle sv == style && svarStopStyle sv == ss -> do
183 -- Here, WE ARE IN THE WORKER/PRODUCER THREAD, we know that because
184 -- the SVar exists. We are running under runOne and the output we
185 -- produce ultimately will be sent to the SVar by runOne.
186 --
187 -- If we came here the worker/runOne is evaluating a `parallel`
188 -- combinator. In this case, we always fork a new worker for the
189 -- first component (m1) in the parallel composition and continue to
190 -- evaluate the second component (m2) in the current worker thread.
191 --
192 -- When m1 is serially composed, the worker would evaluate it
193 -- without any further forks and the resulting output is sent to
194 -- the SVar and the evaluation terminates. If m1 is a `parallel`
195 -- composition of two streams the worker would again recurses here.
196 --
197 -- Similarly, when m2 is serially composed it gets evaluated here
198 -- and the resulting output is sent to the SVar by the runOne
199 -- wrapper. When m2 is composed with `parallel` it will again
200 -- recurse here and so on until it finally terminates.
201 --
202 -- When we create a right associated expression using `parallel`,
203 -- then m1 would always terminate without further forks or
204 -- recursion into this routine, therefore, the worker returns
205 -- immediately after evaluating it. And m2 would continue to
206 -- fork/recurse, therefore, the current thread always recurses and
207 -- forks new workers one after the other. This is a tail recursive
208 -- style execution, m2, the recursive expression always executed at
209 -- the tail.
210 --
211 -- When the expression is left associated, the worker spawned would
212 -- get the forking/recursing responsibility and then again the
213 -- worker spawned by that worker would fork, thus creating layer
214 -- over layer of workers and a chain of threads leading to a very
215 -- inefficient execution.
216 pushWorkerPar sv (runOne st m1)
217 K.foldStreamShared st yld sng stp m2
218 _ ->
219 -- Here WE ARE IN THE CONSUMER THREAD, we create a new SVar, fork
220 -- worker threads to execute m1 and m2 and this thread starts
221 -- pulling the stream from the SVar.
222 K.foldStreamShared st yld sng stp (forkSVarPar ss m1 m2)
224-------------------------------------------------------------------------------
225-- User facing APIs
226-------------------------------------------------------------------------------
228{-# INLINE parallelK #-}
229parallelK :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
230parallelK = joinStreamVarPar ParallelVar StopNone
232-- | XXX we can implement it more efficienty by directly implementing instead
233-- of combining streams using parallel.
234{-# INLINE consM #-}
235{-# SPECIALIZE consM :: IO a -> ParallelT IO a -> ParallelT IO a #-}
236consM :: MonadAsync m => m a -> ParallelT m a -> ParallelT m a
237consM m (ParallelT r) = ParallelT $ parallelK (K.fromEffect m) r
239-- This is a co-parallel like combinator for streams, where first stream is the
240-- main stream and the rest are just supporting it, when the first ends
241-- everything ends.
242--
243-- | Like `parallel` but stops the output as soon as the first stream stops.
244--
245-- /Pre-release/
246{-# INLINE parallelFstK #-}
247parallelFstK :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
248parallelFstK = joinStreamVarPar ParallelVar StopBy
250-- This is a race like combinator for streams.
251--
252-- | Like `parallel` but stops the output as soon as any of the two streams
253-- stops.
254--
255-- /Pre-release/
256{-# INLINE parallelMinK #-}
257parallelMinK :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
258parallelMinK = joinStreamVarPar ParallelVar StopAny
260------------------------------------------------------------------------------
261-- Convert a stream to parallel
262------------------------------------------------------------------------------
264-- | Like 'mkParallel' but uses StreamK internally.
265--
266-- /Pre-release/
267--
268mkParallelK :: MonadAsync m => Stream m a -> Stream m a
269mkParallelK m = K.mkStream $ \st yld sng stp -> do
270 sv <- newParallelVar StopNone (adaptState st)
271 -- pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream m)
272 SVar.toSVarParallel st sv $ D.fromStreamK m
273 K.foldStream st yld sng stp $ Stream.toStreamK $ SVar.fromSVar sv
275-- | Same as 'mkParallel' but for StreamD stream.
276--
277{-# INLINE_NORMAL mkParallelD #-}
278mkParallelD :: MonadAsync m => D.Stream m a -> D.Stream m a
279mkParallelD m = D.Stream step Nothing
280 where
282 step gst Nothing = do
283 sv <- newParallelVar StopNone gst
284 SVar.toSVarParallel gst sv m
285 -- XXX use unfold instead?
286 return $ Skip $ Just $ SVar.fromSVarD sv
288 step gst (Just (D.UnStream step1 st)) = do
289 r <- step1 gst st
290 return $ case r of
291 Yield a s -> Yield a (Just $ D.Stream step1 s)
292 Skip s -> Skip (Just $ D.Stream step1 s)
293 Stop -> Stop
295-------------------------------------------------------------------------------
296-- Concurrent tap
297-------------------------------------------------------------------------------
299-- NOTE: In regular pull style streams, the consumer stream is pulling elements
300-- from the SVar and we have several workers producing elements and pushing to
301-- SVar. In case of folds, we, the parent stream driving the fold, are the
302-- stream producing worker, we start an SVar and start pushing to the SVar, the
303-- fold on the other side of the SVar is the consumer stream.
304--
305-- In the pull stream case exceptions are propagated from the producing workers
306-- to the consumer stream, the exceptions are propagated on the same channel as
307-- the produced stream elements. However, in case of push style folds the
308-- current stream itself is the worker and the fold is the consumer, in this
309-- case we have to propagate the exceptions from the consumer to the producer.
310-- This is reverse of the pull case and we need a reverse direction channel
311-- to propagate the exception.
312--
313-- | Redirect a copy of the stream to a supplied fold and run it concurrently
314-- in an independent thread. The fold may buffer some elements. The buffer size
315-- is determined by the prevailing 'Streamly.Prelude.maxBuffer' setting.
316--
317-- @
318-- Stream m a -> m b
319-- |
320-- -----stream m a ---------------stream m a-----
321--
322-- @
323--
324-- @
325-- > S.drain $ S.tapAsync (S.mapM_ print) (S.enumerateFromTo 1 2)
326-- 1
327-- 2
328-- @
329--
330-- Exceptions from the concurrently running fold are propagated to the current
331-- computation. Note that, because of buffering in the fold, exceptions may be
332-- delayed and may not correspond to the current element being processed in the
333-- parent stream, but we guarantee that before the parent stream stops the tap
334-- finishes and all exceptions from it are drained.
335--
336--
337-- Compare with 'tap'.
338--
339-- /Pre-release/
340{-# INLINE tapAsyncK #-}
341tapAsyncK :: MonadAsync m => (Stream m a -> m b) -> Stream m a -> Stream m a
342tapAsyncK f m = K.mkStream $ \st yld sng stp -> do
343 sv <- SVar.newFoldSVar st (f . Stream.toStreamK)
344 K.foldStreamShared st yld sng stp
345 $ Stream.toStreamK (SVar.teeToSVar sv $ Stream.fromStreamK m)
347data TapState fs st a = TapInit | Tapping !fs st | TapDone st
349-- | Like 'tapAsync' but uses a 'Fold' instead of a fold function.
350--
351{-# INLINE_NORMAL tapAsyncF #-}
352tapAsyncF :: MonadAsync m => Fold m a b -> D.Stream m a -> D.Stream m a
353tapAsyncF f (D.Stream step1 state1) = D.Stream step TapInit
354 where
356 drainFold svr = do
357 -- In general, a Stop event would come equipped with the result
358 -- of the fold. It is not used here but it would be useful in
359 -- applicative and distribute.
360 done <- SVar.fromConsumer svr
361 when (not done) $ do
362 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
363 $ takeMVar (outputDoorBellFromConsumer svr)
364 drainFold svr
366 stopFold svr = do
367 liftIO $ sendStop svr Nothing
368 -- drain/wait until a stop event arrives from the fold.
369 drainFold svr
371 {-# INLINE_LATE step #-}
372 step gst TapInit = do
373 sv <- SVar.newFoldSVarF gst f
374 return $ Skip (Tapping sv state1)
376 step gst (Tapping sv st) = do
377 r <- step1 gst st
378 case r of
379 Yield a s -> do
380 done <- SVar.pushToFold sv a
381 if done
382 then do
383 -- XXX we do not need to wait synchronously here
384 stopFold sv
385 return $ Yield a (TapDone s)
386 else return $ Yield a (Tapping sv s)
387 Skip s -> return $ Skip (Tapping sv s)
388 Stop -> do
389 stopFold sv
390 return Stop
392 step gst (TapDone st) = do
393 r <- step1 gst st
394 return $ case r of
395 Yield a s -> Yield a (TapDone s)
396 Skip s -> Skip (TapDone s)
397 Stop -> Stop
399------------------------------------------------------------------------------
400-- ParallelT
401------------------------------------------------------------------------------
403-- | For 'ParallelT' streams:
404--
405-- @
406-- (<>) = 'Streamly.Prelude.parallel'
407-- (>>=) = flip . 'Streamly.Prelude.concatMapWith' 'Streamly.Prelude.parallel'
408-- @
409--
410-- See 'Streamly.Prelude.AsyncT', 'ParallelT' is similar except that all
411-- iterations are strictly concurrent while in 'AsyncT' it depends on the
412-- consumer demand and available threads. See 'parallel' for more details.
413--
414-- /Since: 0.1.0 ("Streamly")/
415--
416-- /Since: 0.7.0 (maxBuffer applies to ParallelT streams)/
417--
418-- @since 0.8.0
419newtype ParallelT m a = ParallelT {getParallelT :: Stream m a}
420 deriving (MonadTrans)
422-- | A parallely composing IO stream of elements of type @a@.
423-- See 'ParallelT' documentation for more details.
424--
425-- /Since: 0.2.0 ("Streamly")/
426--
427-- @since 0.8.0
428type Parallel = ParallelT IO
430------------------------------------------------------------------------------
431-- Semigroup
432------------------------------------------------------------------------------
434{-# INLINE append #-}
435{-# SPECIALIZE append :: ParallelT IO a -> ParallelT IO a -> ParallelT IO a #-}
436append :: MonadAsync m => ParallelT m a -> ParallelT m a -> ParallelT m a
437append (ParallelT m1) (ParallelT m2) = ParallelT $ parallelK m1 m2
439instance MonadAsync m => Semigroup (ParallelT m a) where
440 (<>) = append
442------------------------------------------------------------------------------
443-- Monoid
444------------------------------------------------------------------------------
446instance MonadAsync m => Monoid (ParallelT m a) where
447 mempty = ParallelT K.nil
448 mappend = (<>)
450------------------------------------------------------------------------------
451-- Applicative
452------------------------------------------------------------------------------
454{-# INLINE apParallel #-}
455{-# SPECIALIZE apParallel ::
456 ParallelT IO (a -> b) -> ParallelT IO a -> ParallelT IO b #-}
457apParallel :: MonadAsync m =>
458 ParallelT m (a -> b) -> ParallelT m a -> ParallelT m b
459apParallel (ParallelT m1) (ParallelT m2) =
460 let f x1 = K.concatMapWith parallelK (pure . x1) m2
461 in ParallelT $ K.concatMapWith parallelK f m1
463instance (Monad m, MonadAsync m) => Applicative (ParallelT m) where
464 {-# INLINE pure #-}
465 pure = ParallelT . K.fromPure
467 {-# INLINE (<*>) #-}
468 (<*>) = apParallel
470------------------------------------------------------------------------------
471-- Monad
472------------------------------------------------------------------------------
474{-# INLINE bind #-}
475{-# SPECIALIZE bind ::
476 ParallelT IO a -> (a -> ParallelT IO b) -> ParallelT IO b #-}
477bind :: MonadAsync m => ParallelT m a -> (a -> ParallelT m b) -> ParallelT m b
478bind (ParallelT m) f = ParallelT $ K.bindWith parallelK m (getParallelT . f)
480instance MonadAsync m => Monad (ParallelT m) where
481 return = pure
483 {-# INLINE (>>=) #-}
484 (>>=) = bind
486------------------------------------------------------------------------------
487-- Other instances
488------------------------------------------------------------------------------
490MONAD_COMMON_INSTANCES(ParallelT, MONADPARALLEL)
492-------------------------------------------------------------------------------
493-- From callback
494-------------------------------------------------------------------------------
496-- Note: we can use another API with two callbacks stop and yield if we want
497-- the callback to be able to indicate end of stream.
498--
499-- | Generates a callback and a stream pair. The callback returned is used to
500-- queue values to the stream. The stream is infinite, there is no way for the
501-- callback to indicate that it is done now.
502--
503-- /Pre-release/
504--
505{-# INLINE_NORMAL newCallbackStream #-}
506newCallbackStream :: MonadAsync m => m (a -> m (), Stream m a)
507newCallbackStream = do
508 sv <- newParallelVar StopNone defState
510 -- XXX Add our own thread-id to the SVar as we can not know the callback's
511 -- thread-id and the callback is not run in a managed worker. We need to
512 -- handle this better.
513 liftIO myThreadId >>= modifyThread sv
515 let callback a = liftIO $ void $ send sv (ChildYield a)
516 -- XXX we can return an SVar and then the consumer can unfold from the
517 -- SVar?
518 return (callback, D.toStreamK (SVar.fromSVarD sv))