@@ -320,23 +320,25 @@ pullDrain ctx = AsyncT $ \_ stp yld -> do
320320 done <- addThread ctx tid
321321 if done then stp else continue
322322
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)
328+ then stp
329+ else (runAsyncT (pullDrain ctx)) Nothing stp yld
330+
323331-- | Split the original computation in a pull-push pair. The original
324332-- computation pulls from a Channel while m1 and m2 push to the channel.
325333{-# NOINLINE pullFork #-}
326334pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
327335pullFork m1 m2 = AsyncT $ \_ stp yld -> do
328336 ctx <- liftIO $ newContext
329− let m = pullDispatch ctx ((m1 <|> m2) <> finalizer ctx) True
337+ let m = pullDispatch ctx ((m1 <|> m2) <> pullDrainStart ctx) True
330338 (runAsyncT m) Nothing stp yld
331339
332340 where
333341
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
339−
340342 newContext = do
341343 channel <- atomically newTChan
342344 running <- liftIO $ newIORef S.empty