on 209e990 help sign in
1046 lines 43 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.Concurrent.Channel.Append
3-- Copyright : (c) 2017 Composewell Technologies
4-- License : BSD-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
9-- The functions in this module are separated from the combinators using
10-- these because of a GHC issue. We need to have newChannel specialized but
11-- not inlined. If we keep it in the same module as its users we cannot achieve
12-- that and the code becomes bloated. But if it is in a separate module we can
13-- use INLINABLE and SPECIALIZE on it which makes it specialized but it is not
14-- actually inlined.
16module Streamly.Internal.Data.Stream.Concurrent.Channel.Append
17 (
18 newChannel
19 )
20where
22import Control.Concurrent (myThreadId)
23import Control.Concurrent.MVar (newEmptyMVar, newMVar, putMVar, takeMVar)
24import Control.Exception (assert)
25import Control.Monad (when, void)
26import Control.Monad.IO.Class (MonadIO(liftIO))
27import Data.Heap (Heap, Entry(..))
28import Data.IORef (IORef, newIORef, readIORef, atomicModifyIORef, writeIORef)
29import Data.Kind (Type)
30import GHC.Exts (inline)
31import Streamly.Internal.Control.Concurrent
32 (MonadRunInIO, RunInIO(..), askRunInIO, restoreM)
33import Streamly.Internal.Data.Atomics
34 (atomicModifyIORefCAS, atomicModifyIORefCAS_)
35import Streamly.Internal.Data.Stream.Channel.Dispatcher (modifyThread)
37import qualified Data.Heap as H
38import qualified Data.Set as Set
39import qualified Streamly.Internal.Data.StreamK as K
41import Streamly.Internal.Data.Stream.Concurrent.Channel.Consumer
42import Streamly.Internal.Data.Stream.Concurrent.Channel.Dispatcher
43import Streamly.Internal.Data.Stream.Concurrent.Channel.Type
44import Streamly.Internal.Data.Stream.Channel.Types
45import Streamly.Internal.Data.Stream.Channel.Worker
47------------------------------------------------------------------------------
48-- Concurrent streams with first-come-first serve results
49------------------------------------------------------------------------------
51-- Note: For purely right associated expressions this queue should have at most
52-- one element. It grows to more than one when we have left associcated
53-- expressions. Large left associated compositions can grow this to a
54-- large size
55{-# INLINE enqueueLIFO #-}
56enqueueLIFO ::
57 Channel m a
58 -> IORef ([(RunInIO m, K.StreamK m a)], [(RunInIO m, K.StreamK m a)])
59 -> Bool
60 -> (RunInIO m, K.StreamK m a)
61 -> IO ()
62enqueueLIFO sv q inner m = do
63 atomicModifyIORefCAS_ q $ \(xs, ys) ->
64 if inner then (xs, m : ys) else (m : xs, ys)
65 ringDoorBell (doorBellOnWorkQ sv) (outputDoorBell sv)
67data QResult a = QEmpty | QOuter a | QInner a
69{-# INLINE dequeue #-}
70dequeue :: MonadIO m =>
71 IORef ([(RunInIO m, K.StreamK m a)], [(RunInIO m, K.StreamK m a)])
72 -> m (QResult (RunInIO m, K.StreamK m a))
73dequeue qref =
74 liftIO
75 $ atomicModifyIORefCAS qref
76 $ \case
77 (xs, y : ys) -> ((xs, ys), QInner y)
78 (x : xs, ys) -> ((xs, ys), QOuter x)
79 x -> (x, QEmpty)
81data WorkerStatus = Continue | Suspend
83{-# INLINE workLoopLIFO #-}
84workLoopLIFO
85 :: MonadRunInIO m
86 => IORef ([(RunInIO m, K.StreamK m a)], [(RunInIO m, K.StreamK m a)])
87 -> Channel m a
88 -> Maybe WorkerInfo
89 -> m ()
90workLoopLIFO qref sv winfo = run
92 where
94 run = do
95 work <- dequeue qref
96 case work of
97 QEmpty ->
98 liftIO $ stop sv winfo
99 QInner (RunInIO runin, m) ->
100 process runin m True
101 QOuter (RunInIO runin, m) ->
102 process runin m False
104 process runin m inner = do
105 -- XXX when we finish we need to send the monadic state back to
106 -- the parent so that the state can be merged back. We capture
107 -- and return the state in the stop continuation.
108 --
109 -- Instead of using the run function we can just restore the
110 -- monad state here. That way it can work easily for
111 -- distributed case as well.
112 r <- liftIO $ runin $
113 K.foldStreamShared
114 undefined
115 yieldk
116 single
117 (return Continue)
118 m
119 res <- restoreM r
120 case res of
121 Continue -> run
122 Suspend -> liftIO $ stop sv winfo
124 where
126 single a = do
127 res <- liftIO $ yield sv winfo a
128 return $ if res then Continue else Suspend
130 yieldk a r = do
131 res <- liftIO $ yield sv winfo a
132 if res
133 then K.foldStreamShared undefined yieldk single (return Continue) r
134 else do
135 runInIO <- askRunInIO
136 liftIO $ enqueueLIFO sv qref inner (runInIO, r)
137 return Suspend
139-- We duplicate workLoop for yield limit and no limit cases because it has
140-- around 40% performance overhead in the worst case.
141--
142-- XXX we can pass yinfo directly as an argument here so that we do not have to
143-- make a check every time.
144{-# INLINE workLoopLIFOLimited #-}
145workLoopLIFOLimited
146 :: forall m a. MonadRunInIO m
147 => IORef ([(RunInIO m, K.StreamK m a)], [(RunInIO m, K.StreamK m a)])
148 -> Channel m a
149 -> Maybe WorkerInfo
150 -> m ()
151workLoopLIFOLimited qref sv winfo = run
153 where
155 incrContinue =
156 liftIO (incrementYieldLimit (remainingWork sv)) >> return Continue
158 run = do
159 work <- dequeue qref
160 case work of
161 QEmpty ->
162 liftIO $ stop sv winfo
163 QInner item ->
164 process item True
165 QOuter item ->
166 process item False
168 process item@(RunInIO runin, m) inner = do
169 -- XXX This is just a best effort minimization of concurrency
170 -- to the yield limit. If the stream is made of concurrent
171 -- streams we do not reserve the yield limit in the constituent
172 -- streams before executing the action. This can be done
173 -- though, by sharing the yield limit ref with downstream
174 -- actions via state passing. Just a todo.
175 yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
176 if yieldLimitOk
177 then do
178 r <- liftIO $ runin $
179 K.foldStreamShared
180 undefined
181 yieldk
182 single
183 incrContinue
184 m
185 res <- restoreM r
186 case res of
187 Continue -> run
188 Suspend -> liftIO $ stop sv winfo
189 -- Avoid any side effects, undo the yield limit decrement if we
190 -- never yielded anything.
191 else liftIO $ do
192 enqueueLIFO sv qref inner item
193 incrementYieldLimit (remainingWork sv)
194 stop sv winfo
196 where
198 single a = do
199 res <- liftIO $ yield sv winfo a
200 return $ if res then Continue else Suspend
202 -- XXX can we pass on the yield limit downstream to limit the
203 -- concurrency of constituent streams.
204 yieldk a r = do
205 res <- liftIO $ yield sv winfo a
206 yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
207 if res && yieldLimitOk
208 then K.foldStreamShared undefined yieldk single incrContinue r
209 else do
210 runInIO <- askRunInIO
211 liftIO $ incrementYieldLimit (remainingWork sv)
212 liftIO $ enqueueLIFO sv qref inner (runInIO, r)
213 return Suspend
215-------------------------------------------------------------------------------
216-- Ahead Channel Data Structures
217-------------------------------------------------------------------------------
219-- XXX Left associated ahead expressions are expensive. We start a new SVar for
220-- each left associative expression. The queue is used only for right
221-- associated expression, we queue the right expression and execute the left.
222-- Therefore the queue never has more than one item in it. However, in case of
223-- parIterateConcatMap the iteration may add more items at the end of the
224-- queue.
225--
226-- XXX we can fix this. When we queue more than one item on the queue we can
227-- mark the previously queued item as not-runnable. The not-runnable item is
228-- not dequeued until the already running one has finished and at that time we
229-- would also know the exact sequence number of the already queued item.
230--
231-- we can even run the already queued items but they will have to be sorted in
232-- layers in the heap. We can use a list of heaps for that.
233{-# ANN enqueueAhead "HLint: ignore" #-}
234{-# INLINE enqueueAhead #-}
235enqueueAhead ::
236 Channel m a
237 -> IORef ([K.StreamK m a], Int)
238 -> (RunInIO m, K.StreamK m a)
239 -> IO ()
240enqueueAhead sv q m = do
241 -- XXX The queue is LIFO. When parConcatIterate queues more than one items
242 -- to the queue it will perform a DFS style traversal. For BFS we will have
243 -- to use a FIFO data structure here. That would require another Config
244 -- option.
245 atomicModifyIORefCAS_ q $ \(xs, n) -> (snd m:xs, n)
246 ringDoorBell (doorBellOnWorkQ sv) (outputDoorBell sv)
248-- Normally the thread that has the token should never go away. The token gets
249-- handed over to another thread, but someone or the other has the token at any
250-- point of time. But if the task that has the token finds that the outputQueue
251-- is full, in that case it can go away without even handing over the token to
252-- another thread. In that case it sets the nextSequence number in the heap its
253-- own sequence number before going away. To handle this case, any task that
254-- does not have the token tries to dequeue from the heap first before
255-- dequeuing from the work queue. If it finds that the task at the top of the
256-- heap is the one that owns the current sequence number then it grabs the
257-- token and starts with that.
258--
259-- XXX instead of queueing just the head element and the remaining computation
260-- on the heap, evaluate as many as we can and place them on the heap. But we
261-- need to give higher priority to the lower sequence numbers so that lower
262-- priority tasks do not fill up the heap making higher priority tasks block
263-- due to full heap. Maybe we can have a weighted space for them in the heap.
264-- The weight is inversely proportional to the sequence number.
265--
266-- XXX review for livelock
268{-# INLINE dequeueAhead #-}
269dequeueAhead :: MonadIO m
270 => IORef ([t m a], Int) -> m (Maybe (t m a, Int))
271dequeueAhead q = liftIO $
272 atomicModifyIORefCAS q $ \case
273 ([], n) -> (([], n), Nothing)
274 (x : xs, n) -> ((xs, n + 1), Just (x, n + 1))
276-- Dequeue only if the seq number matches the expected seq number.
277{-# INLINE dequeueAheadSeqCheck #-}
278dequeueAheadSeqCheck :: MonadIO m
279 => IORef ([t m a], Int) -> Int -> m (Maybe (t m a))
280dequeueAheadSeqCheck q seqNo = liftIO $
281 atomicModifyIORefCAS q $ \case
282 ([], n) -> (([], n), Nothing)
283 (x : xs, n) ->
284 if n + 1 == seqNo
285 then ((xs, n + 1), Just x)
286 else ((x : xs, n), Nothing)
288-------------------------------------------------------------------------------
289-- Heap manipulation
290-------------------------------------------------------------------------------
292withIORef :: IORef a -> (a -> IO b) -> IO b
293withIORef ref f = readIORef ref >>= f
295atomicModifyIORef_ :: IORef a -> (a -> a) -> IO ()
296atomicModifyIORef_ ref f =
297 atomicModifyIORef ref $ \x -> (f x, ())
299data AheadHeapEntry (t :: (Type -> Type) -> Type -> Type) m a =
300 AheadEntryNull
301 | AheadEntryPure a
302 | AheadEntryStream (RunInIO m, t m a)
304data HeapDequeueResult t m a =
305 Clearing
306 | Waiting Int
307 | Ready (Entry Int (AheadHeapEntry t m a))
309{-# INLINE dequeueFromHeap #-}
310dequeueFromHeap
311 :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
312 -> IO (HeapDequeueResult t m a)
313dequeueFromHeap hpVar =
314 atomicModifyIORef hpVar $ \pair@(hp, snum) ->
315 case snum of
316 Nothing -> (pair, Clearing)
317 Just n -> do
318 let r = H.uncons hp
319 case r of
320 Just (ent@(Entry seqNo _ev), hp') ->
321 if seqNo == n
322 then ((hp', Nothing), Ready ent)
323 else assert (seqNo >= n) (pair, Waiting n)
324 Nothing -> (pair, Waiting n)
326{-# INLINE dequeueFromHeapSeq #-}
327dequeueFromHeapSeq
328 :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
329 -> Int
330 -> IO (HeapDequeueResult t m a)
331dequeueFromHeapSeq hpVar i =
332 atomicModifyIORef hpVar $ \(hp, snum) ->
333 case snum of
334 Nothing -> do
335 let r = H.uncons hp
336 case r of
337 Just (ent@(Entry seqNo _ev), hp') ->
338 if seqNo == i
339 then ((hp', Nothing), Ready ent)
340 else assert (seqNo >= i) ((hp, Just i), Waiting i)
341 Nothing -> ((hp, Just i), Waiting i)
342 Just _ -> error "dequeueFromHeapSeq: unreachable"
344heapIsSane :: Maybe Int -> Int -> Bool
345heapIsSane snum seqNo =
346 case snum of
347 Nothing -> True
348 Just n -> seqNo >= n
350{-# INLINE requeueOnHeapTop #-}
351requeueOnHeapTop
352 :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
353 -> Entry Int (AheadHeapEntry t m a)
354 -> Int
355 -> IO ()
356requeueOnHeapTop hpVar ent seqNo =
357 atomicModifyIORef_ hpVar $ \(hp, snum) ->
358 assert (heapIsSane snum seqNo) (H.insert ent hp, Just seqNo)
360{-# INLINE updateHeapSeq #-}
361updateHeapSeq
362 :: IORef (Heap (Entry Int (AheadHeapEntry t m a)), Maybe Int)
363 -> Int
364 -> IO ()
365updateHeapSeq hpVar seqNo =
366 atomicModifyIORef_ hpVar $ \(hp, snum) ->
367 assert (heapIsSane snum seqNo) (hp, Just seqNo)
369------------------------------------------------------------------------------
370-- Ahead: Concurrent streams with ordered results
371------------------------------------------------------------------------------
373-- Lookahead streams can execute multiple tasks concurrently, ahead of time,
374-- but always serve them in the same order as they appear in the stream. To
375-- implement lookahead streams efficiently we assign a sequence number to each
376-- task when the task is picked up for execution. When the task finishes, the
377-- output is tagged with the same sequence number and we rearrange the outputs
378-- in sequence based on that number.
379--
380-- To explain the mechanism imagine that the current task at the head of the
381-- stream has a "token" to yield to the outputQueue. The ownership of the token
382-- is determined by the current sequence number is maintained in outputHeap.
383-- Sequence number is assigned when a task is queued. When a thread dequeues a
384-- task it picks up the sequence number as well and when the output is ready it
385-- uses the sequence number to queue the output to the outputQueue.
386--
387-- The thread with current sequence number sends the output directly to the
388-- outputQueue. Other threads push the output to the outputHeap. When the task
389-- being queued on the heap is a stream of many elements we evaluate only the
390-- first element and keep the rest of the unevaluated computation in the heap.
391-- When such a task gets the "token" for outputQueue it evaluates and directly
392-- yields all the elements to the outputQueue without checking for the
393-- "token".
394--
395-- Note that no two outputs in the heap can have the same sequence numbers and
396-- therefore we do not need a stable heap. We have also separated the buffer
397-- for the current task (outputQueue) and the pending tasks (outputHeap) so
398-- that the pending tasks cannot interfere with the current task. Note that for
399-- a single task just the outputQueue is enough and for the case of many
400-- threads just a heap is good enough. However we balance between these two
401-- cases, so that both are efficient.
402--
403-- For bigger streams it may make sense to have separate buffers for each
404-- stream. However, for singleton streams this may become inefficient. However,
405-- if we do not have separate buffers, then the streams that come later in
406-- sequence may hog the buffer, hindering the streams that are ahead. For this
407-- reason we have a single element buffer limitation for the streams being
408-- executed in advance.
409--
410-- This scheme works pretty efficiently with less than 40% extra overhead
411-- compared to the Async streams where we do not have any kind of sequencing of
412-- the outputs. It is especially devised so that we are most efficient when we
413-- have short tasks and need just a single thread. Also when a thread yields
414-- many items it can hold lockfree access to the outputQueue and do it
415-- efficiently.
416--
417-- XXX Maybe we can start the ahead threads at a lower cpu and IO priority so
418-- that they do not hog the resources and hinder the progress of the threads in
419-- front of them.
421-- Left associated ahead expressions are expensive. We start a new SVar for
422-- each left associative expression. The queue is used only for right
423-- associated expression, we queue the right expression and execute the left.
424-- Thererefore the queue never has more than on item in it.
425--
426-- XXX Also note that limiting concurrency for cases like "take 10" would not
427-- work well with left associative expressions, because we have no visibility
428-- about how much the left side of the expression would yield.
429--
430-- XXX It may be a good idea to increment sequence numbers for each yield,
431-- currently a stream on the left side of the expression may yield many
432-- elements with the same sequene number. We can then use the seq number to
433-- enforce yieldMax and yieldLImit as well.
435-- Invariants:
436--
437-- * A worker should always ensure that it pushes all the consecutive items in
438-- the heap to the outputQueue especially the items on behalf of the workers
439-- that have already left when we were holding the token. This avoids deadlock
440-- conditions when the later workers completion depends on the consumption of
441-- earlier results. For more details see comments in the consumer pull side
442-- code.
444{-# INLINE underMaxHeap #-}
445underMaxHeap ::
446 Channel m a
447 -> Heap (Entry Int (AheadHeapEntry K.StreamK m a))
448 -> IO Bool
449underMaxHeap sv hp = do
450 (_, len) <- readIORef (outputQueue sv)
452 -- XXX simplify this
453 let maxHeap = case maxBufferLimit sv of
454 Limited lim -> Limited $
455 max 0 (lim - fromIntegral len)
456 Unlimited -> Unlimited
458 case maxHeap of
459 Limited lim -> do
460 active <- readIORef (workerCount sv)
461 return $ H.size hp + active <= fromIntegral lim
462 Unlimited -> return True
464-- Return value:
465-- True => stop
466-- False => continue
467preStopCheck ::
468 Channel m a
469 -> IORef (Heap (Entry Int (AheadHeapEntry K.StreamK m a)) , Maybe Int)
470 -> IO Bool
471preStopCheck sv heap =
472 -- check the stop condition under a lock before actually
473 -- stopping so that the whole herd does not stop at once.
474 withIORef heap $ \(hp, _) -> do
475 heapOk <- underMaxHeap sv hp
476 takeMVar (workerStopMVar sv)
477 let stopping = do
478 putMVar (workerStopMVar sv) ()
479 return True
480 continue = do
481 putMVar (workerStopMVar sv) ()
482 return False
483 if heapOk
484 then
485 case yieldRateInfo sv of
486 Nothing -> continue
487 Just yinfo -> do
488 rateOk <-
489 isBeyondMaxRate
490 (maxWorkerLimit sv) (workerCount sv) yinfo
491 if rateOk then continue else stopping
492 else stopping
494abortExecution :: Channel m a -> Maybe WorkerInfo -> IO ()
495abortExecution sv winfo = do
496 incrementYieldLimit (remainingWork sv)
497 stop sv winfo
499-- XXX In absence of a "noyield" primitive (i.e. do not pre-empt inside a
500-- critical section) from GHC RTS, we have a difficult problem. Assume we have
501-- a 100,000 threads producing output and queuing it to the heap for
502-- sequencing. The heap can be drained only by one thread at a time, any thread
503-- that finds that heap can be drained now, takes a lock and starts draining
504-- it, however the thread may get prempted in the middle of it holding the
505-- lock. Since that thread is holding the lock, the other threads cannot pick
506-- up the draining task, therefore they proceed to picking up the next task to
507-- execute. If the draining thread could yield voluntarily at a point where it
508-- has released the lock, then the next threads could pick up the draining
509-- instead of executing more tasks. When there are 100,000 threads the drainer
510-- gets a cpu share to run only 1:100000 of the time. This makes the heap
511-- accumulate a lot of output when we the buffer size is large.
512--
513-- The solutions to this problem are:
514-- 1) make the other threads wait in a queue until the draining finishes
515-- 2) make the other threads queue and go away if draining is in progress
516--
517-- In both cases we give the drainer a chance to run more often.
518--
519processHeap
520 :: MonadRunInIO m
521 => IORef ([K.StreamK m a], Int)
522 -> IORef (Heap (Entry Int (AheadHeapEntry K.StreamK m a)), Maybe Int)
523 -> Channel m a
524 -> Maybe WorkerInfo
525 -> AheadHeapEntry K.StreamK m a
526 -> Int
527 -> Bool -- we are draining the heap before we stop
528 -> m ()
529processHeap q heap sv winfo entry sno stopping = loopHeap sno entry
531 where
533 stopIfNeeded ent seqNo r = do
534 stopIt <- liftIO $ preStopCheck sv heap
535 if stopIt
536 then liftIO $ do
537 -- put the entry back in the heap and stop
538 requeueOnHeapTop heap (Entry seqNo ent) seqNo
539 stop sv winfo
540 else runStreamWithYieldLimit True seqNo r
542 loopHeap seqNo ent =
543 case ent of
544 AheadEntryNull -> nextHeap seqNo
545 AheadEntryPure a -> do
546 -- Use 'send' directly so that we do not account this in worker
547 -- latency as this will not be the real latency.
548 -- Don't stop the worker in this case as we are just
549 -- transferring available results from heap to outputQueue.
550 void
551 $ liftIO
552 $ sendWithDoorBell
553 (outputQueue sv) (outputDoorBell sv) (ChildYield a)
554 nextHeap seqNo
555 AheadEntryStream (RunInIO runin, r) -> do
556 if stopping
557 then stopIfNeeded ent seqNo r
558 else do
559 res <- liftIO $ runin (runStreamWithYieldLimit True seqNo r)
560 restoreM res
562 nextHeap prevSeqNo = do
563 res <- liftIO $ dequeueFromHeapSeq heap (prevSeqNo + 1)
564 case res of
565 Ready (Entry seqNo hent) -> loopHeap seqNo hent
566 Clearing -> liftIO $ stop sv winfo
567 Waiting _ ->
568 if stopping
569 then do
570 r <- liftIO $ preStopCheck sv heap
571 if r
572 then liftIO $ stop sv winfo
573 else processWorkQueue prevSeqNo
574 else inline processWorkQueue prevSeqNo
576 processWorkQueue prevSeqNo = do
577 yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
578 if yieldLimitOk
579 then do
580 work <- dequeueAhead q
581 case work of
582 Nothing -> liftIO $ stop sv winfo
583 Just (m, seqNo) -> do
584 if seqNo == prevSeqNo + 1
585 then processWithToken q heap sv winfo m seqNo
586 else processWithoutToken q heap sv winfo m seqNo
587 else liftIO $ abortExecution sv winfo
589 -- We do not stop the worker on buffer full here as we want to proceed to
590 -- nextHeap anyway so that we can clear any subsequent entries. We stop
591 -- only in yield continuation where we may have a remaining stream to be
592 -- pushed on the heap.
593 singleStreamFromHeap seqNo a = do
594 void $ liftIO $ yield sv winfo a
595 nextHeap seqNo
597 -- XXX when we have an unfinished stream on the heap we cannot account all
598 -- the yields of that stream until it finishes, so if we have picked up
599 -- and executed more actions beyond that in the parent stream and put them
600 -- on the heap then they would eat up some yield limit which is not
601 -- correct, we will think that our yield limit is over even though we have
602 -- to yield items from unfinished stream before them. For this reason, if
603 -- there are pending items in the heap we drain them unconditionally
604 -- without considering the yield limit.
605 runStreamWithYieldLimit continue seqNo r = do
606 _ <- liftIO $ decrementYieldLimit (remainingWork sv)
607 if continue -- see comment above -- && yieldLimitOk
608 then do
609 let stopk = do
610 liftIO (incrementYieldLimit (remainingWork sv))
611 nextHeap seqNo
612 K.foldStreamShared undefined
613 (yieldStreamFromHeap seqNo)
614 (singleStreamFromHeap seqNo)
615 stopk
616 r
617 else do
618 runIn <- askRunInIO
619 let ent = Entry seqNo (AheadEntryStream (runIn, r))
620 liftIO $ do
621 requeueOnHeapTop heap ent seqNo
622 incrementYieldLimit (remainingWork sv)
623 stop sv winfo
625 yieldStreamFromHeap seqNo a r = do
626 continue <- liftIO $ yield sv winfo a
627 runStreamWithYieldLimit continue seqNo r
629{-# NOINLINE drainHeap #-}
630drainHeap
631 :: MonadRunInIO m
632 => IORef ([K.StreamK m a], Int)
633 -> IORef (Heap (Entry Int (AheadHeapEntry K.StreamK m a)), Maybe Int)
634 -> Channel m a
635 -> Maybe WorkerInfo
636 -> m ()
637drainHeap q heap sv winfo = do
638 r <- liftIO $ dequeueFromHeap heap
639 case r of
640 Ready (Entry seqNo hent) ->
641 processHeap q heap sv winfo hent seqNo True
642 _ -> liftIO $ stop sv winfo
644data HeapStatus = HContinue | HStop
646processWithoutToken
647 :: MonadRunInIO m
648 => IORef ([K.StreamK m a], Int)
649 -> IORef (Heap (Entry Int (AheadHeapEntry K.StreamK m a)), Maybe Int)
650 -> Channel m a
651 -> Maybe WorkerInfo
652 -> K.StreamK m a
653 -> Int
654 -> m ()
655processWithoutToken q heap sv winfo m seqNo = do
656 -- we have already decremented the yield limit for m
657 let stopk = do
658 liftIO (incrementYieldLimit (remainingWork sv))
659 -- If the stream stops without yielding anything, and we do not put
660 -- anything on heap, but if heap was waiting for this seq number
661 -- then it will keep waiting forever, because we are never going to
662 -- put it on heap. So we have to put a null entry on heap even when
663 -- we stop.
664 toHeap AheadEntryNull
665 mrun = runInIO $ svarMrun sv
667 r <- liftIO $ mrun $
668 K.foldStreamShared undefined
669 (\a r -> do
670 runIn <- askRunInIO
671 toHeap $ AheadEntryStream (runIn, K.cons a r))
672 (toHeap . AheadEntryPure)
673 stopk
674 m
675 res <- restoreM r
676 case res of
677 Continue -> workLoopAhead q heap sv winfo
678 Suspend -> drainHeap q heap sv winfo
680 where
682 -- XXX to reduce contention each CPU can have its own heap
683 toHeap ent = do
684 -- Heap insertion is an expensive affair so we use a non CAS based
685 -- modification, otherwise contention and retries can make a thread
686 -- context switch and throw it behind other threads which come later in
687 -- sequence.
688 newHp <- liftIO $ atomicModifyIORef heap $ \(hp, snum) ->
689 let hp' = H.insert (Entry seqNo ent) hp
690 in assert (heapIsSane snum seqNo) ((hp', snum), hp')
692 when (svarInspectMode sv) $
693 liftIO $ do
694 maxHp <- readIORef (maxHeapSize $ svarStats sv)
695 when (H.size newHp > maxHp) $
696 writeIORef (maxHeapSize $ svarStats sv) (H.size newHp)
698 heapOk <- liftIO $ underMaxHeap sv newHp
699 status <-
700 case yieldRateInfo sv of
701 Nothing -> return HContinue
702 Just yinfo ->
703 case winfo of
704 Just info -> do
705 rateOk <-
706 liftIO
707 $ workerRateControl
708 (maxWorkerLimit sv)
709 (workerCount sv)
710 yinfo
711 info
712 if rateOk
713 then return HContinue
714 else return HStop
715 Nothing -> return HContinue
717 if heapOk
718 then
719 case status of
720 HContinue -> return Continue
721 HStop -> return Suspend
722 else return Suspend
724data TokenWorkerStatus = TokenContinue Int | TokenSuspend
726processWithToken
727 :: MonadRunInIO m
728 => IORef ([K.StreamK m a], Int)
729 -> IORef (Heap (Entry Int (AheadHeapEntry K.StreamK m a)), Maybe Int)
730 -> Channel m a
731 -> Maybe WorkerInfo
732 -> K.StreamK m a
733 -> Int
734 -> m ()
735processWithToken q heap sv winfo action sno = do
736 -- Note, we enter this function with yield limit already decremented
737 -- XXX deduplicate stop in all invocations
738 let stopk = do
739 liftIO (incrementYieldLimit (remainingWork sv))
740 return $ TokenContinue (sno + 1)
741 mrun = runInIO $ svarMrun sv
743 r <-
744 liftIO
745 $ mrun
746 $ K.foldStreamShared
747 undefined (yieldOutput sno) (singleOutput sno) stopk action
749 res <- restoreM r
750 case res of
751 TokenContinue seqNo -> loopWithToken seqNo
752 TokenSuspend -> drainHeap q heap sv winfo
754 where
756 singleOutput seqNo a = do
757 continue <- liftIO $ yield sv winfo a
758 if continue
759 then return $ TokenContinue (seqNo + 1)
760 else do
761 liftIO $ updateHeapSeq heap (seqNo + 1)
762 return TokenSuspend
764 -- XXX use a wrapper function around stop so that we never miss
765 -- incrementing the yield in a stop continuation. Essentiatlly all
766 -- "unstream" calls in this function must increment yield limit on stop.
767 yieldOutput seqNo a r = do
768 continue <- liftIO $ yield sv winfo a
769 yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
770 if continue && yieldLimitOk
771 then do
772 let stopk = do
773 liftIO (incrementYieldLimit (remainingWork sv))
774 return $ TokenContinue (seqNo + 1)
775 K.foldStreamShared undefined
776 (yieldOutput seqNo)
777 (singleOutput seqNo)
778 stopk
779 r
780 else do
781 runIn <- askRunInIO
782 let ent = Entry seqNo (AheadEntryStream (runIn, r))
783 liftIO $ requeueOnHeapTop heap ent seqNo
784 liftIO $ incrementYieldLimit (remainingWork sv)
785 return TokenSuspend
787 loopWithToken nextSeqNo = do
788 let preExit = liftIO $ do
789 updateHeapSeq heap nextSeqNo
790 incrementYieldLimit (remainingWork sv)
791 yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
792 -- To avoid a race when another thread puts something
793 -- on the heap and goes away, the consumer will not get
794 -- a doorBell and we will not clear the heap before
795 -- executing the next action. If the consumer depends
796 -- on the output that is stuck in the heap then this
797 -- will result in a deadlock. So we always clear the
798 -- heap before executing the next action.
799 if yieldLimitOk
800 then do
801 -- XXX Instead of checking seqno inside dequeue we can dequeue
802 -- unconditionally and if the seqNo is not the same as nextSeqNo
803 -- then release the token and call processWithoutToken. Need
804 -- to check the performance though.
805 work <- dequeueAheadSeqCheck q nextSeqNo
806 case work of
807 Nothing -> preExit >> workLoopAhead q heap sv winfo
808 Just m -> do
809 let stopk = do
810 liftIO (incrementYieldLimit (remainingWork sv))
811 return $ TokenContinue (nextSeqNo + 1)
812 mrun = runInIO $ svarMrun sv
813 r <- liftIO $ mrun $
814 K.foldStreamShared undefined
815 (yieldOutput nextSeqNo)
816 (singleOutput nextSeqNo)
817 stopk
818 m
819 res <- restoreM r
820 case res of
821 TokenContinue seqNo -> loopWithToken seqNo
822 TokenSuspend -> drainHeap q heap sv winfo
823 else preExit >> drainHeap q heap sv winfo
825-- XXX the yield limit changes increased the performance overhead by 30-40%.
826-- Just like AsyncT we can use an implementation without yeidlimit and even
827-- without pacing code to keep the performance higher in the unlimited and
828-- unpaced case.
829--
830-- XXX The yieldLimit stuff is pretty invasive. We can instead do it by using
831-- three hooks, a pre-execute hook, a yield hook and a stop hook. In fact these
832-- hooks can be used for a more general implementation to even check predicates
833-- and not just yield limit.
835workLoopAhead
836 :: MonadRunInIO m
837 => IORef ([K.StreamK m a], Int)
838 -> IORef (Heap (Entry Int (AheadHeapEntry K.StreamK m a)), Maybe Int)
839 -> Channel m a
840 -> Maybe WorkerInfo
841 -> m ()
842workLoopAhead q heap sv winfo = do
843 r <- liftIO $ dequeueFromHeap heap
844 case r of
845 Ready (Entry seqNo hent) ->
846 processHeap q heap sv winfo hent seqNo False
847 Clearing -> liftIO $ stop sv winfo
848 Waiting _ -> do
849 -- Before we execute the next item from the work queue we check
850 -- if we are beyond the yield limit. It is better to check the
851 -- yield limit before we pick up the next item. Otherwise we
852 -- may have already started more tasks even though we may have
853 -- reached the yield limit. We can avoid this by taking active
854 -- workers into account, but that is not as reliable, because
855 -- workers may go away without picking up work and yielding a
856 -- value.
857 --
858 -- Rate control can be done either based on actual yields in
859 -- the output queue or based on any yield either to the heap or
860 -- to the output queue. In both cases we may have one issue or
861 -- the other. We chose to do this based on actual yields to the
862 -- output queue because it makes the code common to both async
863 -- and ahead streams.
864 --
865 yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
866 if yieldLimitOk
867 then do
868 work <- dequeueAhead q
869 case work of
870 Nothing -> liftIO $ stop sv winfo
871 Just (m, seqNo) -> do
872 if seqNo == 0
873 then processWithToken q heap sv winfo m seqNo
874 else processWithoutToken q heap sv winfo m seqNo
875 else liftIO $ abortExecution sv winfo
877-------------------------------------------------------------------------------
878-- SVar creation
879-- This code belongs in SVar.hs but is kept here for perf reasons
880-------------------------------------------------------------------------------
882-- XXX we have this function in this file because passing runStreamLIFO as a
883-- function argument to this function results in a perf degradation of more
884-- than 10%. Need to investigate what the root cause is.
885-- Interestingly, the same thing does not make any difference for Ahead.
886-- {-# INLINABLE getLifoSVar #-}
887getLifoSVar :: forall m a. MonadRunInIO m =>
888 RunInIO m -> Config -> IO (Channel m a)
889getLifoSVar mrun cfg = do
890 outQ <- newIORef ([], 0)
891 -- the second component of the tuple is "Nothing" when heap is being
892 -- cleared, "Just n" when we are expecting sequence number n to arrive
893 -- before we can start clearing the heap.
894 outH <- newIORef (H.empty, Just 0)
895 outQMv <- newEmptyMVar
896 active <- newIORef 0
897 wfw <- newIORef False
898 running <- newIORef Set.empty
899 q <- newIORef
900 ( [] :: [(RunInIO m, K.StreamK m a)]
901 , [] :: [(RunInIO m, K.StreamK m a)]
902 )
903 -- Sequence number is incremented whenever something is de-queued,
904 -- therefore, first sequence number would be 0
905 aheadQ <- newIORef ([], -1)
906 stopMVar <- newMVar ()
907 yl <-
908 case getYieldLimit cfg of
909 Nothing -> return Nothing
910 Just x -> Just <$> newIORef x
911 rateInfo <- newRateInfo cfg
913 stats <- newSVarStats
914 tid <- myThreadId
916 -- We are reading it without lock, the result would be reliable only if no
917 -- worker is pending.
918 let isWorkFinished _ = do
919 (xs, ys) <- readIORef q
920 return (null xs && null ys)
922 let isWorkFinishedLimited sv = do
923 yieldsDone <-
924 case remainingWork sv of
925 Just ref -> do
926 n <- readIORef ref
927 return (n <= 0)
928 Nothing -> return False
929 qEmpty <- isWorkFinished sv
930 return $ qEmpty || yieldsDone
932 let eagerEval = getEagerDispatch cfg
933 inOrder = getOrdered cfg
935 let getSVar :: Channel m a
936 -> (Channel m a -> m [ChildEvent a])
937 -> (Channel m a -> m Bool)
938 -> (Channel m a -> IO Bool)
939 -> (IORef ([(RunInIO m, K.StreamK m a)], [(RunInIO m, K.StreamK m a)])
940 -> Channel m a
941 -> Maybe WorkerInfo
942 -> m())
943 -> Channel m a
944 getSVar sv readOutput postProc workDone wloop = Channel
945 { outputQueue = outQ
946 , remainingWork = yl
947 , maxBufferLimit = getMaxBuffer cfg
948 , maxWorkerLimit = min (getMaxThreads cfg) (getMaxBuffer cfg)
949 , yieldRateInfo = rateInfo
950 , outputDoorBell = outQMv
951 , readOutputQ = readOutput sv
952 , postProcess = postProc sv
953 , workerThreads = running
955 , workLoop =
956 if inOrder
957 then workLoopAhead aheadQ outH sv
958 else wloop q sv
959 , enqueue =
960 \inner ->
961 if inOrder
962 then enqueueAhead sv aheadQ
963 else enqueueLIFO sv q inner
964 , eagerDispatch = when eagerEval $ void $ dispatchWorker 0 sv
965 , isWorkDone =
966 if inOrder
967 then isWorkDoneAhead sv aheadQ outH
968 else workDone sv
969 , isQueueDone =
970 if inOrder
971 then isQueueDoneAhead sv aheadQ
972 else workDone sv
974 , doorBellOnWorkQ = wfw
975 , svarMrun = mrun
976 , workerCount = active
977 -- XXX We can use delThread or modThread based on eager flag.
978 , accountThread = modifyThread running outQMv
979 , workerStopMVar = stopMVar
980 , svarRef = Nothing
981 , svarInspectMode = getInspectMode cfg
982 , svarCreator = tid
983 , svarStats = stats
984 }
986 let sv =
987 case getStreamRate cfg of
988 Nothing ->
989 case getYieldLimit cfg of
990 Nothing -> getSVar sv (readOutputQBounded eagerEval)
991 postProcessBounded
992 isWorkFinished
993 workLoopLIFO
994 Just _ -> getSVar sv (readOutputQBounded eagerEval)
995 postProcessBounded
996 isWorkFinishedLimited
997 workLoopLIFOLimited
998 Just _ ->
999 case getYieldLimit cfg of
1000 Nothing -> getSVar sv readOutputQPaced
1001 postProcessPaced
1002 isWorkFinished
1003 workLoopLIFO
1004 Just _ -> getSVar sv readOutputQPaced
1005 postProcessPaced
1006 isWorkFinishedLimited
1007 workLoopLIFOLimited
1008 in return sv
1010 where
1012 {-# INLINE isQueueDoneAhead #-}
1013 isQueueDoneAhead sv q = do
1014 queueDone <- checkEmpty q
1015 yieldsDone <-
1016 case remainingWork sv of
1017 Just yref -> do
1018 n <- readIORef yref
1019 return (n <= 0)
1020 Nothing -> return False
1021 -- XXX note that yieldsDone can only be authoritative only when there
1022 -- are no workers running. If there are active workers they can
1023 -- later increment the yield count and therefore change the result.
1024 return $ yieldsDone || queueDone
1026 {-# INLINE isWorkDoneAhead #-}
1027 isWorkDoneAhead sv q ref = do
1028 heapDone <- do
1029 (hp, _) <- readIORef ref
1030 return (H.size hp <= 0)
1031 queueDone <- isQueueDoneAhead sv q
1032 return $ heapDone && queueDone
1034 checkEmpty q = do
1035 (xs, _) <- readIORef q
1036 return $ null xs
1038-- | Create a new async style concurrent stream evaluation channel. The monad
1039-- state used to run the stream actions is taken from the call site of
1040-- newChannel.
1041{-# INLINABLE newChannel #-}
1042{-# SPECIALIZE newChannel :: (Config -> Config) -> IO (Channel IO a) #-}
1043newChannel :: MonadRunInIO m => (Config -> Config) -> m (Channel m a)
1044newChannel modifier = do
1045 mrun <- askRunInIO
1046 liftIO $ getLifoSVar mrun (modifier defaultConfig)