@@ -37,10 +37,8 @@ where
3737import Control.Applicative (Alternative (..), liftA2)
3838import Control.Concurrent (ThreadId, forkIO, killThread,
3939 myThreadId, threadDelay)
40−import Control.Concurrent.STM (TBQueue, atomically, newTBQueue,
41− tryReadTBQueue, writeTBQueue,
42− isEmptyTBQueue, isFullTBQueue,
43− peekTBQueue)
40+import Control.Concurrent.MVar (MVar, newEmptyMVar, tryTakeMVar,
41+ tryPutMVar, takeMVar)
4442import Control.Exception (SomeException (..))
4543import qualified Control.Exception.Lifted as EL
4644import Control.Monad (ap, liftM, MonadPlus(..), mzero,
@@ -56,7 +54,9 @@ import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
5654import Data.Functor (void)
5755import Data.IORef (IORef, modifyIORef, newIORef,
5856 writeIORef, readIORef)
59−import Data.Maybe (fromJust)
57+import Data.Atomics (atomicModifyIORefCAS,
58+ atomicModifyIORefCAS_)
59+import Data.Maybe (isNothing)
6060import Data.Monoid ((<>))
6161import Data.Set (Set)
6262import qualified Data.Set as S
@@ -93,9 +93,12 @@ data ChildEvent a =
9393-- State threaded around the monad for thread management
9494------------------------------------------------------------------------------
9595
96+ -- XXX keep just a list of [a] and handle other events out of band as
97+ -- exceptions or using atomicModify
9698data Context m a =
97− Context { outputQueue :: TBQueue (ChildEvent a)
98− , workQueue :: TBQueue (AsyncT m a)
99+ Context { outputQueue :: IORef [ChildEvent a]
100+ , synchOutQ :: MVar Bool -- wakeup mechanism for outQ
101+ , workQueue :: IORef [AsyncT m a]
99102 , runningThreads :: IORef (Set ThreadId)
100103 , doneThreads :: IORef (Set ThreadId)
101104 , fairSched :: Bool
@@ -198,46 +201,40 @@ instance Monad m => Applicative (AsyncT m) where
198201{-# INLINE doFork #-}
199202doFork :: (MonadIO m, MonadBaseControl IO m)
200203 => m ()
201− -> (SomeException -> IO ())
204+ -> (SomeException -> m ())
202205 -> m ThreadId
203206doFork action exHandler =
204207 EL.mask $ \restore ->
205208 liftBaseWith $ \runInIO -> forkIO $ do
206209 -- XXX test the exception handling
207− _ <- runInIO $ EL.catch (restore action) (liftIO . exHandler)
210+ _ <- runInIO $ EL.catch (restore action) exHandler
208211 -- XXX restore state here?
209212 return ()
210213
211−{-# NOINLINE forkWorker #-}
212−forkWorker :: MonadAsync m => Maybe (Context m a) -> m ()
213−forkWorker ctx = do
214− let c = fromJust ctx
215− let q = outputQueue c
216− tid <- doFork (push c) (handleChildException q)
217− liftIO $ atomically $ writeTBQueue q (ChildCreate tid)
214+-- XXX exception safety of all atomic/MVar operations
218215
219−-- Note: This is performance sensitive code.
216+{-# INLINE send #-}
217+send :: MonadIO m => Context m a -> ChildEvent a -> m ()
218+send ctx msg = liftIO $ do
219+ atomicModifyIORefCAS_ (outputQueue ctx) $ \es -> msg : es
220+ -- XXX need a memory barrier? The wake up must happen only after the
221+ -- store has finished otherwise we can have lost wakeup problems.
222+ void $ tryPutMVar (synchOutQ ctx) True
223+
224+-- Note: Left associated operations can grow this queue to a large size
220225{-# INLINE queueWork #-}
221−queueWork :: MonadAsync m
222− => Maybe (Context m a) -> TBQueue (AsyncT m a) -> AsyncT m a -> m ()
223−queueWork ctx q m = do
224− -- To guarantee deadlock avoidance we need to dispatch a worker when the
225− -- workQueue goes full. Otherwise all the worker threads might be waiting
226− -- on the queue and never wakeup.
227− --
228− -- TBD If we run out of threads we can also evaluate the action completely
229− -- right here, disallowing any further child workers and turning the
230− -- parallel composition into interleaved serial composition.
231− workQFull <- liftIO $ atomically $ isFullTBQueue q
232− when (workQFull) $ forkWorker ctx
233− liftIO $ atomically $ writeTBQueue q m
226+queueWork :: MonadAsync m => IORef [AsyncT m a] -> AsyncT m a -> m ()
227+queueWork q m = liftIO $ atomicModifyIORefCAS_ q $ \ ms -> m : ms
234228
235229-- Note: This is performance sensitive code.
236230{-# INLINE dequeueLoop #-}
237231dequeueLoop :: MonadAsync m
238− => Maybe (Context m a) -> TBQueue (AsyncT m a) -> Bool -> AsyncT m a
232+ => Maybe (Context m a) -> IORef [AsyncT m a] -> Bool -> AsyncT m a
239233dequeueLoop ctx q fair = AsyncT $ \_ stp yld -> do
240− work <- liftIO $ atomically $ tryReadTBQueue q
234+ work <- liftIO $ atomicModifyIORefCAS q $ \ ms ->
235+ case ms of
236+ [] -> ([], Nothing)
237+ x : xs -> (xs, Just x)
241238 case work of
242239 Nothing -> stp
243240 Just m -> do
@@ -245,7 +242,7 @@ dequeueLoop ctx q fair = AsyncT $ \_ stp yld -> do
245242 continue a c = yld a c (Just (dequeueLoop ctx q fair))
246243 yield a c Nothing = continue a c
247244 yield a c r | not fair = yld a c r
248− yield a c (Just r) = queueWork ctx q r >> continue a c
245+ yield a c (Just r) = queueWork q r >> continue a c
249246 (runAsyncT m) ctx stop yield
250247
251248{-# NOINLINE push #-}
@@ -255,13 +252,17 @@ push ctx = run (Just ctx) (dequeueLoop (Just ctx) (workQueue ctx)
255252
256253 where
257254
258− send msg = atomically $ writeTBQueue (outputQueue ctx) msg
259255 stop = do
260− workQEmpty <- liftIO $ atomically $ isEmptyTBQueue (workQueue ctx)
261− if (not workQEmpty) then push ctx
262− else liftIO $ myThreadId >>= \tid -> send (ChildStop tid Nothing)
263− yield a _ Nothing = liftIO $ myThreadId >>= \tid -> send (ChildDone tid a)
264− yield a c (Just r) = liftIO (send (ChildYield a)) >> run c r
256+ work <- liftIO $ readIORef (workQueue ctx)
257+ case work of
258+ [] -> do
259+ tid <- liftIO $ myThreadId
260+ send ctx (ChildStop tid Nothing)
261+ _ -> push ctx
262+ yield a _ Nothing = do
263+ tid <- liftIO myThreadId
264+ send ctx (ChildDone tid a)
265+ yield a c (Just r) = send ctx (ChildYield a) >> run c r
265266 run c m = (runAsyncT m) c stop yield
266267
267268-- Thread tracking has a significant performance overhead (~20% on empty
@@ -301,6 +302,13 @@ addThread ctx tid = accountThread tid (doneThreads ctx) (runningThreads ctx)
301302delThread :: MonadIO m => Context m a -> ThreadId -> m Bool
302303delThread ctx tid = accountThread tid (runningThreads ctx) (doneThreads ctx)
303304
305+{-# INLINE allThreadsDone #-}
306+allThreadsDone :: MonadIO m => Context m a -> m Bool
307+allThreadsDone ctx = liftIO $ do
308+ s1 <- readIORef (runningThreads ctx)
309+ s2 <- readIORef (doneThreads ctx)
310+ return $ S.null s1 && S.null s2
311+
304312handleException :: (MonadIO m, MonadThrow m)
305313 => SomeException -> Context m a -> ThreadId -> m r
306314handleException e ctx tid = do
@@ -308,59 +316,64 @@ handleException e ctx tid = do
308316 liftIO $ readIORef (runningThreads ctx) >>= mapM_ killThread
309317 throwM e
310318
311−{-# NOINLINE sendWorkerWait #-}
319+{-# INLINE sendWorkerWait #-}
312320sendWorkerWait :: MonadAsync m => Context m a -> m ()
313321sendWorkerWait ctx = do
314− liftIO $ threadDelay 4
315− let workQ = workQueue ctx
316− outQ = outputQueue ctx
317− workQEmpty <- liftIO $ atomically $ isEmptyTBQueue workQ
318− outQEmpty <- liftIO $ atomically $ isEmptyTBQueue outQ
319− when (not workQEmpty && outQEmpty) $ pushWorker ctx
320− void $ liftIO $ atomically $ peekTBQueue (outputQueue ctx)
322+ liftIO $ threadDelay 200
323+ output <- liftIO $ readIORef (outputQueue ctx)
324+ case output of
325+ [] -> pushWorker ctx
326+ _ -> return ()
327+ void $ liftIO $ takeMVar (synchOutQ ctx)
321328
322329-- Note: This is performance sensitive code.
323330{-# NOINLINE pullWorker #-}
324331pullWorker :: MonadAsync m => Context m a -> AsyncT m a
325332pullWorker ctx = AsyncT $ \pctx stp yld -> do
326− let continue = (runAsyncT (pullWorker ctx)) pctx stp yld
327− yield a = yld a pctx (Just (pullWorker ctx))
328− threadOp tid f finish cont = do
329− done <- f ctx tid
330− if done then finish else cont
331−
332− res <- liftIO $ atomically $ tryReadTBQueue (outputQueue ctx)
333− case res of
334− Nothing -> sendWorkerWait ctx >> continue
335− Just ev ->
336− case ev of
337− ChildYield a -> yield a
338− ChildDone tid a ->
339− threadOp tid delThread (yld a pctx Nothing) (yield a)
340− ChildStop tid e ->
341− case e of
342− Nothing -> threadOp tid delThread stp continue
343− Just ex -> handleException ex ctx tid
344− ChildCreate tid -> threadOp tid addThread stp continue
333+ res <- liftIO $ tryTakeMVar (synchOutQ ctx)
334+ when (isNothing res) $ sendWorkerWait ctx
335+ list <- liftIO $ atomicModifyIORefCAS (outputQueue ctx) $ \x -> ([], x)
336+ (runAsyncT $ processEvents list) pctx stp yld
337+ where
338+
339+ {-# INLINE processEvents #-}
340+ processEvents [] = AsyncT $ \pctx stp yld -> do
341+ done <- allThreadsDone ctx
342+ if not done
343+ then (runAsyncT (pullWorker ctx)) pctx stp yld
344+ else stp
345+
346+ processEvents (ev : es) = AsyncT $ \pctx stp yld -> do
347+ let continue = (runAsyncT (processEvents es)) pctx stp yld
348+ yield a = yld a pctx (Just (processEvents es))
349+
350+ case ev of
351+ ChildYield a -> yield a
352+ ChildDone tid a -> delThread ctx tid >> yield a
353+ ChildStop tid e ->
354+ case e of
355+ Nothing -> delThread ctx tid >> continue
356+ Just ex -> handleException ex ctx tid
357+ ChildCreate tid -> addThread ctx tid >> continue
345358
346359-- If an exception occurs we push it to the channel so that it can handled by
347360-- the parent. 'Paused' exceptions are to be collected at the top level.
348361-- XXX Paused exceptions should only bubble up to the runRecorder computation
349362{-# NOINLINE handleChildException #-}
350−handleChildException :: TBQueue (ChildEvent a) -> SomeException -> IO ()
351−handleChildException pchan e = do
352− tid <- myThreadId
353− atomically $ writeTBQueue pchan (ChildStop tid (Just e))
363+handleChildException :: MonadIO m => Context m a -> SomeException -> m ()
364+handleChildException ctx e = do
365+ tid <- liftIO myThreadId
366+ send ctx (ChildStop tid (Just e))
354367
355368-- This function is different than "forkWorker" because we have to directly
356369-- insert the threadIds here and cannot use the channel to send ChildCreate
357370-- unlike on the push side. If we do that, the first thread's done message
358371-- may arrive even before the second thread is forked, in that case
359372-- pullWorker will falsely detect that all threads are over.
360−{-# INLINE pushWorker #-}
373+{-# NOINLINE pushWorker #-}
361374pushWorker :: MonadAsync m => Context m a -> m ()
362375pushWorker ctx = do
363− tid <- doFork (push ctx) (handleChildException (outputQueue ctx))
376+ tid <- doFork (push ctx) (handleChildException ctx)
364377 liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)
365378
366379-- | Split the original computation in a pull-push pair. The original
@@ -370,17 +383,19 @@ pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a
370383pullFork m1 m2 fair = AsyncT $ \_ stp yld -> do
371384 ctx <- liftIO $ newContext
372385 let q = workQueue ctx
373− queueWork (Just ctx) q m1 >> queueWork (Just ctx) q m2 >> pushWorker ctx
386+ queueWork q m1 >> queueWork q m2 >> pushWorker ctx
374387 (runAsyncT (pullWorker ctx)) Nothing stp yld
375388
376389 where
377390
378391 newContext = do
379− outQ <- atomically $ newTBQueue 32
380− workQ <- atomically $ newTBQueue 32
392+ outQ <- newIORef []
393+ outQMv <- newEmptyMVar
394+ workQ <- newIORef []
381395 running <- newIORef S.empty
382396 done <- newIORef S.empty
383397 return $ Context { outputQueue = outQ
398+ , synchOutQ = outQMv
384399 , runningThreads = running
385400 , doneThreads = done
386401 , workQueue = workQ
@@ -430,7 +445,7 @@ parallel m1 m2 fair = AsyncT $ \ctx stp yld -> do
430445 -- go left to right. If the left one keeps producing results we
431446 -- may or may not run the right one.
432447 let q = workQueue c
433− queueWork ctx q m1 >> queueWork ctx q m2
448+ queueWork q m1 >> queueWork q m2
434449 (runAsyncT (dequeueLoop ctx q fair)) Nothing stp yld
435450
436451instance MonadAsync m => Alternative (AsyncT m) where