@@ -20,17 +20,22 @@ module Streamly.Internal.Data.Stream.Async.Channel.Append
2020where
2121
2222import Control.Concurrent (myThreadId)
23−import Control.Concurrent.MVar (newEmptyMVar)
23+import Control.Concurrent.MVar (newEmptyMVar, newMVar, putMVar, takeMVar)
24+import Control.Exception (assert)
2425import Control.Monad (when, void)
2526import Control.Monad.IO.Class (MonadIO(liftIO))
2627import Control.Monad.Trans.Control (MonadBaseControl, restoreM)
27−import Data.IORef (IORef, newIORef, readIORef)
28+import Data.Heap (Heap, Entry(..))
29+import Data.IORef (IORef, newIORef, readIORef, atomicModifyIORef, writeIORef)
30+import Data.Kind (Type)
31+import GHC.Exts (inline)
2832import Streamly.Internal.Control.Concurrent
2933 (MonadRunInIO, MonadAsync, RunInIO(..), askRunInIO)
3034import Streamly.Internal.Data.Atomics
3135 (atomicModifyIORefCAS, atomicModifyIORefCAS_)
3236import Streamly.Internal.Data.Stream.Channel.Dispatcher (modifyThread)
3337
38+import qualified Data.Heap as H
3439import qualified Data.Set as Set
3540import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
3641
@@ -38,9 +43,10 @@ import Streamly.Internal.Data.Stream.Async.Channel.Consumer
3843import Streamly.Internal.Data.Stream.Async.Channel.Dispatcher
3944import Streamly.Internal.Data.Stream.Async.Channel.Type
4045import Streamly.Internal.Data.Stream.Channel.Types
46+import Streamly.Internal.Data.Stream.Channel.Worker
4147
4248------------------------------------------------------------------------------
43−-- Creating a channel
49+-- Concurrent streams with first-come-first serve results
4450------------------------------------------------------------------------------
4551
4652-- Note: For purely right associated expressions this queue should have at most
@@ -207,6 +213,679 @@ workLoopLIFOLimited qref sv winfo = run
207213 liftIO $ enqueueLIFO sv qref inner (runInIO, r)
208214 return Suspend
209215
216+-------------------------------------------------------------------------------
217+-- Ahead Channel Data Structures
218+-------------------------------------------------------------------------------
219+
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.
224+--
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.
229+--
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 #-}
233+enqueueAhead ::
234+ Channel m a
235+ -> IORef ([K.Stream m a], Int)
236+ -> (RunInIO m, K.Stream m a)
237+ -> IO ()
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)
243+
244+-- enqueue without incrementing the sequence number
245+{-# INLINE reEnqueueAhead #-}
246+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)
253+
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.
264+--
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.
271+--
272+-- XXX review for livelock
273+
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"
282+
283+-------------------------------------------------------------------------------
284+-- Heap manipulation
285+-------------------------------------------------------------------------------
286+
287+withIORef :: IORef a -> (a -> IO b) -> IO b
288+withIORef ref f = readIORef ref >>= f
289+
290+atomicModifyIORef_ :: IORef a -> (a -> a) -> IO ()
291+atomicModifyIORef_ ref f =
292+ atomicModifyIORef ref $ \x -> (f x, ())
293+
294+data AheadHeapEntry (t :: (Type -> Type) -> Type -> Type) m a =
295+ AheadEntryNull
296+ | AheadEntryPure a
297+ | AheadEntryStream (RunInIO m, t m a)
298+
299+data HeapDequeueResult t m a =
300+ Clearing
301+ | Waiting Int
302+ | Ready (Entry Int (AheadHeapEntry t m a))
303+
304+{-# INLINE dequeueFromHeap #-}
305+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) ->
310+ case snum of
311+ Nothing -> (pair, Clearing)
312+ Just n -> do
313+ let r = H.uncons hp
314+ case r of
315+ Just (ent@(Entry seqNo _ev), hp') ->
316+ if seqNo == n
317+ then ((hp', Nothing), Ready ent)
318+ else assert (seqNo >= n) (pair, Waiting n)
319+ Nothing -> (pair, Waiting n)
320+
321+{-# INLINE dequeueFromHeapSeq #-}
322+dequeueFromHeapSeq
323+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
324+ -> Int
325+ -> IO (HeapDequeueResult t m a)
326+dequeueFromHeapSeq hpVar i =
327+ atomicModifyIORef hpVar $ \(hp, snum) ->
328+ case snum of
329+ Nothing -> do
330+ let r = H.uncons hp
331+ case r of
332+ Just (ent@(Entry seqNo _ev), hp') ->
333+ if seqNo == i
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"
338+
339+heapIsSane :: Maybe Int -> Int -> Bool
340+heapIsSane snum seqNo =
341+ case snum of
342+ Nothing -> True
343+ Just n -> seqNo >= n
344+
345+{-# INLINE requeueOnHeapTop #-}
346+requeueOnHeapTop
347+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
348+ -> Entry Int (AheadHeapEntry t m a)
349+ -> Int
350+ -> IO ()
351+requeueOnHeapTop hpVar ent seqNo =
352+ atomicModifyIORef_ hpVar $ \(hp, snum) ->
353+ assert (heapIsSane snum seqNo) (H.insert ent hp, Just seqNo)
354+
355+{-# INLINE updateHeapSeq #-}
356+updateHeapSeq
357+ :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
358+ -> Int
359+ -> IO ()
360+updateHeapSeq hpVar seqNo =
361+ atomicModifyIORef_ hpVar $ \(hp, snum) ->
362+ assert (heapIsSane snum seqNo) (hp, Just seqNo)
363+
364+------------------------------------------------------------------------------
365+-- Ahead: Concurrent streams with ordered results
366+------------------------------------------------------------------------------
367+
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.
374+--
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.
381+--
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
388+-- "token".
389+--
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.
397+--
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.
404+--
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
410+-- efficiently.
411+--
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
414+-- front of them.
415+
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.
420+--
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.
424+--
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.
429+
430+-- Invariants:
431+--
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
437+-- code.
438+
439+{-# INLINE underMaxHeap #-}
440+underMaxHeap ::
441+ Channel m a
442+ -> Heap (Entry Int (AheadHeapEntry K.Stream m a))
443+ -> IO Bool
444+underMaxHeap sv hp = do
445+ (_, len) <- readIORef (outputQueue sv)
446+
447+ -- XXX simplify this
448+ let maxHeap = case maxBufferLimit sv of
449+ Limited lim -> Limited $
450+ max 0 (lim - fromIntegral len)
451+ Unlimited -> Unlimited
452+
453+ case maxHeap of
454+ Limited lim -> do
455+ active <- readIORef (workerCount sv)
456+ return $ H.size hp + active <= fromIntegral lim
457+ Unlimited -> return True
458+
459+-- Return value:
460+-- True => stop
461+-- False => continue
462+preStopCheck ::
463+ Channel m a
464+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)) , Maybe Int)
465+ -> IO Bool
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)
472+ let stopping = do
473+ putMVar (workerStopMVar sv) ()
474+ return True
475+ continue = do
476+ putMVar (workerStopMVar sv) ()
477+ return False
478+ if heapOk
479+ then
480+ case yieldRateInfo sv of
481+ Nothing -> continue
482+ Just yinfo -> do
483+ rateOk <-
484+ isBeyondMaxRate
485+ (maxWorkerLimit sv) (workerCount sv) yinfo
486+ if rateOk then continue else stopping
487+ else stopping
488+
489+abortExecution ::
490+ IORef ([K.Stream m a], Int)
491+ -> Channel m a
492+ -> Maybe WorkerInfo
493+ -> K.Stream m a
494+ -> IO ()
495+abortExecution q sv winfo m = do
496+ reEnqueueAhead sv q m
497+ incrementYieldLimit (remainingWork sv)
498+ stop sv winfo
499+
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.
513+--
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
517+--
518+-- In both cases we give the drainer a chance to run more often.
519+--
520+processHeap
521+ :: MonadRunInIO m
522+ => IORef ([K.Stream m a], Int)
523+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
524+ -> Channel m a
525+ -> Maybe WorkerInfo
526+ -> AheadHeapEntry K.Stream m a
527+ -> Int
528+ -> Bool -- we are draining the heap before we stop
529+ -> m ()
530+processHeap q heap sv winfo entry sno stopping = loopHeap sno entry
531+
532+ where
533+
534+ stopIfNeeded ent seqNo r = do
535+ stopIt <- liftIO $ preStopCheck sv heap
536+ if stopIt
537+ then liftIO $ do
538+ -- put the entry back in the heap and stop
539+ requeueOnHeapTop heap (Entry seqNo ent) seqNo
540+ stop sv winfo
541+ else runStreamWithYieldLimit True seqNo r
542+
543+ loopHeap seqNo ent =
544+ case ent of
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.
551+ void
552+ $ liftIO
553+ $ sendWithDoorBell
554+ (outputQueue sv) (outputDoorBell sv) (ChildYield a)
555+ nextHeap seqNo
556+ AheadEntryStream (RunInIO runin, r) -> do
557+ if stopping
558+ then stopIfNeeded ent seqNo r
559+ else do
560+ res <- liftIO $ runin (runStreamWithYieldLimit True seqNo r)
561+ restoreM res
562+
563+ nextHeap prevSeqNo = do
564+ res <- liftIO $ dequeueFromHeapSeq heap (prevSeqNo + 1)
565+ case res of
566+ Ready (Entry seqNo hent) -> loopHeap seqNo hent
567+ Clearing -> liftIO $ stop sv winfo
568+ Waiting _ ->
569+ if stopping
570+ then do
571+ r <- liftIO $ preStopCheck sv heap
572+ if r
573+ then liftIO $ stop sv winfo
574+ else processWorkQueue prevSeqNo
575+ else inline processWorkQueue prevSeqNo
576+
577+ processWorkQueue prevSeqNo = do
578+ work <- dequeueAhead q
579+ case work of
580+ Nothing -> liftIO $ stop sv winfo
581+ Just (m, seqNo) -> do
582+ yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
583+ if yieldLimitOk
584+ then
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
589+
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
596+ nextHeap seqNo
597+
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
609+ then do
610+ let stopk = do
611+ liftIO (incrementYieldLimit (remainingWork sv))
612+ nextHeap seqNo
613+ K.foldStreamShared undefined
614+ (yieldStreamFromHeap seqNo)
615+ (singleStreamFromHeap seqNo)
616+ stopk
617+ r
618+ else do
619+ runIn <- askRunInIO
620+ let ent = Entry seqNo (AheadEntryStream (runIn, r))
621+ liftIO $ do
622+ requeueOnHeapTop heap ent seqNo
623+ incrementYieldLimit (remainingWork sv)
624+ stop sv winfo
625+
626+ yieldStreamFromHeap seqNo a r = do
627+ continue <- liftIO $ yield sv winfo a
628+ runStreamWithYieldLimit continue seqNo r
629+
630+{-# NOINLINE drainHeap #-}
631+drainHeap
632+ :: MonadRunInIO m
633+ => IORef ([K.Stream m a], Int)
634+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
635+ -> Channel m a
636+ -> Maybe WorkerInfo
637+ -> m ()
638+drainHeap q heap sv winfo = do
639+ r <- liftIO $ dequeueFromHeap heap
640+ case r of
641+ Ready (Entry seqNo hent) ->
642+ processHeap q heap sv winfo hent seqNo True
643+ _ -> liftIO $ stop sv winfo
644+
645+data HeapStatus = HContinue | HStop
646+
647+processWithoutToken
648+ :: MonadRunInIO m
649+ => IORef ([K.Stream m a], Int)
650+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
651+ -> Channel m a
652+ -> Maybe WorkerInfo
653+ -> K.Stream m a
654+ -> Int
655+ -> m ()
656+processWithoutToken q heap sv winfo m seqNo = do
657+ -- we have already decremented the yield limit for m
658+ let stopk = do
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
664+ -- we stop.
665+ toHeap AheadEntryNull
666+ mrun = runInIO $ svarMrun sv
667+
668+ r <- liftIO $ mrun $
669+ K.foldStreamShared undefined
670+ (\a r -> do
671+ runIn <- askRunInIO
672+ toHeap $ AheadEntryStream (runIn, K.cons a r))
673+ (toHeap . AheadEntryPure)
674+ stopk
675+ m
676+ res <- restoreM r
677+ case res of
678+ Continue -> workLoopAhead q heap sv winfo
679+ Suspend -> drainHeap q heap sv winfo
680+
681+ where
682+
683+ -- XXX to reduce contention each CPU can have its own heap
684+ toHeap ent = do
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
688+ -- sequence.
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')
692+
693+ when (svarInspectMode sv) $
694+ liftIO $ do
695+ maxHp <- readIORef (maxHeapSize $ svarStats sv)
696+ when (H.size newHp > maxHp) $
697+ writeIORef (maxHeapSize $ svarStats sv) (H.size newHp)
698+
699+ heapOk <- liftIO $ underMaxHeap sv newHp
700+ status <-
701+ case yieldRateInfo sv of
702+ Nothing -> return HContinue
703+ Just yinfo ->
704+ case winfo of
705+ Just info -> do
706+ rateOk <-
707+ liftIO
708+ $ workerRateControl
709+ (maxWorkerLimit sv)
710+ (workerCount sv)
711+ yinfo
712+ info
713+ if rateOk
714+ then return HContinue
715+ else return HStop
716+ Nothing -> return HContinue
717+
718+ if heapOk
719+ then
720+ case status of
721+ HContinue -> return Continue
722+ HStop -> return Suspend
723+ else return Suspend
724+
725+data TokenWorkerStatus = TokenContinue Int | TokenSuspend
726+
727+processWithToken
728+ :: MonadRunInIO m
729+ => IORef ([K.Stream m a], Int)
730+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
731+ -> Channel m a
732+ -> Maybe WorkerInfo
733+ -> K.Stream m a
734+ -> Int
735+ -> m ()
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
739+ let stopk = do
740+ liftIO (incrementYieldLimit (remainingWork sv))
741+ return $ TokenContinue (sno + 1)
742+ mrun = runInIO $ svarMrun sv
743+
744+ r <-
745+ liftIO
746+ $ mrun
747+ $ K.foldStreamShared
748+ undefined (yieldOutput sno) (singleOutput sno) stopk action
749+
750+ res <- restoreM r
751+ case res of
752+ TokenContinue seqNo -> loopWithToken seqNo
753+ TokenSuspend -> drainHeap q heap sv winfo
754+
755+ where
756+
757+ singleOutput seqNo a = do
758+ continue <- liftIO $ yield sv winfo a
759+ if continue
760+ then return $ TokenContinue (seqNo + 1)
761+ else do
762+ liftIO $ updateHeapSeq heap (seqNo + 1)
763+ return TokenSuspend
764+
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
772+ then do
773+ let stopk = do
774+ liftIO (incrementYieldLimit (remainingWork sv))
775+ return $ TokenContinue (seqNo + 1)
776+ K.foldStreamShared undefined
777+ (yieldOutput seqNo)
778+ (singleOutput seqNo)
779+ stopk
780+ r
781+ else do
782+ runIn <- askRunInIO
783+ let ent = Entry seqNo (AheadEntryStream (runIn, r))
784+ liftIO $ requeueOnHeapTop heap ent seqNo
785+ liftIO $ incrementYieldLimit (remainingWork sv)
786+ return TokenSuspend
787+
788+ loopWithToken nextSeqNo = do
789+ work <- dequeueAhead q
790+ case work of
791+ Nothing -> do
792+ liftIO $ updateHeapSeq heap nextSeqNo
793+ workLoopAhead q heap sv winfo
794+
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)
801+ if yieldLimitOk
802+ then
803+ if seqNo == nextSeqNo
804+ then do
805+ let stopk = do
806+ liftIO (incrementYieldLimit (remainingWork sv))
807+ return $ TokenContinue (seqNo + 1)
808+ mrun = runInIO $ svarMrun sv
809+ r <- liftIO $ mrun $
810+ K.foldStreamShared undefined
811+ (yieldOutput seqNo)
812+ (singleOutput seqNo)
813+ stopk
814+ m
815+ res <- restoreM r
816+ case res of
817+ TokenContinue seqNo1 -> loopWithToken seqNo1
818+ TokenSuspend -> drainHeap q heap sv winfo
819+
820+ else
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
830+
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
834+-- unpaced case.
835+--
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.
840+
841+workLoopAhead
842+ :: MonadRunInIO m
843+ => IORef ([K.Stream m a], Int)
844+ -> IORef (Heap (Entry Int (AheadHeapEntry K.Stream m a)), Maybe Int)
845+ -> Channel m a
846+ -> Maybe WorkerInfo
847+ -> m ()
848+workLoopAhead q heap sv winfo = do
849+ r <- liftIO $ dequeueFromHeap heap
850+ case r of
851+ Ready (Entry seqNo hent) ->
852+ processHeap q heap sv winfo hent seqNo False
853+ Clearing -> liftIO $ stop sv winfo
854+ Waiting _ -> do
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
862+ -- value.
863+ --
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.
870+ --
871+ work <- dequeueAhead q
872+ case work of
873+ Nothing -> liftIO $ stop sv winfo
874+ Just (m, seqNo) -> do
875+ yieldLimitOk <-
876+ liftIO $ decrementYieldLimit (remainingWork sv)
877+ if yieldLimitOk
878+ then
879+ if seqNo == 0
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
888+
210889-------------------------------------------------------------------------------
211890-- SVar creation
212891-- 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) =>
221900 RunInIO m -> Config -> IO (Channel m a)
222901getLifoSVar mrun cfg = do
223902 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)
224907 outQMv <- newEmptyMVar
225908 active <- newIORef 0
226909 wfw <- newIORef False
@@ -229,9 +912,14 @@ getLifoSVar mrun cfg = do
229912 ( [] :: [(RunInIO m, K.Stream m a)]
230913 , [] :: [(RunInIO m, K.Stream m a)]
231914 )
232− yl <- case getYieldLimit cfg of
233− Nothing -> return Nothing
234− Just x -> Just <$> newIORef x
915+ -- Sequence number is incremented whenever something is queued, therefore,
916+ -- first sequence number would be 0
917+ aheadQ <- newIORef ([], -1)
918+ stopMVar <- newMVar ()
919+ yl <-
920+ case getYieldLimit cfg of
921+ Nothing -> return Nothing
922+ Just x -> Just <$> newIORef x
235923 rateInfo <- newRateInfo cfg
236924
237925 stats <- newSVarStats
@@ -254,6 +942,7 @@ getLifoSVar mrun cfg = do
254942 return $ qEmpty || yieldsDone
255943
256944 let eager = getEagerDispatch cfg
945+ inOrder = getOrdered cfg
257946
258947 let getSVar :: Channel m a
259948 -> (Channel m a -> m [ChildEvent a])
@@ -274,17 +963,32 @@ getLifoSVar mrun cfg = do
274963 , readOutputQ = readOutput sv
275964 , postProcess = postProc sv
276965 , workerThreads = running
277− , workLoop = wloop q sv
278− , enqueue = enqueueLIFO sv q
279− , eagerDispatch = when eager $ void $ dispatchWorker 0 sv
280− , isWorkDone = workDone sv
281− , isQueueDone = workDone sv
966+
967+ , workLoop =
968+ if inOrder
969+ then workLoopAhead aheadQ outH sv
970+ else wloop q sv
971+ , enqueue =
972+ \inner ->
973+ if inOrder
974+ then enqueueAhead sv aheadQ
975+ else enqueueLIFO sv q inner
976+ , eagerDispatch = when eager $ void $ dispatchWorker 0 sv
977+ , isWorkDone =
978+ if inOrder
979+ then isWorkDoneAhead sv aheadQ outH
980+ else workDone sv
981+ , isQueueDone =
982+ if inOrder
983+ then isQueueDoneAhead sv aheadQ
984+ else workDone sv
985+
282986 , needDoorBell = wfw
283987 , svarMrun = mrun
284988 , workerCount = active
285989 -- XXX We can use delThread or modThread based on eager flag.
286990 , accountThread = modifyThread running outQMv
287− , workerStopMVar = undefined
991+ , workerStopMVar = stopMVar
288992 , svarRef = Nothing
289993 , svarInspectMode = getInspectMode cfg
290994 , svarCreator = tid
@@ -315,6 +1019,34 @@ getLifoSVar mrun cfg = do
3151019 workLoopLIFOLimited
3161020 in return sv
3171021
1022+ where
1023+
1024+ {-# INLINE isQueueDoneAhead #-}
1025+ isQueueDoneAhead sv q = do
1026+ queueDone <- checkEmpty q
1027+ yieldsDone <-
1028+ case remainingWork sv of
1029+ Just yref -> do
1030+ n <- readIORef yref
1031+ return (n <= 0)
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
1037+
1038+ {-# INLINE isWorkDoneAhead #-}
1039+ isWorkDoneAhead sv q ref = do
1040+ heapDone <- do
1041+ (hp, _) <- readIORef ref
1042+ return (H.size hp <= 0)
1043+ queueDone <- isQueueDoneAhead sv q
1044+ return $ heapDone && queueDone
1045+
1046+ checkEmpty q = do
1047+ (xs, _) <- readIORef q
1048+ return $ null xs
1049+
3181050-- | Create a new async style concurrent stream evaluation channel. The monad
3191051-- state used to run the stream actions is taken from the call site of
3201052-- newChannel.