@@ -127,14 +127,104 @@ import Prelude hiding (map)
127127-- updates the heap count and the other updates the yield count.
128128-- we can pass that as an update function instead.
129129
130−workLoopAhead :: MonadIO m
130+processHeap :: MonadIO m
131131 => IORef ([Stream m a], Int)
132132 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
133133 -> State Stream m a
134134 -> SVar Stream m a
135135 -> WorkerInfo
136+ -> AheadHeapEntry Stream m a
137+ -> Int
136138 -> m ()
137−workLoopAhead q heap st sv winfo = runHeap
139+processHeap q heap st sv winfo entry sno = loopHeap sno entry
140+
141+ where
142+
143+ loopHeap seqNo ent = do
144+#ifdef DIAGNOSTICS
145+ liftIO $ do
146+ maxHp <- readIORef (maxHeapSize sv)
147+ (hp, _) <- readIORef heap
148+ when (H.size hp > maxHp) $ writeIORef (maxHeapSize sv) (H.size hp)
149+#endif
150+ case ent of
151+ AheadEntryPure a -> do
152+ -- Use 'send' directly so that we do not account this in worker
153+ -- latency as this will not be the real latency.
154+ continue <- liftIO $ send sv (ChildYield a)
155+ if continue
156+ then nextHeap seqNo
157+ else stopHeap seqNo
158+ AheadEntryStream r -> runStreamWithYieldLimit True seqNo r
159+
160+ stopHeap seqNo = liftIO $ do
161+ atomicModifyIORef heap $ \(h, _) -> ((h, seqNo + 1), ())
162+ sendStop sv winfo
163+
164+ nextHeap prevSeqNo = do
165+ -- XXX use "dequeueIfSeqential prevSeqNo" instead of always
166+ -- updating the sequence number in heap.
167+ liftIO $ atomicModifyIORef heap $ \(h, _) -> ((h, prevSeqNo + 1), ())
168+ ent <- liftIO $ dequeueFromHeap heap
169+ case ent of
170+ Just (Entry seqNo hent) -> loopHeap seqNo hent
171+ Nothing -> do
172+ work <- dequeueAhead q
173+ case work of
174+ Nothing -> liftIO $ sendStop sv winfo
175+ Just (m, seqNo) -> do
176+ yieldLimitOk <- liftIO $ decrementYieldLimit sv
177+ if yieldLimitOk
178+ then do
179+ if seqNo == prevSeqNo + 1
180+ then processWithToken q heap st sv winfo m seqNo
181+ else processWithoutToken q heap st sv winfo m seqNo
182+ else liftIO $ do
183+ incrementYieldLimit sv
184+ sendStop sv winfo
185+
186+ singleStreamFromHeap seqNo a = do
187+ continue <- liftIO $ sendYield sv winfo (ChildYield a)
188+ if continue
189+ then nextHeap seqNo
190+ else stopHeap seqNo
191+
192+ runStreamWithYieldLimit continue seqNo r = do
193+ yieldLimitOk <- liftIO $ decrementYieldLimit sv
194+ if continue && yieldLimitOk
195+ then do
196+ let stop = do
197+ liftIO (incrementYieldLimit sv)
198+ nextHeap seqNo
199+ unStream r st stop
200+ (singleStreamFromHeap seqNo)
201+ (yieldStreamFromHeap seqNo)
202+ else liftIO $ do
203+ atomicModifyIORef heap $ \(h, _) ->
204+ ((H.insert (Entry seqNo (AheadEntryStream r)) h, seqNo), ())
205+ incrementYieldLimit sv
206+ sendStop sv winfo
207+
208+ yieldStreamFromHeap seqNo a r = do
209+ continue <- liftIO $ sendYield sv winfo (ChildYield a)
210+ runStreamWithYieldLimit continue seqNo r
211+
212+processWithoutToken :: MonadIO m
213+ => IORef ([Stream m a], Int)
214+ -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
215+ -> State Stream m a
216+ -> SVar Stream m a
217+ -> WorkerInfo
218+ -> Stream m a
219+ -> Int
220+ -> m ()
221+processWithoutToken q heap st sv winfo m sno = do
222+ -- we have already decremented the yield limit for m
223+ let stop = do
224+ liftIO (incrementYieldLimit sv)
225+ workLoopAhead q heap st sv winfo
226+
227+ unStream m st stop (singleToHeap sno) (yieldToHeap sno)
138228
139229 where
140230
@@ -170,7 +260,7 @@ workLoopAhead q heap st sv winfo = runHeap
170260 if heapOk
171261 then
172262 case yieldRateInfo sv of
173− Nothing -> runHeap
263+ Nothing -> workLoopAhead q heap st sv winfo
174264 Just yinfo -> do
175265 -- check the stop condition under a lock before actually
176266 -- stopping so that the whole herd does not stop at once.
@@ -179,7 +269,7 @@ workLoopAhead q heap st sv winfo = runHeap
179269 if rateOk
180270 then do
181271 liftIO $ putMVar (workerStopMVar yinfo) ()
182− runHeap
272+ workLoopAhead q heap st sv winfo
183273 else do
184274 liftIO $ sendStop sv winfo
185275 liftIO $ putMVar (workerStopMVar yinfo) ()
@@ -188,125 +278,103 @@ workLoopAhead q heap st sv winfo = runHeap
188278 singleToHeap seqNo a = toHeap seqNo (AheadEntryPure a)
189279 yieldToHeap seqNo a r = toHeap seqNo (AheadEntryStream (a `K.cons` r))
190280
281+processWithToken :: MonadIO m
282+ => IORef ([Stream m a], Int)
283+ -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
284+ -> State Stream m a
285+ -> SVar Stream m a
286+ -> WorkerInfo
287+ -> Stream m a
288+ -> Int
289+ -> m ()
290+processWithToken q heap st sv winfo action sno = do
291+ let stop = do
292+ liftIO (incrementYieldLimit sv)
293+ loopWithToken sno
294+
295+ unStream action st stop
296+ (singleOutput sno)
297+ (yieldOutput sno)
298+
299+ where
300+
191301 singleOutput seqNo a = do
192302 continue <- liftIO $ sendYield sv winfo (ChildYield a)
193303 if continue
194− then runQueueToken seqNo
304+ then loopWithToken seqNo
195305 else liftIO $ do
196306 atomicModifyIORef heap $ \(h, _) -> ((h, seqNo + 1), ())
197307 sendStop sv winfo
198308
199− singleFromHeap seqNo a = do
200− continue <- liftIO $ send sv (ChildYield a)
201− liftIO $ atomicModifyIORef heap $ \(h, _) -> ((h, seqNo + 1), ())
202− if continue
203− then runHeap
204− else liftIO $ sendStop sv winfo
205−
206− singleStreamFromHeap seqNo a = do
207− continue <- liftIO $ sendYield sv winfo (ChildYield a)
208− liftIO $ atomicModifyIORef heap $ \(h, _) -> ((h, seqNo + 1), ())
209− if continue
210− then runHeap
211− else liftIO $ sendStop sv winfo
212−
213309 -- XXX use a wrapper function around stop so that we never miss
214310 -- incrementing the yield in a stop continuation. Essentiatlly all
215311 -- "unstream" calls in this function must increment yield limit on stop.
216312 yieldOutput seqNo a r = do
217− continue <- liftIO $ sendYield sv winfo (ChildYield a)
218− yieldLimitOk <- liftIO $ decrementYieldLimit sv
219− if continue && yieldLimitOk
220− then do
221− let stop = liftIO (incrementYieldLimit sv) >> runQueueToken seqNo
222− unStream r st stop
223− (singleOutput seqNo)
224− (yieldOutput seqNo)
225− else liftIO $ do
226− atomicModifyIORef heap $ \(h, _) ->
227− ((H.insert (Entry seqNo (AheadEntryStream r)) h, seqNo), ())
228− incrementYieldLimit sv
229− sendStop sv winfo
230−
231− yieldStreamFromHeap seqNo a r = do
232313 continue <- liftIO $ sendYield sv winfo (ChildYield a)
233314 yieldLimitOk <- liftIO $ decrementYieldLimit sv
234315 if continue && yieldLimitOk
235316 then do
236317 let stop = do
237− liftIO $ atomicModifyIORef heap $ \(h, _) ->
238− ((h, seqNo + 1), ())
239318 liftIO (incrementYieldLimit sv)
240− runHeap
319+ loopWithToken seqNo
241320 unStream r st stop
242− (singleStreamFromHeap seqNo)
243− (yieldStreamFromHeap seqNo)
321+ (singleOutput seqNo)
322+ (yieldOutput seqNo)
244323 else liftIO $ do
245324 atomicModifyIORef heap $ \(h, _) ->
246− ((H.insert (Entry seqNo (AheadEntryStream r)) h, seqNo), ())
325+ ((H.insert (Entry seqNo (AheadEntryStream r)) h, seqNo), ())
247326 incrementYieldLimit sv
248327 sendStop sv winfo
249328
250− {-# INLINE runQueueToken #-}
251− runQueueToken prevSeqNo = do
329+ loopWithToken prevSeqNo = do
252330 work <- dequeueAhead q
253331 case work of
254332 Nothing -> do
255333 liftIO $ atomicModifyIORef heap $ \(h, _) ->
256334 ((h, prevSeqNo + 1), ())
257− runHeap
335+
336+ ent <- liftIO $ dequeueFromHeap heap
337+ case ent of
338+ Just (Entry seqNo hent) ->
339+ processHeap q heap st sv winfo hent seqNo
340+ Nothing -> liftIO $ sendStop sv winfo
341+
258342 Just (m, seqNo) -> do
259343 yieldLimitOk <- liftIO $ decrementYieldLimit sv
260344 if yieldLimitOk
261345 then do
262346 if seqNo == prevSeqNo + 1
263347 then do
264− let stop = liftIO (incrementYieldLimit sv)
265− >> runQueueToken seqNo
348+ let stop = do
349+ liftIO (incrementYieldLimit sv)
350+ loopWithToken seqNo
266351 unStream m st stop
267352 (singleOutput seqNo)
268353 (yieldOutput seqNo)
269354 else do
270355 liftIO $ atomicModifyIORef heap $ \(h, _) ->
271− ((h, prevSeqNo + 1), ())
272− let stop = liftIO (incrementYieldLimit sv)
273− >> runHeap
274− unStream m st stop
275− (singleToHeap seqNo)
276− (yieldToHeap seqNo)
356+ ((h, prevSeqNo + 1), ())
357+ -- XXX put it back on the queue and start with the Heap?
358+ processWithoutToken q heap st sv winfo m seqNo
277359 else liftIO $ do
278360 atomicModifyIORef heap $ \(h, _) ->
279− ((h, prevSeqNo + 1), ())
361+ ((h, prevSeqNo + 1), ())
280362 incrementYieldLimit sv
281363 sendStop sv winfo
282364
283− runQueueNoToken = do
284− work <- dequeueAhead q
285− case work of
286− Nothing -> runHeap
287− Just (m, seqNo) -> do
288− yieldLimitOk <- liftIO $ decrementYieldLimit sv
289− if yieldLimitOk
290− then do
291− if seqNo == 0
292− then do
293− let stop = liftIO (incrementYieldLimit sv)
294− >> runQueueToken seqNo
295− unStream m st stop
296− (singleOutput seqNo)
297− (yieldOutput seqNo)
298− else do
299− let stop = liftIO (incrementYieldLimit sv)
300− >> runHeap
301− unStream m st stop
302− (singleToHeap seqNo)
303− (yieldToHeap seqNo)
304− else liftIO $ do
305− incrementYieldLimit sv
306− sendStop sv winfo
307−
308− {-# NOINLINE runHeap #-}
309− runHeap = do
365+-- XXX the yield limit changes increased the performance overhead by 30-40%.
366+-- Just like AsyncT we can use an implementation without yeidlimit and even
367+-- without pacing code to keep the performance higher in the unlimited and
368+-- unpaced case.
369+{-# INLINE workLoopAhead #-}
370+workLoopAhead :: MonadIO m
371+ => IORef ([Stream m a], Int)
372+ -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
373+ -> State Stream m a
374+ -> SVar Stream m a
375+ -> WorkerInfo
376+ -> m ()
377+workLoopAhead q heap st sv winfo = do
310378#ifdef DIAGNOSTICS
311379 liftIO $ do
312380 maxHp <- readIORef (maxHeapSize sv)
@@ -316,7 +384,7 @@ workLoopAhead q heap st sv winfo = runHeap
316384 ent <- liftIO $ dequeueFromHeap heap
317385 case ent of
318386 Nothing -> do
319− -- Before we pick up the next item from the work queue we check
387+ -- Before we execute the next item from the work queue we check
320388 -- if we are beyond the yield limit. It is better to check the
321389 -- yield limit before we pick up the next item. Otherwise we
322390 -- may have already started more tasks even though we may have
@@ -325,8 +393,6 @@ workLoopAhead q heap st sv winfo = runHeap
325393 -- workers may go away without picking up work and yielding a
326394 -- value.
327395 --
328− -- XXX do the same for Async stream as well.
329− --
330396 -- Rate control can be done either based on actual yields in
331397 -- the output queue or based on any yield either to the heap or
332398 -- to the output queue. In both cases we may have one issue or
@@ -334,26 +400,20 @@ workLoopAhead q heap st sv winfo = runHeap
334400 -- output queue because it makes the code common to both async
335401 -- and ahead streams.
336402 --
337− done <- queueEmptyAhead q
338− if done
339− -- XXX need to run heap here if not empty
340− then liftIO $ sendStop sv winfo
341− else runQueueNoToken
342− Just (Entry seqNo hent) -> do
343− case hent of
344− -- Do not account this in worker latency as this will not
345− -- be the real latency.
346− AheadEntryPure a -> singleFromHeap seqNo a
347− AheadEntryStream r -> do
348− let stp = do
349− liftIO $ atomicModifyIORef heap $ \(h, _) ->
350− ((h, seqNo + 1), ())
351− liftIO (incrementYieldLimit sv)
352− runHeap
353−
354− unStream r st stp
355− (singleStreamFromHeap seqNo)
356− (yieldStreamFromHeap seqNo)
403+ work <- dequeueAhead q
404+ case work of
405+ Nothing -> liftIO $ sendStop sv winfo
406+ Just (m, seqNo) -> do
407+ yieldLimitOk <- liftIO $ decrementYieldLimit sv
408+ if yieldLimitOk
409+ then do
410+ if seqNo == 0
411+ then processWithToken q heap st sv winfo m seqNo
412+ else processWithoutToken q heap st sv winfo m seqNo
413+ else liftIO $ do
414+ incrementYieldLimit sv
415+ sendStop sv winfo
416+ Just (Entry seqNo hent) -> processHeap q heap st sv winfo hent seqNo
357417
358418-------------------------------------------------------------------------------
359419-- WAhead