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