on b8178cd help sign in
567 lines 22 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 -- * Parallel Stream Type
25 ParallelT
26 , Parallel
27 , parallely
29 -- * Merge Concurrently
30 , parallel
31 , parallelFst
32 , parallelMin
34 -- * Evaluate Concurrently
35 , mkParallel
36 , applyParallel
37 , foldParallel
39 -- * Tap Concurrently
40 , tapAsync
41 , distributeAsync_
43 )
44where
46import Control.Concurrent (myThreadId, takeMVar)
47import Control.Monad (when)
48import Control.Monad.Base (MonadBase(..), liftBaseDefault)
49import Control.Monad.Catch (MonadThrow, throwM)
50-- import Control.Monad.Error.Class (MonadError(..))
51import Control.Monad.IO.Class (MonadIO(..))
52import Control.Monad.Reader.Class (MonadReader(..))
53import Control.Monad.State.Class (MonadState(..))
54import Control.Monad.Trans.Class (MonadTrans(lift))
55import Data.Functor (void)
56import Data.IORef (readIORef, writeIORef)
57import Data.Maybe (fromJust)
58#if __GLASGOW_HASKELL__ < 808
59import Data.Semigroup (Semigroup(..))
60#endif
61import Prelude hiding (map)
63import qualified Data.Set as Set
65import Streamly.Internal.Data.Stream.SVar
66 (fromSVar, fromProducer, fromConsumer, pushToFold)
67import Streamly.Internal.Data.Stream.StreamK
68 (IsStream(..), Stream, mkStream, foldStream, foldStreamShared, adapt)
70import Streamly.Internal.Data.SVar
72import qualified Streamly.Internal.Data.Stream.StreamK as K
73import qualified Streamly.Internal.Data.Stream.StreamD as D
75#include "Instances.hs"
77-------------------------------------------------------------------------------
78-- Parallel
79-------------------------------------------------------------------------------
81-------------------------------------------------------------------------------
82-- StreamK based worker routines
83-------------------------------------------------------------------------------
85{-# NOINLINE runOne #-}
86runOne
87 :: MonadIO m
88 => State Stream m a -> Stream m a -> Maybe WorkerInfo -> m ()
89runOne st m0 winfo =
90 case getYieldLimit st of
91 Nothing -> go m0
92 Just _ -> runOneLimited st m0 winfo
94 where
96 go m = do
97 liftIO $ decrementBufferLimit sv
98 foldStreamShared st yieldk single stop m
100 sv = fromJust $ streamVar st
102 stop = liftIO $ do
103 incrementBufferLimit sv
104 sendStop sv winfo
105 sendit a = liftIO $ void $ send sv (ChildYield a)
106 single a = sendit a >> (liftIO $ sendStop sv winfo)
107 yieldk a r = sendit a >> go r
109runOneLimited
110 :: MonadIO m
111 => State Stream m a -> Stream m a -> Maybe WorkerInfo -> m ()
112runOneLimited st m0 winfo = go m0
114 where
116 go m = do
117 yieldLimitOk <- liftIO $ decrementYieldLimit sv
118 if yieldLimitOk
119 then do
120 liftIO $ decrementBufferLimit sv
121 foldStreamShared st yieldk single stop m
122 else do
123 liftIO $ cleanupSVarFromWorker sv
124 liftIO $ sendStop sv winfo
126 sv = fromJust $ streamVar st
128 stop = liftIO $ do
129 incrementBufferLimit sv
130 incrementYieldLimit sv
131 sendStop sv winfo
132 sendit a = liftIO $ void $ send sv (ChildYield a)
133 single a = sendit a >> (liftIO $ sendStop sv winfo)
134 yieldk a r = sendit a >> go r
136-------------------------------------------------------------------------------
137-- Consing and appending a stream in parallel style
138-------------------------------------------------------------------------------
140-- Note that consing and appending requires StreamK as it would not scale well
141-- with StreamD unless we are only consing a very small number of streams or
142-- elements in a stream. StreamK allows us to manipulate control flow in a way
143-- which StreamD cannot allow. StreamK can make a jump without having to
144-- remember the past state.
146{-# NOINLINE forkSVarPar #-}
147forkSVarPar :: (IsStream t, MonadAsync m)
148 => SVarStopStyle -> t m a -> t m a -> t m a
149forkSVarPar ss m r = mkStream $ \st yld sng stp -> do
150 sv <- newParallelVar ss st
151 pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream m)
152 case ss of
153 StopBy -> liftIO $ do
154 set <- readIORef (workerThreads sv)
155 writeIORef (svarStopBy sv) $ Set.elemAt 0 set
156 _ -> return ()
157 pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream r)
158 foldStream st yld sng stp (fromSVar sv)
160{-# INLINE joinStreamVarPar #-}
161joinStreamVarPar :: (IsStream t, MonadAsync m)
162 => SVarStyle -> SVarStopStyle -> t m a -> t m a -> t m a
163joinStreamVarPar style ss m1 m2 = mkStream $ \st yld sng stp ->
164 case streamVar st of
165 Just sv | svarStyle sv == style && svarStopStyle sv == ss -> do
166 -- Here, WE ARE IN THE WORKER/PRODUCER THREAD, we know that because
167 -- the SVar exists. We are running under runOne and the output we
168 -- produce ultimately will be sent to the SVar by runOne.
169 --
170 -- If we came here the worker/runOne is evaluating a `parallel`
171 -- combinator. In this case, we always fork a new worker for the
172 -- first component (m1) in the parallel composition and continue to
173 -- evaluate the second component (m2) in the current worker thread.
174 --
175 -- When m1 is serially composed, the worker would evaluate it
176 -- without any further forks and the resulting output is sent to
177 -- the SVar and the evaluation terminates. If m1 is a `parallel`
178 -- composition of two streams the worker would again recurses here.
179 --
180 -- Similarly, when m2 is serially composed it gets evaluated here
181 -- and the resulting output is sent to the SVar by the runOne
182 -- wrapper. When m2 is composed with `parallel` it will again
183 -- recurse here and so on until it finally terminates.
184 --
185 -- When we create a right associated expression using `parallel`,
186 -- then m1 would always terminate without further forks or
187 -- recursion into this routine, therefore, the worker returns
188 -- immediately after evaluating it. And m2 would continue to
189 -- fork/recurse, therefore, the current thread always recurses and
190 -- forks new workers one after the other. This is a tail recursive
191 -- style execution, m2, the recursive expression always executed at
192 -- the tail.
193 --
194 -- When the expression is left associated, the worker spawned would
195 -- get the forking/recursing responsibility and then again the
196 -- worker spawned by that worker would fork, thus creating layer
197 -- over layer of workers and a chain of threads leading to a very
198 -- inefficient execution.
199 pushWorkerPar sv (runOne st $ toStream m1)
200 foldStreamShared st yld sng stp m2
201 _ ->
202 -- Here WE ARE IN THE CONSUMER THREAD, we create a new SVar, fork
203 -- worker threads to execute m1 and m2 and this thread starts
204 -- pulling the stream from the SVar.
205 foldStreamShared st yld sng stp (forkSVarPar ss m1 m2)
207-------------------------------------------------------------------------------
208-- User facing APIs
209-------------------------------------------------------------------------------
211-- | XXX we can implement it more efficienty by directly implementing instead
212-- of combining streams using parallel.
213{-# INLINE consMParallel #-}
214{-# SPECIALIZE consMParallel :: IO a -> ParallelT IO a -> ParallelT IO a #-}
215consMParallel :: MonadAsync m => m a -> ParallelT m a -> ParallelT m a
216consMParallel m r = fromStream $ K.yieldM m `parallel` (toStream r)
218-- | Polymorphic version of the 'Semigroup' operation '<>' of 'ParallelT'
219-- Merges two streams concurrently.
220--
221-- @since 0.2.0
222{-# INLINE parallel #-}
223parallel :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
224parallel = joinStreamVarPar ParallelVar StopNone
226-- This is a co-parallel like combinator for streams, where first stream is the
227-- main stream and the rest are just supporting it, when the first ends
228-- everything ends.
229--
230-- | Like `parallel` but stops the output as soon as the first stream stops.
231--
232-- @since 0.7.0
233{-# INLINE parallelFst #-}
234parallelFst :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
235parallelFst = joinStreamVarPar ParallelVar StopBy
237-- This is a race like combinator for streams.
238--
239-- | Like `parallel` but stops the output as soon as any of the two streams
240-- stops.
241--
242-- @since 0.7.0
243{-# INLINE parallelMin #-}
244parallelMin :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a
245parallelMin = joinStreamVarPar ParallelVar StopAny
247------------------------------------------------------------------------------
248-- Convert a stream to parallel
249------------------------------------------------------------------------------
251mkParallel :: (IsStream t, MonadAsync m)
252 => State Stream m a -> t m a -> m (t m a)
253mkParallel st m = do
254 sv <- newParallelVar StopNone defState
255 -- pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream m)
256 D.toSVarParallel st sv $ D.toStreamD m
257 return $ fromSVar sv
259------------------------------------------------------------------------------
260-- Stream to stream concurrent function application
261------------------------------------------------------------------------------
263{-# INLINE applyParallel #-}
264applyParallel :: (IsStream t, MonadAsync m)
265 => (t m a -> t m b) -> t m a -> t m b
266applyParallel f m = mkStream $ \st yld sng stp -> do
267 sv <- newParallelVar StopNone (adaptState st)
268 -- pushWorkerPar sv (runOne st{streamVar = Just sv} (toStream m))
269 D.toSVarParallel (adaptState st) sv $ D.toStreamD m
270 foldStream st yld sng stp $ f $ fromSVar sv
272------------------------------------------------------------------------------
273-- Stream runner concurrent function application
274------------------------------------------------------------------------------
276{-# INLINE foldParallel #-}
277foldParallel :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
278foldParallel f m = do
279 sv <- newParallelVar StopNone defState
280 -- pushWorkerPar sv (runOne defState{streamVar = Just sv} $ toStream m)
281 D.toSVarParallel defState sv $ D.toStreamD m
282 f $ fromSVar sv
284------------------------------------------------------------------------------
285-- Clone and distribute a stream in parallel
286------------------------------------------------------------------------------
288-- Tap a stream and send the elements to the specified SVar in addition to
289-- yielding them again.
290--
291-- XXX this could be written in StreamD style for better efficiency with fusion.
292--
293{-# INLINE teeToSVar #-}
294teeToSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> t m a
295teeToSVar svr m = mkStream $ \st yld sng stp -> do
296 foldStreamShared st yld sng stp (go False m)
298 where
300 go False m0 = mkStream $ \st yld _ stp -> do
301 let drain = do
302 -- In general, a Stop event would come equipped with the result
303 -- of the fold. It is not used here but it would be useful in
304 -- applicative and distribute.
305 done <- fromConsumer svr
306 when (not done) $ do
307 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
308 $ takeMVar (outputDoorBellFromConsumer svr)
309 drain
311 stopFold = do
312 liftIO $ sendStop svr Nothing
313 -- drain/wait until a stop event arrives from the fold.
314 drain
316 stop = stopFold >> stp
317 single a = do
318 done <- pushToFold svr a
319 yld a (go done (K.nilM stopFold))
320 yieldk a r = pushToFold svr a >>= \done -> yld a (go done r)
321 in foldStreamShared st yieldk single stop m0
323 go True m0 = m0
325-- In case of folds the roles of worker and parent on an SVar are reversed. The
326-- parent stream pushes values to an SVar instead of pulling from it and a
327-- worker thread running the fold pulls from the SVar and folds the stream. We
328-- keep a separate channel for pushing exceptions in the reverse direction i.e.
329-- from the fold to the parent stream.
330--
331-- Note: If we terminate due to an exception, we do not actively terminate the
332-- fold. It gets cleaned up by the GC.
334-- | Create an SVar with a fold consumer that will fold any elements sent to it
335-- using the supplied fold function.
336{-# INLINE newFoldSVar #-}
337newFoldSVar :: (IsStream t, MonadAsync m)
338 => State Stream m a -> (t m a -> m b) -> m (SVar Stream m a)
339newFoldSVar stt f = do
340 -- Buffer size for the SVar is derived from the current state
341 sv <- newParallelVar StopAny (adaptState stt)
343 -- Add the producer thread-id to the SVar.
344 liftIO myThreadId >>= modifyThread sv
346 void $ doFork (void $ f $ fromStream $ fromProducer sv)
347 (svarMrun sv)
348 (handleFoldException sv)
349 return sv
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 '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-- @since 0.7.0
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 <- newFoldSVar st f
396 foldStreamShared st yld sng stp (teeToSVar sv m)
398-- | Concurrently distribute a stream to a collection of fold functions,
399-- discarding the outputs of the folds.
400--
401-- >>> S.drain $ distributeAsync_ [S.mapM_ print, S.mapM_ print] (S.enumerateFromTo 1 2)
402--
403-- @
404-- distributeAsync_ = flip (foldr tapAsync)
405-- @
406--
407-- /Internal/
408--
409{-# INLINE distributeAsync_ #-}
410distributeAsync_ :: (Foldable f, IsStream t, MonadAsync m)
411 => f (t m a -> m b) -> t m a -> t m a
412distributeAsync_ = flip (foldr tapAsync)
414------------------------------------------------------------------------------
415-- ParallelT
416------------------------------------------------------------------------------
418-- | Async composition with strict concurrent execution of all streams.
419--
420-- The 'Semigroup' instance of 'ParallelT' executes both the streams
421-- concurrently without any delay or without waiting for the consumer demand
422-- and /merges/ the results as they arrive. If the consumer does not consume
423-- the results, they are buffered upto a configured maximum, controlled by the
424-- 'maxBuffer' primitive. If the buffer becomes full the concurrent tasks will
425-- block until there is space in the buffer.
426--
427-- Both 'WAsyncT' and 'ParallelT', evaluate the constituent streams fairly in a
428-- round robin fashion. The key difference is that 'WAsyncT' might wait for the
429-- consumer demand before it executes the tasks whereas 'ParallelT' starts
430-- executing all the tasks immediately without waiting for the consumer demand.
431-- For 'WAsyncT' the 'maxThreads' limit applies whereas for 'ParallelT' it does
432-- not apply. In other words, 'WAsyncT' can be lazy whereas 'ParallelT' is
433-- strict.
434--
435-- 'ParallelT' is useful for cases when the streams are required to be
436-- evaluated simultaneously irrespective of how the consumer consumes them e.g.
437-- when we want to race two tasks and want to start both strictly at the same
438-- time or if we have timers in the parallel tasks and our results depend on
439-- the timers being started at the same time. If we do not have such
440-- requirements then 'AsyncT' or 'AheadT' are recommended as they can be more
441-- efficient than 'ParallelT'.
442--
443-- @
444-- main = ('toList' . 'parallely' $ (fromFoldable [1,2]) \<> (fromFoldable [3,4])) >>= print
445-- @
446-- @
447-- [1,3,2,4]
448-- @
449--
450-- When streams with more than one element are merged, it yields whichever
451-- stream yields first without any bias, unlike the 'Async' style streams.
452--
453-- Any exceptions generated by a constituent stream are propagated to the
454-- output stream. The output and exceptions from a single stream are guaranteed
455-- to arrive in the same order in the resulting stream as they were generated
456-- in the input stream. However, the relative ordering of elements from
457-- different streams in the resulting stream can vary depending on scheduling
458-- and generation delays.
459--
460-- Similarly, the 'Monad' instance of 'ParallelT' runs /all/ iterations
461-- of the loop concurrently.
462--
463-- @
464-- import "Streamly"
465-- import qualified "Streamly.Prelude" as S
466-- import Control.Concurrent
467--
468-- main = 'drain' . 'parallely' $ do
469-- n <- return 3 \<\> return 2 \<\> return 1
470-- S.yieldM $ do
471-- threadDelay (n * 1000000)
472-- myThreadId >>= \\tid -> putStrLn (show tid ++ ": Delay " ++ show n)
473-- @
474-- @
475-- ThreadId 40: Delay 1
476-- ThreadId 39: Delay 2
477-- ThreadId 38: Delay 3
478-- @
479--
480-- Note that parallel composition can only combine a finite number of
481-- streams as it needs to retain state for each unfinished stream.
482--
483-- /Since: 0.7.0 (maxBuffer applies to ParallelT streams)/
484--
485-- /Since: 0.1.0/
486newtype ParallelT m a = ParallelT {getParallelT :: Stream m a}
487 deriving (MonadTrans)
489-- | A parallely composing IO stream of elements of type @a@.
490-- See 'ParallelT' documentation for more details.
491--
492-- @since 0.2.0
493type Parallel = ParallelT IO
495-- | Fix the type of a polymorphic stream as 'ParallelT'.
496--
497-- @since 0.1.0
498parallely :: IsStream t => ParallelT m a -> t m a
499parallely = adapt
501instance IsStream ParallelT where
502 toStream = getParallelT
503 fromStream = ParallelT
505 {-# INLINE consM #-}
506 {-# SPECIALIZE consM :: IO a -> ParallelT IO a -> ParallelT IO a #-}
507 consM = consMParallel
509 {-# INLINE (|:) #-}
510 {-# SPECIALIZE (|:) :: IO a -> ParallelT IO a -> ParallelT IO a #-}
511 (|:) = consM
513------------------------------------------------------------------------------
514-- Semigroup
515------------------------------------------------------------------------------
517{-# INLINE mappendParallel #-}
518{-# SPECIALIZE mappendParallel :: ParallelT IO a -> ParallelT IO a -> ParallelT IO a #-}
519mappendParallel :: MonadAsync m => ParallelT m a -> ParallelT m a -> ParallelT m a
520mappendParallel m1 m2 = fromStream $ parallel (toStream m1) (toStream m2)
522instance MonadAsync m => Semigroup (ParallelT m a) where
523 (<>) = mappendParallel
525------------------------------------------------------------------------------
526-- Monoid
527------------------------------------------------------------------------------
529instance MonadAsync m => Monoid (ParallelT m a) where
530 mempty = K.nil
531 mappend = (<>)
533------------------------------------------------------------------------------
534-- Applicative
535------------------------------------------------------------------------------
537{-# INLINE apParallel #-}
538{-# SPECIALIZE apParallel :: ParallelT IO (a -> b) -> ParallelT IO a -> ParallelT IO b #-}
539apParallel :: MonadAsync m => ParallelT m (a -> b) -> ParallelT m a -> ParallelT m b
540apParallel (ParallelT m1) (ParallelT m2) =
541 let f x1 = K.concatMapBy parallel (pure . x1) m2
542 in ParallelT $ K.concatMapBy parallel f m1
544instance (Monad m, MonadAsync m) => Applicative (ParallelT m) where
545 {-# INLINE pure #-}
546 pure = ParallelT . K.yield
547 {-# INLINE (<*>) #-}
548 (<*>) = apParallel
550------------------------------------------------------------------------------
551-- Monad
552------------------------------------------------------------------------------
554{-# INLINE bindParallel #-}
555{-# SPECIALIZE bindParallel :: ParallelT IO a -> (a -> ParallelT IO b) -> ParallelT IO b #-}
556bindParallel :: MonadAsync m => ParallelT m a -> (a -> ParallelT m b) -> ParallelT m b
557bindParallel m f = fromStream $ K.bindWith parallel (K.adapt m) (\a -> K.adapt $ f a)
559instance MonadAsync m => Monad (ParallelT m) where
560 return = pure
561 (>>=) = bindParallel
563------------------------------------------------------------------------------
564-- Other instances
565------------------------------------------------------------------------------
567MONAD_COMMON_INSTANCES(ParallelT, MONADPARALLEL)