Make scanChannelRaw a non-terminating scan to remove deadlock

Author: adithyaov

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

@@ -394,6 +394,7 @@ _scanToChannel chan (Scanl step initial extract final) =
 data SendChannelRaw s a
     = SCREmptyBuffer s
     | SCRBuffered s a
+    | SCRDrain
 
 {-# INLINE scanToChannelRaw #-}
 scanToChannelRaw ::
@@ -423,15 +424,27 @@ scanToChannelRaw chan (Scanl step initial extract final) =
                 b <- extract s
                 void $ sendPartialToDriver chan b
                 return $ Fold.Partial (SCRBuffered s (ChildYield x1))
-            Fold.Done b ->
-                Fold.Done <$> void (sendYieldToDriver chan b)
+            Fold.Done b -> do
+                -- 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.
+                --
+                void (sendYieldToDriver chan b)
+                pure $ Fold.Partial SCRDrain
     step1 (SCRBuffered st (ChildYield x)) ChildStopChannel = do
         r <- step st x
         b <-
             case r of
                 Fold.Partial s -> extract s
                 Fold.Done b0 -> pure b0
         Fold.Done <$> void (sendYieldToDriver chan b)
+    step1 SCRDrain _ = pure $ Fold.Partial SCRDrain
     step1 _ _ = error "scanToChannelRaw: Unsupported constructor"
 
     extract1 _ = return ()
@@ -441,6 +454,7 @@ scanToChannelRaw chan (Scanl step initial extract final) =
     -- XXX We are losing the input here.
     -- XXX Should we consume the input and finalize it instead?
     final1 (SCRBuffered st _val) = void (final st)
+    final1 SCRDrain = pure ()
 
 {-# INLINABLE newChannelWithScan #-}
 {-# SPECIALIZE newChannelWithScan ::

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

@@ -369,6 +369,9 @@ parDistributeScan cfg getFolds (Stream sstep state) =
         res <- sstep (adaptState gst) st
         next <- case res of
             Yield x s -> do
+                -- UDPATE: forever block does not occur anymore as
+                -- "scanToChannelRaw" is not a terminating fold anymore.
+                --
                 -- 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

test/Streamly/Test/Data/Scanl/Concurrent.hs

@@ -39,9 +39,11 @@ oddScan =
     Scanl.filtering odd
         & Scanl.lmapM (\x -> threadDelay 100 >> pure x)
 
-parDistributeScanTest :: (Stream.Config -> Stream.Config) -> IO ()
-parDistributeScanTest concOpts = do
-    ref <- newIORef [evenScan, oddScan]
+parDistributeScanTestScanEnd :: (Stream.Config -> Stream.Config) -> IO ()
+parDistributeScanTestScanEnd concOpts = do
+    let streamLen = 1000
+        evenLen = 100
+    ref <- newIORef [Scanl.take evenLen evenScan, oddScan]
     let gen = atomicModifyIORef' ref (\xs -> ([], xs))
         inpList = [1..1_000]
         inpStream = Stream.fromList inpList
@@ -50,6 +52,20 @@ parDistributeScanTest concOpts = do
             & Stream.concatMap Stream.fromList
             & Stream.catMaybes
             & Stream.fold Fold.toList
+    sort res1 `shouldBe` [1..evenLen] ++ filter odd [(evenLen+1)..streamLen]
+
+parDistributeScanTestStreamEnd :: (Stream.Config -> Stream.Config) -> IO ()
+parDistributeScanTestStreamEnd concOpts = do
+    let streamLen = 1000
+    ref <- newIORef [evenScan, oddScan]
+    let gen = atomicModifyIORef' ref (\xs -> ([], xs))
+        inpList = [1..streamLen]
+        inpStream = Stream.fromList inpList
+    res1 <-
+        Scanl.parDistributeScan concOpts gen inpStream
+            & Stream.concatMap Stream.fromList
+            & Stream.catMaybes
+            & Stream.fold Fold.toList
     sort res1 `shouldBe` inpList
 
 main :: IO ()
@@ -59,5 +75,7 @@ main = hspec
   $ modifyMaxSuccess (const 10)
 #endif
   $ describe moduleName $ do
-        it "parDistributeScan (maxBuffer 1)"
-            $ parDistributeScanTest (Stream.maxBuffer 1)
+        it "parDistributeScan (stream end) (maxBuffer 1)"
+            $ parDistributeScanTestStreamEnd (Stream.maxBuffer 1)
+        it "parDistributeScan (scan end)"
+            $ parDistributeScanTestScanEnd (Stream.maxBuffer 1)