git.iko.soy/streamly
on 0780dbb help sign in
streamly/src/Streamly/Internal/Data/Stream/Parallel.hs
522 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.Data.Fold (Fold)
68import Streamly.Internal.Control.Concurrent (MonadAsync)
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
73 (foldStreamShared, mkStream, foldStream, fromEffect
74 , nil, concatMapWith, fromPure, bindWith, withLocal)
75import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
76 (Stream(..), mapM, toStreamK, fromStreamK)
77import qualified Streamly.Internal.Data.Stream.SVar.Generate as SVar
78import qualified Streamly.Internal.Data.Stream.SVar.Eliminate as SVar
79import qualified Streamly.Internal.Data.Stream as Stream
80 (toStreamK, fromStreamK)
82import Streamly.Internal.Data.SVar
84#include "inline.hs"
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 K.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 K.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 :: MonadAsync m
169 => SVarStopStyle -> Stream m a -> Stream m a -> Stream m a
170forkSVarPar ss m r = K.mkStream $ \st yld sng stp -> do
171 sv <- newParallelVar ss st
172 pushWorkerPar sv (runOne st{streamVar = Just sv} 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} r)
179 K.foldStream st yld sng stp $ Stream.toStreamK (SVar.fromSVar sv)
181{-# INLINE joinStreamVarPar #-}
182joinStreamVarPar :: MonadAsync m
183 => SVarStyle -> SVarStopStyle -> Stream m a -> Stream m a -> Stream m a
184joinStreamVarPar style ss m1 m2 = K.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 m1)
221 K.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 K.foldStreamShared st yld sng stp (forkSVarPar ss m1 m2)
228-------------------------------------------------------------------------------
229-- User facing APIs
230-------------------------------------------------------------------------------
232{-# INLINE parallelK #-}
233parallelK :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
234parallelK = joinStreamVarPar ParallelVar StopNone
236-- | XXX we can implement it more efficienty by directly implementing instead
237-- of combining streams using parallel.
238{-# INLINE consM #-}
239{-# SPECIALIZE consM :: IO a -> ParallelT IO a -> ParallelT IO a #-}
240consM :: MonadAsync m => m a -> ParallelT m a -> ParallelT m a
241consM m (ParallelT r) = ParallelT $ parallelK (K.fromEffect m) r
243-- This is a co-parallel like combinator for streams, where first stream is the
244-- main stream and the rest are just supporting it, when the first ends
245-- everything ends.
246--
247-- | Like `parallel` but stops the output as soon as the first stream stops.
248--
249-- /Pre-release/
250{-# INLINE parallelFstK #-}
251parallelFstK :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
252parallelFstK = joinStreamVarPar ParallelVar StopBy
254-- This is a race like combinator for streams.
255--
256-- | Like `parallel` but stops the output as soon as any of the two streams
257-- stops.
258--
259-- /Pre-release/
260{-# INLINE parallelMinK #-}
261parallelMinK :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
262parallelMinK = joinStreamVarPar ParallelVar StopAny
264------------------------------------------------------------------------------
265-- Convert a stream to parallel
266------------------------------------------------------------------------------
268-- | Like 'mkParallel' but uses StreamK internally.
269--
270-- /Pre-release/
271--
272mkParallelK :: MonadAsync m => Stream m a -> Stream m a
273mkParallelK m = K.mkStream $ \st yld sng stp -> do
274 sv <- newParallelVar StopNone (adaptState st)
275 -- pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream m)
276 SVar.toSVarParallel st sv $ D.fromStreamK m
277 K.foldStream st yld sng stp $ Stream.toStreamK $ SVar.fromSVar sv
279-- | Same as 'mkParallel' but for StreamD stream.
280--
281{-# INLINE_NORMAL mkParallelD #-}
282mkParallelD :: MonadAsync m => D.Stream m a -> D.Stream m a
283mkParallelD m = D.Stream step Nothing
284 where
286 step gst Nothing = do
287 sv <- newParallelVar StopNone gst
288 SVar.toSVarParallel gst sv m
289 -- XXX use unfold instead?
290 return $ Skip $ Just $ SVar.fromSVarD sv
292 step gst (Just (D.UnStream step1 st)) = do
293 r <- step1 gst st
294 return $ case r of
295 Yield a s -> Yield a (Just $ D.Stream step1 s)
296 Skip s -> Skip (Just $ D.Stream step1 s)
297 Stop -> Stop
299-------------------------------------------------------------------------------
300-- Concurrent tap
301-------------------------------------------------------------------------------
303-- NOTE: In regular pull style streams, the consumer stream is pulling elements
304-- from the SVar and we have several workers producing elements and pushing to
305-- SVar. In case of folds, we, the parent stream driving the fold, are the
306-- stream producing worker, we start an SVar and start pushing to the SVar, the
307-- fold on the other side of the SVar is the consumer stream.
308--
309-- In the pull stream case exceptions are propagated from the producing workers
310-- to the consumer stream, the exceptions are propagated on the same channel as
311-- the produced stream elements. However, in case of push style folds the
312-- current stream itself is the worker and the fold is the consumer, in this
313-- case we have to propagate the exceptions from the consumer to the producer.
314-- This is reverse of the pull case and we need a reverse direction channel
315-- to propagate the exception.
316--
317-- | Redirect a copy of the stream to a supplied fold and run it concurrently
318-- in an independent thread. The fold may buffer some elements. The buffer size
319-- is determined by the prevailing 'Streamly.Prelude.maxBuffer' setting.
320--
321-- @
322-- Stream m a -> m b
323-- |
324-- -----stream m a ---------------stream m a-----
325--
326-- @
327--
328-- @
329-- > S.drain $ S.tapAsync (S.mapM_ print) (S.enumerateFromTo 1 2)
330-- 1
331-- 2
332-- @
333--
334-- Exceptions from the concurrently running fold are propagated to the current
335-- computation. Note that, because of buffering in the fold, exceptions may be
336-- delayed and may not correspond to the current element being processed in the
337-- parent stream, but we guarantee that before the parent stream stops the tap
338-- finishes and all exceptions from it are drained.
339--
340--
341-- Compare with 'tap'.
342--
343-- /Pre-release/
344{-# INLINE tapAsyncK #-}
345tapAsyncK :: MonadAsync m => (Stream m a -> m b) -> Stream m a -> Stream m a
346tapAsyncK f m = K.mkStream $ \st yld sng stp -> do
347 sv <- SVar.newFoldSVar st (f . Stream.toStreamK)
348 K.foldStreamShared st yld sng stp
349 $ Stream.toStreamK (SVar.teeToSVar sv $ Stream.fromStreamK m)
351data TapState fs st a = TapInit | Tapping !fs st | TapDone st
353-- | Like 'tapAsync' but uses a 'Fold' instead of a fold function.
354--
355{-# INLINE_NORMAL tapAsyncF #-}
356tapAsyncF :: MonadAsync m => Fold m a b -> D.Stream m a -> D.Stream m a
357tapAsyncF f (D.Stream step1 state1) = D.Stream step TapInit
358 where
360 drainFold svr = do
361 -- In general, a Stop event would come equipped with the result
362 -- of the fold. It is not used here but it would be useful in
363 -- applicative and distribute.
364 done <- SVar.fromConsumer svr
365 when (not done) $ do
366 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
367 $ takeMVar (outputDoorBellFromConsumer svr)
368 drainFold svr
370 stopFold svr = do
371 liftIO $ sendStop svr Nothing
372 -- drain/wait until a stop event arrives from the fold.
373 drainFold svr
375 {-# INLINE_LATE step #-}
376 step gst TapInit = do
377 sv <- SVar.newFoldSVarF gst f
378 return $ Skip (Tapping sv state1)
380 step gst (Tapping sv st) = do
381 r <- step1 gst st
382 case r of
383 Yield a s -> do
384 done <- SVar.pushToFold sv a
385 if done
386 then do
387 -- XXX we do not need to wait synchronously here
388 stopFold sv
389 return $ Yield a (TapDone s)
390 else return $ Yield a (Tapping sv s)
391 Skip s -> return $ Skip (Tapping sv s)
392 Stop -> do
393 stopFold sv
394 return Stop
396 step gst (TapDone st) = do
397 r <- step1 gst st
398 return $ case r of
399 Yield a s -> Yield a (TapDone s)
400 Skip s -> Skip (TapDone s)
401 Stop -> Stop
403------------------------------------------------------------------------------
404-- ParallelT
405------------------------------------------------------------------------------
407-- | For 'ParallelT' streams:
408--
409-- @
410-- (<>) = 'Streamly.Prelude.parallel'
411-- (>>=) = flip . 'Streamly.Prelude.concatMapWith' 'Streamly.Prelude.parallel'
412-- @
413--
414-- See 'Streamly.Prelude.AsyncT', 'ParallelT' is similar except that all
415-- iterations are strictly concurrent while in 'AsyncT' it depends on the
416-- consumer demand and available threads. See 'parallel' for more details.
417--
418-- /Since: 0.1.0 ("Streamly")/
419--
420-- /Since: 0.7.0 (maxBuffer applies to ParallelT streams)/
421--
422-- @since 0.8.0
423newtype ParallelT m a = ParallelT {getParallelT :: Stream m a}
424 deriving (MonadTrans)
426-- | A parallely composing IO stream of elements of type @a@.
427-- See 'ParallelT' documentation for more details.
428--
429-- /Since: 0.2.0 ("Streamly")/
430--
431-- @since 0.8.0
432type Parallel = ParallelT IO
434------------------------------------------------------------------------------
435-- Semigroup
436------------------------------------------------------------------------------
438{-# INLINE append #-}
439{-# SPECIALIZE append :: ParallelT IO a -> ParallelT IO a -> ParallelT IO a #-}
440append :: MonadAsync m => ParallelT m a -> ParallelT m a -> ParallelT m a
441append (ParallelT m1) (ParallelT m2) = ParallelT $ parallelK m1 m2
443instance MonadAsync m => Semigroup (ParallelT m a) where
444 (<>) = append
446------------------------------------------------------------------------------
447-- Monoid
448------------------------------------------------------------------------------
450instance MonadAsync m => Monoid (ParallelT m a) where
451 mempty = ParallelT K.nil
452 mappend = (<>)
454------------------------------------------------------------------------------
455-- Applicative
456------------------------------------------------------------------------------
458{-# INLINE apParallel #-}
459{-# SPECIALIZE apParallel ::
460 ParallelT IO (a -> b) -> ParallelT IO a -> ParallelT IO b #-}
461apParallel :: MonadAsync m =>
462 ParallelT m (a -> b) -> ParallelT m a -> ParallelT m b
463apParallel (ParallelT m1) (ParallelT m2) =
464 let f x1 = K.concatMapWith parallelK (pure . x1) m2
465 in ParallelT $ K.concatMapWith parallelK f m1
467instance (Monad m, MonadAsync m) => Applicative (ParallelT m) where
468 {-# INLINE pure #-}
469 pure = ParallelT . K.fromPure
471 {-# INLINE (<*>) #-}
472 (<*>) = apParallel
474------------------------------------------------------------------------------
475-- Monad
476------------------------------------------------------------------------------
478{-# INLINE bind #-}
479{-# SPECIALIZE bind ::
480 ParallelT IO a -> (a -> ParallelT IO b) -> ParallelT IO b #-}
481bind :: MonadAsync m => ParallelT m a -> (a -> ParallelT m b) -> ParallelT m b
482bind (ParallelT m) f = ParallelT $ K.bindWith parallelK m (getParallelT . f)
484instance MonadAsync m => Monad (ParallelT m) where
485 return = pure
487 {-# INLINE (>>=) #-}
488 (>>=) = bind
490------------------------------------------------------------------------------
491-- Other instances
492------------------------------------------------------------------------------
494MONAD_COMMON_INSTANCES(ParallelT, MONADPARALLEL)
496-------------------------------------------------------------------------------
497-- From callback
498-------------------------------------------------------------------------------
500-- Note: we can use another API with two callbacks stop and yield if we want
501-- the callback to be able to indicate end of stream.
502--
503-- | Generates a callback and a stream pair. The callback returned is used to
504-- queue values to the stream. The stream is infinite, there is no way for the
505-- callback to indicate that it is done now.
506--
507-- /Pre-release/
508--
509{-# INLINE_NORMAL newCallbackStream #-}
510newCallbackStream :: MonadAsync m => m (a -> m (), Stream m a)
511newCallbackStream = do
512 sv <- newParallelVar StopNone defState
514 -- XXX Add our own thread-id to the SVar as we can not know the callback's
515 -- thread-id and the callback is not run in a managed worker. We need to
516 -- handle this better.
517 liftIO myThreadId >>= modifyThread sv
519 let callback a = liftIO $ void $ send sv (ChildYield a)
520 -- XXX we can return an SVar and then the consumer can unfold from the
521 -- SVar?
522 return (callback, D.toStreamK (SVar.fromSVarD sv))