@@ -36,6 +36,7 @@ module Streamly.Core
3636 -- * Semigroup Style Composition
3737 , serial
3838 , wSerial
39+ , ahead
3940 , async
4041 , wAsync
4142 , parallel
@@ -66,7 +67,7 @@ where
6667
6768import Control.Concurrent (ThreadId, myThreadId, threadDelay)
6869import Control.Concurrent.MVar (MVar, newEmptyMVar, tryTakeMVar,
69− tryPutMVar, takeMVar)
70+ tryPutMVar, takeMVar, putMVar)
7071import Control.Exception (SomeException (..))
7172import qualified Control.Exception.Lifted as EL
7273import Control.Monad (when)
@@ -81,6 +82,8 @@ import Data.Atomics (casIORef, readForCAS, peekTicket
8182import Data.Concurrent.Queue.MichaelScott (LinkedQueue, newQ, pushL,
8283 tryPopR, nullQ)
8384import Data.Functor (void)
85+import Data.Heap (Heap, Entry(..))
86+import qualified Data.Heap as H
8487import Data.IORef (IORef, modifyIORef, newIORef,
8588 readIORef, atomicModifyIORef)
8689import Data.Maybe (isNothing, fromJust)
@@ -102,19 +105,27 @@ data ChildEvent a =
102105 ChildYield a
103106 | ChildStop ThreadId (Maybe SomeException)
104107
108+-- | Sorting out of turn outputs in a heap for Ahead style streams
109+data AheadHeapEntry a =
110+ AheadYieldSingle a
111+ | AheadYieldMany a (MVar ()) -- value, MVar to wakeup for more values
112+
105113------------------------------------------------------------------------------
106114-- State threaded around the monad for thread management
107115------------------------------------------------------------------------------
108116
117+-- XXX this is redundant as we always start a new SVar on bind
109118-- | Conjunction is used for monadic/product style composition. Disjunction is
110119-- used for fold/sum style composition. We need to distinguish the two types of
111120-- SVars so that the scheduling of the two is independent.
112121data SVarTag = Conjunction | Disjunction deriving Eq
113122
123+-- XXX use a separate data structure for each type of SVar
114124data SVarSched =
115125 LIFO -- depth first concurrent
116126 | FIFO -- breadth first concurrent
117127 | Par -- all parallel
128+ | SVarAhead
118129 deriving Eq
119130
120131-- | Identify the type of the SVar. Two computations using the same style can
@@ -155,7 +166,18 @@ data SVar m a =
155166 , doorBell :: MVar () -- signal the consumer about output
156167 , siren :: MVar () -- hooter for workers to begin work
157168
169+ -- Output synchronization mechanism for Ahead streams (Ahead and
170+ -- wAhead). We maintain a heap of out of sequence ahead of time
171+ -- generated outputs and the sequence number of the task that is
172+ -- currently at the head of the stream. Concurrent execute ahead
173+ -- tasks that have a sequence number greater than the task at the
174+ -- head should add their output to the heap.
175+ , outputHeap :: IORef (Heap (Entry Int (AheadHeapEntry a)), Int)
176+ , heapSiren :: MVar () -- hooter for workers to continue
177+ -- adding to the heap
178+
158179 -- Shared work queue
180+ , workQueue :: IORef ([Stream m a], Int)
159181 , enqueue :: Stream m a -> IO ()
160182 , queueEmpty :: m Bool
161183 , runqueue :: m ()
@@ -308,6 +330,9 @@ doFork action exHandler =
308330
309331-- TBD Each worker can have their own queue and the consumer can empty one
310332-- queue at a time, that way contention can be reduced.
333+
334+-- | This function is used by the producer threads to queue output for the
335+-- consumer thread to consume.
311336{-# INLINE send #-}
312337send :: SVar m a -> ChildEvent a -> IO ()
313338send sv msg = do
@@ -340,6 +365,10 @@ send sv msg = do
340365sendStop :: SVar m a -> IO ()
341366sendStop sv = myThreadId >>= \tid -> send sv (ChildStop tid Nothing)
342367
368+-------------------------------------------------------------------------------
369+-- Async
370+-------------------------------------------------------------------------------
371+
343372-- Note: For purely right associated expressions this queue should have at most
344373-- one element. It grows to more than one when we have left associcated
345374-- expressions. Large left associated compositions can grow this to a
@@ -367,6 +396,10 @@ runqueueLIFO sv q = run
367396 [] -> ([], Nothing)
368397 x : xs -> (xs, Just x)
369398
399+-------------------------------------------------------------------------------
400+-- WAsync
401+-------------------------------------------------------------------------------
402+
370403{-# INLINE enqueueFIFO #-}
371404enqueueFIFO :: LinkedQueue (Stream m a) -> Stream m a -> IO ()
372405enqueueFIFO = pushL
@@ -387,6 +420,10 @@ runqueueFIFO sv q = run
387420 single a = sendit a >> run
388421 yield a r = sendit a >> liftIO (enqueueFIFO q r) >> run
389422
423+-------------------------------------------------------------------------------
424+-- Parallel
425+-------------------------------------------------------------------------------
426+
390427{-# INLINE runOne #-}
391428runOne :: MonadIO m => SVar m a -> Stream m a -> m ()
392429runOne sv m = (runStream m) (Just sv) stop single yield
@@ -398,6 +435,226 @@ runOne sv m = (runStream m) (Just sv) stop single yield
398435 single a = sendit a >> stop
399436 yield a r = sendit a >> runOne sv r
400437
438+-------------------------------------------------------------------------------
439+-- Ahead
440+-------------------------------------------------------------------------------
441+
442+-- Lookahead streams can execute multiple tasks concurrently, ahead of time but
443+-- always serves them in the same order as they appear in the stream. To
444+-- implement lookahead streams efficiently we assign a sequence number to each
445+-- task when the task is picked up for execution. When the task finishes the
446+-- output is tagged with the same sequence number and we rearrange the outputs
447+-- in sequence based on that number.
448+--
449+-- To explain the mechanism imagine that the current task at the head of the
450+-- stream has a "token" to yield to the output. The ownership of the token is
451+-- determined by the current sequence number maintained in outputHeap. Sequence
452+-- number is assigned when a task is queued. When a thread dequeues a task it
453+-- picks up the sequence number as well and when the output is ready it uses
454+-- the sequence number to queue the output to the outputQueue.
455+--
456+-- The thread with current sequence number sends the output directly to the
457+-- outputQueue. Other threads push the output to the outputHeap. When the task
458+-- being queued on the heap is a stream of many elements we yield the first
459+-- element on the heap and an MVar to wakeup for more elements. When such a
460+-- task gets the "token" for outputQueue it can directly keep yielding all the
461+-- elements to the outputQueue without checking for the "token".
462+--
463+-- Note that no two outputs in the heap can have the same sequence numbers and
464+-- therefore we do not need a stable heap. We have also separated the buffer
465+-- for the current task (outputQueue) and the pending tasks (outputHeap) so
466+-- that the pending tasks cannot interfere with the current task. Note that for
467+-- a single task just the outputQueue is enough and for the case of many
468+-- threads just a heap is good enough. However we balance between these two
469+-- cases, so that both are efficient.
470+--
471+-- For bigger streams it may make sense to have separate buffers for each
472+-- stream. However, for singleton streams this may become inefficient. However,
473+-- if we do not have separate buffers, then the streams that come later in
474+-- sequence may hog the buffer, hindering the streams that are ahead. For this
475+-- reason we have a single element buffer limitation for the streams being
476+-- executed in advance.
477+--
478+-- This scheme works pretty efficiently with only 50% more overhead compared to
479+-- the Async streams where we do not have any kind of sequencing of the
480+-- outputs. It is especially devised so that we are most efficient when we have
481+-- short tasks and need just a single thread. Also when a thread yields many
482+-- items it can hold lockfree access to the outputQueue and do it efficiently.
483+--
484+-- XXX Maybe we can start the ahead threads at a lower cpu and IO priority so
485+-- that they do not hog the resources and hinder the progress of the threads in
486+-- front of them.
487+
488+{-# INLINE sendToHeap #-}
489+sendToHeap :: SVar m a -> AheadHeapEntry a -> Int -> IO Bool
490+sendToHeap sv msg seqNo = do
491+ let maxHeapSize = 1500
492+
493+ -- Maybe we can use a lock instead of using CAS
494+ -- insertions in the heap may be expensive if the heap is big
495+ -- Note: len is valid only when added is True
496+ (len, added) <- atomicModifyIORefCAS (outputHeap sv) $
497+ \ref@(h, snum) ->
498+ if seqNo > snum
499+ then ((H.insert (Entry seqNo msg) h, snum), (H.size h, True))
500+ else (ref, (0, False))
501+
502+ -- this is only for the case when the consumer generated siren can get
503+ -- missed due to race.
504+ when (len < maxHeapSize) $ void $ tryPutMVar (heapSiren sv) ()
505+
506+ if added
507+ then do
508+ when (len + 1 >= maxHeapSize) $ do
509+ takeMVar (heapSiren sv)
510+ (h, _) <- readIORef (outputHeap sv)
511+ when (H.size h <= maxHeapSize) $ void $ tryPutMVar (siren sv) ()
512+ return False
513+ else do
514+ case msg of
515+ AheadYieldSingle a -> send sv (ChildYield a)
516+ AheadYieldMany a mv -> do
517+ send sv (ChildYield a)
518+ putMVar mv ()
519+ return True
520+
521+popFromHeap
522+ :: IORef (Heap (Entry Int (AheadHeapEntry a)), Int)
523+ -> Int
524+ -> IO (Maybe (AheadHeapEntry a))
525+popFromHeap hp nextSeqNo =
526+ atomicModifyIORefCAS hp $ \(h, snum) -> do
527+ let r = H.uncons h
528+ let n = max nextSeqNo snum
529+ case r of
530+ Nothing -> ((h, n), Nothing)
531+ Just (Entry seqNo ev, hp') ->
532+ if (seqNo <= n)
533+ then ((hp', seqNo), Just ev)
534+ else ((h, n), Nothing)
535+
536+-- | As soon as we set the nextSeqNo in outputHeap we have transferred the
537+-- output "token" to the thread executing the task with that sequence number.
538+-- We have to make sure that there is no entry in the heap that has nextSeqNo
539+-- before we set nextSeqNo as the expected sequence number in the heap. It is
540+-- possible that when we are clearing the heap entry of the current sequence
541+-- number the next sequence number gets added and then we have to clear that,
542+-- this can keep going on and on.
543+drainHeap :: SVar m a -> Int -> IO ()
544+drainHeap sv nextSeqNo = do
545+ r <- popFromHeap (outputHeap sv) nextSeqNo
546+ case r of
547+ Nothing -> return ()
548+ Just ev -> do
549+ case ev of
550+ AheadYieldSingle a -> do
551+ send sv (ChildYield a)
552+ atomicModifyIORefCAS_ (outputHeap sv) $ \(h, snum) ->
553+ (h, snum + 1)
554+ drainHeap sv nextSeqNo
555+ AheadYieldMany a mv -> do
556+ send sv (ChildYield a)
557+ putMVar mv ()
558+
559+-- | The thread that has the token has the responsibility to move any pending
560+-- entries in the heap to the outputQueue. When the token is assigned to a
561+-- sequence number we must not have any entries for that sequence number in the
562+-- heap.
563+releaseToken :: SVar m a -> Int -> IO ()
564+releaseToken sv nextSeqNo = do
565+ drainHeap sv nextSeqNo
566+ void $ tryPutMVar (heapSiren sv) ()
567+
568+-- | When we have the output "token" we are allowed to append to the
569+-- outputQueue otherwise we have to insert the output in the outputHeap.
570+-- Returns whether we still have the token.
571+handleOutput :: SVar m a
572+ -> Bool
573+ -> Int
574+ -> Int
575+ -> a
576+ -> AheadHeapEntry a
577+ -> IO Bool
578+handleOutput sv haveToken prevSeqNo seqNo a ev = do
579+ if haveToken
580+ then do
581+ let nextSeqNo = prevSeqNo + 1
582+ if seqNo == nextSeqNo
583+ then send sv (ChildYield a) >> return True
584+ else do
585+ releaseToken sv nextSeqNo
586+ toHeap
587+ else toHeap
588+ where toHeap = sendToHeap sv ev seqNo
589+
590+-- Left associated ahead expressions are expensive. We start a new SVar for
591+-- each left associative expression. The queue is used only for right
592+-- associated expression, we queue the right expression and execute the left.
593+-- Thererefore the queue never has more than on item in it.
594+{-# INLINE enqueueAhead #-}
595+enqueueAhead :: IORef ([Stream m a], Int) -> Stream m a -> IO ()
596+enqueueAhead q m = do
597+ atomicModifyIORefCAS_ q $ \ case
598+ ([], n) -> ([m], n + 1) -- increment sequence
599+ _ -> error "not empty"
600+
601+runqueueAhead :: MonadIO m => SVar m a -> IORef ([Stream m a], Int) -> m ()
602+runqueueAhead sv q = run False (-1)
603+
604+ where
605+
606+ run haveToken prevSeqNo = do
607+ work <- dequeue
608+ case work of
609+ Nothing -> liftIO $ do
610+ when haveToken $ releaseToken sv (prevSeqNo + 1)
611+ sendStop sv
612+ Just (m, seqNo) -> do
613+ let tStatus = haveToken || seqNo == 0
614+ let stop = do
615+ let nextSeqNo = prevSeqNo + 1
616+ if tStatus && seqNo /= nextSeqNo
617+ then do
618+ liftIO $ releaseToken sv nextSeqNo
619+ run False seqNo
620+ else run tStatus seqNo
621+ let single a = do
622+ tokenStatus <- liftIO $
623+ handleOutput sv tStatus prevSeqNo seqNo a
624+ (AheadYieldSingle a)
625+ run tokenStatus seqNo
626+
627+ -- If we are yielding more than one item and we do not have the
628+ -- token then we insert the current item in the pending
629+ -- outputHeap and wait for our turn.
630+ let
631+ stopNext = run True seqNo
632+ singleNext a = do
633+ liftIO $ send sv (ChildYield a)
634+ run True seqNo
635+ yieldNext a r = do
636+ liftIO $ send sv (ChildYield a)
637+ (runStream r) (Just sv) stopNext singleNext yieldNext
638+
639+ yield a r = do
640+ -- XXX we can allocate the MVar when actually inserting
641+ -- in the heap and then return it.
642+ liftIO $ do
643+ mv <- newEmptyMVar
644+ gotToken <-
645+ handleOutput sv tStatus prevSeqNo seqNo a
646+ (AheadYieldMany a mv)
647+ when (not gotToken) $ void $ takeMVar mv
648+ (runStream r) (Just sv) stopNext singleNext yieldNext
649+
650+ (runStream m) (Just sv) stop single yield
651+
652+ dequeue = liftIO $ do
653+ atomicModifyIORefCAS q $ \case
654+ ([], n) -> (([], n), Nothing)
655+ (x : [], n) -> (([], n), Just (x, n))
656+ _ -> error "more than one item on queue"
657+
401658-- Thread tracking is needed for two reasons:
402659--
403660-- 1) Killing threads on exceptions. Threads may not be left to go away by
@@ -481,7 +738,9 @@ sendWorkerWait sv = do
481738 (_, n) <- liftIO $ readIORef (outputQueue sv)
482739 when (n <= 0) $ do
483740 -- XXX the queue may be empty temporarily if the worker has dequeued
484− -- the work item but has not enqueued the remaining part yet.
741+ -- the work item but has not enqueued the remaining part yet. For the
742+ -- same reason, a worker may come back if it tries to dequeue when the
743+ -- queue is empty, even though the whole work has not finished yet.
485744 done <- queueEmpty sv
486745 if not done
487746 then do
@@ -553,7 +812,10 @@ getFifoSVar ctype = do
553812 SVar { outputQueue = outQ
554813 , doorBell = outQMv
555814 , siren = hooter
815+ , outputHeap = undefined
816+ , heapSiren = undefined
556817 , runningThreads = running
818+ , workQueue = undefined
557819 , runqueue = runqueueFIFO sv q
558820 , enqueue = pushL q
559821 , queueEmpty = liftIO $ nullQ q
@@ -575,7 +837,10 @@ getLifoSVar ctype = do
575837 SVar { outputQueue = outQ
576838 , doorBell = outQMv
577839 , siren = hooter
840+ , outputHeap = undefined
841+ , heapSiren = undefined
578842 , runningThreads = running
843+ , workQueue = undefined
579844 , runqueue = runqueueLIFO sv q
580845 , enqueue = enqueueLIFO q
581846 , queueEmpty = checkEmpty
@@ -595,7 +860,10 @@ getParSVar style = do
595860 SVar { outputQueue = outQ
596861 , doorBell = outQMv
597862 , siren = hooter
863+ , outputHeap = undefined
864+ , heapSiren = undefined
598865 , runningThreads = running
866+ , workQueue = undefined
599867 , runqueue = undefined
600868 , enqueue = undefined
601869 , queueEmpty = undefined
@@ -604,6 +872,35 @@ getParSVar style = do
604872 }
605873 in return sv
606874
875+getAheadSVar :: MonadIO m => SVarStyle -> IO (SVar m a)
876+getAheadSVar style = do
877+ outQ <- newIORef ([], 0)
878+ outH <- newIORef (H.empty, 0)
879+ outQMv <- newEmptyMVar
880+ hooter <- newEmptyMVar
881+ hsiren <- newEmptyMVar
882+ active <- newIORef 0
883+ running <- newIORef S.empty
884+ q <- newIORef ([], -1)
885+ let checkEmpty = liftIO $ do
886+ (xs, _) <- readIORef q
887+ return $ null xs
888+ let sv =
889+ SVar { outputQueue = outQ
890+ , doorBell = outQMv
891+ , siren = hooter
892+ , outputHeap = outH
893+ , heapSiren = hsiren
894+ , runningThreads = running
895+ , workQueue = q
896+ , runqueue = runqueueAhead sv q
897+ , enqueue = undefined
898+ , queueEmpty = checkEmpty
899+ , svarStyle = style
900+ , activeWorkers = active
901+ }
902+ in return sv
903+
607904-- | Create a new empty SVar.
608905newEmptySVar :: MonadAsync m => SVarStyle -> m (SVar m a)
609906newEmptySVar style = do
@@ -612,6 +909,7 @@ newEmptySVar style = do
612909 SVarStyle _ FIFO -> getFifoSVar style
613910 SVarStyle _ LIFO -> getLifoSVar style
614911 SVarStyle _ Par -> getParSVar style
912+ SVarStyle _ SVarAhead -> getAheadSVar style
615913
616914-- | Create a new SVar and enqueue one stream computation on it.
617915newStreamVar1 :: MonadAsync m => SVarStyle -> Stream m a -> m (SVar m a)
@@ -624,6 +922,17 @@ newStreamVar1 style m = do
624922 pushWorker sv
625923 return sv
626924
925+-- | Create a new SVar and enqueue one stream computation on it.
926+newStreamVarAhead :: MonadAsync m => Stream m a -> m (SVar m a)
927+newStreamVarAhead m = do
928+ sv <- newEmptySVar (SVarStyle Disjunction SVarAhead)
929+ -- Note: We must have all the work on the queue before sending the
930+ -- pushworker, otherwise the pushworker may exit before we even get a
931+ -- chance to push.
932+ liftIO $ enqueueAhead (workQueue sv) m
933+ pushWorker sv
934+ return sv
935+
627936-- | Write a stream to an 'SVar' in a non-blocking manner. The stream can then
628937-- be read back from the SVar using 'fromSVar'.
629938toStreamVar :: MonadAsync m => SVar m a -> Stream m a -> m ()
@@ -727,6 +1036,25 @@ joinStreamVarPar style m1 m2 = Stream $ \svr stp sng yld ->
7271036 pushWorkerPar sv m2
7281037 (runStream (fromStreamVar sv)) Nothing stp sng yld
7291038
1039+{-# INLINE ahead #-}
1040+ahead :: MonadAsync m => Stream m a -> Stream m a -> Stream m a
1041+ahead m1 m2 = Stream $ \svr stp sng yld -> do
1042+ let style = SVarStyle Disjunction SVarAhead
1043+ case svr of
1044+ Just sv | svarStyle sv == style -> do
1045+ liftIO $ enqueueAhead (workQueue sv) m2
1046+ -- Always run the left side on a new SVar to avoid complexity in
1047+ -- sequencing results. This means the left side cannot further
1048+ -- split into more ahead computations on the same SVar.
1049+ (runStream m1) Nothing stp sng yld
1050+ _ -> do
1051+ sv <- newStreamVarAhead (concurrently m1 m2)
1052+ (runStream (fromStreamVar sv)) Nothing stp sng yld
1053+ where
1054+ concurrently ma mb = Stream $ \svr stp sng yld -> do
1055+ liftIO $ enqueueAhead (workQueue (fromJust svr)) mb
1056+ (runStream ma) Nothing stp sng yld
1057+
7301058------------------------------------------------------------------------------
7311059-- Semigroup and Monoid style compositions for parallel actions
7321060------------------------------------------------------------------------------