on master help sign in
Stop the channel when a particular stream ends
6c4b361Harendra Kumar
src/Streamly/Internal/Data/Stream/Async.hs+32−1
@@ -29,6 +29,8 @@ module Streamly.Internal.Data.Stream.Async
29 , maxThreads
29 , maxThreads
30 , maxBuffer
30 , maxBuffer
31 , eagerEval
31 , eagerEval
32+ , StopWhen (..)
33+ , stopWhen
32 -- maxYields
34 -- maxYields
33 , Rate(..)
35 , Rate(..)
34 , rate
36 , rate
@@ -41,7 +43,6 @@ module Streamly.Internal.Data.Stream.Async
41 -- * Combinators
43 -- * Combinators
42 -- | Stream combinators using Async channel
44 -- | Stream combinators using Async channel
43
45
44− -- XXX Move to parallel module
45 , eval
46 , eval
46 , evalWith
47 , evalWith
47
48
@@ -55,6 +56,9 @@ module Streamly.Internal.Data.Stream.Async
55 , appendWith
56 , appendWith
56 , interleave
57 , interleave
57 , interleaveWith
58 , interleaveWith
59+ , parallel
60+ , parallelFst
61+ , parallelMin
58
62
59 , apply
63 , apply
60 , applyWith
64 , applyWith
@@ -173,6 +177,33 @@ append = appendWith id
173interleave :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
177interleave :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
174interleave = interleaveWith id
178interleave = interleaveWith id
175
179
180+-- | Like 'append' but with 'eagerEval' on.
181+--
182+-- >>> parallel = Async.appendWith Async.eagerEval
183+--
184+{-# INLINE parallel #-}
185+parallel :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
186+parallel = appendWith eagerEval
187+
188+-- | Like 'parallel' but stops the output as soon as the first stream stops.
189+--
190+-- >>> parallelFst = Async.appendWith (Async.eagerEval . Async.stopWhen Async.FirstStops)
191+--
192+-- /Pre-release/
193+{-# INLINE parallelFst #-}
194+parallelFst :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
195+parallelFst = appendWith (eagerEval . stopWhen FirstStops)
196+
197+-- | Like 'parallel' but stops the output as soon as any of the two streams
198+-- stops.
199+--
200+-- >>> parallelMin = Async.appendWith (Async.eagerEval . Async.stopWhen Async.AnyStops)
201+--
202+-- /Pre-release/
203+{-# INLINE parallelMin #-}
204+parallelMin :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
205+parallelMin = appendWith (eagerEval . stopWhen AnyStops)
206+
176-- concatMapWith modifier f stream = concatWith modifier $ fmap f stream
207-- concatMapWith modifier f stream = concatWith modifier $ fmap f stream
177
208
178-- | Like 'concatMap' but we can also specify the concurrent channel's
209-- | Like 'concatMap' but we can also specify the concurrent channel's
src/Streamly/Internal/Data/Stream/Async/Channel.hs+73−24
@@ -23,6 +23,8 @@ module Streamly.Internal.Data.Stream.Async.Channel
23 , maxRate
23 , maxRate
24 , constRate
24 , constRate
25 , inspectMode
25 , inspectMode
26+ , StopWhen (..)
27+ , stopWhen
26
28
27 -- * Operations
29 -- * Operations
28 , fromChannel
30 , fromChannel
@@ -129,35 +131,46 @@ evalWithD modifier m = D.Stream step Nothing
129-- appending two streams
131-- appending two streams
130-------------------------------------------------------------------------------
132-------------------------------------------------------------------------------
131
133
132−-- | Create a new channel and add both the streams to it for async evaluation.
134+{-# INLINE appendGeneric #-}
133−-- The output stream is the result of the evaluation.
134−{-# INLINE appendWithK #-}
135−appendWithK :: MonadAsync m =>
135+appendGeneric :: MonadAsync m =>
136+ ((Config -> Config) -> m (Channel m a))
136− (Config -> Config) -> K.Stream m a -> K.Stream m a -> K.Stream m a
137+ -> (Config -> Config)
138+ -> K.Stream m a
139+ -> K.Stream m a
140+ -> K.Stream m a
137−appendWithK modifier stream1 stream2 = K.concatEffect action
141+appendGeneric newChan modifier stream1 stream2 = K.concatEffect action
138
142
139 where
143 where
140
144
141 action = do
145 action = do
142− chan <- Append.newChannel modifier
146+ chan <- newChan modifier
147+ let cfg = modifier defaultConfig
148+ done = K.nilM (stopChannel chan)
149+ case getStopWhen cfg of
150+ AllStop -> do
151+ toChannelK chan stream2
152+ toChannelK chan stream1
153+ FirstStops -> do
143− toChannelK chan stream2
154+ toChannelK chan stream2
144− toChannelK chan stream1
155+ toChannelK chan (K.serial stream1 done)
156+ AnyStops -> do
157+ toChannelK chan (K.serial stream2 done)
158+ toChannelK chan (K.serial stream1 done)
145 return $ Stream.toStreamK $ fromChannel chan
159 return $ Stream.toStreamK $ fromChannel chan
146
160
161+-- | Create a new channel and add both the streams to it for async evaluation.
162+-- The output stream is the result of the evaluation.
163+{-# INLINE appendWithK #-}
164+appendWithK :: MonadAsync m =>
165+ (Config -> Config) -> K.Stream m a -> K.Stream m a -> K.Stream m a
166+appendWithK = appendGeneric Append.newChannel
167+
147-- | Create a new channel and add both the streams to it for async evaluation.
168-- | Create a new channel and add both the streams to it for async evaluation.
148-- The output stream is the result of the evaluation.
169-- The output stream is the result of the evaluation.
149{-# INLINE interleaveWithK #-}
170{-# INLINE interleaveWithK #-}
150interleaveWithK :: MonadAsync m =>
171interleaveWithK :: MonadAsync m =>
151 (Config -> Config) -> K.Stream m a -> K.Stream m a -> K.Stream m a
172 (Config -> Config) -> K.Stream m a -> K.Stream m a -> K.Stream m a
152−interleaveWithK modifier stream1 stream2 = K.concatEffect action
153−
154− where
155−
156− action = do
157− chan <- Interleave.newChannel modifier
173+interleaveWithK = appendGeneric Interleave.newChannel
158− toChannelK chan stream2
159− toChannelK chan stream1
160− return $ Stream.toStreamK $ fromChannel chan
161
174
162-- | Evaluate the first stream in the current thread and add the second stream
175-- | Evaluate the first stream in the current thread and add the second stream
163-- to the supplied channel. This is to be used by a worker thread.
176-- to the supplied channel. This is to be used by a worker thread.
@@ -219,26 +232,62 @@ concatMapWithChanK chan f stream =
219 in K.concatMapEffect (`run` stream) (mkEnqueue chan run)
232 in K.concatMapEffect (`run` stream) (mkEnqueue chan run)
220 -- K.concatMapWith (_appendWithChanK chan) f stream
233 -- K.concatMapWith (_appendWithChanK chan) f stream
221
234
235+{-# INLINE concatMapWithChanKAny #-}
236+concatMapWithChanKAny :: MonadAsync m =>
237+ Channel m b -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
238+concatMapWithChanKAny chan f stream =
239+ let done = K.nilM (stopChannel chan)
240+ run q = concatMapDivK q (\x -> K.serial (f x) done)
241+ in K.concatMapEffect (`run` stream) (mkEnqueue chan run)
242+
243+{-# INLINE concatMapWithChanKFirst #-}
244+concatMapWithChanKFirst :: MonadAsync m =>
245+ Channel m b -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
246+concatMapWithChanKFirst chan f stream =
247+ let done = K.nilM (stopChannel chan)
248+ run q = concatMapDivK q f
249+ in K.concatEffect $ do
250+ res <- K.uncons stream
251+ case res of
252+ Nothing -> return K.nil
253+ Just (h, t) -> do
254+ q <- mkEnqueue chan run
255+ q t
256+ return $ K.serial (f h) done
257+
258+{-# INLINE concatMapWithChanKGeneric #-}
259+concatMapWithChanKGeneric :: MonadAsync m =>
260+ (Config -> Config)
261+ -> Channel m b
262+ -> (a -> K.Stream m b)
263+ -> K.Stream m a
264+ -> K.Stream m b
265+concatMapWithChanKGeneric modifier chan f stream = do
266+ let cfg = modifier defaultConfig
267+ case getStopWhen cfg of
268+ AllStop -> concatMapWithChanK chan f stream
269+ FirstStops -> concatMapWithChanKFirst chan f stream
270+ AnyStops -> concatMapWithChanKAny chan f stream
271+
222-- XXX Add a deep evaluation variant that evaluates individual elements in the
272-- XXX Add a deep evaluation variant that evaluates individual elements in the
223-- generated streams in parallel.
273-- generated streams in parallel.
224
274
225-- | Allocate a channel and use it to concurrently evaluate the streams
275-- | Allocate a channel and use it to concurrently evaluate the streams
226-- generated by the mapped function.
276-- generated by the mapped function.
227--
277--
228−-- >> concatMapWithK modifier f = withChannelK (`concatMapWithChanK` f) modifier
229−--
230{-# INLINE concatMapWithK #-}
278{-# INLINE concatMapWithK #-}
231concatMapWithK :: MonadAsync m =>
279concatMapWithK :: MonadAsync m =>
232 (Config -> Config) -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
280 (Config -> Config) -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
233−concatMapWithK modifier f = withChannelK (`concatMapWithChanK` f) modifier
281+concatMapWithK modifier f =
282+ let g = concatMapWithChanKGeneric modifier
283+ in withChannelK (`g` f) modifier
234
284
235-- | Allocate a channel and use it to concurrently evaluate the streams
285-- | Allocate a channel and use it to concurrently evaluate the streams
236-- generated by the mapped function.
286-- generated by the mapped function.
237--
287--
238−-- >> concatMapInterleaveWithK modifier f = withInterleaveChannelK (`concatMapWithChanK` f) modifier
239−--
240{-# INLINE concatMapInterleaveWithK #-}
288{-# INLINE concatMapInterleaveWithK #-}
241concatMapInterleaveWithK :: MonadAsync m =>
289concatMapInterleaveWithK :: MonadAsync m =>
242 (Config -> Config) -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
290 (Config -> Config) -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
243concatMapInterleaveWithK modifier f =
291concatMapInterleaveWithK modifier f =
292+ let g = concatMapWithChanKGeneric modifier
244− withInterleaveChannelK (`concatMapWithChanK` f) modifier
293+ in withInterleaveChannelK (`g` f) modifier
src/Streamly/Internal/Data/Stream/Async/Channel/Operations.hs+22−11
@@ -129,18 +129,17 @@ fromChannelRaw sv = K.MkStream $ \st yld sng stp -> do
129
129
130 where
130 where
131
131
132− allDone stp = do
132+ cleanup = do
133− when (svarInspectMode sv) $ do
133+ when (svarInspectMode sv) $ liftIO $ do
134− t <- liftIO $ getTime Monotonic
134+ t <- getTime Monotonic
135− liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
135+ writeIORef (svarStopTime (svarStats sv)) (Just t)
136− liftIO $ printSVar (dumpSVar sv) "SVar Done"
136+ printSVar (dumpSVar sv) "SVar Done"
137− stp
138
137
139 {-# INLINE processEvents #-}
138 {-# INLINE processEvents #-}
140 processEvents [] = K.MkStream $ \st yld sng stp -> do
139 processEvents [] = K.MkStream $ \st yld sng stp -> do
141 done <- postProcess sv
140 done <- postProcess sv
142 if done
141 if done
143− then allDone stp
142+ then cleanup >> stp
144 else K.foldStream st yld sng stp $ fromChannelRaw sv
143 else K.foldStream st yld sng stp $ fromChannelRaw sv
145
144
146 processEvents (ev : es) = K.MkStream $ \st yld sng stp -> do
145 processEvents (ev : es) = K.MkStream $ \st yld sng stp -> do
@@ -153,11 +152,23 @@ fromChannelRaw sv = K.MkStream $ \st yld sng stp -> do
153 Nothing -> K.foldStream st yld sng stp rest
152 Nothing -> K.foldStream st yld sng stp rest
154 Just ex ->
153 Just ex ->
155 case fromException ex of
154 case fromException ex of
156− Just ThreadAbort ->
155+ Just ChannelStop -> do
157− K.foldStream st yld sng stp rest
158− Nothing -> do
159 liftIO (cleanupSVar (workerThreads sv))
156 liftIO (cleanupSVar (workerThreads sv))
157+ cleanup >> stp
158+ Nothing -> do
159+ case fromException ex of
160+ Just ThreadAbort ->
161+ -- We terminate the loop after sending
162+ -- ThreadAbort to workers so we should
163+ -- never get it unless it is thrown
164+ -- from inside a worker thread or by
165+ -- someone else to our thread.
166+ error "processEvents: got ThreadAbort"
167+ -- K.foldStream st yld sng stp rest
168+ Nothing -> do
169+ -- Throw ThreadAbort to all and exit.
170+ liftIO (cleanupSVar (workerThreads sv))
160− throwM ex
171+ cleanup >> throwM ex
161
172
162#if __GLASGOW_HASKELL__ < 810
173#if __GLASGOW_HASKELL__ < 810
163#ifdef INSPECTION
174#ifdef INSPECTION
src/Streamly/Internal/Data/Stream/Async/Channel/Type.hs+19−3
@@ -11,19 +11,24 @@ module Streamly.Internal.Data.Stream.Async.Channel.Type
11 Channel(..)
11 Channel(..)
12 , yield
12 , yield
13 , stop
13 , stop
14+ , stopChannel
14 , dumpSVar
15 , dumpSVar
15 )
16 )
16where
17where
17
18
18−import Control.Concurrent (ThreadId)
19+import Control.Concurrent (ThreadId, myThreadId)
19import Control.Concurrent.MVar (MVar)
20import Control.Concurrent.MVar (MVar)
21+import Control.Exception (toException)
22+import Control.Monad (void)
23+import Control.Monad.IO.Class (MonadIO(..))
20import Data.IORef (IORef)
24import Data.IORef (IORef)
21import Data.List (intersperse)
25import Data.List (intersperse)
22import Data.Set (Set)
26import Data.Set (Set)
23−
24import Streamly.Internal.Control.Concurrent (RunInIO)
27import Streamly.Internal.Control.Concurrent (RunInIO)
28+import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
25import Streamly.Internal.Data.Stream.Channel.Dispatcher (dumpSVarStats)
29import Streamly.Internal.Data.Stream.Channel.Dispatcher (dumpSVarStats)
26−import Streamly.Internal.Data.Stream.Channel.Worker (sendYield, sendStop)
30+import Streamly.Internal.Data.Stream.Channel.Worker
31+ (sendYield, sendStop, sendWithDoorBell)
27import Streamly.Internal.Data.Stream.StreamK.Type (Stream)
32import Streamly.Internal.Data.Stream.StreamK.Type (Stream)
28
33
29import Streamly.Internal.Data.Stream.Channel.Types
34import Streamly.Internal.Data.Stream.Channel.Types
@@ -137,6 +142,17 @@ stop sv winfo =
137 (outputQueue sv)
142 (outputQueue sv)
138 (outputDoorBell sv)
143 (outputDoorBell sv)
139
144
145+{-# INLINABLE stopChannel #-}
146+stopChannel :: MonadIO m => Channel m a -> m ()
147+stopChannel chan = liftIO $ do
148+ atomicModifyIORefCAS_ (workerCount chan) $ \n -> n - 1
149+ myThreadId >>= \tid ->
150+ void
151+ $ sendWithDoorBell
152+ (outputQueue chan)
153+ (outputDoorBell chan)
154+ (ChildStop tid (Just (toException ChannelStop)))
155+
140{-# NOINLINE dumpSVar #-}
156{-# NOINLINE dumpSVar #-}
141dumpSVar :: Channel m a -> IO String
157dumpSVar :: Channel m a -> IO String
142dumpSVar sv = do
158dumpSVar sv = do
src/Streamly/Internal/Data/Stream/Channel/Types.hs+39−4
@@ -37,6 +37,7 @@ module Streamly.Internal.Data.Stream.Channel.Types
37 Count (..)
37 Count (..)
38 , Limit (..)
38 , Limit (..)
39 , ThreadAbort (..)
39 , ThreadAbort (..)
40+ , ChannelStop (..)
40 , ChildEvent (..)
41 , ChildEvent (..)
41
42
42 -- * Stats
43 -- * Stats
@@ -60,6 +61,7 @@ module Streamly.Internal.Data.Stream.Channel.Types
60
61
61 -- * Configuration
62 -- * Configuration
62 , Rate (..)
63 , Rate (..)
64+ , StopWhen (..)
63 , Config
65 , Config
64
66
65 -- ** Default config
67 -- ** Default config
@@ -72,6 +74,7 @@ module Streamly.Internal.Data.Stream.Channel.Types
72 , maxYields
74 , maxYields
73 , inspectMode
75 , inspectMode
74 , eagerEval
76 , eagerEval
77+ , stopWhen
75
78
76 , rate
79 , rate
77 , avgRate
80 , avgRate
@@ -88,6 +91,7 @@ module Streamly.Internal.Data.Stream.Channel.Types
88 , getYieldLimit
91 , getYieldLimit
89 , getInspectMode
92 , getInspectMode
90 , getEagerDispatch
93 , getEagerDispatch
94+ , getStopWhen
91
95
92 -- * Cleanup
96 -- * Cleanup
93 , cleanupSVar
97 , cleanupSVar
@@ -107,7 +111,7 @@ module Streamly.Internal.Data.Stream.Channel.Types
107 )
111 )
108where
112where
109
113
110−import Control.Concurrent (ThreadId, myThreadId, throwTo, MVar, tryReadMVar)
114+import Control.Concurrent (ThreadId, throwTo, MVar, tryReadMVar)
111import Control.Concurrent.MVar (tryPutMVar)
115import Control.Concurrent.MVar (tryPutMVar)
112import Control.Exception
116import Control.Exception
113 ( SomeException(..), Exception, catches, throwIO, Handler(..)
117 ( SomeException(..), Exception, catches, throwIO, Handler(..)
@@ -163,10 +167,19 @@ instance Ord Limit where
163-- Parent child thread communication type
167-- Parent child thread communication type
164------------------------------------------------------------------------------
168------------------------------------------------------------------------------
165
169
170+-- | Channel driver throws this exception to all active workers to clean up
171+-- the channel.
166data ThreadAbort = ThreadAbort deriving Show
172data ThreadAbort = ThreadAbort deriving Show
167
173
168instance Exception ThreadAbort
174instance Exception ThreadAbort
169
175
176+-- | If any stream in the channel throws this exception, the channel stops
177+-- immediately.
178+data ChannelStop = ChannelStop deriving Show
179+
180+instance Exception ChannelStop
181+
182+-- XXX Use a ChildSingle event to speed up mapM?
170-- | Events that a child thread may send to a parent thread.
183-- | Events that a child thread may send to a parent thread.
171data ChildEvent a =
184data ChildEvent a =
172 ChildYield a
185 ChildYield a
@@ -298,6 +311,14 @@ data Rate = Rate
298 , rateBuffer :: Int -- ^ Maximum slack from the goal
311 , rateBuffer :: Int -- ^ Maximum slack from the goal
299 }
312 }
300
313
314+-- | Specify when the 'Channel' should stop. Note that for ordered streams only
315+-- 'FirstStops' and 'AnyStop' would be the same, and usually only 'AllStop'
316+-- would make sense. This option is most useful for eager unordered streams.
317+data StopWhen =
318+ FirstStops -- ^ Stop when the first stream/action ends.
319+ | AllStop -- ^ Stop when all the streams end.
320+ | AnyStops -- ^ Stop when any one stream ends.
321+
301-- XXX we can put the resettable fields in a oneShotConfig field and others in
322-- XXX we can put the resettable fields in a oneShotConfig field and others in
302-- a persistentConfig field. That way reset would be fast and scalable
323-- a persistentConfig field. That way reset would be fast and scalable
303-- irrespective of the number of fields.
324-- irrespective of the number of fields.
@@ -323,6 +344,7 @@ data Config = Config
323 , _maxStreamRate :: Maybe Rate
344 , _maxStreamRate :: Maybe Rate
324 , _inspectMode :: Bool
345 , _inspectMode :: Bool
325 , _eagerDispatch :: Bool
346 , _eagerDispatch :: Bool
347+ , _stopWhen :: StopWhen
326 }
348 }
327
349
328-------------------------------------------------------------------------------
350-------------------------------------------------------------------------------
@@ -354,6 +376,7 @@ defaultConfig = Config
354 , _inspectMode = False
376 , _inspectMode = False
355 -- XXX Set it to True when Rate is not set?
377 -- XXX Set it to True when Rate is not set?
356 , _eagerDispatch = False
378 , _eagerDispatch = False
379+ , _stopWhen = AllStop
357 }
380 }
358
381
359-------------------------------------------------------------------------------
382-------------------------------------------------------------------------------
@@ -463,6 +486,8 @@ setStreamLatency n st =
463getStreamLatency :: Config -> Maybe NanoSecond64
486getStreamLatency :: Config -> Maybe NanoSecond64
464getStreamLatency = _streamLatency
487getStreamLatency = _streamLatency
465
488
489+-- XXX Rename to "inspect"
490+
466-- | Print debug information about the 'Channel' when the stream ends.
491-- | Print debug information about the 'Channel' when the stream ends.
467--
492--
468-- /Pre-release/
493-- /Pre-release/
@@ -473,6 +498,8 @@ inspectMode st = st { _inspectMode = True }
473getInspectMode :: Config -> Bool
498getInspectMode :: Config -> Bool
474getInspectMode = _inspectMode
499getInspectMode = _inspectMode
475
500
501+-- XXX rename to "eager"
502+
476-- | By default the stream is evaluated on demand, new evaluation may be
503-- | By default the stream is evaluated on demand, new evaluation may be
477-- deferred until the consumer consumes the evaluated stream. When 'eagerEval'
504-- deferred until the consumer consumes the evaluated stream. When 'eagerEval'
478-- option is on the evaluation is aggressive and does not wait for consumption
505-- option is on the evaluation is aggressive and does not wait for consumption
@@ -487,6 +514,13 @@ eagerEval st = st { _eagerDispatch = True }
487getEagerDispatch :: Config -> Bool
514getEagerDispatch :: Config -> Bool
488getEagerDispatch = _eagerDispatch
515getEagerDispatch = _eagerDispatch
489
516
517+-- | Specify when the 'Channel' should stop.
518+stopWhen :: StopWhen -> Config -> Config
519+stopWhen cond st = st { _stopWhen = cond }
520+
521+getStopWhen :: Config -> StopWhen
522+getStopWhen = _stopWhen
523+
490-------------------------------------------------------------------------------
524-------------------------------------------------------------------------------
491-- Initialization
525-- Initialization
492-------------------------------------------------------------------------------
526-------------------------------------------------------------------------------
@@ -740,13 +774,14 @@ printSVar dump how = do
740-- Cleanup
774-- Cleanup
741-------------------------------------------------------------------------------
775-------------------------------------------------------------------------------
742
776
743−-- Can be called from a worker as well, avoids throwing to itself
777+-- | Never called from a worker thread.
744cleanupSVar :: IORef (Set ThreadId) -> IO ()
778cleanupSVar :: IORef (Set ThreadId) -> IO ()
745cleanupSVar workerSet = do
779cleanupSVar workerSet = do
746 workers <- readIORef workerSet
780 workers <- readIORef workerSet
747− self <- myThreadId
781+ -- self <- myThreadId
748 Prelude.mapM_ (`throwTo` ThreadAbort)
782 Prelude.mapM_ (`throwTo` ThreadAbort)
749− (Prelude.filter (/= self) $ Set.toList workers)
783+ -- (Prelude.filter (/= self) $ Set.toList workers)
784+ (Set.toList workers)
750
785
751-------------------------------------------------------------------------------
786-------------------------------------------------------------------------------
752-- Evaluator
787-- Evaluator