Add tests for Concurrent.parDemuxScan

Author: adithyaov

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

@@ -345,20 +345,6 @@ parDemuxScan cfg getKey getFold (Stream sstep state) =
                             r@(chan, _) <- newChannelWithScan q db cfg (fmap (k,) fld)
                             return (Map.insert k r keyToChan1, chan)
                         Just (chan, _) -> return (keyToChan1, chan)
-                -- 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.
                 sendToWorker_ ch x
                 return $ DemuxGo s q db keyToChan2
             Skip s ->

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

@@ -39,8 +39,8 @@ oddScan =
     Scanl.filtering odd
         & Scanl.lmapM (\x -> threadDelay 100 >> pure x)
 
-parDistributeScanTestScanEnd :: (Stream.Config -> Stream.Config) -> IO ()
-parDistributeScanTestScanEnd concOpts = do
+parDistributeScan_ScanEnd :: (Stream.Config -> Stream.Config) -> IO ()
+parDistributeScan_ScanEnd concOpts = do
     let streamLen = 10000
         evenLen = 100
     ref <- newIORef [Scanl.take evenLen evenScan, oddScan]
@@ -54,8 +54,28 @@ parDistributeScanTestScanEnd concOpts = do
             & Stream.fold Fold.toList
     sort res1 `shouldBe` [1..evenLen] ++ filter odd [(evenLen+1)..streamLen]
 
-parDistributeScanTestStreamEnd :: (Stream.Config -> Stream.Config) -> IO ()
-parDistributeScanTestStreamEnd concOpts = do
+parDemuxScan_ScanEnd :: (Stream.Config -> Stream.Config) -> IO ()
+parDemuxScan_ScanEnd concOpts = do
+    let streamLen = 10000
+        evenLen = 100
+        demuxer i = even (i :: Int)
+    ref <- newIORef (Scanl.take evenLen $ Scanl.mkScanl1 (\_ x -> x))
+    let gen True =
+            atomicModifyIORef' ref (\xs -> (fmap (const Nothing) Scanl.drain, xs))
+        gen False = pure $ Scanl.mkScanl1 (\_ x -> x)
+        inpList = [1..streamLen]
+        inpStream = Stream.fromList inpList
+    res <-
+        Scanl.parDemuxScan concOpts demuxer gen inpStream
+            & Stream.concatMap Stream.fromList
+            & fmap (\x -> (fst x,) <$> snd x)
+            & Stream.catMaybes
+            & Stream.fold Fold.toList
+    map snd (filter fst res) `shouldBe` take evenLen [2, 4 ..]
+    map snd (filter (not . fst) res) `shouldBe` filter odd [1..streamLen]
+
+parDistributeScan_StreamEnd :: (Stream.Config -> Stream.Config) -> IO ()
+parDistributeScan_StreamEnd concOpts = do
     let streamLen = 10000
     ref <- newIORef [evenScan, oddScan]
     let gen = atomicModifyIORef' ref (\xs -> ([], xs))
@@ -68,6 +88,22 @@ parDistributeScanTestStreamEnd concOpts = do
             & Stream.fold Fold.toList
     sort res1 `shouldBe` inpList
 
+parDemuxScan_StreamEnd :: (Stream.Config -> Stream.Config) -> IO ()
+parDemuxScan_StreamEnd concOpts = do
+    let streamLen = 10000
+        demuxer i = even (i :: Int)
+        gen _ = pure $ Scanl.mkScanl1 (\_ x -> x)
+        inpList = [1..streamLen]
+        inpStream = Stream.fromList inpList
+    res <-
+        Scanl.parDemuxScan concOpts demuxer gen inpStream
+            & Stream.concatMap Stream.fromList
+            & fmap (\x -> (fst x,) <$> snd x)
+            & Stream.catMaybes
+            & Stream.fold Fold.toList
+    map snd (filter fst res) `shouldBe` filter even [1..streamLen]
+    map snd (filter (not . fst) res) `shouldBe` filter odd [1..streamLen]
+
 main :: IO ()
 main = hspec
   $ H.parallel
@@ -76,6 +112,10 @@ main = hspec
 #endif
   $ describe moduleName $ do
         it "parDistributeScan (stream end) (maxBuffer 1)"
-            $ parDistributeScanTestStreamEnd (Stream.maxBuffer 1)
-        it "parDistributeScan (scan end)"
-            $ parDistributeScanTestScanEnd (Stream.maxBuffer 1)
+            $ parDistributeScan_StreamEnd (Stream.maxBuffer 1)
+        it "parDistributeScan (scan end) (maxBuffer 1)"
+            $ parDistributeScan_ScanEnd (Stream.maxBuffer 1)
+        it "parDemuxScan (stream end) (maxBuffer 1)"
+            $ parDemuxScan_StreamEnd (Stream.maxBuffer 1)
+        it "parDemuxScan (scan end) (maxBuffer 1)"
+            $ parDemuxScan_ScanEnd (Stream.maxBuffer 1)