@@ -82,10 +82,9 @@ data ChildEvent a =
8282-- State threaded around the monad for thread management
8383------------------------------------------------------------------------------
8484
85−data Context m a =
85+data Context a =
8686 Context { childChannel :: TChan (ChildEvent a)
8787 , pullSide :: Bool
88− , pendingWork :: IORef (AsyncT m a)
8988 }
9089
9190-- The 'Maybe (AsyncT m a)' is redundant as we can use 'stop' value for the
@@ -99,9 +98,9 @@ data Context m a =
9998newtype AsyncT m a =
10099 AsyncT {
101100 runAsyncT :: forall r.
102− Maybe (Context m a) -- state
101+ Maybe (Context a) -- state
103102 -> m r -- stop
104− -> (a -> Maybe (Context m a) -> Maybe (AsyncT m a) -> m r) -- yield
103+ -> (a -> Maybe (Context a) -> Maybe (AsyncT m a) -> m r) -- yield
105104 -> m r
106105 }
107106
@@ -178,7 +177,7 @@ doFork action preExit =
178177 return ()
179178
180179{-# NOINLINE push #-}
181−push :: MonadIO m => Context m a -> AsyncT m a -> m ()
180+push :: MonadIO m => Context a -> AsyncT m a -> m ()
182181push context action = run (Just context) action
183182
184183 where
@@ -211,7 +210,7 @@ handlePushException pchan e = do
211210
212211{-# NOINLINE pushSideDispatch #-}
213212pushSideDispatch :: MonadAsync m
214− => Context m a -> AsyncT m a -> AsyncT m a -> AsyncT m a
213+ => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a
215214pushSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
216215 let chan = childChannel ctx
217216 _ <- doFork (push ctx m1) (handlePushException chan)
@@ -223,15 +222,14 @@ pushSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
223222-- there.
224223{-# NOINLINE pullDispatch #-}
225224pullDispatch :: (MonadIO m, MonadThrow m)
226− => Context m a -> Bool -> AsyncT m a
227−pullDispatch ctx dispatch = AsyncT $ \_ stp yld -> do
225+ => Context a -> AsyncT m a -> Bool -> AsyncT m a
226+pullDispatch ctx m dispatch = AsyncT $ \_ stp yld -> do
228227 let chan = childChannel ctx
229228 eHandler e = handleException e stp
230229 evHandler ev = handleEvent ev stp yld
231230 maybeEvHandler ev = maybe (continue stp yld) evHandler ev
232231
233232 if dispatch then do
234− m <- liftIO $ readIORef (pendingWork ctx)
235233 (runAsyncT m) (Just ctx) stp yld
236234 else do
237235 res <- liftIO $ try $ atomically $ tryReadTChan chan
@@ -245,13 +243,13 @@ pullDispatch ctx dispatch = AsyncT $ \_ stp yld -> do
245243 Nothing -> throwM e
246244
247245 {-# INLINE continue #-}
248− continue stp yld = (runAsyncT (pullDispatch ctx True)) (Just ctx) stp yld
246+ continue stp yld = (runAsyncT (pullDispatch ctx m True)) (Just ctx) stp yld
249247
250248 {-# INLINE handleEvent #-}
251249 handleEvent ev stp yld =
252250 case ev of
253− ChildYield a -> yld a Nothing (Just (pullDispatch ctx False))
254− ChildDone _ a -> yld a Nothing (Just (pullDispatch ctx True))
251+ ChildYield a -> yld a Nothing (Just (pullDispatch ctx m False))
252+ ChildDone _ a -> yld a Nothing (Just (pullDispatch ctx m True))
255253 -- XXX kill threads
256254 ChildStop _ e -> maybe (continue stp yld) throwM e
257255
@@ -259,13 +257,12 @@ pullDispatch ctx dispatch = AsyncT $ \_ stp yld -> do
259257-- m2 in the parent thread. Any exceptions are also pushed to the channel.
260258{-# INLINE pullSideDispatch #-}
261259pullSideDispatch :: MonadAsync m
262− => Context m a -> AsyncT m a -> AsyncT m a -> AsyncT m a
260+ => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a
263261pullSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
264262 let chan = childChannel ctx
265263 _ <- doFork (push (ctx {pullSide = False}) m1) (handlePushException chan)
266264 -- liftIO $ threadDelay 0
267− liftIO $ writeIORef (pendingWork ctx) m2
268− (runAsyncT (pullDispatch ctx False)) Nothing stp yld
265+ (runAsyncT (pullDispatch ctx m2 False)) Nothing stp yld
269266
270267{-# NOINLINE pullDrain #-}
271268pullDrain :: (MonadIO m, MonadThrow m) => TChan (ChildEvent a) -> AsyncT m a
@@ -300,17 +297,16 @@ pullDrain chan = AsyncT $ \_ stp yld -> do
300297{-# NOINLINE pullFork #-}
301298pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
302299pullFork m1 m2 = AsyncT $ \_ stp yld -> do
303− ctx <- liftIO $ newContext (m1 <|> m2)
304− (runAsyncT (pullDispatch ctx True)) Nothing stp yld
300+ ctx <- liftIO $ newContext
301+ let m = pullDispatch ctx ((m1 <|> m2) <> pullDrain (childChannel ctx)) True
302+ (runAsyncT m) Nothing stp yld
305303
306304 where
307305
308− newContext m = do
306+ newContext = do
309307 channel <- atomically newTChan
310− work <- liftIO $ newIORef (m <> pullDrain channel)
311308 return $ Context { childChannel = channel
312309 , pullSide = True
313− , pendingWork = work
314310 }
315311
316312-- Concurrency rate control. Our objective is to create more threads on