Introduce FoldEOF and tackle the cleanup after the scan terminates

Author: adithyaov

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

@@ -66,6 +66,7 @@ data OutEvent b =
       FoldException ThreadId SomeException
     | FoldPartial b
     | FoldDone ThreadId b
+    | FoldEOF ThreadId
 
 -- | The fold driver thread queues the input of the fold in the 'inputQueue'
 -- The driver rings the doorbell when the queue transitions from empty to
@@ -212,6 +213,11 @@ sendPartialToDriver :: MonadIO m => Channel m a b -> b -> m ()
 sendPartialToDriver sv res = liftIO $ do
     void $ sendToDriver sv (FoldPartial res)
 
+sendEOFToDriver :: MonadIO m => Channel m a b -> m ()
+sendEOFToDriver sv = liftIO $ do
+    tid <- myThreadId
+    void $ sendToDriver sv (FoldEOF tid)
+
 {-# NOINLINE sendExceptionToDriver #-}
 sendExceptionToDriver :: Channel m a b -> SomeException -> IO ()
 sendExceptionToDriver sv e = do
@@ -331,9 +337,9 @@ newChannelWith outq outqDBell modifier f = do
          in D.fold f1 $ fromInputQueue chan
 
 {-# INLINE scanToChannel #-}
-scanToChannel :: MonadIO m => Channel m a b -> Scanl m a b -> Scanl m a ()
+scanToChannel :: MonadIO m => Channel m a b -> Scanl m a b -> Fold m a ()
 scanToChannel chan (Scanl step initial extract final) =
-    Scanl step1 initial1 extract1 final1
+    Fold step1 initial1 extract1 final1
 
     where
 
@@ -344,8 +350,7 @@ scanToChannel chan (Scanl step initial extract final) =
                 b <- extract s
                 void $ sendPartialToDriver chan b
                 return $ Fold.Partial s
-            Fold.Done b ->
-                Fold.Done <$> void (sendYieldToDriver chan b)
+            Fold.Done b -> Fold.Done <$> sendYieldToDriver chan b
 
     step1 st x = do
         r <- step st x
@@ -354,13 +359,12 @@ scanToChannel chan (Scanl step initial extract final) =
                 b <- extract s
                 void $ sendPartialToDriver chan b
                 return $ Fold.Partial s
-            Fold.Done b ->
-                Fold.Done <$> void (sendYieldToDriver chan b)
+            Fold.Done b -> Fold.Done <$> sendYieldToDriver chan b
 
-    extract1 _ = return ()
+    extract1 _ = pure ()
 
     -- XXX Should we not discard the result?
-    final1 st = void (final st)
+    final1 st = void $ final st
 
 {-# INLINABLE newChannelWithScan #-}
 {-# SPECIALIZE newChannelWithScan ::
@@ -386,7 +390,10 @@ newChannelWithScan outq outqDBell modifier f = do
     where
 
     {-# NOINLINE work #-}
-    work chan = D.drain $ D.scanl (scanToChannel chan f) $ fromInputQueue chan
+    work chan = do
+        (_, next) <- D.foldBreak (scanToChannel chan f) (fromInputQueue chan)
+        sendEOFToDriver chan
+        D.drain next
 
 {-# INLINABLE newChannel #-}
 {-# SPECIALIZE newChannel ::
@@ -441,6 +448,7 @@ checkFoldStatus sv = do
         case ev of
             FoldException _ e -> throwM e
             FoldDone _ b -> return (Just b)
+            FoldEOF _ -> return Nothing
             FoldPartial _ -> undefined
 
 {-# INLINE isBufferAvailable #-}

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

@@ -190,6 +190,13 @@ parDistributeScan cfg getFolds (Stream sstep state) =
                             Stream.fold (spanChans tid) (Stream.fromList chans)
                         Prelude.mapM_ (finalize . fst) chToClose
                         processOutputs ch xs (b:done)
+                    FoldEOF tid -> do
+                        -- We have to send ChildStopChannel to all the folds
+                        -- that are done to stop the manager fold.
+                        (ch, chToClose) <-
+                            Stream.fold (spanChans tid) (Stream.fromList chans)
+                        Prelude.mapM_ (finalize . fst) chToClose
+                        processOutputs ch xs done
                     FoldPartial b ->
                          processOutputs chans xs (b:done)