Make concurrent channel cleanup robust and reliable

Author: harendra-kumar

benchmark/bench-runner/Main.hs

@@ -101,7 +101,7 @@ rtsOpts exeName benchName0 = unwords [general, exeSpecific, benchSpecific]
             "-M64M"
 
         | "Data.Stream.ConcurrentEager/o-n-heap.monad-outer-product.toNullAp"
-            `isPrefixOf` benchName = "-M1024M"
+            `isPrefixOf` benchName = "-M1500M"
         | "Data.Stream.ConcurrentEager/o-1-space."
             `isPrefixOf` benchName = "-M128M"
 

src/Streamly/Internal/Control/ForkLifted.hs

@@ -18,7 +18,8 @@ where
 import Control.Concurrent (ThreadId, forkIO, forkOS)
 import Control.Exception (SomeException(..), catch, mask)
 import Data.Functor (void)
-import Streamly.Internal.Control.Concurrent (MonadRunInIO, RunInIO(..), withRunInIO, withRunInIONoRestore)
+import Streamly.Internal.Control.Concurrent
+    (MonadRunInIO, RunInIO(..), withRunInIO, withRunInIONoRestore)
 import Streamly.Internal.Control.ForkIO (rawForkIO, forkManagedWith)
 
 -- | Fork a thread to run the given computation, installing the provided
@@ -47,9 +48,8 @@ doForkWith :: MonadRunInIO m
 doForkWith bound action (RunInIO mrun) exHandler =
     withRunInIO $ \run ->
         mask $ \restore -> do
-                tid <- (if bound then forkOS else rawForkIO) $
-                    catch (restore $ void $ mrun action)
-                          exHandler
+                let frk = if bound then forkOS else rawForkIO
+                tid <- frk $ catch (restore $ void $ mrun action) exHandler
                 run (return tid)
 
 -- | 'fork' lifted to any monad with 'MonadBaseControl IO m' capability.

src/Streamly/Internal/Data/Channel/Dispatcher.hs

@@ -32,7 +32,8 @@ import Control.Exception (assert)
 import Control.Monad (when, void)
 import Control.Monad.IO.Class (MonadIO(liftIO))
 import Data.IORef (IORef, modifyIORef, readIORef, writeIORef)
-import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS, writeBarrier)
+import Streamly.Internal.Data.Atomics
+    (atomicModifyIORefCAS, writeBarrier, atomicModifyIORefCAS_)
 import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
 import Streamly.Internal.Data.Time.Units
        ( AbsTime, NanoSecond64(..), diffAbsTime64, showNanoSecond64
@@ -214,8 +215,8 @@ dumpSVarStats inspecting rateInfo ss = do
     stopTime <- readIORef (svarStopTime ss)
     let stopReason =
             case stopTime of
-                Nothing -> "on GC"
-                Just _ -> "normal"
+                Nothing -> "stream abandoned"
+                Just _ -> "stream finished"
     (svarCnt, svarGainLossCnt, svarLat, interval) <- case rateInfo of
         Nothing -> return (0, 0, 0, 0)
         Just yinfo -> do
@@ -277,15 +278,15 @@ dumpSVarStats inspecting rateInfo ss = do
 {-# NOINLINE addThread #-}
 addThread :: MonadIO m => IORef (Set ThreadId) -> ThreadId -> m ()
 addThread workerSet tid =
-    liftIO $ modifyIORef workerSet (S.insert tid)
+    liftIO $ atomicModifyIORefCAS_ workerSet (S.insert tid)
 
 -- This is cheaper than modifyThread because we do not have to send a
 -- outputDoorBell This can make a difference when more workers are being
 -- dispatched.
 {-# INLINE delThread #-}
 delThread :: MonadIO m => IORef (Set ThreadId) -> ThreadId -> m ()
 delThread workerSet tid =
-    liftIO $ modifyIORef workerSet (S.delete tid)
+    liftIO $ atomicModifyIORefCAS_ workerSet (S.delete tid)
 
 -- If present then delete else add. This takes care of out of order add and
 -- delete i.e. a delete arriving before we even added a thread.

src/Streamly/Internal/Data/Channel/Types.hs

@@ -62,9 +62,6 @@ module Streamly.Internal.Data.Channel.Types
     , StopWhen (..)
     , magicMaxBuffer
 
-    -- ** Cleanup
-    , cleanupSVar
-
     -- ** Diagnostics
     , dumpCreator
     , dumpOutputQ
@@ -78,22 +75,19 @@ module Streamly.Internal.Data.Channel.Types
     )
 where
 
-import Control.Concurrent (ThreadId, throwTo, MVar, tryReadMVar)
+import Control.Concurrent (ThreadId, MVar, tryReadMVar)
 import Control.Concurrent.MVar (tryPutMVar)
 import Control.Exception
     ( SomeException(..), Exception, catches, throwIO, Handler(..)
     , BlockedIndefinitelyOnMVar(..), BlockedIndefinitelyOnSTM(..))
 import Control.Monad (void, when)
 import Data.Int (Int64)
 import Data.IORef (IORef, newIORef, readIORef, writeIORef)
-import Data.Set (Set)
 import Streamly.Internal.Data.Atomics
     (atomicModifyIORefCAS, atomicModifyIORefCAS_, storeLoadBarrier)
 import Streamly.Internal.Data.Time.Units (AbsTime, NanoSecond64(..))
 import System.IO (hPutStrLn, stderr)
 
-import qualified Data.Set as Set
-
 ------------------------------------------------------------------------------
 -- Common types
 ------------------------------------------------------------------------------
@@ -462,17 +456,3 @@ printSVar :: IO String -> String -> IO ()
 printSVar dump how = do
     svInfo <- dump
     hPutStrLn stderr $ "\n" <> how <> "\n" <> svInfo
-
--------------------------------------------------------------------------------
--- Cleanup
--------------------------------------------------------------------------------
-
--- | Never called from a worker thread.
-cleanupSVar :: IORef (Set ThreadId) -> IO ()
-cleanupSVar workerSet = do
-    workers <- readIORef workerSet
-    -- self <- myThreadId
-    Prelude.mapM_ (`throwTo` ThreadAbort)
-          -- (Prelude.filter (/= self) $ Set.toList workers)
-          (Set.toList workers)
-    writeIORef workerSet Set.empty

src/Streamly/Internal/Data/Channel/Worker.hs

@@ -423,18 +423,20 @@ sendStop workerCount rateInfo q bell workerInfo = do
           workerStopUpdate winfo rinfo
       _ ->
           return ()
-    myThreadId >>= \tid ->
-        void $ sendEvent q bell (ChildStop tid Nothing)
+    tid <- myThreadId
+    void $ sendEvent q bell (ChildStop tid Nothing)
 
 -- XXX Shouldn't we perform a workerStopUpdate even in this case?
 
 -- | Add a 'ChildStop' event with exception to the channel's output queue.
 {-# NOINLINE sendException #-}
 sendException ::
-       IORef ([ChildEvent a], Int) -- ^ Queue where the exception event is added
+       IORef Int -- ^ Channel's current worker count
+    -> IORef ([ChildEvent a], Int) -- ^ Queue where the event is added
     -> MVar () -- ^ Door bell to ring
     -> SomeException -- ^ The exception to send
     -> IO ()
-sendException q bell e = do
+sendException workerCount q bell e = do
+    atomicModifyIORefCAS_ workerCount $ \n -> n - 1
     tid <- myThreadId
     void $ sendEvent q bell (ChildStop tid (Just e))

src/Streamly/Internal/Data/Stream/Channel/Append.hs

@@ -84,7 +84,7 @@ workLoopLIFO qref sv winfo = run
     run = do
         work <- dequeue qref
         case work of
-            Nothing -> liftIO $ stopWith winfo sv
+            Nothing -> return ()
             Just (RunInIO runin, m) -> process runin m
 
     process runin m = do
@@ -105,7 +105,7 @@ workLoopLIFO qref sv winfo = run
         res <- restoreM r
         case res of
             Continue -> run
-            Suspend -> liftIO $ stopWith winfo sv
+            Suspend -> return ()
 
         where
 
@@ -143,8 +143,9 @@ workLoopLIFOLimited qref sv winfo = run
 
     run = do
         work <- dequeue qref
+        {- HLINT ignore "Use forM_" -}
         case work of
-            Nothing -> liftIO $ stopWith winfo sv
+            Nothing -> return ()
             Just item -> process item
 
     process item@(RunInIO runin, m) = do
@@ -167,13 +168,12 @@ workLoopLIFOLimited qref sv winfo = run
             res <- restoreM r
             case res of
                 Continue -> run
-                Suspend -> liftIO $ stopWith winfo sv
+                Suspend -> return ()
         -- Avoid any side effects, undo the yield limit decrement if we
         -- never yielded anything.
         else liftIO $ do
             enqueueLIFO sv qref item
             incrementYieldLimit (remainingWork sv)
-            stopWith winfo sv
 
         where
 
@@ -497,11 +497,6 @@ preStopCheck sv heap =
                     if beyondRate then stopping else continue
         else stopping
 
-abortExecution :: Channel m a -> Maybe WorkerInfo -> IO ()
-abortExecution sv winfo = do
-    incrementYieldLimit (remainingWork sv)
-    stopWith winfo sv
-
 -- XXX In absence of a "noyield" primitive (i.e. do not pre-empt inside a
 -- critical section) from GHC RTS, we have a difficult problem. Assume we have
 -- a 100,000 threads producing output and queuing it to the heap for
@@ -580,32 +575,31 @@ processHeap q heap sv winfo entry sno stopping = loopHeap sno entry
             liftIO $ do
                 requeueOnHeapTop heap ent seqNo
                 incrementYieldLimit (remainingWork sv)
-                stopWith winfo sv
 
     processWorkQueue prevSeqNo = do
         yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
         if yieldLimitOk
         then do
             work <- dequeueAhead q
             case work of
-                Nothing -> liftIO $ stopWith winfo sv
+                Nothing -> return ()
                 Just (m, seqNo) -> do
                     if seqNo == prevSeqNo + 1
                     then processWithToken q heap sv winfo m seqNo
                     else processWithoutToken q heap sv winfo m seqNo
-        else liftIO $ abortExecution sv winfo
+        else liftIO $ incrementYieldLimit (remainingWork sv)
 
     nextHeap prevSeqNo = do
         res <- liftIO $ dequeueFromHeapSeq heap (prevSeqNo + 1)
         case res of
             Ready (Entry seqNo hent) -> loopHeap seqNo hent
-            Clearing -> liftIO $ stopWith winfo sv
+            Clearing -> return ()
             Waiting _ ->
                 if stopping
                 then do
                     r <- liftIO $ preStopCheck sv heap
                     if r
-                    then liftIO $ stopWith winfo sv
+                    then return ()
                     else processWorkQueue prevSeqNo
                 else inline processWorkQueue prevSeqNo
 
@@ -643,7 +637,6 @@ processHeap q heap sv winfo entry sno stopping = loopHeap sno entry
                     then liftIO $ do
                         -- put the entry back in the heap and stop
                         requeueOnHeapTop heap (Entry seqNo ent) seqNo
-                        stopWith winfo sv
                     else go
                 else go
             AheadEntryStream (RunInIO runin, Nothing, r) -> do
@@ -660,7 +653,6 @@ processHeap q heap sv winfo entry sno stopping = loopHeap sno entry
                     then liftIO $ do
                         -- put the entry back in the heap and stop
                         requeueOnHeapTop heap (Entry seqNo ent) seqNo
-                        stopWith winfo sv
                     else go
                 else go
 
@@ -677,7 +669,7 @@ drainHeap q heap sv winfo = do
     case r of
         Ready (Entry seqNo hent) ->
             processHeap q heap sv winfo hent seqNo True
-        _ -> liftIO $ stopWith winfo sv
+        _ -> return ()
 
 data HeapStatus = HContinue | HStop
 
@@ -902,7 +894,7 @@ workLoopAhead q heap sv winfo = do
         case r of
             Ready (Entry seqNo hent) ->
                 processHeap q heap sv winfo hent seqNo False
-            Clearing -> liftIO $ stopWith winfo sv
+            Clearing -> return ()
             Waiting _ -> do
                 -- Before we execute the next item from the work queue we check
                 -- if we are beyond the yield limit. It is better to check the
@@ -925,12 +917,12 @@ workLoopAhead q heap sv winfo = do
                 then do
                     work <- dequeueAhead q
                     case work of
-                        Nothing -> liftIO $ stopWith winfo sv
+                        Nothing -> return ()
                         Just (m, seqNo) -> do
                             if seqNo == 0
                             then processWithToken q heap sv winfo m seqNo
                             else processWithoutToken q heap sv winfo m seqNo
-                else liftIO $ abortExecution sv winfo
+                else liftIO $ incrementYieldLimit (remainingWork sv)
 
 -------------------------------------------------------------------------------
 -- SVar creation
@@ -967,6 +959,7 @@ getLifoSVar mrun cfg = do
         case getYieldLimit cfg of
             Nothing -> return Nothing
             Just x -> Just <$> newIORef x
+    stopRef <- newIORef False
     rateInfo <- newRateInfo cfg
 
     stats <- newSVarStats
@@ -1015,6 +1008,7 @@ getLifoSVar mrun cfg = do
                 if inOrder
                 then workLoopAhead aheadQ outH sv
                 else wloop q sv
+            , channelStopped = stopRef
             , enqueue =
                     if inOrder
                     then enqueueAhead sv aheadQ

src/Streamly/Internal/Data/Stream/Channel/Consumer.hs

@@ -20,7 +20,6 @@ where
 import Control.Monad (when, void)
 import Control.Monad.IO.Class (MonadIO(liftIO))
 import Data.IORef (readIORef)
-import Streamly.Internal.Control.Concurrent (MonadRunInIO)
 
 import Streamly.Internal.Data.Channel.Dispatcher
 import Streamly.Internal.Data.Channel.Types
@@ -42,7 +41,7 @@ readOutputQChan sv = do
 -- there is at least one outstanding worker.
 --
 -- To be used as 'readOutputQ' function for the channel.
-readOutputQBounded :: MonadRunInIO m => Bool -> Channel m a -> m [ChildEvent a]
+readOutputQBounded :: MonadIO m => Bool -> Channel m a -> m [ChildEvent a]
 readOutputQBounded eagerEval sv = do
     (list, len) <- liftIO $ readOutputQChan sv
     -- When there is no output seen we dispatch more workers to help
@@ -75,7 +74,7 @@ readOutputQBounded eagerEval sv = do
 --
 -- To be used as 'readOutputQ' function for the channel when rate control is
 -- on.
-readOutputQPaced :: MonadRunInIO m => Channel m a -> m [ChildEvent a]
+readOutputQPaced :: MonadIO m => Channel m a -> m [ChildEvent a]
 readOutputQPaced sv = do
     (list, len) <- liftIO $ readOutputQChan sv
     if len <= 0
@@ -98,7 +97,7 @@ readOutputQPaced sv = do
 --
 -- To be used as 'postProcess' function for the channel when rate control is
 -- enabled.
-postProcessPaced :: MonadRunInIO m => Channel m a -> m Bool
+postProcessPaced :: MonadIO m => Channel m a -> m Bool
 postProcessPaced sv = do
     workersDone <- allThreadsDone (workerThreads sv)
     -- XXX If during consumption we figure out we are getting delayed then we
@@ -120,7 +119,7 @@ postProcessPaced sv = do
 -- | If there is work to do ensure that we have at least one worker disptached.
 --
 -- To be used as 'postProcess' function for the channel.
-postProcessBounded :: MonadRunInIO m => Channel m a -> m Bool
+postProcessBounded :: MonadIO m => Channel m a -> m Bool
 postProcessBounded sv = do
     workersDone <- allThreadsDone (workerThreads sv)
     -- There may still be work pending even if there are no workers pending