@@ -20,17 +20,22 @@ module Streamly.Internal.Data.Stream.Async.Channel.Append
22import Control.Concurrent (myThreadId)
22import Control.Concurrent (myThreadId)
23−import Control.Concurrent.MVar (newEmptyMVar)
23+import Control.Concurrent.MVar (newEmptyMVar, newMVar, putMVar, takeMVar)
24+import Control.Exception (assert)
24import Control.Monad (when, void)
25import Control.Monad (when, void)
25import Control.Monad.IO.Class (MonadIO(liftIO))
26import Control.Monad.IO.Class (MonadIO(liftIO))
26import Control.Monad.Trans.Control (MonadBaseControl, restoreM)
27import Control.Monad.Trans.Control (MonadBaseControl, restoreM)
28+import Data.Heap (Heap, Entry(..))
27−import Data.IORef (IORef, newIORef, readIORef)
29+import Data.IORef (IORef, newIORef, readIORef, atomicModifyIORef, writeIORef)
30+import Data.Kind (Type)
31+import GHC.Exts (inline)
28import Streamly.Internal.Control.Concurrent
32import Streamly.Internal.Control.Concurrent
29 (MonadRunInIO, MonadAsync, RunInIO(..), askRunInIO)
33 (MonadRunInIO, MonadAsync, RunInIO(..), askRunInIO)
30import Streamly.Internal.Data.Atomics
34import Streamly.Internal.Data.Atomics
31 (atomicModifyIORefCAS, atomicModifyIORefCAS_)
35 (atomicModifyIORefCAS, atomicModifyIORefCAS_)
32import Streamly.Internal.Data.Stream.Channel.Dispatcher (modifyThread)
36import Streamly.Internal.Data.Stream.Channel.Dispatcher (modifyThread)
38+import qualified Data.Heap as H
34import qualified Data.Set as Set
39import qualified Data.Set as Set
35import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
40import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
@@ -38,9 +43,10 @@ import Streamly.Internal.Data.Stream.Async.Channel.Consumer
38import Streamly.Internal.Data.Stream.Async.Channel.Dispatcher
43import Streamly.Internal.Data.Stream.Async.Channel.Dispatcher
39import Streamly.Internal.Data.Stream.Async.Channel.Type
44import Streamly.Internal.Data.Stream.Async.Channel.Type
40import Streamly.Internal.Data.Stream.Channel.Types
45import Streamly.Internal.Data.Stream.Channel.Types
46+import Streamly.Internal.Data.Stream.Channel.Worker
42------------------------------------------------------------------------------
48------------------------------------------------------------------------------
43−-- Creating a channel
49+-- Concurrent streams with first-come-first serve results
44------------------------------------------------------------------------------
50------------------------------------------------------------------------------
46-- Note: For purely right associated expressions this queue should have at most
52-- Note: For purely right associated expressions this queue should have at most
@@ -207,6 +213,679 @@ workLoopLIFOLimited qref sv winfo = run
207 liftIO $ enqueueLIFO sv qref inner (runInIO, r)
213 liftIO $ enqueueLIFO sv qref inner (runInIO, r)
208 return Suspend
214 return Suspend
216+-------------------------------------------------------------------------------
217+-- Ahead Channel Data Structures
218+-------------------------------------------------------------------------------
220+-- XXX Left associated ahead expressions are expensive. We start a new SVar for
221+-- each left associative expression. The queue is used only for right
222+-- associated expression, we queue the right expression and execute the left.
223+-- Thererefore the queue never has more than one item in it.
225+-- XXX we can fix this. When we queue more than one item on the queue we can
226+-- mark the previously queued item as not-runnable. The not-runnable item is
227+-- not dequeued until the already running one has finished and at that time we
228+-- would also know the exact sequence number of the already queued item.
230+-- we can even run the already queued items but they will have to be sorted in
231+-- layers in the heap. We can use a list of heaps for that.
232+{-# INLINE enqueueAhead #-}
235+ -> IORef ([K.Stream m a], Int)
236+ -> (RunInIO m, K.Stream m a)
238+enqueueAhead sv q m = do
239+ atomicModifyIORefCAS_ q $ \ case
240+ ([], n) -> ([snd m], n + 1) -- increment sequence
241+ _ -> error "enqueueAhead: queue is not empty"
242+ ringDoorBell (needDoorBell sv) (outputDoorBell sv)
244+-- enqueue without incrementing the sequence number
245+{-# INLINE reEnqueueAhead #-}
247+ Channel m a -> IORef ([K.Stream m a], Int) -> K.Stream m a -> IO ()
248+reEnqueueAhead sv q m = do
249+ atomicModifyIORefCAS_ q $ \ case
250+ ([], n) -> ([m], n) -- DO NOT increment sequence
251+ _ -> error "reEnqueueAhead: queue is not empty"
252+ ringDoorBell (needDoorBell sv) (outputDoorBell sv)
254+-- Normally the thread that has the token should never go away. The token gets
255+-- handed over to another thread, but someone or the other has the token at any
256+-- point of time. But if the task that has the token finds that the outputQueue
257+-- is full, in that case it can go away without even handing over the token to
258+-- another thread. In that case it sets the nextSequence number in the heap its
259+-- own sequence number before going away. To handle this case, any task that
260+-- does not have the token tries to dequeue from the heap first before
261+-- dequeuing from the work queue. If it finds that the task at the top of the
262+-- heap is the one that owns the current sequence number then it grabs the
263+-- token and starts with that.
265+-- XXX instead of queueing just the head element and the remaining computation
266+-- on the heap, evaluate as many as we can and place them on the heap. But we
267+-- need to give higher priority to the lower sequence numbers so that lower
268+-- priority tasks do not fill up the heap making higher priority tasks block
269+-- due to full heap. Maybe we can have a weighted space for them in the heap.
270+-- The weight is inversely proportional to the sequence number.
272+-- XXX review for livelock
274+{-# INLINE dequeueAhead #-}
275+dequeueAhead :: MonadIO m
276+ => IORef ([t m a], Int) -> m (Maybe (t m a, Int))
277+dequeueAhead q = liftIO $
278+ atomicModifyIORefCAS q $ \case
279+ ([], n) -> (([], n), Nothing)
280+ (x : [], n) -> (([], n), Just (x, n))
281+ _ -> error "more than one item on queue"
283+-------------------------------------------------------------------------------
285+-------------------------------------------------------------------------------
287+withIORef :: IORef a -> (a -> IO b) -> IO b
288+withIORef ref f = readIORef ref >>= f
290+atomicModifyIORef_ :: IORef a -> (a -> a) -> IO ()
291+atomicModifyIORef_ ref f =
292+ atomicModifyIORef ref $ \x -> (f x, ())
294+data AheadHeapEntry (t :: (Type -> Type) -> Type -> Type) m a =
297+ | AheadEntryStream (RunInIO m, t m a)
299+data HeapDequeueResult t m a =
302+ | Ready (Entry Int (AheadHeapEntry t m a))
304+{-# INLINE dequeueFromHeap #-}
306+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
307+ -> IO (HeapDequeueResult t m a)
308+dequeueFromHeap hpVar =
309+ atomicModifyIORef hpVar $ \pair@(hp, snum) ->
311+ Nothing -> (pair, Clearing)
315+ Just (ent@(Entry seqNo _ev), hp') ->
317+ then ((hp', Nothing), Ready ent)
318+ else assert (seqNo >= n) (pair, Waiting n)
319+ Nothing -> (pair, Waiting n)
321+{-# INLINE dequeueFromHeapSeq #-}
323+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
325+ -> IO (HeapDequeueResult t m a)
326+dequeueFromHeapSeq hpVar i =
327+ atomicModifyIORef hpVar $ \(hp, snum) ->
332+ Just (ent@(Entry seqNo _ev), hp') ->
334+ then ((hp', Nothing), Ready ent)
335+ else assert (seqNo >= i) ((hp, Just i), Waiting i)
336+ Nothing -> ((hp, Just i), Waiting i)
337+ Just _ -> error "dequeueFromHeapSeq: unreachable"
339+heapIsSane :: Maybe Int -> Int -> Bool
340+heapIsSane snum seqNo =
343+ Just n -> seqNo >= n
345+{-# INLINE requeueOnHeapTop #-}
347+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
348+ -> Entry Int (AheadHeapEntry t m a)
351+requeueOnHeapTop hpVar ent seqNo =
352+ atomicModifyIORef_ hpVar $ \(hp, snum) ->
353+ assert (heapIsSane snum seqNo) (H.insert ent hp, Just seqNo)
355+{-# INLINE updateHeapSeq #-}
357+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
360+updateHeapSeq hpVar seqNo =
361+ atomicModifyIORef_ hpVar $ \(hp, snum) ->
362+ assert (heapIsSane snum seqNo) (hp, Just seqNo)
364+------------------------------------------------------------------------------
365+-- Ahead: Concurrent streams with ordered results
366+------------------------------------------------------------------------------
368+-- Lookahead streams can execute multiple tasks concurrently, ahead of time,
369+-- but always serve them in the same order as they appear in the stream. To
370+-- implement lookahead streams efficiently we assign a sequence number to each
371+-- task when the task is picked up for execution. When the task finishes, the
372+-- output is tagged with the same sequence number and we rearrange the outputs
373+-- in sequence based on that number.
375+-- To explain the mechanism imagine that the current task at the head of the
376+-- stream has a "token" to yield to the outputQueue. The ownership of the token
377+-- is determined by the current sequence number is maintained in outputHeap.
378+-- Sequence number is assigned when a task is queued. When a thread dequeues a
379+-- task it picks up the sequence number as well and when the output is ready it
380+-- uses the sequence number to queue the output to the outputQueue.
382+-- The thread with current sequence number sends the output directly to the
383+-- outputQueue. Other threads push the output to the outputHeap. When the task
384+-- being queued on the heap is a stream of many elements we evaluate only the
385+-- first element and keep the rest of the unevaluated computation in the heap.
386+-- When such a task gets the "token" for outputQueue it evaluates and directly
387+-- yields all the elements to the outputQueue without checking for the
390+-- Note that no two outputs in the heap can have the same sequence numbers and
391+-- therefore we do not need a stable heap. We have also separated the buffer
392+-- for the current task (outputQueue) and the pending tasks (outputHeap) so
393+-- that the pending tasks cannot interfere with the current task. Note that for
394+-- a single task just the outputQueue is enough and for the case of many
395+-- threads just a heap is good enough. However we balance between these two
396+-- cases, so that both are efficient.
398+-- For bigger streams it may make sense to have separate buffers for each
399+-- stream. However, for singleton streams this may become inefficient. However,
400+-- if we do not have separate buffers, then the streams that come later in
401+-- sequence may hog the buffer, hindering the streams that are ahead. For this
402+-- reason we have a single element buffer limitation for the streams being
403+-- executed in advance.
405+-- This scheme works pretty efficiently with less than 40% extra overhead
406+-- compared to the Async streams where we do not have any kind of sequencing of
407+-- the outputs. It is especially devised so that we are most efficient when we
408+-- have short tasks and need just a single thread. Also when a thread yields
409+-- many items it can hold lockfree access to the outputQueue and do it
412+-- XXX Maybe we can start the ahead threads at a lower cpu and IO priority so
413+-- that they do not hog the resources and hinder the progress of the threads in
416+-- Left associated ahead expressions are expensive. We start a new SVar for
417+-- each left associative expression. The queue is used only for right
418+-- associated expression, we queue the right expression and execute the left.
419+-- Thererefore the queue never has more than on item in it.
421+-- XXX Also note that limiting concurrency for cases like "take 10" would not
422+-- work well with left associative expressions, because we have no visibility
423+-- about how much the left side of the expression would yield.
425+-- XXX It may be a good idea to increment sequence numbers for each yield,
426+-- currently a stream on the left side of the expression may yield many
427+-- elements with the same sequene number. We can then use the seq number to
428+-- enforce yieldMax and yieldLImit as well.
432+-- * A worker should always ensure that it pushes all the consecutive items in
433+-- the heap to the outputQueue especially the items on behalf of the workers
434+-- that have already left when we were holding the token. This avoids deadlock
435+-- conditions when the later workers completion depends on the consumption of
436+-- earlier results. For more details see comments in the consumer pull side
439+{-# INLINE underMaxHeap #-}
442+ -> Heap (Entry Int (AheadHeapEntry K.Stream m a))
444+underMaxHeap sv hp = do
445+ (_, len) <- readIORef (outputQueue sv)
447+ -- XXX simplify this
448+ let maxHeap = case maxBufferLimit sv of
449+ Limited lim -> Limited $
450+ max 0 (lim - fromIntegral len)
451+ Unlimited -> Unlimited
455+ active <- readIORef (workerCount sv)
456+ return $ H.size hp + active <= fromIntegral lim
457+ Unlimited -> return True
464+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)) , Maybe Int)
466+preStopCheck sv heap =
467+ -- check the stop condition under a lock before actually
468+ -- stopping so that the whole herd does not stop at once.
469+ withIORef heap $ \(hp, _) -> do
470+ heapOk <- underMaxHeap sv hp
471+ takeMVar (workerStopMVar sv)
473+ putMVar (workerStopMVar sv) ()
476+ putMVar (workerStopMVar sv) ()
480+ case yieldRateInfo sv of
485+ (maxWorkerLimit sv) (workerCount sv) yinfo
486+ if rateOk then continue else stopping
490+ IORef ([K.Stream m a], Int)
495+abortExecution q sv winfo m = do
496+ reEnqueueAhead sv q m
497+ incrementYieldLimit (remainingWork sv)
500+-- XXX In absence of a "noyield" primitive (i.e. do not pre-empt inside a
501+-- critical section) from GHC RTS, we have a difficult problem. Assume we have
502+-- a 100,000 threads producing output and queuing it to the heap for
503+-- sequencing. The heap can be drained only by one thread at a time, any thread
504+-- that finds that heap can be drained now, takes a lock and starts draining
505+-- it, however the thread may get prempted in the middle of it holding the
506+-- lock. Since that thread is holding the lock, the other threads cannot pick
507+-- up the draining task, therefore they proceed to picking up the next task to
508+-- execute. If the draining thread could yield voluntarily at a point where it
509+-- has released the lock, then the next threads could pick up the draining
510+-- instead of executing more tasks. When there are 100,000 threads the drainer
511+-- gets a cpu share to run only 1:100000 of the time. This makes the heap
512+-- accumulate a lot of output when we the buffer size is large.
514+-- The solutions to this problem are:
515+-- 1) make the other threads wait in a queue until the draining finishes
516+-- 2) make the other threads queue and go away if draining is in progress
518+-- In both cases we give the drainer a chance to run more often.
522+ => IORef ([K.Stream m a], Int)
523+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
526+ -> AheadHeapEntry K.Stream m a
528+ -> Bool -- we are draining the heap before we stop
530+processHeap q heap sv winfo entry sno stopping = loopHeap sno entry
534+ stopIfNeeded ent seqNo r = do
535+ stopIt <- liftIO $ preStopCheck sv heap
538+ -- put the entry back in the heap and stop
539+ requeueOnHeapTop heap (Entry seqNo ent) seqNo
541+ else runStreamWithYieldLimit True seqNo r
543+ loopHeap seqNo ent =
545+ AheadEntryNull -> nextHeap seqNo
546+ AheadEntryPure a -> do
547+ -- Use 'send' directly so that we do not account this in worker
548+ -- latency as this will not be the real latency.
549+ -- Don't stop the worker in this case as we are just
550+ -- transferring available results from heap to outputQueue.
554+ (outputQueue sv) (outputDoorBell sv) (ChildYield a)
556+ AheadEntryStream (RunInIO runin, r) -> do
558+ then stopIfNeeded ent seqNo r
560+ res <- liftIO $ runin (runStreamWithYieldLimit True seqNo r)
563+ nextHeap prevSeqNo = do
564+ res <- liftIO $ dequeueFromHeapSeq heap (prevSeqNo + 1)
566+ Ready (Entry seqNo hent) -> loopHeap seqNo hent
567+ Clearing -> liftIO $ stop sv winfo
571+ r <- liftIO $ preStopCheck sv heap
573+ then liftIO $ stop sv winfo
574+ else processWorkQueue prevSeqNo
575+ else inline processWorkQueue prevSeqNo
577+ processWorkQueue prevSeqNo = do
578+ work <- dequeueAhead q
580+ Nothing -> liftIO $ stop sv winfo
581+ Just (m, seqNo) -> do
582+ yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
585+ if seqNo == prevSeqNo + 1
586+ then processWithToken q heap sv winfo m seqNo
587+ else processWithoutToken q heap sv winfo m seqNo
588+ else liftIO $ abortExecution q sv winfo m
590+ -- We do not stop the worker on buffer full here as we want to proceed to
591+ -- nextHeap anyway so that we can clear any subsequent entries. We stop
592+ -- only in yield continuation where we may have a remaining stream to be
593+ -- pushed on the heap.
594+ singleStreamFromHeap seqNo a = do
595+ void $ liftIO $ yield sv winfo a
598+ -- XXX when we have an unfinished stream on the heap we cannot account all
599+ -- the yields of that stream until it finishes, so if we have picked up
600+ -- and executed more actions beyond that in the parent stream and put them
601+ -- on the heap then they would eat up some yield limit which is not
602+ -- correct, we will think that our yield limit is over even though we have
603+ -- to yield items from unfinished stream before them. For this reason, if
604+ -- there are pending items in the heap we drain them unconditionally
605+ -- without considering the yield limit.
606+ runStreamWithYieldLimit continue seqNo r = do
607+ _ <- liftIO $ decrementYieldLimit (remainingWork sv)
608+ if continue -- see comment above -- && yieldLimitOk
611+ liftIO (incrementYieldLimit (remainingWork sv))
613+ K.foldStreamShared undefined
614+ (yieldStreamFromHeap seqNo)
615+ (singleStreamFromHeap seqNo)
620+ let ent = Entry seqNo (AheadEntryStream (runIn, r))
622+ requeueOnHeapTop heap ent seqNo
623+ incrementYieldLimit (remainingWork sv)
626+ yieldStreamFromHeap seqNo a r = do
627+ continue <- liftIO $ yield sv winfo a
628+ runStreamWithYieldLimit continue seqNo r
630+{-# NOINLINE drainHeap #-}
633+ => IORef ([K.Stream m a], Int)
634+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
638+drainHeap q heap sv winfo = do
639+ r <- liftIO $ dequeueFromHeap heap
641+ Ready (Entry seqNo hent) ->
642+ processHeap q heap sv winfo hent seqNo True
643+ _ -> liftIO $ stop sv winfo
645+data HeapStatus = HContinue | HStop
649+ => IORef ([K.Stream m a], Int)
650+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
656+processWithoutToken q heap sv winfo m seqNo = do
657+ -- we have already decremented the yield limit for m
659+ liftIO (incrementYieldLimit (remainingWork sv))
660+ -- If the stream stops without yielding anything, and we do not put
661+ -- anything on heap, but if heap was waiting for this seq number
662+ -- then it will keep waiting forever, because we are never going to
663+ -- put it on heap. So we have to put a null entry on heap even when
665+ toHeap AheadEntryNull
666+ mrun = runInIO $ svarMrun sv
668+ r <- liftIO $ mrun $
669+ K.foldStreamShared undefined
672+ toHeap $ AheadEntryStream (runIn, K.cons a r))
673+ (toHeap . AheadEntryPure)
678+ Continue -> workLoopAhead q heap sv winfo
679+ Suspend -> drainHeap q heap sv winfo
683+ -- XXX to reduce contention each CPU can have its own heap
685+ -- Heap insertion is an expensive affair so we use a non CAS based
686+ -- modification, otherwise contention and retries can make a thread
687+ -- context switch and throw it behind other threads which come later in
689+ newHp <- liftIO $ atomicModifyIORef heap $ \(hp, snum) ->
690+ let hp' = H.insert (Entry seqNo ent) hp
691+ in assert (heapIsSane snum seqNo) ((hp', snum), hp')
693+ when (svarInspectMode sv) $
695+ maxHp <- readIORef (maxHeapSize $ svarStats sv)
696+ when (H.size newHp > maxHp) $
697+ writeIORef (maxHeapSize $ svarStats sv) (H.size newHp)
699+ heapOk <- liftIO $ underMaxHeap sv newHp
701+ case yieldRateInfo sv of
702+ Nothing -> return HContinue
714+ then return HContinue
716+ Nothing -> return HContinue
721+ HContinue -> return Continue
722+ HStop -> return Suspend
725+data TokenWorkerStatus = TokenContinue Int | TokenSuspend
729+ => IORef ([K.Stream m a], Int)
730+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
736+processWithToken q heap sv winfo action sno = do
737+ -- Note, we enter this function with yield limit already decremented
738+ -- XXX deduplicate stop in all invocations
740+ liftIO (incrementYieldLimit (remainingWork sv))
741+ return $ TokenContinue (sno + 1)
742+ mrun = runInIO $ svarMrun sv
747+ $ K.foldStreamShared
748+ undefined (yieldOutput sno) (singleOutput sno) stopk action
752+ TokenContinue seqNo -> loopWithToken seqNo
753+ TokenSuspend -> drainHeap q heap sv winfo
757+ singleOutput seqNo a = do
758+ continue <- liftIO $ yield sv winfo a
760+ then return $ TokenContinue (seqNo + 1)
762+ liftIO $ updateHeapSeq heap (seqNo + 1)
765+ -- XXX use a wrapper function around stop so that we never miss
766+ -- incrementing the yield in a stop continuation. Essentiatlly all
767+ -- "unstream" calls in this function must increment yield limit on stop.
768+ yieldOutput seqNo a r = do
769+ continue <- liftIO $ yield sv winfo a
770+ yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
771+ if continue && yieldLimitOk
774+ liftIO (incrementYieldLimit (remainingWork sv))
775+ return $ TokenContinue (seqNo + 1)
776+ K.foldStreamShared undefined
778+ (singleOutput seqNo)
783+ let ent = Entry seqNo (AheadEntryStream (runIn, r))
784+ liftIO $ requeueOnHeapTop heap ent seqNo
785+ liftIO $ incrementYieldLimit (remainingWork sv)
788+ loopWithToken nextSeqNo = do
789+ work <- dequeueAhead q
792+ liftIO $ updateHeapSeq heap nextSeqNo
793+ workLoopAhead q heap sv winfo
795+ Just (m, seqNo) -> do
796+ yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
797+ let undo = liftIO $ do
798+ updateHeapSeq heap nextSeqNo
799+ reEnqueueAhead sv q m
800+ incrementYieldLimit (remainingWork sv)
803+ if seqNo == nextSeqNo
806+ liftIO (incrementYieldLimit (remainingWork sv))
807+ return $ TokenContinue (seqNo + 1)
808+ mrun = runInIO $ svarMrun sv
809+ r <- liftIO $ mrun $
810+ K.foldStreamShared undefined
812+ (singleOutput seqNo)
817+ TokenContinue seqNo1 -> loopWithToken seqNo1
818+ TokenSuspend -> drainHeap q heap sv winfo
821+ -- To avoid a race when another thread puts something
822+ -- on the heap and goes away, the consumer will not get
823+ -- a doorBell and we will not clear the heap before
824+ -- executing the next action. If the consumer depends
825+ -- on the output that is stuck in the heap then this
826+ -- will result in a deadlock. So we always clear the
827+ -- heap before executing the next action.
828+ undo >> workLoopAhead q heap sv winfo
829+ else undo >> drainHeap q heap sv winfo
831+-- XXX the yield limit changes increased the performance overhead by 30-40%.
832+-- Just like AsyncT we can use an implementation without yeidlimit and even
833+-- without pacing code to keep the performance higher in the unlimited and
836+-- XXX The yieldLimit stuff is pretty invasive. We can instead do it by using
837+-- three hooks, a pre-execute hook, a yield hook and a stop hook. In fact these
838+-- hooks can be used for a more general implementation to even check predicates
839+-- and not just yield limit.
843+ => IORef ([K.Stream m a], Int)
844+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
848+workLoopAhead q heap sv winfo = do
849+ r <- liftIO $ dequeueFromHeap heap
851+ Ready (Entry seqNo hent) ->
852+ processHeap q heap sv winfo hent seqNo False
853+ Clearing -> liftIO $ stop sv winfo
855+ -- Before we execute the next item from the work queue we check
856+ -- if we are beyond the yield limit. It is better to check the
857+ -- yield limit before we pick up the next item. Otherwise we
858+ -- may have already started more tasks even though we may have
859+ -- reached the yield limit. We can avoid this by taking active
860+ -- workers into account, but that is not as reliable, because
861+ -- workers may go away without picking up work and yielding a
864+ -- Rate control can be done either based on actual yields in
865+ -- the output queue or based on any yield either to the heap or
866+ -- to the output queue. In both cases we may have one issue or
867+ -- the other. We chose to do this based on actual yields to the
868+ -- output queue because it makes the code common to both async
869+ -- and ahead streams.
871+ work <- dequeueAhead q
873+ Nothing -> liftIO $ stop sv winfo
874+ Just (m, seqNo) -> do
876+ liftIO $ decrementYieldLimit (remainingWork sv)
880+ then processWithToken q heap sv winfo m seqNo
881+ else processWithoutToken q heap sv winfo m seqNo
882+ -- If some worker decremented the yield limit but then
883+ -- did not yield anything and therefore incremented it
884+ -- later, then if we did not requeue m here we may find
885+ -- the work queue empty and therefore miss executing
886+ -- the remaining action.
887+ else liftIO $ abortExecution q sv winfo m
210-------------------------------------------------------------------------------
889-------------------------------------------------------------------------------
211-- SVar creation
890-- SVar creation
212-- This code belongs in SVar.hs but is kept here for perf reasons
891-- This code belongs in SVar.hs but is kept here for perf reasons
@@ -221,6 +900,10 @@ getLifoSVar :: forall m a. (MonadIO m, MonadBaseControl IO m) =>
221 RunInIO m -> Config -> IO (Channel m a)
900 RunInIO m -> Config -> IO (Channel m a)
222getLifoSVar mrun cfg = do
901getLifoSVar mrun cfg = do
223 outQ <- newIORef ([], 0)
902 outQ <- newIORef ([], 0)
903+ -- the second component of the tuple is "Nothing" when heap is being
904+ -- cleared, "Just n" when we are expecting sequence number n to arrive
905+ -- before we can start clearing the heap.
906+ outH <- newIORef (H.empty, Just 0)
224 outQMv <- newEmptyMVar
907 outQMv <- newEmptyMVar
225 active <- newIORef 0
908 active <- newIORef 0
226 wfw <- newIORef False
909 wfw <- newIORef False
@@ -229,9 +912,14 @@ getLifoSVar mrun cfg = do
229 ( [] :: [(RunInIO m, K.Stream m a)]
912 ( [] :: [(RunInIO m, K.Stream m a)]
230 , [] :: [(RunInIO m, K.Stream m a)]
913 , [] :: [(RunInIO m, K.Stream m a)]
915+ -- Sequence number is incremented whenever something is queued, therefore,
916+ -- first sequence number would be 0
917+ aheadQ <- newIORef ([], -1)
918+ stopMVar <- newMVar ()
232− yl <- case getYieldLimit cfg of
920+ case getYieldLimit cfg of
233− Nothing -> return Nothing
921+ Nothing -> return Nothing
234− Just x -> Just <$> newIORef x
922+ Just x -> Just <$> newIORef x
235 rateInfo <- newRateInfo cfg
923 rateInfo <- newRateInfo cfg
237 stats <- newSVarStats
925 stats <- newSVarStats
@@ -254,6 +942,7 @@ getLifoSVar mrun cfg = do
254 return $ qEmpty || yieldsDone
942 return $ qEmpty || yieldsDone
256 let eager = getEagerDispatch cfg
944 let eager = getEagerDispatch cfg
945+ inOrder = getOrdered cfg
258 let getSVar :: Channel m a
947 let getSVar :: Channel m a
259 -> (Channel m a -> m [ChildEvent a])
948 -> (Channel m a -> m [ChildEvent a])
@@ -274,17 +963,32 @@ getLifoSVar mrun cfg = do
274 , readOutputQ = readOutput sv
963 , readOutputQ = readOutput sv
275 , postProcess = postProc sv
964 , postProcess = postProc sv
276 , workerThreads = running
965 , workerThreads = running
277− , workLoop = wloop q sv
967+ , workLoop =
969+ then workLoopAhead aheadQ outH sv
974+ then enqueueAhead sv aheadQ
278− , enqueue = enqueueLIFO sv q
975+ else enqueueLIFO sv q inner
279− , eagerDispatch = when eager $ void $ dispatchWorker 0 sv
976+ , eagerDispatch = when eager $ void $ dispatchWorker 0 sv
280− , isWorkDone = workDone sv
977+ , isWorkDone =
979+ then isWorkDoneAhead sv aheadQ outH
281− , isQueueDone = workDone sv
981+ , isQueueDone =
983+ then isQueueDoneAhead sv aheadQ
282 , needDoorBell = wfw
986 , needDoorBell = wfw
283 , svarMrun = mrun
987 , svarMrun = mrun
284 , workerCount = active
988 , workerCount = active
285 -- XXX We can use delThread or modThread based on eager flag.
989 -- XXX We can use delThread or modThread based on eager flag.
286 , accountThread = modifyThread running outQMv
990 , accountThread = modifyThread running outQMv
287− , workerStopMVar = undefined
991+ , workerStopMVar = stopMVar
288 , svarRef = Nothing
992 , svarRef = Nothing
289 , svarInspectMode = getInspectMode cfg
993 , svarInspectMode = getInspectMode cfg
290 , svarCreator = tid
994 , svarCreator = tid
@@ -315,6 +1019,34 @@ getLifoSVar mrun cfg = do
315 workLoopLIFOLimited
1019 workLoopLIFOLimited
316 in return sv
1020 in return sv
1024+ {-# INLINE isQueueDoneAhead #-}
1025+ isQueueDoneAhead sv q = do
1026+ queueDone <- checkEmpty q
1028+ case remainingWork sv of
1030+ n <- readIORef yref
1032+ Nothing -> return False
1033+ -- XXX note that yieldsDone can only be authoritative only when there
1034+ -- are no workers running. If there are active workers they can
1035+ -- later increment the yield count and therefore change the result.
1036+ return $ yieldsDone || queueDone
1038+ {-# INLINE isWorkDoneAhead #-}
1039+ isWorkDoneAhead sv q ref = do
1041+ (hp, _) <- readIORef ref
1042+ return (H.size hp <= 0)
1043+ queueDone <- isQueueDoneAhead sv q
1044+ return $ heapDone && queueDone
1047+ (xs, _) <- readIORef q
318-- | Create a new async style concurrent stream evaluation channel. The monad
1050-- | Create a new async style concurrent stream evaluation channel. The monad
319-- state used to run the stream actions is taken from the call site of
1051-- state used to run the stream actions is taken from the call site of
320-- newChannel.
1052-- newChannel.