@@ -26,10 +26,10 @@ where
2626
2727import Control.Applicative (Alternative (..))
2828import Control.Concurrent (ThreadId, forkIO, killThread,
29− myThreadId)
29+ myThreadId, newQSem, QSem,
30+ signalQSem, waitQSem)
3031import Control.Concurrent.STM (TChan, atomically, newTChan,
31− readTChan, tryReadTChan,
32− writeTChan)
32+ tryReadTChan, writeTChan)
3333import Control.Exception (SomeException (..))
3434import qualified Control.Exception.Lifted as EL
3535import Control.Monad (ap, liftM, MonadPlus(..), mzero)
@@ -85,10 +85,10 @@ data ChildEvent a =
8585------------------------------------------------------------------------------
8686
8787data Context a =
88− Context { childChannel :: TChan (ChildEvent a)
89− , pullSide :: Bool
90− , runningThreads :: IORef (Set ThreadId)
91− , doneThreads :: IORef (Set ThreadId)
88+ Context { childChannel :: TChan (ChildEvent a)
89+ , dispatchReq :: QSem
90+ , runningThreads :: IORef (Set ThreadId)
91+ , doneThreads :: IORef (Set ThreadId)
9292 }
9393
9494-- The 'Maybe (AsyncT m a)' is redundant as we can use 'stop' value for the
@@ -182,7 +182,9 @@ doFork action exHandler =
182182
183183{-# NOINLINE push #-}
184184push :: MonadIO m => Context a -> AsyncT m a -> m ()
185−push context action = run (Just context) action
185+push context action = do
186+ liftIO $ waitQSem (dispatchReq context)
187+ run (Just context) action
186188
187189 where
188190
@@ -203,24 +205,17 @@ push context action = run (Just context) action
203205 continue a ctx m = channelYield a >> run ctx m
204206 yielder a ctx r = maybe (done a) (\rx -> continue a ctx rx) r
205207
206−-- If an exception occurs we push it to the channel so that it can handled by
207−-- the parent. 'Paused' exceptions are to be collected at the top level.
208−-- XXX Paused exceptions should only bubble up to the runRecorder computation
209−{-# NOINLINE handleChildException #-}
210−handleChildException :: TChan (ChildEvent a) -> SomeException -> IO ()
211−handleChildException pchan e = do
212− tid <- myThreadId
213− atomically $ writeTChan pchan (ChildStop tid (Just e))
214−
215−{-# NOINLINE pushSideDispatch #-}
216−pushSideDispatch :: MonadAsync m
217− => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a
218−pushSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
219− let chan = childChannel ctx
220− tid <- doFork (push ctx m1) (handleChildException chan)
221− liftIO $ atomically $ writeTChan chan (ChildCreate tid)
222− (runAsyncT m2) (Just ctx) stp yld
223−
208+-- Thread tracking has a significant performance overhead (~20% on empty
209+-- threads, it will be lower for heavy threads). It is needed for two reasons:
210+--
211+-- 1) Killing threads on exceptions. Threads may not be allowed to go away by
212+-- themselves because they may run for significant times before going away or
213+-- worse they may be stuck in IO and never go away.
214+--
215+-- 2) To know when all threads are done. This can be acheived by detecting a
216+-- BlockedIndefinitelyOnSTM exception too. But we will have to trigger a GC to
217+-- make sure that we detect it promptly.
218+--
224219-- This is a bit messy because ChildCreate and ChildDone events can arrive out
225220-- of order in case of pushSideDispatch. Returns whether we are done draining
226221-- threads.
@@ -258,75 +253,49 @@ handleException e ctx tid = do
258253-- We re-raise any exceptions received from the child threads, that way
259254-- exceptions get propagated to the top level computation and can be handled
260255-- there.
261−{-# NOINLINE pullDispatch #-}
262−pullDispatch :: (MonadIO m, MonadThrow m)
263− => Context a -> AsyncT m a -> Bool -> AsyncT m a
264−pullDispatch ctx m dispatch = AsyncT $ \_ stp yld -> do
265− if dispatch then do
266− (runAsyncT m) (Just ctx) stp yld
267− else do
268− res <- liftIO $ atomically $ tryReadTChan (childChannel ctx)
269− maybe (continue stp yld) (\ev -> handleEvent ev stp yld) res
256+{-# NOINLINE pullDrain #-}
257+pullDrain :: (MonadIO m, MonadThrow m) => Context a -> AsyncT m a
258+pullDrain ctx = AsyncT $ \_ stp yld -> do
259+ res <- liftIO $ atomically $ tryReadTChan (childChannel ctx)
260+ maybe (dispatch >> continue stp yld)
261+ (\ev -> handleEvent ev stp yld) res
270262
271263 where
272264
273265 {-# INLINE continue #-}
274− continue stp yld = (runAsyncT (pullDispatch ctx m True)) (Just ctx) stp yld
266+ continue stp yld = (runAsyncT (pullDrain ctx)) Nothing stp yld
267+
268+ {-# INLINE dispatch #-}
269+ dispatch = liftIO $ signalQSem (dispatchReq ctx)
275270
276271 {-# INLINE handleEvent #-}
277− handleEvent ev stp yld =
278− case ev of
279− ChildYield a -> yld a Nothing (Just (pullDispatch ctx m False))
272+ handleEvent ev stp yld = do
273+ let yielder a = yld a Nothing (Just (pullDrain ctx))
274+ case ev of
275+ ChildYield a -> yielder a
280276 ChildDone tid a -> do
281− void $ delThread ctx tid
282− yld a Nothing (Just (pullDispatch ctx m True))
277+ dispatch
278+ done <- delThread ctx tid
279+ if done then (yld a Nothing Nothing) else (yielder a)
283280 ChildStop tid e -> do
284281 case e of
285− Nothing -> void (delThread ctx tid) >> continue stp yld
282+ Nothing -> do
283+ dispatch
284+ done <- delThread ctx tid
285+ if done then stp else continue stp yld
286286 Just x -> handleException x ctx tid
287− ChildCreate tid -> void (addThread ctx tid) >> continue stp yld
288−
289−-- | run m1 in a new thread, pushing its results to a pull channel and then run
290−-- m2 in the parent thread. Any exceptions are also pushed to the channel.
291−{-# INLINE pullSideDispatch #-}
292−pullSideDispatch :: MonadAsync m
293− => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a
294−pullSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
295− let chan = childChannel ctx
296− tid <- doFork (push (ctx {pullSide = False}) m1)
297− (handleChildException chan)
298− liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)
299− (runAsyncT (pullDispatch ctx m2 False)) Nothing stp yld
287+ ChildCreate tid -> do
288+ done <- addThread ctx tid
289+ if done then stp else continue stp yld
300290
301−{-# NOINLINE pullDrain #-}
302−pullDrain :: (MonadIO m, MonadThrow m) => Context a -> AsyncT m a
303−pullDrain ctx = AsyncT $ \_ stp yld -> do
304− let yielder a = yld a Nothing (Just (pullDrain ctx))
305− continue = (runAsyncT (pullDrain ctx)) Nothing stp yld
306−
307− res <- liftIO $ atomically $ readTChan (childChannel ctx)
308− case res of
309− ChildYield a -> yielder a
310− ChildDone tid a -> do
311− done <- delThread ctx tid
312− if done then (yld a Nothing Nothing) else (yielder a)
313− ChildStop tid e -> do
314− case e of
315− Nothing -> do
316− done <- delThread ctx tid
317− if done then stp else continue
318− Just x -> handleException x ctx tid
319− ChildCreate tid -> do
320− done <- addThread ctx tid
321− if done then stp else continue
322−
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
291+-- If an exception occurs we push it to the channel so that it can handled by
292+-- the parent. 'Paused' exceptions are to be collected at the top level.
293+-- XXX Paused exceptions should only bubble up to the runRecorder computation
294+{-# NOINLINE handleChildException #-}
295+handleChildException :: TChan (ChildEvent a) -> SomeException -> IO ()
296+handleChildException pchan e = do
297+ tid <- myThreadId
298+ atomically $ writeTChan pchan (ChildStop tid (Just e))
330299
331300-- | Split the original computation in a pull-push pair. The original
332301-- computation pulls from a Channel while m1 and m2 push to the channel.
@@ -334,19 +303,30 @@ pullDrainStart ctx = AsyncT $ \_ stp yld -> do
334303pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
335304pullFork m1 m2 = AsyncT $ \_ stp yld -> do
336305 ctx <- liftIO $ newContext
337− let m = pullDispatch ctx ((m1 <|> m2) <> pullDrainStart ctx) True
338− (runAsyncT m) Nothing stp yld
306+ initialFork ctx m1 >> initialFork ctx m2
307+ (runAsyncT (pullDrain ctx)) Nothing stp yld
339308
340309 where
341310
311+ -- This function is different than "fork" because we have to directly
312+ -- insert the threadIds here and cannot use the channel to send ChildCreate
313+ -- unlike on the push side. If we do that, the first thread's done message
314+ -- may arrive even before the second thread is forked, in that case
315+ -- pullDrain will falsely detect that all threads are over.
316+ initialFork ctx m = do
317+ let chan = childChannel ctx
318+ tid <- doFork (push ctx m) (handleChildException chan)
319+ liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)
320+
342321 newContext = do
343322 channel <- atomically newTChan
344323 running <- liftIO $ newIORef S.empty
345− done <- liftIO $ newIORef S.empty
346− return $ Context { childChannel = channel
347− , pullSide = True
324+ done <- liftIO $ newIORef S.empty
325+ count <- liftIO $ newQSem 1
326+ return $ Context { childChannel = channel
327+ , dispatchReq = count
348328 , runningThreads = running
349− , doneThreads = done
329+ , doneThreads = done
350330 }
351331
352332-- Concurrency rate control. Our objective is to create more threads on
@@ -379,7 +359,14 @@ pullFork m1 m2 = AsyncT $ \_ stp yld -> do
379359-- XXX to rate control left folded structrues we will have to return the
380360-- residual work back to the dispatcher. It will also consume a lot of
381361-- memory due to queueing of all the work before execution starts.
382−--
362+
363+{-# INLINE fork #-}
364+fork :: MonadAsync m => Context a -> AsyncT m a -> m ()
365+fork ctx m = do
366+ let chan = childChannel ctx
367+ tid <- doFork (push ctx m) (handleChildException chan)
368+ liftIO $ atomically $ writeTChan chan (ChildCreate tid)
369+
383370instance MonadAsync m => Alternative (AsyncT m) where
384371 empty = mempty
385372
@@ -389,12 +376,7 @@ instance MonadAsync m => Alternative (AsyncT m) where
389376 m1 <|> m2 = AsyncT $ \ctx stp yld -> do
390377 case ctx of
391378 Nothing -> (runAsyncT (pullFork m1 m2)) ctx stp yld
392− Just c -> do
393− if pullSide c then
394− (runAsyncT (pullSideDispatch c m1 m2)) ctx stp yld
395− else
396− -- for left associated compositions
397− (runAsyncT (pushSideDispatch c m1 m2)) ctx stp yld
379+ Just c -> fork c m2 >> (runAsyncT m1) ctx stp yld
398380
399381instance MonadAsync m => MonadPlus (AsyncT m) where
400382 mzero = empty