on d2c4df4 help sign in
618 lines 25 kB history raw
1{-# LANGUAGE ConstraintKinds #-}
2{-# LANGUAGE FlexibleContexts #-}
3{-# LANGUAGE FlexibleInstances #-}
4{-# LANGUAGE LambdaCase #-}
5{-# LANGUAGE MultiParamTypeClasses #-}
6{-# LANGUAGE RankNTypes #-}
7{-# LANGUAGE UndecidableInstances #-} -- XXX
9-- |
10-- Module : Streamly.Core
11-- Copyright : (c) 2017 Harendra Kumar
12--
13-- License : BSD3
14-- Maintainer : harendra.kumar@gmail.com
15-- Stability : experimental
16-- Portability : GHC
17--
18--
19module Streamly.Core
20 (
21 MonadParallel
22 , MonadAsync -- deprecated
24 -- * Streams
25 , Stream (..)
27 -- * Construction
28 , singleton
29 , cons
30 , repeat
31 , nil
33 -- * Semigroup Style Composition
34 , serial
35 , coserial
36 , coparallel
37 , parallel
39 -- * Alternative
40 , alt
42 -- * Transformers
43 , withLocal
44 , withCatchError
46 -- * Concurrent Stream Vars (SVars)
47 , SVar
48 , SVarSched (..)
49 , SVarTag (..)
50 , SVarStyle (..)
51 , newEmptySVar
52 , newStreamVar1
53 , newStreamVar2
54 , joinStreamVar2
55 , fromStreamVar
56 , toStreamVar
57 )
58where
60import Control.Concurrent (ThreadId, forkIO,
61 myThreadId, threadDelay)
62import Control.Concurrent.MVar (MVar, newEmptyMVar, tryTakeMVar,
63 tryPutMVar, takeMVar)
64import Control.Exception (SomeException (..))
65import qualified Control.Exception.Lifted as EL
66import Control.Monad (when)
67import Control.Monad.Catch (MonadThrow, throwM)
68import Control.Monad.Error.Class (MonadError(..))
69import Control.Monad.IO.Class (MonadIO(..))
70import Control.Monad.Reader.Class (MonadReader(..))
71import Control.Monad.Trans.Class (MonadTrans (lift))
72import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
73import Data.Atomics (atomicModifyIORefCAS,
74 atomicModifyIORefCAS_)
75import Data.Concurrent.Queue.MichaelScott (LinkedQueue, newQ, pushL,
76 tryPopR, nullQ)
77import Data.Functor (void)
78import Data.IORef (IORef, modifyIORef, newIORef,
79 readIORef)
80import Data.Maybe (isNothing)
81import Data.Semigroup (Semigroup(..))
82import Data.Set (Set)
83import qualified Data.Set as S
84import Prelude hiding (repeat)
86------------------------------------------------------------------------------
87-- Parent child thread communication type
88------------------------------------------------------------------------------
90-- | Events that a child thread may send to a parent thread.
91data ChildEvent a =
92 ChildYield a
93 | ChildStop ThreadId (Maybe SomeException)
95------------------------------------------------------------------------------
96-- State threaded around the monad for thread management
97------------------------------------------------------------------------------
99-- | Conjunction is used for monadic/product style composition. Disjunction is
100-- used for fold/sum style composition. We need to distinguish the two types of
101-- SVars so that the scheduling of the two is independent.
102data SVarTag = Conjunction | Disjunction deriving Eq
104-- | For fairly interleaved parallel composition the sched policy is FIFO
105-- whereas for left biased parallel composition it is LIFO.
106data SVarSched = LIFO | FIFO deriving Eq
108-- | Identify the type of the SVar. Two computations using the same style can
109-- be scheduled on the same SVar.
110data SVarStyle = SVarStyle SVarTag SVarSched deriving Eq
112-- | An SVar or a Stream Var is a conduit to the output from multiple streams
113-- running concurrently and asynchronously. An SVar can be thought of as an
114-- asynchronous IO handle. We can write any number of streams to an SVar in a
115-- non-blocking manner and then read them back at any time at any pace. The
116-- SVar would run the streams asynchronously and accumulate results. An SVar
117-- may not really execute the stream completely and accumulate all the results.
118-- However, it ensures that the reader can read the results at whatever paces
119-- it wants to read. The SVar monitors and adapts to the consumer's pace.
120--
121-- An SVar is a mini scheduler, it has an associated runqueue that holds the
122-- stream tasks to be picked and run by a pool of worker threads. It has an
123-- associated output queue where the output stream elements are placed by the
124-- worker threads. A doorBell is used by the worker threads to intimate the
125-- consumer thread about availability of new results in the output queue. More
126-- workers are added to the SVar by 'fromStreamVar' on demand if the output
127-- produced is not keeping pace with the consumer. On bounded SVars, workers
128-- block on the output queue to provide throttling of the producer when the
129-- consumer is not pulling fast enough. The number of workers may even get
130-- reduced depending on the consuming pace.
131--
132-- New work is enqueued either at the time of creation of the SVar or as a
133-- result of executing the parallel combinators i.e. '<|' and '<|>' when the
134-- already enqueued computations get evaluated. See 'joinStreamVar2'.
135--
136data SVar m a =
137 SVar { outputQueue :: IORef [ChildEvent a]
138 , doorBell :: MVar Bool -- wakeup mechanism for outQ
139 , enqueue :: Stream m a -> IO ()
140 , runqueue :: m ()
141 , runningThreads :: IORef (Set ThreadId)
142 , queueEmpty :: m Bool
143 , svarStyle :: SVarStyle
144 }
146------------------------------------------------------------------------------
147-- The stream type
148------------------------------------------------------------------------------
150-- | The type 'Stream m a' represents a monadic stream of values of type 'a'
151-- constructed using actions in monad 'm'. It uses stop, singleton and yield
152-- continuations equivalent to the following direct style type:
153--
154-- data Stream m a = Stop | Singleton a | Yield a (Stream m a)
155--
156-- To facilitate parallel composition we maintain a local state in an SVar that
157-- is shared across and is used for synchronization of the streams being
158-- composed.
159--
160-- The singleton case can be expressed in terms of stop and yield but we have
161-- it as a separate case to optimize composition operations for streams with
162-- single element. We build singleton streams in the implementation of 'pure'
163-- for Applicative and Monad, and in 'lift' for MonadTrans.
164--
165newtype Stream m a =
166 Stream {
167 runStream :: forall r.
168 Maybe (SVar m a) -- local state
169 -> m r -- stop
170 -> (a -> m r) -- singleton
171 -> (a -> Stream m a -> m r) -- yield
172 -> m r
173 }
175nil :: Stream m a
176nil = Stream $ \_ stp _ _ -> stp
178{-# INLINE singleton #-}
179singleton :: a -> Stream m a
180singleton a = Stream $ \_ _ single _ -> single a
182cons :: a -> Stream m a -> Stream m a
183cons a r = Stream $ \_ _ _ yld -> yld a r
185repeat :: a -> Stream m a
186repeat a = let x = cons a x in x
188------------------------------------------------------------------------------
189-- Semigroup
190------------------------------------------------------------------------------
192-- | Concatenates two streams sequentially i.e. the first stream is
193-- exhausted completely before yielding any element from the second stream.
194serial :: Stream m a -> Stream m a -> Stream m a
195serial m1 m2 = go m1
196 where
197 go (Stream m) = Stream $ \_ stp sng yld ->
198 let stop = (runStream m2) Nothing stp sng yld
199 single a = yld a m2
200 yield a r = yld a (go r)
201 in m Nothing stop single yield
203instance Semigroup (Stream m a) where
204 (<>) = serial
206------------------------------------------------------------------------------
207-- Monoid
208------------------------------------------------------------------------------
210instance Monoid (Stream m a) where
211 mempty = nil
212 mappend = (<>)
214------------------------------------------------------------------------------
215-- Interleave
216------------------------------------------------------------------------------
218coserial :: Stream m a -> Stream m a -> Stream m a
219coserial m1 m2 = Stream $ \_ stp sng yld -> do
220 let stop = (runStream m2) Nothing stp sng yld
221 single a = yld a m2
222 yield a r = yld a (coserial m2 r)
223 (runStream m1) Nothing stop single yield
225------------------------------------------------------------------------------
226-- Spawning threads and collecting result in streamed fashion
227------------------------------------------------------------------------------
229-- | A monad that can perform concurrent or parallel IO operations. Streams
230-- that can be composed concurrently require the underlying monad to be
231-- 'MonadParallel'.
232type MonadParallel m = (MonadIO m, MonadBaseControl IO m, MonadThrow m)
234{-# DEPRECATED MonadAsync "Please use MonadParallel instead." #-}
235type MonadAsync m = MonadParallel m
237{-# INLINE doFork #-}
238doFork :: MonadBaseControl IO m
239 => m ()
240 -> (SomeException -> m ())
241 -> m ThreadId
242doFork action exHandler =
243 EL.mask $ \restore ->
244 liftBaseWith $ \runInIO -> forkIO $ do
245 -- XXX test the exception handling
246 _ <- runInIO $ EL.catch (restore action) exHandler
247 -- XXX restore state here?
248 return ()
250-- XXX exception safety of all atomic/MVar operations
252{-# INLINE send #-}
253send :: MonadIO m => SVar m a -> ChildEvent a -> m ()
254send sv msg = liftIO $ do
255 atomicModifyIORefCAS_ (outputQueue sv) $ \es -> msg : es
256 -- XXX need a memory barrier? The wake up must happen only after the
257 -- store has finished otherwise we can have lost wakeup problems.
258 void $ tryPutMVar (doorBell sv) True
260{-# INLINE sendStop #-}
261sendStop :: MonadIO m => SVar m a -> m ()
262sendStop sv = liftIO myThreadId >>= \tid -> send sv (ChildStop tid Nothing)
264-- Note: Left associated compositions can grow this queue to a large size
265{-# INLINE enqueueLIFO #-}
266enqueueLIFO :: IORef [Stream m a] -> Stream m a -> IO ()
267enqueueLIFO q m = atomicModifyIORefCAS_ q $ \ ms -> m : ms
269runqueueLIFO :: MonadIO m => SVar m a -> IORef [Stream m a] -> m ()
270runqueueLIFO sv q = run
272 where
274 run = do
275 work <- dequeue
276 case work of
277 Nothing -> sendStop sv
278 Just m -> (runStream m) (Just sv) run single yield
280 sendit a = send sv (ChildYield a)
281 single a = sendit a >> run
282 yield a r = sendit a >> (runStream r) (Just sv) run single yield
284 dequeue = liftIO $ atomicModifyIORefCAS q $ \case
285 [] -> ([], Nothing)
286 x : xs -> (xs, Just x)
288{-# INLINE enqueueFIFO #-}
289enqueueFIFO :: LinkedQueue (Stream m a) -> Stream m a -> IO ()
290enqueueFIFO = pushL
292runqueueFIFO :: MonadIO m => SVar m a -> LinkedQueue (Stream m a) -> m ()
293runqueueFIFO sv q = run
295 where
297 run = do
298 work <- dequeue
299 case work of
300 Nothing -> sendStop sv
301 Just m -> (runStream m) (Just sv) run single yield
303 dequeue = liftIO $ tryPopR q
304 sendit a = send sv (ChildYield a)
305 single a = sendit a >> run
306 yield a r = sendit a >> liftIO (enqueueFIFO q r) >> run
308-- Thread tracking is needed for two reasons:
309--
310-- 1) Killing threads on exceptions. Threads may not be allowed to go away by
311-- themselves because they may run for significant times before going away or
312-- worse they may be stuck in IO and never go away.
313--
314-- 2) To know when all threads are done.
316{-# NOINLINE addThread #-}
317addThread :: MonadIO m => SVar m a -> ThreadId -> m ()
318addThread sv tid =
319 liftIO $ modifyIORef (runningThreads sv) (S.insert tid)
321{-# INLINE delThread #-}
322delThread :: MonadIO m => SVar m a -> ThreadId -> m ()
323delThread sv tid =
324 liftIO $ modifyIORef (runningThreads sv) $ (\s -> S.delete tid s)
326{-# INLINE allThreadsDone #-}
327allThreadsDone :: MonadIO m => SVar m a -> m Bool
328allThreadsDone sv = liftIO $ S.null <$> readIORef (runningThreads sv)
330{-# NOINLINE handleChildException #-}
331handleChildException :: MonadIO m => SVar m a -> SomeException -> m ()
332handleChildException sv e = do
333 tid <- liftIO myThreadId
334 send sv (ChildStop tid (Just e))
336{-# NOINLINE pushWorker #-}
337pushWorker :: MonadParallel m => SVar m a -> m ()
338pushWorker sv =
339 doFork (runqueue sv) (handleChildException sv) >>= addThread sv
341-- XXX When the queue is LIFO we can put a limit on the number of dispatches.
342-- Also, if a worker blocks on the output queue we can decide if we want to
343-- block or make it go away entirely, depending on the number of workers and
344-- the type of the queue.
345{-# INLINE sendWorkerWait #-}
346sendWorkerWait :: MonadParallel m => SVar m a -> m ()
347sendWorkerWait sv = do
348 case svarStyle sv of
349 SVarStyle _ LIFO -> liftIO $ threadDelay 200
350 SVarStyle _ FIFO -> liftIO $ threadDelay 0
352 output <- liftIO $ readIORef (outputQueue sv)
353 when (null output) $ do
354 done <- queueEmpty sv
355 if not done
356 then pushWorker sv >> sendWorkerWait sv
357 else void (liftIO $ takeMVar (doorBell sv))
359-- | Pull a stream from an SVar.
360{-# NOINLINE fromStreamVar #-}
361fromStreamVar :: MonadParallel m => SVar m a -> Stream m a
362fromStreamVar sv = Stream $ \_ stp sng yld -> do
363 -- XXX if reading the IORef is costly we can use a flag in the SVar to
364 -- indicate we are done.
365 done <- allThreadsDone sv
366 if done
367 then stp
368 else do
369 res <- liftIO $ tryTakeMVar (doorBell sv)
370 when (isNothing res) $ sendWorkerWait sv
371 list <- liftIO $ atomicModifyIORefCAS (outputQueue sv) $ \x -> ([], x)
372 -- To avoid lock overhead we read all events at once instead of reading
373 -- one at a time. We just reverse the list to process the events in the
374 -- order they arrived. Maybe we can use a queue instead?
375 (runStream $ processEvents (reverse list)) Nothing stp sng yld
377 where
379 handleException e tid = do
380 delThread sv tid
381 -- XXX implement kill async exception handling
382 -- liftIO $ readIORef (runningThreads sv) >>= mapM_ killThread
383 throwM e
385 {-# INLINE processEvents #-}
386 processEvents [] = Stream $ \_ stp sng yld -> do
387 done <- allThreadsDone sv
388 if not done
389 then (runStream (fromStreamVar sv)) Nothing stp sng yld
390 else stp
392 processEvents (ev : es) = Stream $ \_ stp sng yld -> do
393 let continue = (runStream (processEvents es)) Nothing stp sng yld
394 yield a = yld a (processEvents es)
396 case ev of
397 ChildYield a -> yield a
398 ChildStop tid e ->
399 case e of
400 Nothing -> delThread sv tid >> continue
401 Just ex -> handleException ex tid
403getFifoSVar :: MonadIO m => SVarStyle -> IO (SVar m a)
404getFifoSVar ctype = do
405 outQ <- newIORef []
406 outQMv <- newEmptyMVar
407 running <- newIORef S.empty
408 q <- newQ
409 let sv =
410 SVar { outputQueue = outQ
411 , doorBell = outQMv
412 , runningThreads = running
413 , runqueue = runqueueFIFO sv q
414 , enqueue = pushL q
415 , queueEmpty = liftIO $ nullQ q
416 , svarStyle = ctype
417 }
418 in return sv
420getLifoSVar :: MonadIO m => SVarStyle -> IO (SVar m a)
421getLifoSVar ctype = do
422 outQ <- newIORef []
423 outQMv <- newEmptyMVar
424 running <- newIORef S.empty
425 q <- newIORef []
426 let checkEmpty = null <$> liftIO (readIORef q)
427 let sv =
428 SVar { outputQueue = outQ
429 , doorBell = outQMv
430 , runningThreads = running
431 , runqueue = runqueueLIFO sv q
432 , enqueue = enqueueLIFO q
433 , queueEmpty = checkEmpty
434 , svarStyle = ctype
435 }
436 in return sv
438-- | Create a new empty SVar.
439newEmptySVar :: MonadParallel m => SVarStyle -> m (SVar m a)
440newEmptySVar style = do
441 liftIO $
442 case style of
443 SVarStyle _ FIFO -> getFifoSVar style
444 SVarStyle _ LIFO -> getLifoSVar style
446-- | Create a new SVar and enqueue one stream computation on it.
447newStreamVar1 :: MonadParallel m => SVarStyle -> Stream m a -> m (SVar m a)
448newStreamVar1 style m = do
449 sv <- newEmptySVar style
450 -- Note: We must have all the work on the queue before sending the
451 -- pushworker, otherwise the pushworker may exit before we even get a
452 -- chance to push.
453 liftIO $ (enqueue sv) m
454 pushWorker sv
455 return sv
457-- | Create a new SVar and enqueue two stream computations on it.
458newStreamVar2 :: MonadParallel m
459 => SVarStyle -> Stream m a -> Stream m a -> m (SVar m a)
460newStreamVar2 style m1 m2 = do
461 -- Note: We must have all the work on the queue before sending the
462 -- pushworker, otherwise the pushworker may exit before we even get a
463 -- chance to push.
464 sv <- liftIO $
465 case style of
466 SVarStyle _ FIFO -> do
467 c <- getFifoSVar style
468 (enqueue c) m1 >> (enqueue c) m2
469 return c
470 SVarStyle _ LIFO -> do
471 c <- getLifoSVar style
472 (enqueue c) m2 >> (enqueue c) m1
473 return c
474 pushWorker sv
475 return sv
477-- | Write a stream to an 'SVar' in a non-blocking manner. The stream can then
478-- be read back from the SVar using 'fromSVar'.
479toStreamVar :: MonadParallel m => SVar m a -> Stream m a -> m ()
480toStreamVar sv m = do
481 liftIO $ (enqueue sv) m
482 done <- allThreadsDone sv
483 -- XXX there may be a race here unless we are running in the consumer
484 -- thread. This is safe only when called from the consumer thread or when
485 -- no consumer is present.
486 when done $ pushWorker sv
488------------------------------------------------------------------------------
489-- Running streams concurrently
490------------------------------------------------------------------------------
492-- Concurrency rate control. Our objective is to create more threads on demand
493-- if the consumer is running faster than us. As soon as we encounter an
494-- Alternative composition we create a push pull pair of threads. We use a
495-- channel for communication between the consumer pulling from the channel and
496-- the producer who is pushing to the channel. The producer creates more threads
497-- if no output is seen on the channel, that is the consumer is running faster.
498-- However this mechanism can be problematic if the initial production latency
499-- is high, we may end up creating too many threads. So we need some way to
500-- monitor and use the latency as well.
501--
502-- TBD We may run computations at the lower level of the composition tree
503-- serially even if they are composed using a parallel combinator. We can use
504-- <> in place of <| and <=> in place of <|>. If we find that a parallel
505-- channel immediately above a computation becomes empty we can switch to
506-- parallelizing the computation. For that we can use a state flag to fork the
507-- rest of the computation at any point of time inside the Monad bind operation
508-- if the consumer is running at a faster speed.
509--
510-- TBD the alternative composition allows us to dispatch a chunkSize of only 1.
511-- If we have to dispatch in arbitrary chunksizes we will need to compose the
512-- parallel actions using a data constructor (Free Alternative) instead so that
513-- we can divide it in chunks of arbitrary size before dispatching. If the stream
514-- is composed of hierarchically composed grains of different sizes then we can
515-- always switch to a desired granularity depending on the consumer speed.
516--
517-- TBD for pure work (when we are not in the IO monad) we can divide it into
518-- just the number of CPUs.
520{-# NOINLINE withNewSVar2 #-}
521withNewSVar2 :: MonadParallel m
522 => SVarStyle -> Stream m a -> Stream m a -> Stream m a
523withNewSVar2 style m1 m2 = Stream $ \_ stp sng yld -> do
524 sv <- newStreamVar2 style m1 m2
525 (runStream (fromStreamVar sv)) Nothing stp sng yld
527-- | Join two computations on the currently running 'SVar' queue for concurrent
528-- execution. The 'SVarStyle' required by the current composition context is
529-- passed as one of the parameters. If the style does not match with the style
530-- of the current 'SVar' we create a new 'SVar' and schedule the computations
531-- on that. The newly created SVar joins as one of the computations on the
532-- current SVar queue.
533--
534-- When we are using parallel composition, an SVar is passed around as a state
535-- variable. We try to schedule a new parallel computation on the SVar passed
536-- to us. The first time, when no SVar exists, a new SVar is created.
537-- Subsequently, 'joinStreamVar2' may get called when a computation already
538-- scheduled on the SVar is further evaluated. For example, when (a \<|> b) is
539-- evaluated it calls a 'joinStreamVar2' to put 'a' and 'b' on the current scheduler
540-- queue. However, if the scheduling and composition style of the new
541-- computation being scheduled is different than the style of the current SVar,
542-- then we create a new SVar and schedule it on that.
543--
544-- For example:
545--
546-- * (x \<|> y) \<|> (t \<|> u) -- all of them get scheduled on the same SVar
547-- * (x \<|> y) \<|> (t \<| u) -- @t@ and @u@ get scheduled on a new child SVar
548-- because of the scheduling policy change.
549-- * if we 'adapt' a stream of type 'AsyncT' to a stream of type
550-- 'ParallelT', we create a new SVar at the transitioning bind.
551-- * When the stream is switching from disjunctive composition to conjunctive
552-- composition and vice-versa we create a new SVar to isolate the scheduling
553-- of the two.
554--
555{-# INLINE joinStreamVar2 #-}
556joinStreamVar2 :: MonadParallel m
557 => SVarStyle -> Stream m a -> Stream m a -> Stream m a
558joinStreamVar2 style m1 m2 = Stream $ \svr stp sng yld ->
559 case svr of
560 Just sv | svarStyle sv == style ->
561 liftIO ((enqueue sv) m2) >> (runStream m1) svr stp sng yld
562 _ -> (runStream (withNewSVar2 style m1 m2)) Nothing stp sng yld
564------------------------------------------------------------------------------
565-- Semigroup and Monoid style compositions for parallel actions
566------------------------------------------------------------------------------
568{-# INLINE coparallel #-}
569coparallel :: MonadParallel m => Stream m a -> Stream m a -> Stream m a
570coparallel = joinStreamVar2 (SVarStyle Disjunction LIFO)
572{-# INLINE parallel #-}
573parallel :: MonadParallel m => Stream m a -> Stream m a -> Stream m a
574parallel = joinStreamVar2 (SVarStyle Disjunction FIFO)
576-------------------------------------------------------------------------------
577-- Functor instace is the same for all types
578-------------------------------------------------------------------------------
580instance Monad m => Functor (Stream m) where
581 fmap f m = Stream $ \_ stp sng yld ->
582 let single = sng . f
583 yield a r = yld (f a) (fmap f r)
584 in (runStream m) Nothing stp single yield
586------------------------------------------------------------------------------
587-- Alternative & MonadPlus
588------------------------------------------------------------------------------
590alt :: Stream m a -> Stream m a -> Stream m a
591alt m1 m2 = Stream $ \_ stp sng yld ->
592 let stop = runStream m2 Nothing stp sng yld
593 in runStream m1 Nothing stop sng yld
595-------------------------------------------------------------------------------
596-- Transformers
597-------------------------------------------------------------------------------
599instance MonadTrans Stream where
600 lift mx = Stream $ \_ _ single _ -> mx >>= single
602withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
603withLocal f m =
604 Stream $ \svr stp sng yld ->
605 let single = local f . sng
606 yield a r = local f $ yld a (withLocal f r)
607 in (runStream m) svr (local f stp) single yield
609-- XXX handle and test cross thread state transfer
610withCatchError
611 :: MonadError e m
612 => Stream m a -> (e -> Stream m a) -> Stream m a
613withCatchError m h =
614 Stream $ \svr stp sng yld ->
615 let run x = runStream x svr stp sng yield
616 handle r = r `catchError` \e -> run $ h e
617 yield a r = yld a (withCatchError r h)
618 in handle $ run m