Change the edge-case behaviour of parDistributeScan

- Add a new type of concurrent fold output indicating EOF
- Add a reference in the channel: closedForInput
- Check on closedForInput in sendToWorker_

Author: adithyaov

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

@@ -45,6 +45,7 @@ import Data.List (intersperse)
 import Streamly.Internal.Control.Concurrent
     (MonadAsync, MonadRunInIO, askRunInIO)
 import Streamly.Internal.Control.ForkLifted (doForkWith)
+import Streamly.Internal.Data.Atomics (writeBarrier)
 import Streamly.Internal.Data.Fold (Fold(..))
 import Streamly.Internal.Data.Scanl (Scanl(..))
 import Streamly.Internal.Data.Channel.Dispatcher (dumpSVarStats)
@@ -66,6 +67,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
@@ -107,6 +109,7 @@ data Channel m a b = Channel
     --
     -- [LOCKING] Infrequent, MVar.
     , inputItemDoorBell :: MVar ()
+    , closedForInput :: IORef Bool
 
     -- | Doorbell to tell the driver that there is now space available in the
     -- 'inputQueue' and more items can be queued.
@@ -212,6 +215,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
@@ -281,6 +289,7 @@ mkNewChannelWith outQRev outQMvRev cfg = do
     outQ <- newIORef ([], 0)
     outQMv <- newEmptyMVar
     bufferMv <- newEmptyMVar
+    ref <- newIORef False
 
     stats <- newSVarStats
     tid <- myThreadId
@@ -292,6 +301,7 @@ mkNewChannelWith outQRev outQMvRev cfg = do
             , outputQueue = outQRev
             , outputDoorBell = outQMvRev
             , inputSpaceDoorBell = bufferMv
+            , closedForInput = ref
             , maxInputBuffer   = getMaxBuffer cfg
             , readInputQ      = liftIO $ fmap fst (readInputQWithDB sv)
             , svarRef          = Nothing
@@ -330,10 +340,12 @@ newChannelWith outq outqDBell modifier f = do
         let f1 = Fold.rmapM (void . sendYieldToDriver chan) f
          in D.fold f1 $ fromInputQueue chan
 
+-- | Returns True if the fold terminated due to completion and False when due
+-- to end-of-stream.
 {-# 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 Bool
 scanToChannel chan (Scanl step initial extract final) =
-    Scanl step1 initial1 extract1 final1
+    Fold step1 initial1 extract1 final1
 
     where
 
@@ -344,8 +356,9 @@ 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 -> do
+                sendYieldToDriver chan b
+                return $ Fold.Done True
 
     step1 st x = do
         r <- step st x
@@ -354,13 +367,16 @@ 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 -> do
+                sendYieldToDriver chan b
+                return $ Fold.Done True
 
-    extract1 _ = return ()
+    extract1 _ = error "extract: not supported by folds"
 
     -- XXX Should we not discard the result?
-    final1 st = void (final st)
+    final1 st = do
+        void (final st)
+        return False
 
 {-# INLINABLE newChannelWithScan #-}
 {-# SPECIALIZE newChannelWithScan ::
@@ -386,7 +402,15 @@ newChannelWithScan outq outqDBell modifier f = do
     where
 
     {-# NOINLINE work #-}
-    work chan = D.drain $ D.scanl (scanToChannel chan f) $ fromInputQueue chan
+    work chan = do
+        completed <- D.fold (scanToChannel chan f) (fromInputQueue chan)
+        -- We check for only one item in the outputqueue, for example in
+        -- parTeeWith, multiple messages can make that complicated. Therefore,
+        -- we first check if we already sent a FoldDone.
+        when (not completed) $ sendEOFToDriver chan
+        liftIO $ writeIORef (closedForInput chan) True
+        liftIO writeBarrier
+        void $ liftIO $ tryPutMVar (inputSpaceDoorBell chan) ()
 
 {-# INLINABLE newChannel #-}
 {-# SPECIALIZE newChannel ::
@@ -441,7 +465,10 @@ checkFoldStatus sv = do
         case ev of
             FoldException _ e -> throwM e
             FoldDone _ b -> return (Just b)
-            FoldPartial _ -> undefined
+            FoldPartial _ ->
+                error "checkFoldStatus: FoldPartial can occur only for scans"
+            FoldEOF _ ->
+                error "checkFoldStatus: FoldEOF can occur only for scans"
 
 {-# INLINE isBufferAvailable #-}
 isBufferAvailable :: MonadIO m => Channel m a b -> m Bool
@@ -510,10 +537,10 @@ sendToWorker_ chan a = go
                     (inputItemDoorBell chan)
                     (ChildYield a)
         else do
-            error "sendToWorker_: No space available in the buffer"
             -- Block for space
-            -- () <- liftIO $ takeMVar (inputSpaceDoorBell chan)
-            -- go
+            () <- liftIO $ takeMVar (inputSpaceDoorBell chan)
+            closed <- liftIO $ readIORef (closedForInput chan)
+            when (not closed) go
 
 -- XXX Cleanup the fold if the stream is interrupted. Add a GC hook.
 

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

@@ -340,7 +340,11 @@ parDistributeScan cfg getFolds (Stream sstep state) =
                     FoldDone tid b ->
                         let ch = filter (\(_, t) -> t /= tid) chans
                          in processOutputs ch xs (b:done)
-                    FoldPartial _ -> undefined
+                    FoldPartial _ ->
+                        error "parDistributeScan: cannot occur for folds"
+                    FoldEOF _ ->
+                        error
+                            "parDistributeScan: FoldEOF cannot occur for folds"
 
     collectOutputs qref chans = do
         (_, n) <- liftIO $ readIORef qref
@@ -464,7 +468,10 @@ parDemuxScan cfg getKey getFold (Stream sstep state) =
                     FoldDone _tid o@(k, _) ->
                         let ch = Map.delete k keyToChan
                          in processOutputs ch xs (o:done)
-                    FoldPartial _ -> undefined
+                    FoldPartial _ ->
+                        error "parDemuxScan: cannot occur for folds"
+                    FoldEOF _ ->
+                        error "parDemuxScan: FoldEOF cannot occur for folds"
 
     collectOutputs qref keyToChan = do
         (_, n) <- liftIO $ readIORef qref

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

@@ -89,6 +89,7 @@ parTeeWith cfg f c1 c2 = Scanl step initial extract final
                         liftIO $ throwM ex
                     FoldDone _tid b -> return (Left b)
                     FoldPartial b -> return (Right b)
+                    FoldEOF _ -> error "parTeeWith: FoldEOF cannot occur here"
             _ -> error "parTeeWith: not expecting more than one msg in q"
 
     processResponses ch1 ch2 r1 r2 =
@@ -145,12 +146,18 @@ data ScanState s q db f =
 -- XXX We can use a one way mailbox type abstraction instead of using an IORef
 -- for adding new folds dynamically.
 
--- | Evaluate a stream and scan its outputs using zero or more dynamically
--- generated parallel scans. It checks for any new folds at each input
--- generation step. Any new fold is added to the list of folds which are
--- currently running. If there are no folds available, the input is discarded.
--- If a fold completes its output is emitted in the output of the scan. The
--- outputs of the parallel scans are merged in the output stream.
+-- | Evaluate a stream and scan its outputs using zero or more parallel scans,
+-- which can be generated dynamically. It takes an action for producing new
+-- scans which is run before processing each input. The list of scans produced
+-- is added to the currently running scans. If you do not want the same scan
+-- added every time then the action should generate it only once (see the
+-- example below). If there are no scans available, the input is discarded. The
+-- outputs of all the scans are merged in the output stream.
+--
+-- If the input buffer (see maxBuffer) is limited then a scan may block until
+-- space becomes available in the input buffer. If a scan blocks then input is
+-- not provided to any of the scans, input is distributed to scans only when
+-- all scans have input buffer available.
 --
 -- >>> import Data.IORef
 -- >>> ref <- newIORef [Scanl.take 5 Scanl.sum, Scanl.take 5 Scanl.length :: Scanl.Scanl IO Int Int]
@@ -180,6 +187,9 @@ parDistributeScan cfg getFolds (Stream sstep state) =
                     FoldDone tid b ->
                         let ch = filter (\(_, t) -> t /= tid) chans
                          in processOutputs ch xs (b:done)
+                    FoldEOF tid -> do
+                        let ch = filter (\(_, t) -> t /= tid) chans
+                         in processOutputs ch xs done
                     FoldPartial b ->
                          processOutputs chans xs (b:done)
 
@@ -209,20 +219,8 @@ parDistributeScan cfg getFolds (Stream sstep state) =
         res <- sstep (adaptState gst) st
         next <- case res of
             Yield x s -> do
-                -- XXX We might block forever if some folds are already
-                -- done but we have not read the output queue yet. To
-                -- avoid that we have to either (1) precheck if space
-                -- is available in the input queues of all folds so
-                -- that this does not block, or (2) we have to use a
-                -- non-blocking read and track progress so that we can
-                -- restart from where we left.
-                --
-                -- If there is no space available then we should block
-                -- on doorbell db or inputSpaceDoorBell of the relevant
-                -- channel. To avoid deadlock the output space can be
-                -- kept unlimited. However, the blocking will delay the
-                -- processing of outputs. We should yield the outputs
-                -- before blocking.
+                -- XXX The blocking will delay the processing of outputs.
+                -- Should we yield the outputs before blocking?
                 Prelude.mapM_ (`sendToWorker_` x) (fmap fst running)
                 return $ ScanGo s q db running
             Skip s -> do
@@ -305,6 +303,10 @@ parDemuxScan cfg getKey getFold (Stream sstep state) =
                     FoldDone _tid o@(k, _) ->
                         let ch = Map.delete k keyToChan
                          in processOutputs ch xs (o:done)
+                    FoldEOF tid ->
+                        let chans = Map.toList keyToChan
+                            ch = filter (\(_, (_, t)) -> t /= tid) chans
+                         in processOutputs (Map.fromList ch) xs done
                     FoldPartial b ->
                          processOutputs keyToChan xs (b:done)