@@ -320,23 +320,25 @@ pullDrain ctx = AsyncT $ \_ stp yld -> do
320 done <- addThread ctx tid
320 done <- addThread ctx tid
321 if done then stp else continue
321 if done then stp else continue
323+pullDrainStart :: (MonadIO m, MonadThrow m) => Context a -> AsyncT m a
324+pullDrainStart ctx = AsyncT $ \_ stp yld -> do
325+ r <- liftIO $ readIORef (runningThreads ctx)
326+ d <- liftIO $ readIORef (doneThreads ctx)
327+ if (S.null r && S.null d)
329+ else (runAsyncT (pullDrain ctx)) Nothing stp yld
323-- | Split the original computation in a pull-push pair. The original
331-- | Split the original computation in a pull-push pair. The original
324-- computation pulls from a Channel while m1 and m2 push to the channel.
332-- computation pulls from a Channel while m1 and m2 push to the channel.
325{-# NOINLINE pullFork #-}
333{-# NOINLINE pullFork #-}
326pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
334pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
327pullFork m1 m2 = AsyncT $ \_ stp yld -> do
335pullFork m1 m2 = AsyncT $ \_ stp yld -> do
328 ctx <- liftIO $ newContext
336 ctx <- liftIO $ newContext
329− let m = pullDispatch ctx ((m1 <|> m2) <> finalizer ctx) True
337+ let m = pullDispatch ctx ((m1 <|> m2) <> pullDrainStart ctx) True
330 (runAsyncT m) Nothing stp yld
338 (runAsyncT m) Nothing stp yld
334− finalizer ctx = AsyncT $ \_ stp yld -> do
335− running <- liftIO $ readIORef (runningThreads ctx)
336− done <- liftIO $ readIORef (runningThreads ctx)
337− if (S.null running && S.null done) then stp
338− else (runAsyncT (pullDrain ctx)) Nothing stp yld
340 newContext = do
342 newContext = do
341 channel <- atomically newTChan
343 channel <- atomically newTChan
342 running <- liftIO $ newIORef S.empty
344 running <- liftIO $ newIORef S.empty