@@ -26,10 +26,11 @@ where
2626
2727import Control.Applicative (Alternative (..))
2828import Control.Concurrent (ThreadId, forkIO, killThread,
29− myThreadId)
29+ myThreadId, threadDelay)
3030import Control.Concurrent.STM (TBQueue, atomically, newTBQueue,
31− readTBQueue, tryReadTBQueue,
32− writeTBQueue, isFullTBQueue)
31+ tryReadTBQueue, writeTBQueue,
32+ isEmptyTBQueue, isFullTBQueue,
33+ peekTBQueue)
3334import Control.Exception (SomeException (..))
3435import qualified Control.Exception.Lifted as EL
3536import Control.Monad (ap, liftM, MonadPlus(..), mzero,
@@ -184,10 +185,14 @@ push ctx = run (Just ctx) (dequeueLoop ctx)
184185
185186 where
186187
187− send msg = atomically $ writeTBQueue (childChannel ctx) msg
188+ send msg = atomically $ writeTBQueue (childChannel ctx) msg
189+ stop = do
190+ workQEmpty <- liftIO $ atomically $ isEmptyTBQueue (pendingWork ctx)
191+ if (not workQEmpty) then push ctx
192+ else liftIO $ myThreadId >>= \tid -> send (ChildStop tid Nothing)
188193 yield a _ Nothing = liftIO $ myThreadId >>= \tid -> send (ChildDone tid a)
189194 yield a c (Just r) = liftIO (send (ChildYield a)) >> run c r
190− run c m = (runAsyncT m) c (push ctx) yield
195+ run c m = (runAsyncT m) c stop yield
191196
192197-- Thread tracking has a significant performance overhead (~20% on empty
193198-- threads, it will be lower for heavy threads). It is needed for two reasons:
@@ -237,7 +242,7 @@ handleException e ctx tid = do
237242-- exceptions get propagated to the top level computation and can be handled
238243-- there.
239244{-# NOINLINE pullWorker #-}
240−pullWorker :: (MonadIO m, MonadThrow m) => Context m a -> AsyncT m a
245+pullWorker :: MonadAsync m => Context m a -> AsyncT m a
241246pullWorker ctx = AsyncT $ \pctx stp yld -> do
242247 let continue = (runAsyncT (pullWorker ctx)) pctx stp yld
243248 yield a = yld a pctx (Just (pullWorker ctx))
@@ -245,16 +250,27 @@ pullWorker ctx = AsyncT $ \pctx stp yld -> do
245250 done <- f ctx tid
246251 if done then finish else cont
247252
248− ev <- liftIO $ atomically $ readTBQueue (childChannel ctx)
249− case ev of
250− ChildYield a -> yield a
251− ChildDone tid a ->
252− threadOp tid delThread (yld a pctx Nothing) (yield a)
253− ChildStop tid e ->
254− case e of
255− Nothing -> threadOp tid delThread stp continue
256− Just ex -> handleException ex ctx tid
257− ChildCreate tid -> threadOp tid addThread stp continue
253+ res <- liftIO $ atomically $ tryReadTBQueue (childChannel ctx)
254+ case res of
255+ Nothing -> do
256+ liftIO $ threadDelay 4
257+ let workQ = pendingWork ctx
258+ outQ = childChannel ctx
259+ workQEmpty <- liftIO $ atomically $ isEmptyTBQueue workQ
260+ outQEmpty <- liftIO $ atomically $ isEmptyTBQueue outQ
261+ when (not workQEmpty && outQEmpty) $ pushWorker ctx
262+ void $ liftIO $ atomically $ peekTBQueue (childChannel ctx)
263+ continue
264+ Just ev ->
265+ case ev of
266+ ChildYield a -> yield a
267+ ChildDone tid a ->
268+ threadOp tid delThread (yld a pctx Nothing) (yield a)
269+ ChildStop tid e ->
270+ case e of
271+ Nothing -> threadOp tid delThread stp continue
272+ Just ex -> handleException ex ctx tid
273+ ChildCreate tid -> threadOp tid addThread stp continue
258274
259275-- If an exception occurs we push it to the channel so that it can handled by
260276-- the parent. 'Paused' exceptions are to be collected at the top level.
@@ -265,6 +281,18 @@ handleChildException pchan e = do
265281 tid <- myThreadId
266282 atomically $ writeTBQueue pchan (ChildStop tid (Just e))
267283
284+-- This function is different than "forkWorker" because we have to directly
285+-- insert the threadIds here and cannot use the channel to send ChildCreate
286+-- unlike on the push side. If we do that, the first thread's done message
287+-- may arrive even before the second thread is forked, in that case
288+-- pullWorker will falsely detect that all threads are over.
289+{-# INLINE pushWorker #-}
290+pushWorker :: MonadAsync m => Context m a -> m ()
291+pushWorker ctx = do
292+ let chan = childChannel ctx
293+ tid <- doFork (push ctx) (handleChildException chan)
294+ liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)
295+
268296-- | Split the original computation in a pull-push pair. The original
269297-- computation pulls from a Channel while m1 and m2 push to the channel.
270298{-# NOINLINE pullFork #-}
@@ -276,19 +304,9 @@ pullFork m1 m2 = AsyncT $ \_ stp yld -> do
276304
277305 where
278306
279− -- This function is different than "forkWorker" because we have to directly
280− -- insert the threadIds here and cannot use the channel to send ChildCreate
281− -- unlike on the push side. If we do that, the first thread's done message
282− -- may arrive even before the second thread is forked, in that case
283− -- pullWorker will falsely detect that all threads are over.
284− pushWorker ctx = do
285− let chan = childChannel ctx
286− tid <- doFork (push ctx) (handleChildException chan)
287− liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)
288−
289307 newContext = do
290− channel <- atomically $ newTBQueue 16
291− work <- atomically $ newTBQueue 16
308+ channel <- atomically $ newTBQueue 32
309+ work <- atomically $ newTBQueue 32
292310 running <- newIORef S.empty
293311 done <- newIORef S.empty
294312 return $ Context { childChannel = channel
@@ -354,11 +372,7 @@ dequeueLoop :: MonadAsync m => Context m a -> AsyncT m a
354372dequeueLoop ctx = AsyncT $ \_ stp yld -> do
355373 work <- liftIO $ atomically $ tryReadTBQueue (pendingWork ctx)
356374 case work of
357− Nothing -> do
358− let chan = childChannel ctx
359− tid <- liftIO myThreadId
360− liftIO $ atomically $ writeTBQueue chan (ChildStop tid Nothing)
361− case () of {} -- keep the typechecker happy
375+ Nothing -> stp
362376 Just m -> do
363377 let stop = (runAsyncT (dequeueLoop ctx)) Nothing stp yld
364378 yield a c Nothing = yld a c (Just (dequeueLoop ctx))