@@ -169,23 +169,23 @@ preStopCheck ::
169169preStopCheck sv heap = do
170170 -- check the stop condition under a lock before actually
171171 -- stopping so that the whole herd does not stop at once.
172− takeMVar (workerStopMVar sv)
173− let stop = do
174− putMVar (workerStopMVar sv) ()
175− return True
176− continue = do
177− putMVar (workerStopMVar sv) ()
178− return False
179− (hp, _) <- readIORef heap
180− heapOk <- underMaxHeap sv hp
181− if heapOk
182− then
183− case yieldRateInfo sv of
184− Nothing -> continue
185− Just yinfo -> do
186− rateOk <- isBeyondMaxRate sv yinfo
187− if rateOk then continue else stop
188− else stop
172+ withIORef heap $ \(hp, _) -> do
173+ heapOk <- underMaxHeap sv hp
174+ takeMVar (workerStopMVar sv)
175+ let stop = do
176+ putMVar (workerStopMVar sv) ()
177+ return True
178+ continue = do
179+ putMVar (workerStopMVar sv) ()
180+ return False
181+ if heapOk
182+ then
183+ case yieldRateInfo sv of
184+ Nothing -> continue
185+ Just yinfo -> do
186+ rateOk <- isBeyondMaxRate sv yinfo
187+ if rateOk then continue else stop
188+ else stop
189189
190190processHeap :: MonadIO m
191191 => IORef ([Stream m a], Int)
@@ -206,19 +206,11 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry
206206 if stopIt
207207 then liftIO $ do
208208 -- put the entry back in the heap and stop
209− atomicModifyIORef heap $ \(h, _) ->
210− ((H.insert (Entry seqNo ent) h, seqNo), ())
209+ liftIO $ requeueOnHeapTop heap (Entry seqNo ent) seqNo
211210 sendStop sv winfo
212211 else runStreamWithYieldLimit True seqNo r
213212
214213 loopHeap seqNo ent = do
215−#ifdef DIAGNOSTICS
216− liftIO $ do
217− maxHp <- readIORef (maxHeapSize $ svarStats sv)
218− (hp, _) <- readIORef heap
219− when (H.size hp > maxHp) $ writeIORef (maxHeapSize $ svarStats sv)
220− (H.size hp)
221−#endif
222214 case ent of
223215 AheadEntryPure a -> do
224216 -- Use 'send' directly so that we do not account this in worker
@@ -233,10 +225,7 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry
233225 else runStreamWithYieldLimit True seqNo r
234226
235227 nextHeap prevSeqNo = do
236− -- XXX use "dequeueIfSeqential prevSeqNo" instead of always
237− -- updating the sequence number in heap.
238− liftIO $ atomicModifyIORef heap $ \(h, _) -> ((h, prevSeqNo + 1), ())
239− ent <- liftIO $ dequeueFromHeap heap
228+ ent <- liftIO $ dequeueFromHeapSeq heap (prevSeqNo + 1)
240229 case ent of
241230 Just (Entry seqNo hent) -> loopHeap seqNo hent
242231 Nothing -> do
@@ -291,8 +280,8 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry
291280 (singleStreamFromHeap seqNo)
292281 (yieldStreamFromHeap seqNo)
293282 else liftIO $ do
294− atomicModifyIORef heap $ \(h, _) ->
295− ((H.insert (Entry seqNo (AheadEntryStream r)) h, seqNo), ())
283+ let ent = Entry seqNo (AheadEntryStream r)
284+ liftIO $ requeueOnHeapTop heap ent seqNo
296285 incrementYieldLimit sv
297286 sendStop sv winfo
298287
@@ -315,6 +304,8 @@ drainHeap q heap st sv winfo = do
315304 Just (Entry seqNo hent) ->
316305 processHeap q heap st sv winfo hent seqNo True
317306
307+data HeapStatus = HContinue | HStop
308+
318309processWithoutToken :: MonadIO m
319310 => IORef ([Stream m a], Int)
320311 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
@@ -340,25 +331,37 @@ processWithoutToken q heap st sv winfo m sno = do
340331 -- modification, otherwise contention and retries can make a thread
341332 -- context switch and throw it behind other threads which come later in
342333 -- sequence.
343− hp <- liftIO $ atomicModifyIORef heap $ \(h, snum) ->
344− ((H.insert (Entry seqNo ent) h, snum), h)
334+ newHp <- liftIO $ atomicModifyIORef heap $ \(hp, snum) ->
335+ let hp' = H.insert (Entry seqNo ent) hp
336+ in ((hp', snum), hp')
345337
346− heapOk <- liftIO $ underMaxHeap sv hp
347− let keepDraining = drainHeap q heap st sv winfo
338+#ifdef DIAGNOSTICS
339+ liftIO $ do
340+ maxHp <- readIORef (maxHeapSize $ svarStats sv)
341+ when (H.size newHp > maxHp) $
342+ writeIORef (maxHeapSize $ svarStats sv) (H.size newHp)
343+#endif
344+ heapOk <- liftIO $ underMaxHeap sv newHp
345+ let drainAndStop = drainHeap q heap st sv winfo
348346 mainLoop = workLoopAhead q heap st sv winfo
349− if heapOk
350− then
347+ status <-
351348 case yieldRateInfo sv of
352− Nothing -> mainLoop
349+ Nothing -> return HContinue
353350 Just yinfo -> do
354351 case winfo of
355352 Just info -> do
356353 rateOk <- liftIO $ workerRateControl sv yinfo info
357354 if rateOk
358− then mainLoop
359− else keepDraining
360− Nothing -> mainLoop
361− else keepDraining
355+ then return HContinue
356+ else return HStop
357+ Nothing -> return HContinue
358+
359+ if heapOk
360+ then
361+ case status of
362+ HContinue -> mainLoop
363+ HStop -> drainAndStop
364+ else drainAndStop
362365
363366 singleToHeap seqNo a = toHeap seqNo (AheadEntryPure a)
364367 yieldToHeap seqNo a r = toHeap seqNo (AheadEntryStream (a `K.cons` r))
@@ -388,7 +391,7 @@ processWithToken q heap st sv winfo action sno = do
388391 if continue
389392 then loopWithToken seqNo
390393 else do
391− liftIO $ atomicModifyIORef heap $ \(h, _) -> ((h, seqNo + 1), ())
394+ liftIO $ updateHeapSeq heap (seqNo + 1)
392395 drainHeap q heap st sv winfo
393396
394397 -- XXX use a wrapper function around stop so that we never miss
@@ -406,8 +409,8 @@ processWithToken q heap st sv winfo action sno = do
406409 (singleOutput seqNo)
407410 (yieldOutput seqNo)
408411 else do
409− liftIO $ atomicModifyIORef heap $ \(h, _) ->
410− ((H.insert (Entry seqNo (AheadEntryStream r)) h, seqNo), ())
412+ let ent = Entry seqNo (AheadEntryStream r)
413+ liftIO $ requeueOnHeapTop heap ent seqNo
411414 liftIO $ incrementYieldLimit sv
412415 drainHeap q heap st sv winfo
413416
@@ -415,8 +418,7 @@ processWithToken q heap st sv winfo action sno = do
415418 work <- dequeueAhead q
416419 case work of
417420 Nothing -> do
418− liftIO $ atomicModifyIORef heap $ \(h, _) ->
419− ((h, prevSeqNo + 1), ())
421+ liftIO $ updateHeapSeq heap (prevSeqNo + 1)
420422 workLoopAhead q heap st sv winfo
421423
422424 Just (m, seqNo) -> do
@@ -432,8 +434,7 @@ processWithToken q heap st sv winfo action sno = do
432434 (singleOutput seqNo)
433435 (yieldOutput seqNo)
434436 else do
435− liftIO $ atomicModifyIORef heap $ \(h, _) ->
436− ((h, prevSeqNo + 1), ())
437+ liftIO $ updateHeapSeq heap (prevSeqNo + 1)
437438 liftIO (incrementYieldLimit sv)
438439 -- To avoid a race when another thread puts something
439440 -- on the heap and goes away, the consumer will not get
@@ -445,8 +446,7 @@ processWithToken q heap st sv winfo action sno = do
445446 liftIO $ reEnqueueAhead sv q m
446447 workLoopAhead q heap st sv winfo
447448 else do
448− liftIO $ atomicModifyIORef heap $ \(h, _) ->
449− ((h, prevSeqNo + 1), ())
449+ liftIO $ updateHeapSeq heap (prevSeqNo + 1)
450450 liftIO $ reEnqueueAhead sv q m
451451 liftIO $ incrementYieldLimit sv
452452 drainHeap q heap st sv winfo