@@ -219,7 +219,9 @@ workLoopLIFOLimited qref sv winfo = run
219219-- XXX Left associated ahead expressions are expensive. We start a new SVar for
220220-- each left associative expression. The queue is used only for right
221221-- associated expression, we queue the right expression and execute the left.
222−-- Thererefore the queue never has more than one item in it.
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.
223225--
224226-- XXX we can fix this. When we queue more than one item on the queue we can
225227-- mark the previously queued item as not-runnable. The not-runnable item is
@@ -228,6 +230,7 @@ workLoopLIFOLimited qref sv winfo = run
228230--
229231-- we can even run the already queued items but they will have to be sorted in
230232-- layers in the heap. We can use a list of heaps for that.
233+{-# ANN enqueueAhead "HLint: ignore" #-}
231234{-# INLINE enqueueAhead #-}
232235enqueueAhead ::
233236 Channel m a
@@ -235,19 +238,11 @@ enqueueAhead ::
235238 -> (RunInIO m, K.Stream m a)
236239 -> IO ()
237240enqueueAhead sv q m = do
238− atomicModifyIORefCAS_ q $ \ case
239− ([], n) -> ([snd m], n + 1) -- increment sequence
240− _ -> error "enqueueAhead: queue is not empty"
241− ringDoorBell (doorBellOnWorkQ sv) (outputDoorBell sv)
242−
243−-- enqueue without incrementing the sequence number
244−{-# INLINE reEnqueueAhead #-}
245−reEnqueueAhead ::
246− Channel m a -> IORef ([K.Stream m a], Int) -> K.Stream m a -> IO ()
247−reEnqueueAhead sv q m = do
248− atomicModifyIORefCAS_ q $ \ case
249− ([], n) -> ([m], n) -- DO NOT increment sequence
250− _ -> error "reEnqueueAhead: queue is not empty"
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)
251246 ringDoorBell (doorBellOnWorkQ sv) (outputDoorBell sv)
252247
253248-- Normally the thread that has the token should never go away. The token gets
@@ -276,8 +271,19 @@ dequeueAhead :: MonadIO m
276271dequeueAhead q = liftIO $
277272 atomicModifyIORefCAS q $ \case
278273 ([], n) -> (([], n), Nothing)
279− (x : [], n) -> (([], n), Just (x, n))
280− _ -> error "more than one item on queue"
274+ (x : xs, n) -> ((xs, n + 1), Just (x, n + 1))
275+
276+-- Dequeue only if the seq number matches the expected seq number.
277+{-# INLINE dequeueAheadSeqCheck #-}
278+dequeueAheadSeqCheck :: MonadIO m
279+ => IORef ([t m a], Int) -> Int -> m (Maybe (t m a))
280+dequeueAheadSeqCheck 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)
281287
282288-------------------------------------------------------------------------------
283289-- Heap manipulation
@@ -485,14 +491,8 @@ preStopCheck sv heap =
485491 if rateOk then continue else stopping
486492 else stopping
487493
488−abortExecution ::
489− IORef ([K.Stream m a], Int)
490− -> Channel m a
491− -> Maybe WorkerInfo
492− -> K.Stream m a
493− -> IO ()
494−abortExecution q sv winfo m = do
495− reEnqueueAhead sv q m
494+abortExecution :: Channel m a -> Maybe WorkerInfo -> IO ()
495+abortExecution sv winfo = do
496496 incrementYieldLimit (remainingWork sv)
497497 stop sv winfo
498498
@@ -574,17 +574,17 @@ processHeap q heap sv winfo entry sno stopping = loopHeap sno entry
574574 else inline processWorkQueue prevSeqNo
575575
576576 processWorkQueue prevSeqNo = do
577− work <- dequeueAhead q
578− case work of
579− Nothing -> liftIO $ stop sv winfo
580− Just (m, seqNo) -> do
581− yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
582− if yieldLimitOk
583− then
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
584584 if seqNo == prevSeqNo + 1
585585 then processWithToken q heap sv winfo m seqNo
586586 else processWithoutToken q heap sv winfo m seqNo
587− else liftIO $ abortExecution q sv winfo m
587+ else liftIO $ abortExecution sv winfo
588588
589589 -- We do not stop the worker on buffer full here as we want to proceed to
590590 -- nextHeap anyway so that we can clear any subsequent entries. We stop
@@ -785,47 +785,42 @@ processWithToken q heap sv winfo action sno = do
785785 return TokenSuspend
786786
787787 loopWithToken nextSeqNo = do
788− work <- dequeueAhead q
789− case work of
790− Nothing -> do
791− liftIO $ updateHeapSeq heap nextSeqNo
792− workLoopAhead q heap sv winfo
793−
794− Just (m, seqNo) -> do
795− yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
796− let undo = liftIO $ do
797− updateHeapSeq heap nextSeqNo
798− reEnqueueAhead sv q m
799− incrementYieldLimit (remainingWork sv)
800− if yieldLimitOk
801− then
802− if seqNo == nextSeqNo
803− then do
804− let stopk = do
805− liftIO (incrementYieldLimit (remainingWork sv))
806− return $ TokenContinue (seqNo + 1)
807− mrun = runInIO $ svarMrun sv
808− r <- liftIO $ mrun $
809− K.foldStreamShared undefined
810− (yieldOutput seqNo)
811− (singleOutput seqNo)
812− stopk
813− m
814− res <- restoreM r
815− case res of
816− TokenContinue seqNo1 -> loopWithToken seqNo1
817− TokenSuspend -> drainHeap q heap sv winfo
818−
819− else
820− -- To avoid a race when another thread puts something
821− -- on the heap and goes away, the consumer will not get
822− -- a doorBell and we will not clear the heap before
823− -- executing the next action. If the consumer depends
824− -- on the output that is stuck in the heap then this
825− -- will result in a deadlock. So we always clear the
826− -- heap before executing the next action.
827− undo >> workLoopAhead q heap sv winfo
828− else undo >> drainHeap q heap sv winfo
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
829824
830825-- XXX the yield limit changes increased the performance overhead by 30-40%.
831826-- Just like AsyncT we can use an implementation without yeidlimit and even
@@ -867,23 +862,17 @@ workLoopAhead q heap sv winfo = do
867862 -- output queue because it makes the code common to both async
868863 -- and ahead streams.
869864 --
870− work <- dequeueAhead q
871− case work of
872− Nothing -> liftIO $ stop sv winfo
873− Just (m, seqNo) -> do
874− yieldLimitOk <-
875− liftIO $ decrementYieldLimit (remainingWork sv)
876− if yieldLimitOk
877− then
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
878872 if seqNo == 0
879873 then processWithToken q heap sv winfo m seqNo
880874 else processWithoutToken q heap sv winfo m seqNo
881− -- If some worker decremented the yield limit but then
882− -- did not yield anything and therefore incremented it
883− -- later, then if we did not requeue m here we may find
884− -- the work queue empty and therefore miss executing
885− -- the remaining action.
886− else liftIO $ abortExecution q sv winfo m
875+ else liftIO $ abortExecution sv winfo
887876
888877-------------------------------------------------------------------------------
889878-- SVar creation
@@ -911,8 +900,8 @@ getLifoSVar mrun cfg = do
911900 ( [] :: [(RunInIO m, K.Stream m a)]
912901 , [] :: [(RunInIO m, K.Stream m a)]
913902 )
914− -- Sequence number is incremented whenever something is queued, therefore,
915− -- first sequence number would be 0
903+ -- Sequence number is incremented whenever something is de-queued,
904+ -- therefore, first sequence number would be 0
916905 aheadQ <- newIORef ([], -1)
917906 stopMVar <- newMVar ()
918907 yl <-