Add a way to stop the drain fold

Author: adithyaov

src/Streamly/Internal/Data/Fold/Channel/Type.hs

@@ -428,12 +428,8 @@ scanToChannelRaw chan (Scanl step initial extract final) =
                 -- We don't end the fold here so it can drain any input in the
                 -- input buffer.
                 --
-                -- As a consequence this thread is kept alive until the stream
-                -- ends.
-                --
-                -- One way to schedule a stop is: Producer can choose to send a
-                -- special event to stop this fold. That way the responsibility
-                -- of the deadlock is on the producer.
+                -- The caller should send a ChildStopChannel to the fold again
+                -- after the fold is done to complete this scan.
                 --
                 void (sendYieldToDriver chan b)
                 pure $ Fold.Partial SCRDrain
@@ -444,6 +440,7 @@ scanToChannelRaw chan (Scanl step initial extract final) =
                 Fold.Partial s -> extract s
                 Fold.Done b0 -> pure b0
         Fold.Done <$> void (sendYieldToDriver chan b)
+    step1 SCRDrain ChildStopChannel = pure $ Fold.Done ()
     step1 SCRDrain _ = pure $ Fold.Partial SCRDrain
     step1 _ _ = error "scanToChannelRaw: Unsupported constructor"
 
@@ -481,7 +478,7 @@ newChannelWithScan outq outqDBell modifier f = do
 
     {-# NOINLINE work #-}
     work chan =
-        D.drain $ D.scanl (scanToChannelRaw chan f) $ fromInputQueueRaw chan
+        D.drain $ D.scanl (scanToChannelRaw chan f) $  fromInputQueueRaw chan
 
 {-# INLINABLE newChannel #-}
 {-# SPECIALIZE newChannel ::

src/Streamly/Internal/Data/Scanl/Concurrent.hs

@@ -28,7 +28,8 @@ import Streamly.Internal.Data.Scanl (Scanl(..))
 import Streamly.Internal.Data.Stream (Stream(..), Step(..))
 import Streamly.Internal.Data.SVar.Type (adaptState)
 import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))
-
+import qualified Streamly.Data.Fold as Fold
+import qualified Streamly.Data.Stream as Stream
 import qualified Data.Map.Strict as Map
 
 import Streamly.Internal.Data.Fold.Channel.Type
@@ -166,6 +167,11 @@ parDistributeScan cfg getFolds (Stream sstep state) =
 
     where
 
+    spanChans tid =
+        Fold.tee
+            (Fold.filter (\(_, t) -> t /= tid) Fold.toList)
+            (Fold.filter (\(_, t) -> t == tid) Fold.toList)
+
     -- XXX can be written as a fold
     processOutputs chans events done = do
         case events of
@@ -177,9 +183,11 @@ parDistributeScan cfg getFolds (Stream sstep state) =
                         liftIO $ mapM_ (`throwTo` ThreadAbort) (fmap snd chans)
                         mapM_ cleanup (fmap fst chans)
                         liftIO $ throwM ex
-                    FoldDone tid b ->
-                        let ch = filter (\(_, t) -> t /= tid) chans
-                         in processOutputs ch xs (b:done)
+                    FoldDone tid b -> do
+                        (ch, chToClose) <-
+                            Stream.fold (spanChans tid) (Stream.fromList chans)
+                        Prelude.mapM_ (finalize . fst) chToClose
+                        processOutputs ch xs (b:done)
                     FoldPartial b ->
                          processOutputs chans xs (b:done)