@@ -7,6 +7,8 @@
77{-# LANGUAGE MultiParamTypeClasses #-}
88{-# LANGUAGE UndecidableInstances #-} -- XXX
99
10+#include "inline.hs"
11+
1012-- |
1113-- Module : Streamly.Streams.Parallel
1214-- Copyright : (c) 2017 Harendra Kumar
@@ -47,7 +49,7 @@ import Control.Monad.Reader.Class (MonadReader(..))
4749import Control.Monad.State.Class (MonadState(..))
4850import Control.Monad.Trans.Class (MonadTrans(lift))
4951import Data.Functor (void)
50−import Data.IORef (readIORef, writeIORef)
52+import Data.IORef (readIORef, writeIORef, newIORef)
5153import Data.Maybe (fromJust)
5254#if __GLASGOW_HASKELL__ < 808
5355import Data.Semigroup (Semigroup(..))
@@ -58,10 +60,17 @@ import qualified Data.Set as Set
5860
5961import Streamly.Streams.SVar (fromSVar, fromStreamVar)
6062import Streamly.Streams.Serial (map)
61−import Streamly.Internal.Data.SVar
6263import Streamly.Streams.StreamK (IsStream(..), Stream, mkStream, foldStream,
6364 foldStreamShared, adapt)
65+
66+import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
67+import Streamly.Internal.Data.Time.Clock (Clock(..), getTime)
68+import Streamly.Internal.Data.SVar hiding (handleChildException)
69+
6470import qualified Streamly.Streams.StreamK as K
71+import qualified Streamly.Streams.StreamD as D
72+import qualified Streamly.Internal.Data.SVar as SVar
73+import qualified Streamly.Internal.Data.Fold as FL
6574
6675#include "Instances.hs"
6776
@@ -69,6 +78,10 @@ import qualified Streamly.Streams.StreamK as K
6978-- Parallel
7079-------------------------------------------------------------------------------
7180
81+-------------------------------------------------------------------------------
82+-- StreamK based worker routines
83+-------------------------------------------------------------------------------
84+
7285{-# NOINLINE runOne #-}
7386runOne
7487 :: MonadIO m
@@ -120,6 +133,81 @@ runOneLimited st m0 winfo = go m0
120133 single a = sendit a >> (liftIO $ sendStop sv winfo)
121134 yieldk a r = sendit a >> go r
122135
136+-------------------------------------------------------------------------------
137+-- StreamD based worker routines
138+-------------------------------------------------------------------------------
139+
140+-- Using StreamD the worker stream producing code can fuse with the code to
141+-- queue output to the SVar giving some perf boost.
142+--
143+-- Note that StreamD can only be used in limited situations, specifically, we
144+-- cannot implement joinStreamVarPar using this.
145+
146+{-# INLINE_NORMAL pushWorkerParD #-}
147+pushWorkerParD :: MonadAsync m
148+ => State t m a -> SVar t m a -> D.Stream m a -> m ()
149+pushWorkerParD st sv xs =
150+ if svarInspectMode sv
151+ then forkWithDiag
152+ else do
153+ tid <-
154+ case getYieldLimit st of
155+ Nothing -> doFork (work Nothing)
156+ (svarMrun sv)
157+ (SVar.handleChildException sv)
158+ Just _ -> doFork (workLim Nothing)
159+ (svarMrun sv)
160+ (SVar.handleChildException sv)
161+ modifyThread sv tid
162+
163+ where
164+
165+ {-# NOINLINE work #-}
166+ work info = (D.runFold (FL.toParallelSVar sv info) xs)
167+
168+ {-# NOINLINE workLim #-}
169+ workLim info = D.runFold (FL.toParallelSVarLimited sv info) xs
170+
171+ {-# NOINLINE forkWithDiag #-}
172+ forkWithDiag = do
173+ -- We do not use workerCount in case of ParallelVar but still there is
174+ -- no harm in maintaining it correctly.
175+ liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
176+ recordMaxWorkers sv
177+ -- This allocation matters when significant number of workers are being
178+ -- sent. We allocate it only when needed. The overhead increases by 4x.
179+ winfo <-
180+ case yieldRateInfo sv of
181+ Nothing -> return Nothing
182+ Just _ -> liftIO $ do
183+ cntRef <- newIORef 0
184+ t <- getTime Monotonic
185+ lat <- newIORef (0, t)
186+ return $ Just WorkerInfo
187+ { workerYieldMax = 0
188+ , workerYieldCount = cntRef
189+ , workerLatencyStart = lat
190+ }
191+ tid <-
192+ case getYieldLimit st of
193+ Nothing -> doFork (work winfo)
194+ (svarMrun sv)
195+ (SVar.handleChildException sv)
196+ Just _ -> doFork (workLim winfo)
197+ (svarMrun sv)
198+ (SVar.handleChildException sv)
199+ modifyThread sv tid
200+
201+-------------------------------------------------------------------------------
202+-- Consing and appending a stream in parallel style
203+-------------------------------------------------------------------------------
204+
205+-- Note that consing and appending requires StreamK as it would not scale well
206+-- with StreamD unless we are only consing a very small number of streams or
207+-- elements in a stream. StreamK allows us to manipulate control flow in a way
208+-- which StreamD cannot allow. StreamK can make a jump without having to
209+-- remember the past state.
210+
123211{-# NOINLINE forkSVarPar #-}
124212forkSVarPar :: (IsStream t, MonadAsync m)
125213 => SVarStopStyle -> t m a -> t m a -> t m a
@@ -144,6 +232,10 @@ joinStreamVarPar style ss m1 m2 = mkStream $ \st yld sng stp ->
144232 foldStreamShared st yld sng stp m2
145233 _ -> foldStreamShared st yld sng stp (forkSVarPar ss m1 m2)
146234
235+-------------------------------------------------------------------------------
236+-- User facing APIs
237+-------------------------------------------------------------------------------
238+
147239-- | XXX we can implement it more efficienty by directly implementing instead
148240-- of combining streams using parallel.
149241{-# INLINE consMParallel #-}
@@ -194,24 +286,57 @@ mkParallel m = do
194286-- Stream to stream concurrent function application
195287------------------------------------------------------------------------------
196288
197−{-# INLINE applyParallel #-}
198−applyParallel :: (IsStream t, MonadAsync m) => (t m a -> t m b) -> t m a -> t m b
199−applyParallel f m = mkStream $ \st yld sng stp -> do
289+{-# INLINE _applyParallelK #-}
290+_applyParallelK :: (IsStream t, MonadAsync m)
291+ => (t m a -> t m b) -> t m a -> t m b
292+_applyParallelK f m = mkStream $ \st yld sng stp -> do
200293 sv <- newParallelVar StopNone (adaptState st)
201294 pushWorkerPar sv (runOne st{streamVar = Just sv} (toStream m))
202295 foldStream st yld sng stp $ f $ fromSVar sv
203296
297+-- XXX need to implement fromSVar in direct style for fused generation from
298+-- SVar.
299+{-# INLINE applyParallelD #-}
300+applyParallelD :: (IsStream t, MonadAsync m)
301+ => (t m a -> t m b) -> t m a -> t m b
302+applyParallelD f m = mkStream $ \st yld sng stp -> do
303+ sv <- newParallelVar StopNone (adaptState st)
304+ pushWorkerParD (adaptState st) sv (D.toStreamD m)
305+ foldStream st yld sng stp $ f $ fromSVar sv
306+
307+{-# INLINE applyParallel #-}
308+applyParallel :: (IsStream t, MonadAsync m)
309+ => (t m a -> t m b) -> t m a -> t m b
310+applyParallel = applyParallelD
311+
204312------------------------------------------------------------------------------
205313-- Stream runner concurrent function application
206314------------------------------------------------------------------------------
207315
208−{-# INLINE foldParallel #-}
209−foldParallel :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
210−foldParallel f m = do
316+{-# INLINE _foldParallelK #-}
317+_foldParallelK :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
318+_foldParallelK f m = do
211319 sv <- newParallelVar StopNone defState
212320 pushWorkerPar sv (runOne defState{streamVar = Just sv} $ toStream m)
213321 f $ fromSVar sv
214322
323+-- XXX need to implement fromSVar in direct style for fused generation from
324+-- SVar.
325+{-# INLINE foldParallelD #-}
326+foldParallelD :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
327+foldParallelD f m = do
328+ sv <- newParallelVar StopNone defState
329+ pushWorkerParD defState sv (D.toStreamD m)
330+ f $ fromSVar sv
331+
332+{-# INLINE foldParallel #-}
333+foldParallel :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
334+foldParallel = foldParallelD
335+
336+------------------------------------------------------------------------------
337+-- Parallel tee
338+------------------------------------------------------------------------------
339+
215340{-# INLINE teeToSVar #-}
216341teeToSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> t m a
217342teeToSVar svr m = mkStream $ \st yld sng stp -> do