@@ -28,7 +28,7 @@ module Streamly.Streams.Ahead
2828where
2929
3030import Control.Concurrent.MVar (putMVar, takeMVar)
31−import Control.Monad (ap)
31+import Control.Monad (ap, void)
3232import Control.Monad.Base (MonadBase(..), liftBaseDefault)
3333import Control.Monad.Catch (MonadThrow, throwM)
3434-- import Control.Monad.Error.Class (MonadError(..))
@@ -40,6 +40,7 @@ import Data.Heap (Heap, Entry(..))
4040import Data.IORef (IORef, readIORef, atomicModifyIORef)
4141import Data.Maybe (fromJust)
4242import Data.Semigroup (Semigroup(..))
43+import GHC.Exts (inline)
4344
4445import qualified Data.Heap as H
4546
@@ -127,6 +128,65 @@ import Prelude hiding (map)
127128-- updates the heap count and the other updates the yield count.
128129-- we can pass that as an update function instead.
129130
131+-- Invariants:
132+--
133+-- * A worker should always ensure that it pushes all the consecutive items in
134+-- the heap to the outputQueue especially the items on behalf of the workers
135+-- that have already left when we were holding the token. This avoids deadlock
136+-- conditions when the later workers completion depends on the consumption of
137+-- earlier results. For more details see comments in the consumer pull side
138+-- code.
139+
140+{-# INLINE underMaxHeap #-}
141+underMaxHeap ::
142+ SVar Stream m a
143+ -> Heap (Entry Int (AheadHeapEntry Stream m a))
144+ -> IO Bool
145+underMaxHeap sv hp = do
146+ (_, len) <- readIORef (outputQueue sv)
147+
148+ -- XXX simplify this
149+ let maxHeap = case maxBufferLimit sv of
150+ Limited lim -> Limited $
151+ if (fromIntegral lim) >= len
152+ then lim - (fromIntegral len)
153+ else 0
154+ Unlimited -> Unlimited
155+
156+ case maxHeap of
157+ Limited lim -> do
158+ active <- readIORef (workerCount sv)
159+ return $ H.size hp + active <= (fromIntegral lim)
160+ Unlimited -> return True
161+
162+-- Return value:
163+-- True => stop
164+-- False => continue
165+preStopCheck ::
166+ SVar Stream m a
167+ -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
168+ -> IO Bool
169+preStopCheck sv heap = do
170+ -- check the stop condition under a lock before actually
171+ -- 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
189+
130190processHeap :: MonadIO m
131191 => IORef ([Stream m a], Int)
132192 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
@@ -135,11 +195,22 @@ processHeap :: MonadIO m
135195 -> WorkerInfo
136196 -> AheadHeapEntry Stream m a
137197 -> Int
198+ -> Bool -- we are draining the heap before we stop
138199 -> m ()
139−processHeap q heap st sv winfo entry sno = loopHeap sno entry
200+processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry
140201
141202 where
142203
204+ stopIfNeeded ent seqNo r = do
205+ stopIt <- liftIO $ preStopCheck sv heap
206+ if stopIt
207+ then liftIO $ do
208+ -- put the entry back in the heap and stop
209+ atomicModifyIORef heap $ \(h, _) ->
210+ ((H.insert (Entry seqNo ent) h, seqNo), ())
211+ sendStop sv winfo
212+ else runStreamWithYieldLimit True seqNo r
213+
143214 loopHeap seqNo ent = do
144215#ifdef DIAGNOSTICS
145216 liftIO $ do
@@ -151,15 +222,14 @@ processHeap q heap st sv winfo entry sno = loopHeap sno entry
151222 AheadEntryPure a -> do
152223 -- Use 'send' directly so that we do not account this in worker
153224 -- 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
225+ -- Don't stop the worker in this case as we are just
226+ -- transferring available results from heap to outputQueue.
227+ void $ liftIO $ send sv (ChildYield a)
228+ nextHeap seqNo
229+ AheadEntryStream r -> do
230+ if stopping
231+ then stopIfNeeded ent seqNo r
232+ else runStreamWithYieldLimit True seqNo r
163233
164234 nextHeap prevSeqNo = do
165235 -- XXX use "dequeueIfSeqential prevSeqNo" instead of always
@@ -169,25 +239,36 @@ processHeap q heap st sv winfo entry sno = loopHeap sno entry
169239 case ent of
170240 Just (Entry seqNo hent) -> loopHeap seqNo hent
171241 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
242+ if stopping
243+ then do
244+ r <- liftIO $ preStopCheck sv heap
245+ if r
246+ then liftIO $ sendStop sv winfo
247+ else processWorkQueue prevSeqNo
248+ else (inline processWorkQueue) prevSeqNo
185249
250+ processWorkQueue prevSeqNo = do
251+ work <- dequeueAhead q
252+ case work of
253+ Nothing -> liftIO $ sendStop sv winfo
254+ Just (m, seqNo) -> do
255+ yieldLimitOk <- liftIO $ decrementYieldLimit sv
256+ if yieldLimitOk
257+ then do
258+ if seqNo == prevSeqNo + 1
259+ then processWithToken q heap st sv winfo m seqNo
260+ else processWithoutToken q heap st sv winfo m seqNo
261+ else liftIO $ do
262+ incrementYieldLimit sv
263+ sendStop sv winfo
264+
265+ -- We do not stop the worker on buffer full here as we want to proceed to
266+ -- nextHeap anyway so that we can clear any subsequent entries. We stop
267+ -- only in yield continuation where we may have a remaining stream to be
268+ -- pushed on the heap.
186269 singleStreamFromHeap seqNo a = do
187− continue <- liftIO $ sendYield sv winfo (ChildYield a)
188− if continue
189− then nextHeap seqNo
190− else stopHeap seqNo
270+ void $ liftIO $ sendYield sv winfo (ChildYield a)
271+ nextHeap seqNo
191272
192273 runStreamWithYieldLimit continue seqNo r = do
193274 yieldLimitOk <- liftIO $ decrementYieldLimit sv
@@ -209,6 +290,21 @@ processHeap q heap st sv winfo entry sno = loopHeap sno entry
209290 continue <- liftIO $ sendYield sv winfo (ChildYield a)
210291 runStreamWithYieldLimit continue seqNo r
211292
293+{-# NOINLINE drainHeap #-}
294+drainHeap :: MonadIO m
295+ => IORef ([Stream m a], Int)
296+ -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
297+ -> State Stream m a
298+ -> SVar Stream m a
299+ -> WorkerInfo
300+ -> m ()
301+drainHeap q heap st sv winfo = do
302+ ent <- liftIO $ dequeueFromHeap heap
303+ case ent of
304+ Nothing -> liftIO $ sendStop sv winfo
305+ Just (Entry seqNo hent) ->
306+ processHeap q heap st sv winfo hent seqNo True
307+
212308processWithoutToken :: MonadIO m
213309 => IORef ([Stream m a], Int)
214310 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
@@ -228,25 +324,6 @@ processWithoutToken q heap st sv winfo m sno = do
228324
229325 where
230326
231− maxBuf = maxBufferLimit sv
232−
233− underMaxHeap hp = do
234− (_, len) <- liftIO $ readIORef (outputQueue sv)
235−
236− -- XXX simplify this
237− let maxHeap = case maxBuf of
238− Limited lim -> Limited $
239− if (fromIntegral lim) >= len
240− then lim - (fromIntegral len)
241− else 0
242− Unlimited -> Unlimited
243−
244− case maxHeap of
245− Limited lim -> do
246− active <- liftIO $ readIORef (workerCount sv)
247− return $ H.size hp + active <= (fromIntegral lim)
248− Unlimited -> return True
249−
250327 -- XXX to reduce contention each CPU can have its own heap
251328 toHeap seqNo ent = do
252329 -- Heap insertion is an expensive affair so we use a non CAS based
@@ -255,25 +332,18 @@ processWithoutToken q heap st sv winfo m sno = do
255332 -- sequence.
256333 hp <- liftIO $ atomicModifyIORef heap $ \(h, snum) ->
257334 ((H.insert (Entry seqNo ent) h, snum), h)
258− -- XXX heap check should also be under MVar
259− heapOk <- underMaxHeap hp
335+
336+ heapOk <- liftIO $ underMaxHeap sv hp
260337 if heapOk
261338 then
262339 case yieldRateInfo sv of
263340 Nothing -> workLoopAhead q heap st sv winfo
264341 Just yinfo -> do
265− -- check the stop condition under a lock before actually
266− -- stopping so that the whole herd does not stop at once.
267− liftIO $ takeMVar (workerStopMVar yinfo)
268342 rateOk <- liftIO $ workerRateControl sv yinfo winfo
269343 if rateOk
270− then do
271− liftIO $ putMVar (workerStopMVar yinfo) ()
272− workLoopAhead q heap st sv winfo
273− else do
274− liftIO $ sendStop sv winfo
275− liftIO $ putMVar (workerStopMVar yinfo) ()
276− else liftIO $ sendStop sv winfo
344+ then workLoopAhead q heap st sv winfo
345+ else drainHeap q heap st sv winfo
346+ else drainHeap q heap st sv winfo
277347
278348 singleToHeap seqNo a = toHeap seqNo (AheadEntryPure a)
279349 yieldToHeap seqNo a r = toHeap seqNo (AheadEntryStream (a `K.cons` r))
@@ -302,9 +372,9 @@ processWithToken q heap st sv winfo action sno = do
302372 continue <- liftIO $ sendYield sv winfo (ChildYield a)
303373 if continue
304374 then loopWithToken seqNo
305− else liftIO $ do
306− atomicModifyIORef heap $ \(h, _) -> ((h, seqNo + 1), ())
307− sendStop sv winfo
375+ else do
376+ liftIO $ atomicModifyIORef heap $ \(h, _) -> ((h, seqNo + 1), ())
377+ drainHeap q heap st sv winfo
308378
309379 -- XXX use a wrapper function around stop so that we never miss
310380 -- incrementing the yield in a stop continuation. Essentiatlly all
@@ -320,11 +390,11 @@ processWithToken q heap st sv winfo action sno = do
320390 unStream r st stop
321391 (singleOutput seqNo)
322392 (yieldOutput seqNo)
323− else liftIO $ do
324− atomicModifyIORef heap $ \(h, _) ->
393+ else do
394+ liftIO $ atomicModifyIORef heap $ \(h, _) ->
325395 ((H.insert (Entry seqNo (AheadEntryStream r)) h, seqNo), ())
326− incrementYieldLimit sv
327− sendStop sv winfo
396+ liftIO $ incrementYieldLimit sv
397+ drainHeap q heap st sv winfo
328398
329399 loopWithToken prevSeqNo = do
330400 work <- dequeueAhead q
@@ -332,12 +402,7 @@ processWithToken q heap st sv winfo action sno = do
332402 Nothing -> do
333403 liftIO $ atomicModifyIORef heap $ \(h, _) ->
334404 ((h, prevSeqNo + 1), ())
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
405+ workLoopAhead q heap st sv winfo
341406
342407 Just (m, seqNo) -> do
343408 yieldLimitOk <- liftIO $ decrementYieldLimit sv
@@ -354,19 +419,26 @@ processWithToken q heap st sv winfo action sno = do
354419 else do
355420 liftIO $ atomicModifyIORef heap $ \(h, _) ->
356421 ((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
359− else liftIO $ do
360− atomicModifyIORef heap $ \(h, _) ->
422+ -- To avoid a race when another thread puts something
423+ -- on the heap and goes away, the consumer will not get
424+ -- a doorBell and we will not clear the heap before
425+ -- executing the next action. If the consumer depends
426+ -- on the output that is stuck in the heap then this
427+ -- will result in a deadlock. So we always clear the
428+ -- heap before executing the next action.
429+ -- processWithoutToken q heap st sv winfo m seqNo
430+ liftIO $ reEnqueueAhead sv q m
431+ workLoopAhead q heap st sv winfo
432+ else do
433+ liftIO $ atomicModifyIORef heap $ \(h, _) ->
361434 ((h, prevSeqNo + 1), ())
362− incrementYieldLimit sv
363− sendStop sv winfo
435+ liftIO $ incrementYieldLimit sv
436+ drainHeap q heap st sv winfo
364437
365438-- XXX the yield limit changes increased the performance overhead by 30-40%.
366439-- Just like AsyncT we can use an implementation without yeidlimit and even
367440-- without pacing code to keep the performance higher in the unlimited and
368441-- unpaced case.
369−{-# INLINE workLoopAhead #-}
370442workLoopAhead :: MonadIO m
371443 => IORef ([Stream m a], Int)
372444 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)) , Int)
@@ -413,7 +485,8 @@ workLoopAhead q heap st sv winfo = do
413485 else liftIO $ do
414486 incrementYieldLimit sv
415487 sendStop sv winfo
416− Just (Entry seqNo hent) -> processHeap q heap st sv winfo hent seqNo
488+ Just (Entry seqNo hent) ->
489+ processHeap q heap st sv winfo hent seqNo False
417490
418491-------------------------------------------------------------------------------
419492-- WAhead