@@ -25,14 +25,12 @@ module Asyncly.AsyncT
2525where
2626
2727import Control.Applicative (Alternative (..))
28−import Control.Concurrent (ThreadId, forkIO,
29− myThreadId, threadDelay)
28+import Control.Concurrent (ThreadId, forkIO, killThread,
29+ myThreadId)
3030import Control.Concurrent.STM (TChan, atomically, newTChan,
3131 readTChan, tryReadTChan,
3232 writeTChan)
33−import Control.Exception (fromException, try,
34− SomeException (..),
35− BlockedIndefinitelyOnSTM(..))
33+import Control.Exception (SomeException (..))
3634import qualified Control.Exception.Lifted as EL
3735import Control.Monad (ap, liftM, MonadPlus(..), mzero)
3836--import Control.Monad.Base (MonadBase (..), liftBaseDefault)
@@ -46,10 +44,13 @@ import Control.Monad.Trans.Control (ComposeSt, MonadBaseControl (..),
4644 defaultLiftBaseWith,
4745 defaultRestoreM, liftBaseWith)
4846 -}
49−import Data.IORef (IORef, writeIORef, newIORef,
50− readIORef)
47+import Data.Functor (void)
48+import Data.IORef (IORef, modifyIORef, newIORef,
49+ writeIORef, readIORef)
5150import Data.Maybe (maybe)
5251import Data.Monoid ((<>))
52+import Data.Set (Set)
53+import qualified Data.Set as S
5354
5455import Control.Monad.Trans.Recorder (MonadRecorder(..))
5556
@@ -77,6 +78,7 @@ data ChildEvent a =
7778 ChildYield a
7879 | ChildDone ThreadId a
7980 | ChildStop ThreadId (Maybe SomeException)
81+ | ChildCreate ThreadId
8082
8183------------------------------------------------------------------------------
8284-- State threaded around the monad for thread management
@@ -85,6 +87,8 @@ data ChildEvent a =
8587data Context a =
8688 Context { childChannel :: TChan (ChildEvent a)
8789 , pullSide :: Bool
90+ , runningThreads :: IORef (Set ThreadId)
91+ , doneThreads :: IORef (Set ThreadId)
8892 }
8993
9094-- The 'Maybe (AsyncT m a)' is redundant as we can use 'stop' value for the
@@ -98,7 +102,7 @@ data Context a =
98102newtype AsyncT m a =
99103 AsyncT {
100104 runAsyncT :: forall r.
101− Maybe (Context a) -- state
105+ Maybe (Context a) -- state
102106 -> m r -- stop
103107 -> (a -> Maybe (Context a) -> Maybe (AsyncT m a) -> m r) -- yield
104108 -> m r
@@ -168,11 +172,11 @@ doFork :: (MonadIO m, MonadBaseControl IO m)
168172 => m ()
169173 -> (SomeException -> IO ())
170174 -> m ThreadId
171−doFork action preExit =
175+doFork action exHandler =
172176 EL.mask $ \restore ->
173177 liftBaseWith $ \runInIO -> forkIO $ do
174178 -- XXX test the exception handling
175− _ <- runInIO $ EL.catch (restore action) (liftIO . preExit)
179+ _ <- runInIO $ EL.catch (restore action) (liftIO . exHandler)
176180 -- XXX restore state here?
177181 return ()
178182
@@ -202,9 +206,9 @@ push context action = run (Just context) action
202206-- If an exception occurs we push it to the channel so that it can handled by
203207-- the parent. 'Paused' exceptions are to be collected at the top level.
204208-- XXX Paused exceptions should only bubble up to the runRecorder computation
205−{-# NOINLINE handlePushException #-}
206−handlePushException :: TChan (ChildEvent a) -> SomeException -> IO ()
207−handlePushException pchan e = do
209+{-# NOINLINE handleChildException #-}
210+handleChildException :: TChan (ChildEvent a) -> SomeException -> IO ()
211+handleChildException pchan e = do
208212 tid <- myThreadId
209213 atomically $ writeTChan pchan (ChildStop tid (Just e))
210214
@@ -213,9 +217,43 @@ pushSideDispatch :: MonadAsync m
213217 => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a
214218pushSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
215219 let chan = childChannel ctx
216− _ <- doFork (push ctx m1) (handlePushException chan)
220+ tid <- doFork (push ctx m1) (handleChildException chan)
221+ liftIO $ atomically $ writeTChan chan (ChildCreate tid)
217222 (runAsyncT m2) (Just ctx) stp yld
218223
224+-- This is a bit messy because ChildCreate and ChildDone events can arrive out
225+-- of order in case of pushSideDispatch. Returns whether we are done draining
226+-- threads.
227+{-# INLINE accountThread #-}
228+accountThread :: MonadIO m
229+ => ThreadId -> IORef (Set ThreadId) -> IORef (Set ThreadId) -> m Bool
230+accountThread tid ref1 ref2 = liftIO $ do
231+ s1 <- readIORef ref1
232+ s2 <- readIORef ref2
233+
234+ if (S.member tid s1) then do
235+ let r = S.delete tid s1
236+ writeIORef ref1 r
237+ return $ S.null r && S.null s2
238+ else do
239+ liftIO $ writeIORef ref2 (S.insert tid s2)
240+ return False
241+
242+{-# NOINLINE addThread #-}
243+addThread :: MonadIO m => Context a -> ThreadId -> m Bool
244+addThread ctx tid = accountThread tid (doneThreads ctx) (runningThreads ctx)
245+
246+{-# INLINE delThread #-}
247+delThread :: MonadIO m => Context a -> ThreadId -> m Bool
248+delThread ctx tid = accountThread tid (runningThreads ctx) (doneThreads ctx)
249+
250+handleException :: (MonadIO m, MonadThrow m)
251+ => SomeException -> Context a -> ThreadId -> m r
252+handleException e ctx tid = do
253+ void (delThread ctx tid)
254+ liftIO $ readIORef (runningThreads ctx) >>= mapM_ killThread
255+ throwM e
256+
219257-- XXX the TBQueue size should be proportional to the pending threads.
220258-- We re-raise any exceptions received from the child threads, that way
221259-- exceptions get propagated to the top level computation and can be handled
@@ -224,24 +262,14 @@ pushSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
224262pullDispatch :: (MonadIO m, MonadThrow m)
225263 => Context a -> AsyncT m a -> Bool -> AsyncT m a
226264pullDispatch ctx m dispatch = AsyncT $ \_ stp yld -> do
227− let chan = childChannel ctx
228− eHandler e = handleException e stp
229− evHandler ev = handleEvent ev stp yld
230− maybeEvHandler ev = maybe (continue stp yld) evHandler ev
231−
232265 if dispatch then do
233266 (runAsyncT m) (Just ctx) stp yld
234267 else do
235− res <- liftIO $ try $ atomically $ tryReadTChan chan
236− either eHandler maybeEvHandler res
268+ res <- liftIO $ atomically $ tryReadTChan (childChannel ctx)
269+ maybe (continue stp yld) (\ev -> handleEvent ev stp yld) res
237270
238271 where
239272
240− handleException e stp =
241− case fromException e of
242− Just BlockedIndefinitelyOnSTM -> stp
243− Nothing -> throwM e
244−
245273 {-# INLINE continue #-}
246274 continue stp yld = (runAsyncT (pullDispatch ctx m True)) (Just ctx) stp yld
247275
@@ -249,9 +277,14 @@ pullDispatch ctx m dispatch = AsyncT $ \_ stp yld -> do
249277 handleEvent ev stp yld =
250278 case ev of
251279 ChildYield a -> yld a Nothing (Just (pullDispatch ctx m False))
252− ChildDone _ a -> yld a Nothing (Just (pullDispatch ctx m True))
253− -- XXX kill threads
254− ChildStop _ e -> maybe (continue stp yld) throwM e
280+ ChildDone tid a -> do
281+ void $ delThread ctx tid
282+ yld a Nothing (Just (pullDispatch ctx m True))
283+ ChildStop tid e -> do
284+ case e of
285+ Nothing -> void (delThread ctx tid) >> continue stp yld
286+ Just x -> handleException x ctx tid
287+ ChildCreate tid -> void (addThread ctx tid) >> continue stp yld
255288
256289-- | run m1 in a new thread, pushing its results to a pull channel and then run
257290-- m2 in the parent thread. Any exceptions are also pushed to the channel.
@@ -260,37 +293,32 @@ pullSideDispatch :: MonadAsync m
260293 => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a
261294pullSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
262295 let chan = childChannel ctx
263− _ <- doFork (push (ctx {pullSide = False}) m1) (handlePushException chan)
264− -- liftIO $ threadDelay 0
296+ tid <- doFork (push (ctx {pullSide = False}) m1)
297+ (handleChildException chan)
298+ liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)
265299 (runAsyncT (pullDispatch ctx m2 False)) Nothing stp yld
266300
267301{-# NOINLINE pullDrain #-}
268−pullDrain :: (MonadIO m, MonadThrow m) => TChan (ChildEvent a) -> AsyncT m a
269−pullDrain chan = AsyncT $ \_ stp yld -> do
270− let eHandler e = handleException e stp
271− evHandler ev = handleEvent ev stp yld
272−
273− res <- liftIO $ try (atomically $ readTChan chan)
274− either eHandler evHandler res
275−
276− where
277−
278− handleException e stp =
279− case fromException e of
280− Just BlockedIndefinitelyOnSTM -> stp
281− Nothing -> throwM e
282−
283− {-# INLINE continue #-}
284− continue stp yld = (runAsyncT (pullDrain chan)) Nothing stp yld
285−
286− {-# INLINE handleEvent #-}
287− handleEvent ev stp yld =
288− let yielder a = yld a Nothing (Just (pullDrain chan))
289− in case ev of
290− ChildYield a -> yielder a
291− ChildDone _ a -> yielder a
292− -- XXX kill threads
293− ChildStop _ e -> maybe (continue stp yld) throwM e
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
294322
295323-- | Split the original computation in a pull-push pair. The original
296324-- computation pulls from a Channel while m1 and m2 push to the channel.
@@ -298,15 +326,25 @@ pullDrain chan = AsyncT $ \_ stp yld -> do
298326pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
299327pullFork m1 m2 = AsyncT $ \_ stp yld -> do
300328 ctx <- liftIO $ newContext
301− let m = pullDispatch ctx ((m1 <|> m2) <> pullDrain (childChannel ctx)) True
329+ let m = pullDispatch ctx ((m1 <|> m2) <> finalizer ctx) True
302330 (runAsyncT m) Nothing stp yld
303331
304332 where
305333
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+
306340 newContext = do
307341 channel <- atomically newTChan
342+ running <- liftIO $ newIORef S.empty
343+ done <- liftIO $ newIORef S.empty
308344 return $ Context { childChannel = channel
309345 , pullSide = True
346+ , runningThreads = running
347+ , doneThreads = done
310348 }
311349
312350-- Concurrency rate control. Our objective is to create more threads on