on 33b00a6 help sign in
587 lines 23 kB history raw
1{-# LANGUAGE CPP #-}
2{-# LANGUAGE ConstraintKinds #-}
3{-# LANGUAGE FlexibleContexts #-}
4{-# LANGUAGE FlexibleInstances #-}
5{-# LANGUAGE GeneralizedNewtypeDeriving#-}
6{-# LANGUAGE InstanceSigs #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE UndecidableInstances #-} -- XXX
10#include "inline.hs"
12-- |
13-- Module : Streamly.Internal.Data.Stream.Parallel
14-- Copyright : (c) 2017 Harendra Kumar
15--
16-- License : BSD3
17-- Maintainer : streamly@composewell.com
18-- Stability : experimental
19-- Portability : GHC
20--
21--
22module Streamly.Internal.Data.Stream.Parallel
23 (
24 ParallelT
25 , Parallel
26 , parallely
27 , parallel
28 , parallelFst
29 , parallelMin
30 , tapAsync
31 , mkParallel
32 , applyParallel
33 , foldParallel
34 )
35where
37import Control.Concurrent (myThreadId, takeMVar)
38import Control.Exception (SomeException(..))
39import Control.Monad (when)
40import Control.Monad.Base (MonadBase(..), liftBaseDefault)
41import Control.Monad.Catch (MonadThrow, throwM)
42-- import Control.Monad.Error.Class (MonadError(..))
43import Control.Monad.IO.Class (MonadIO(..))
44import Control.Monad.Reader.Class (MonadReader(..))
45import Control.Monad.State.Class (MonadState(..))
46import Control.Monad.Trans.Class (MonadTrans(lift))
47import Data.Functor (void)
48import Data.IORef (readIORef, writeIORef)
49import Data.Maybe (fromJust)
50#if __GLASGOW_HASKELL__ < 808
51import Data.Semigroup (Semigroup(..))
52#endif
53import Prelude hiding (map)
55import qualified Data.Set as Set
57import Streamly.Internal.Data.Stream.SVar
58 (fromSVar, fromProducer, fromConsumer)
59import Streamly.Internal.Data.Stream.StreamK
60 (IsStream(..), Stream, mkStream, foldStream, foldStreamShared, adapt)
62import Streamly.Internal.Data.SVar
64import qualified Streamly.Internal.Data.Stream.StreamK as K
65import qualified Streamly.Internal.Data.Stream.StreamD as D
67#include "Instances.hs"
69-------------------------------------------------------------------------------
70-- Parallel
71-------------------------------------------------------------------------------
73-------------------------------------------------------------------------------
74-- StreamK based worker routines
75-------------------------------------------------------------------------------
77{-# NOINLINE runOne #-}
78runOne
79 :: MonadIO m
80 => State Stream m a -> Stream m a -> Maybe WorkerInfo -> m ()
81runOne st m0 winfo =
82 case getYieldLimit st of
83 Nothing -> go m0
84 Just _ -> runOneLimited st m0 winfo
86 where
88 go m = do
89 liftIO $ decrementBufferLimit sv
90 foldStreamShared st yieldk single stop m
92 sv = fromJust $ streamVar st
94 stop = liftIO $ do
95 incrementBufferLimit sv
96 sendStop sv winfo
97 sendit a = liftIO $ void $ send sv (ChildYield a)
98 single a = sendit a >> (liftIO $ sendStop sv winfo)
99 yieldk a r = sendit a >> go r
101runOneLimited
102 :: MonadIO m
103 => State Stream m a -> Stream m a -> Maybe WorkerInfo -> m ()
104runOneLimited st m0 winfo = go m0
106 where
108 go m = do
109 yieldLimitOk <- liftIO $ decrementYieldLimit sv
110 if yieldLimitOk
111 then do
112 liftIO $ decrementBufferLimit sv
113 foldStreamShared st yieldk single stop m
114 else do
115 liftIO $ cleanupSVarFromWorker sv
116 liftIO $ sendStop sv winfo
118 sv = fromJust $ streamVar st
120 stop = liftIO $ do
121 incrementBufferLimit sv
122 incrementYieldLimit 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
128-------------------------------------------------------------------------------
129-- Consing and appending a stream in parallel style
130-------------------------------------------------------------------------------
132-- Note that consing and appending requires StreamK as it would not scale well
133-- with StreamD unless we are only consing a very small number of streams or
134-- elements in a stream. StreamK allows us to manipulate control flow in a way
135-- which StreamD cannot allow. StreamK can make a jump without having to
136-- remember the past state.
138{-# NOINLINE forkSVarPar #-}
139forkSVarPar :: (IsStream t, MonadAsync m)
140 => SVarStopStyle -> t m a -> t m a -> t m a
141forkSVarPar ss m r = mkStream $ \st yld sng stp -> do
142 sv <- newParallelVar ss st
143 pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream m)
144 case ss of
145 StopBy -> liftIO $ do
146 set <- readIORef (workerThreads sv)
147 writeIORef (svarStopBy sv) $ Set.elemAt 0 set
148 _ -> return ()
149 pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream r)
150 foldStream st yld sng stp (fromSVar sv)
152{-# INLINE joinStreamVarPar #-}
153joinStreamVarPar :: (IsStream t, MonadAsync m)
154 => SVarStyle -> SVarStopStyle -> t m a -> t m a -> t m a
155joinStreamVarPar style ss m1 m2 = mkStream $ \st yld sng stp ->
156 case streamVar st of
157 Just sv | svarStyle sv == style && svarStopStyle sv == ss -> do
158 -- Here, WE ARE IN THE WORKER/PRODUCER THREAD, we know that because
159 -- the SVar exists. We are running under runOne and the output we
160 -- produce ultimately will be sent to the SVar by runOne.
161 --
162 -- If we came here the worker/runOne is evaluating a `parallel`
163 -- combinator. In this case, we always fork a new worker for the
164 -- first component (m1) in the parallel composition and continue to
165 -- evaluate the second component (m2) in the current worker thread.
166 --
167 -- When m1 is serially composed, the worker would evaluate it
168 -- without any further forks and the resulting output is sent to
169 -- the SVar and the evaluation terminates. If m1 is a `parallel`
170 -- composition of two streams the worker would again recurses here.
171 --
172 -- Similarly, when m2 is serially composed it gets evaluated here
173 -- and the resulting output is sent to the SVar by the runOne
174 -- wrapper. When m2 is composed with `parallel` it will again
175 -- recurse here and so on until it finally terminates.
176 --
177 -- When we create a right associated expression using `parallel`,
178 -- then m1 would always terminate without further forks or
179 -- recursion into this routine, therefore, the worker returns
180 -- immediately after evaluating it. And m2 would continue to
181 -- fork/recurse, therefore, the current thread always recurses and
182 -- forks new workers one after the other. This is a tail recursive
183 -- style execution, m2, the recursive expression always executed at
184 -- the tail.
185 --
186 -- When the expression is left associated, the worker spawned would
187 -- get the forking/recursing responsibility and then again the
188 -- worker spawned by that worker would fork, thus creating layer
189 -- over layer of workers and a chain of threads leading to a very
190 -- inefficient execution.
191 pushWorkerPar sv (runOne st $ toStream m1)
192 foldStreamShared st yld sng stp m2
193 _ ->
194 -- Here WE ARE IN THE CONSUMER THREAD, we create a new SVar, fork
195 -- worker threads to execute m1 and m2 and this thread starts
196 -- pulling the stream from the SVar.
197 foldStreamShared st yld sng stp (forkSVarPar ss m1 m2)
199-------------------------------------------------------------------------------
200-- User facing APIs
201-------------------------------------------------------------------------------
203-- | XXX we can implement it more efficienty by directly implementing instead
204-- of combining streams using parallel.
205{-# INLINE consMParallel #-}
206{-# SPECIALIZE consMParallel :: IO a -> ParallelT IO a -> ParallelT IO a #-}
207consMParallel :: MonadAsync m => m a -> ParallelT m a -> ParallelT m a
208consMParallel m r = fromStream $ K.yieldM m `parallel` (toStream r)
210-- | Polymorphic version of the 'Semigroup' operation '<>' of 'ParallelT'
211-- Merges two streams concurrently.
212--
213-- @since 0.2.0
214{-# INLINE parallel #-}
215parallel :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
216parallel = joinStreamVarPar ParallelVar StopNone
218-- This is a co-parallel like combinator for streams, where first stream is the
219-- main stream and the rest are just supporting it, when the first ends
220-- everything ends.
221--
222-- | Like `parallel` but stops the output as soon as the first stream stops.
223--
224-- @since 0.7.0
225{-# INLINE parallelFst #-}
226parallelFst :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
227parallelFst = joinStreamVarPar ParallelVar StopBy
229-- This is a race like combinator for streams.
230--
231-- | Like `parallel` but stops the output as soon as any of the two streams
232-- stops.
233--
234-- @since 0.7.0
235{-# INLINE parallelMin #-}
236parallelMin :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
237parallelMin = joinStreamVarPar ParallelVar StopAny
239------------------------------------------------------------------------------
240-- Convert a stream to parallel
241------------------------------------------------------------------------------
243mkParallel :: (IsStream t, MonadAsync m)
244 => State Stream m a -> t m a -> m (t m a)
245mkParallel st m = do
246 sv <- newParallelVar StopNone defState
247 -- pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream m)
248 D.toSVarParallel st sv $ D.toStreamD m
249 return $ fromSVar sv
251------------------------------------------------------------------------------
252-- Stream to stream concurrent function application
253------------------------------------------------------------------------------
255{-# INLINE applyParallel #-}
256applyParallel :: (IsStream t, MonadAsync m)
257 => (t m a -> t m b) -> t m a -> t m b
258applyParallel f m = mkStream $ \st yld sng stp -> do
259 sv <- newParallelVar StopNone (adaptState st)
260 -- pushWorkerPar sv (runOne st{streamVar = Just sv} (toStream m))
261 D.toSVarParallel (adaptState st) sv $ D.toStreamD m
262 foldStream st yld sng stp $ f $ fromSVar sv
264------------------------------------------------------------------------------
265-- Stream runner concurrent function application
266------------------------------------------------------------------------------
268{-# INLINE foldParallel #-}
269foldParallel :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
270foldParallel f m = do
271 sv <- newParallelVar StopNone defState
272 -- pushWorkerPar sv (runOne defState{streamVar = Just sv} $ toStream m)
273 D.toSVarParallel defState sv $ D.toStreamD m
274 f $ fromSVar sv
276------------------------------------------------------------------------------
277-- Clone and distribute a stream in parallel
278------------------------------------------------------------------------------
280-- push values to a fold worker via an SVar. Returns whether the fold is done.
281{-# INLINE pushToFold #-}
282pushToFold :: MonadAsync m => SVar Stream m a -> a -> m Bool
283pushToFold sv a = do
284 -- Check for exceptions before decrement so that we do not
285 -- block forever if the child already exited with an exception.
286 --
287 -- We avoid a race between the consumer fold sending an event and we
288 -- blocking on decrementBufferLimit by waking up the producer thread in
289 -- sendToProducer before any event is sent by the fold to the producer
290 -- stream.
291 let qref = outputQueueFromConsumer sv
292 done <- do
293 (_, n) <- liftIO $ readIORef qref
294 if (n > 0)
295 then fromConsumer sv
296 else return False
297 if done
298 then return True
299 else liftIO $ do
300 decrementBufferLimit sv
301 void $ send sv (ChildYield a)
302 return False
304-- Tap a stream and send the elements to the specified SVar in addition to
305-- yielding them again.
306--
307-- XXX this could be written in StreamD style for better efficiency with fusion.
308--
309{-# INLINE teeToSVar #-}
310teeToSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> t m a
311teeToSVar svr m = mkStream $ \st yld sng stp -> do
312 foldStreamShared st yld sng stp (go False m)
314 where
316 go False m0 = mkStream $ \st yld _ stp -> do
317 let drain = do
318 -- In general, a Stop event would come equipped with the result
319 -- of the fold. It is not used here but it would be useful in
320 -- applicative and distribute.
321 done <- fromConsumer svr
322 when (not done) $ do
323 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
324 $ takeMVar (outputDoorBellFromConsumer svr)
325 drain
327 stopFold = do
328 liftIO $ sendStop svr Nothing
329 -- drain/wait until a stop event arrives from the fold.
330 drain
332 stop = stopFold >> stp
333 single a = do
334 done <- pushToFold svr a
335 yld a (go done (K.nilM stopFold))
336 yieldk a r = pushToFold svr a >>= \done -> yld a (go done r)
337 in foldStreamShared st yieldk single stop m0
339 go True m0 = m0
341-- In case of folds the roles of worker and parent on an SVar are reversed. The
342-- parent stream pushes values to an SVar instead of pulling from it and a
343-- worker thread running the fold pulls from the SVar and folds the stream. We
344-- keep a separate channel for pushing exceptions in the reverse direction i.e.
345-- from the fold to the parent stream.
346--
347-- Note: If we terminate due to an exception, we do not actively terminate the
348-- fold. It gets cleaned up by the GC.
350{-# NOINLINE handleFoldException #-}
351handleFoldException :: SVar t m a -> SomeException -> IO ()
352handleFoldException sv e = do
353 tid <- myThreadId
354 void $ sendToProducer sv (ChildStop tid (Just e))
356{-# NOINLINE sendStopToConsumer #-}
357sendStopToConsumer :: MonadIO m => SVar t m a -> m ()
358sendStopToConsumer sv = liftIO $ do
359 tid <- myThreadId
360 void $ sendToProducer sv (ChildStop tid Nothing)
362-- | Create an SVar with a fold consumer that will fold any elements sent to it
363-- using the supplied fold function.
364{-# INLINE newFoldSVar #-}
365newFoldSVar :: (IsStream t, MonadAsync m)
366 => State Stream m a -> (t m a -> m b) -> m (SVar Stream m a)
367newFoldSVar stt f = do
368 -- Buffer size for the SVar is derived from the current state
369 sv <- newParallelVar StopAny (adaptState stt)
371 -- Add the producer thread-id to the SVar.
372 liftIO myThreadId >>= modifyThread sv
374 -- A wrapper to send a Stop event back to the producer when the fold
375 -- receives a stop from the producer.
376 let pull m = mkStream $ \st yld _ stp -> do
377 let stop = sendStopToConsumer sv >> stp
378 single a = yld a (K.nilM (sendStopToConsumer sv))
379 yieldk a r = yld a (pull r)
380 in foldStreamShared st yieldk single stop m
382 void $ doFork (void $ f $ fromStream $ pull $ fromProducer sv)
383 (svarMrun sv)
384 (handleFoldException sv)
385 return sv
387-- NOTE: In regular pull style streams, the consumer stream is pulling elements
388-- from the SVar and we have several workers producing elements and pushing to
389-- SVar. In case of folds, we, the parent stream driving the fold, are the
390-- stream producing worker, we start an SVar and start pushing to the SVar, the
391-- fold on the other side of the SVar is the consumer stream.
392--
393-- In the pull stream case exceptions are propagated from the producing workers
394-- to the consumer stream, the exceptions are propagated on the same channel as
395-- the produced stream elements. However, in case of push style folds the
396-- current stream itself is the worker and the fold is the consumer, in this
397-- case we have to propagate the exceptions from the consumer to the producer.
398-- This is reverse of the pull case and we need a reverse direction channel
399-- to propagate the exception.
400--
401-- | Redirect a copy of the stream to a supplied fold and run it concurrently
402-- in an independent thread. The fold may buffer some elements. The buffer size
403-- is determined by the prevailing 'maxBuffer' setting.
404--
405-- @
406-- Stream m a -> m b
407-- |
408-- -----stream m a ---------------stream m a-----
409--
410-- @
411--
412-- @
413-- > S.drain $ S.tapAsync (S.mapM_ print) (S.enumerateFromTo 1 2)
414-- 1
415-- 2
416-- @
417--
418-- Exceptions from the concurrently running fold are propagated to the current
419-- computation. Note that, because of buffering in the fold, exceptions may be
420-- delayed and may not correspond to the current element being processed in the
421-- parent stream, but we guarantee that before the parent stream stops the tap
422-- finishes and all exceptions from it are drained.
423--
424--
425-- Compare with 'tap'.
426--
427-- @since 0.7.0
428{-# INLINE tapAsync #-}
429tapAsync :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> t m a
430tapAsync f m = mkStream $ \st yld sng stp -> do
431 sv <- newFoldSVar st f
432 foldStreamShared st yld sng stp (teeToSVar sv m)
434------------------------------------------------------------------------------
435-- ParallelT
436------------------------------------------------------------------------------
438-- | Async composition with strict concurrent execution of all streams.
439--
440-- The 'Semigroup' instance of 'ParallelT' executes both the streams
441-- concurrently without any delay or without waiting for the consumer demand
442-- and /merges/ the results as they arrive. If the consumer does not consume
443-- the results, they are buffered upto a configured maximum, controlled by the
444-- 'maxBuffer' primitive. If the buffer becomes full the concurrent tasks will
445-- block until there is space in the buffer.
446--
447-- Both 'WAsyncT' and 'ParallelT', evaluate the constituent streams fairly in a
448-- round robin fashion. The key difference is that 'WAsyncT' might wait for the
449-- consumer demand before it executes the tasks whereas 'ParallelT' starts
450-- executing all the tasks immediately without waiting for the consumer demand.
451-- For 'WAsyncT' the 'maxThreads' limit applies whereas for 'ParallelT' it does
452-- not apply. In other words, 'WAsyncT' can be lazy whereas 'ParallelT' is
453-- strict.
454--
455-- 'ParallelT' is useful for cases when the streams are required to be
456-- evaluated simultaneously irrespective of how the consumer consumes them e.g.
457-- when we want to race two tasks and want to start both strictly at the same
458-- time or if we have timers in the parallel tasks and our results depend on
459-- the timers being started at the same time. If we do not have such
460-- requirements then 'AsyncT' or 'AheadT' are recommended as they can be more
461-- efficient than 'ParallelT'.
462--
463-- @
464-- main = ('toList' . 'parallely' $ (fromFoldable [1,2]) \<> (fromFoldable [3,4])) >>= print
465-- @
466-- @
467-- [1,3,2,4]
468-- @
469--
470-- When streams with more than one element are merged, it yields whichever
471-- stream yields first without any bias, unlike the 'Async' style streams.
472--
473-- Any exceptions generated by a constituent stream are propagated to the
474-- output stream. The output and exceptions from a single stream are guaranteed
475-- to arrive in the same order in the resulting stream as they were generated
476-- in the input stream. However, the relative ordering of elements from
477-- different streams in the resulting stream can vary depending on scheduling
478-- and generation delays.
479--
480-- Similarly, the 'Monad' instance of 'ParallelT' runs /all/ iterations
481-- of the loop concurrently.
482--
483-- @
484-- import "Streamly"
485-- import qualified "Streamly.Prelude" as S
486-- import Control.Concurrent
487--
488-- main = 'drain' . 'parallely' $ do
489-- n <- return 3 \<\> return 2 \<\> return 1
490-- S.yieldM $ do
491-- threadDelay (n * 1000000)
492-- myThreadId >>= \\tid -> putStrLn (show tid ++ ": Delay " ++ show n)
493-- @
494-- @
495-- ThreadId 40: Delay 1
496-- ThreadId 39: Delay 2
497-- ThreadId 38: Delay 3
498-- @
499--
500-- Note that parallel composition can only combine a finite number of
501-- streams as it needs to retain state for each unfinished stream.
502--
503-- /Since: 0.7.0 (maxBuffer applies to ParallelT streams)/
504--
505-- /Since: 0.1.0/
506newtype ParallelT m a = ParallelT {getParallelT :: Stream m a}
507 deriving (MonadTrans)
509-- | A parallely composing IO stream of elements of type @a@.
510-- See 'ParallelT' documentation for more details.
511--
512-- @since 0.2.0
513type Parallel = ParallelT IO
515-- | Fix the type of a polymorphic stream as 'ParallelT'.
516--
517-- @since 0.1.0
518parallely :: IsStream t => ParallelT m a -> t m a
519parallely = adapt
521instance IsStream ParallelT where
522 toStream = getParallelT
523 fromStream = ParallelT
525 {-# INLINE consM #-}
526 {-# SPECIALIZE consM :: IO a -> ParallelT IO a -> ParallelT IO a #-}
527 consM = consMParallel
529 {-# INLINE (|:) #-}
530 {-# SPECIALIZE (|:) :: IO a -> ParallelT IO a -> ParallelT IO a #-}
531 (|:) = consM
533------------------------------------------------------------------------------
534-- Semigroup
535------------------------------------------------------------------------------
537{-# INLINE mappendParallel #-}
538{-# SPECIALIZE mappendParallel :: ParallelT IO a -> ParallelT IO a -> ParallelT IO a #-}
539mappendParallel :: MonadAsync m => ParallelT m a -> ParallelT m a -> ParallelT m a
540mappendParallel m1 m2 = fromStream $ parallel (toStream m1) (toStream m2)
542instance MonadAsync m => Semigroup (ParallelT m a) where
543 (<>) = mappendParallel
545------------------------------------------------------------------------------
546-- Monoid
547------------------------------------------------------------------------------
549instance MonadAsync m => Monoid (ParallelT m a) where
550 mempty = K.nil
551 mappend = (<>)
553------------------------------------------------------------------------------
554-- Applicative
555------------------------------------------------------------------------------
557{-# INLINE apParallel #-}
558{-# SPECIALIZE apParallel :: ParallelT IO (a -> b) -> ParallelT IO a -> ParallelT IO b #-}
559apParallel :: MonadAsync m => ParallelT m (a -> b) -> ParallelT m a -> ParallelT m b
560apParallel (ParallelT m1) (ParallelT m2) =
561 let f x1 = K.concatMapBy parallel (pure . x1) m2
562 in ParallelT $ K.concatMapBy parallel f m1
564instance (Monad m, MonadAsync m) => Applicative (ParallelT m) where
565 {-# INLINE pure #-}
566 pure = ParallelT . K.yield
567 {-# INLINE (<*>) #-}
568 (<*>) = apParallel
570------------------------------------------------------------------------------
571-- Monad
572------------------------------------------------------------------------------
574{-# INLINE bindParallel #-}
575{-# SPECIALIZE bindParallel :: ParallelT IO a -> (a -> ParallelT IO b) -> ParallelT IO b #-}
576bindParallel :: MonadAsync m => ParallelT m a -> (a -> ParallelT m b) -> ParallelT m b
577bindParallel m f = fromStream $ K.bindWith parallel (K.adapt m) (\a -> K.adapt $ f a)
579instance MonadAsync m => Monad (ParallelT m) where
580 return = pure
581 (>>=) = bindParallel
583------------------------------------------------------------------------------
584-- Other instances
585------------------------------------------------------------------------------
587MONAD_COMMON_INSTANCES(ParallelT, MONADPARALLEL)