@@ -26,13 +26,14 @@ where
2626
2727import Control.Applicative (Alternative (..))
2828import Control.Concurrent (ThreadId, forkIO, killThread,
29− myThreadId, newQSem, QSem,
30− signalQSem, waitQSem)
29+ myThreadId)
3130import Control.Concurrent.STM (TBQueue, atomically, newTBQueue,
32− tryReadTBQueue, writeTBQueue)
31+ readTBQueue, tryReadTBQueue,
32+ writeTBQueue, isFullTBQueue)
3333import Control.Exception (SomeException (..))
3434import qualified Control.Exception.Lifted as EL
35−import Control.Monad (ap, liftM, MonadPlus(..), mzero)
35+import Control.Monad (ap, liftM, MonadPlus(..), mzero,
36+ when)
3637--import Control.Monad.Base (MonadBase (..), liftBaseDefault)
3738import Control.Monad.Catch (MonadThrow, throwM)
3839import Control.Monad.IO.Class (MonadIO(..))
@@ -84,9 +85,9 @@ data ChildEvent a =
8485-- State threaded around the monad for thread management
8586------------------------------------------------------------------------------
8687
87−data Context a =
88+data Context m a =
8889 Context { childChannel :: TBQueue (ChildEvent a)
89− , dispatchReq :: QSem
90+ , pendingWork :: TBQueue (AsyncT m a)
9091 , runningThreads :: IORef (Set ThreadId)
9192 , doneThreads :: IORef (Set ThreadId)
9293 }
@@ -102,9 +103,9 @@ data Context a =
102103newtype AsyncT m a =
103104 AsyncT {
104105 runAsyncT :: forall r.
105− Maybe (Context a) -- state
106+ Maybe (Context m a) -- state
106107 -> m r -- stop
107− -> (a -> Maybe (Context a) -> Maybe (AsyncT m a) -> m r) -- yield
108+ -> (a -> Maybe (Context m a) -> Maybe (AsyncT m a) -> m r) -- yield
108109 -> m r
109110 }
110111
@@ -181,17 +182,13 @@ doFork action exHandler =
181182 return ()
182183
183184{-# NOINLINE push #-}
184−push :: MonadIO m => Context a -> AsyncT m a -> m ()
185−push context action = do
186− liftIO $ waitQSem (dispatchReq context)
187− run (Just context) action
185+push :: MonadAsync m => Context m a -> m ()
186+push context = run (Just context) (dequeueLoop context)
188187
189188 where
190189
191190 run ctx m = (runAsyncT m) ctx channelStop yielder
192191
193− -- XXX make this a bounded channel so that we block if the previous value
194− -- is not consumed yet.
195192 chan = childChannel context
196193 channelYield a = liftIO $ atomically $ writeTBQueue chan (ChildYield a)
197194 channelDone a = do
@@ -235,15 +232,15 @@ accountThread tid ref1 ref2 = liftIO $ do
235232 return False
236233
237234{-# NOINLINE addThread #-}
238−addThread :: MonadIO m => Context a -> ThreadId -> m Bool
235+addThread :: MonadIO m => Context m a -> ThreadId -> m Bool
239236addThread ctx tid = accountThread tid (doneThreads ctx) (runningThreads ctx)
240237
241238{-# INLINE delThread #-}
242−delThread :: MonadIO m => Context a -> ThreadId -> m Bool
239+delThread :: MonadIO m => Context m a -> ThreadId -> m Bool
243240delThread ctx tid = accountThread tid (runningThreads ctx) (doneThreads ctx)
244241
245242handleException :: (MonadIO m, MonadThrow m)
246− => SomeException -> Context a -> ThreadId -> m r
243+ => SomeException -> Context m a -> ThreadId -> m r
247244handleException e ctx tid = do
248245 void (delThread ctx tid)
249246 liftIO $ readIORef (runningThreads ctx) >>= mapM_ killThread
@@ -252,34 +249,27 @@ handleException e ctx tid = do
252249-- We re-raise any exceptions received from the child threads, that way
253250-- exceptions get propagated to the top level computation and can be handled
254251-- there.
255−{-# NOINLINE pullDrain #-}
256−pullDrain :: (MonadIO m, MonadThrow m) => Context a -> AsyncT m a
257−pullDrain ctx = AsyncT $ \_ stp yld -> do
258− res <- liftIO $ atomically $ tryReadTBQueue (childChannel ctx)
259− maybe (dispatch >> continue stp yld)
260− (\ev -> handleEvent ev stp yld) res
252+{-# NOINLINE pullWorker #-}
253+pullWorker :: (MonadIO m, MonadThrow m) => Context m a -> AsyncT m a
254+pullWorker ctx = AsyncT $ \_ stp yld -> do
255+ ev <- liftIO $ atomically $ readTBQueue (childChannel ctx)
256+ handleEvent ev stp yld
261257
262258 where
263259
264− {-# INLINE continue #-}
265− continue stp yld = (runAsyncT (pullDrain ctx)) Nothing stp yld
266−
267− {-# INLINE dispatch #-}
268− dispatch = liftIO $ signalQSem (dispatchReq ctx)
260+ continue stp yld = (runAsyncT (pullWorker ctx)) Nothing stp yld
269261
270262 {-# INLINE handleEvent #-}
271263 handleEvent ev stp yld = do
272− let yielder a = yld a Nothing (Just (pullDrain ctx))
264+ let yielder a = yld a Nothing (Just (pullWorker ctx))
273265 case ev of
274266 ChildYield a -> yielder a
275267 ChildDone tid a -> do
276− dispatch
277268 done <- delThread ctx tid
278269 if done then (yld a Nothing Nothing) else (yielder a)
279270 ChildStop tid e -> do
280271 case e of
281272 Nothing -> do
282− dispatch
283273 done <- delThread ctx tid
284274 if done then stp else continue stp yld
285275 Just x -> handleException x ctx tid
@@ -302,30 +292,30 @@ handleChildException pchan e = do
302292pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
303293pullFork m1 m2 = AsyncT $ \_ stp yld -> do
304294 ctx <- liftIO $ newContext
305− initialFork ctx m1 >> initialFork ctx m2
306− (runAsyncT (pullDrain ctx)) Nothing stp yld
295+ queueWork ctx m1 >> queueWork ctx m2 >> pushWorker ctx
296+ (runAsyncT (pullWorker ctx)) Nothing stp yld
307297
308298 where
309299
310− -- This function is different than "fork" because we have to directly
300+ -- This function is different than "forkWorker" because we have to directly
311301 -- insert the threadIds here and cannot use the channel to send ChildCreate
312302 -- unlike on the push side. If we do that, the first thread's done message
313303 -- may arrive even before the second thread is forked, in that case
314− -- pullDrain will falsely detect that all threads are over.
315− initialFork ctx m = do
304+ -- pullWorker will falsely detect that all threads are over.
305+ pushWorker ctx = do
316306 let chan = childChannel ctx
317− tid <- doFork (push ctx m) (handleChildException chan)
307+ tid <- doFork (push ctx) (handleChildException chan)
318308 liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)
319309
320310 newContext = do
321311 channel <- atomically $ newTBQueue 16
312+ work <- atomically $ newTBQueue 16
322313 running <- newIORef S.empty
323314 done <- newIORef S.empty
324− count <- newQSem 1
325315 return $ Context { childChannel = channel
326− , dispatchReq = count
327316 , runningThreads = running
328317 , doneThreads = done
318+ , pendingWork = work
329319 }
330320
331321-- Concurrency rate control. Our objective is to create more threads on
@@ -359,13 +349,39 @@ pullFork m1 m2 = AsyncT $ \_ stp yld -> do
359349-- residual work back to the dispatcher. It will also consume a lot of
360350-- memory due to queueing of all the work before execution starts.
361351
362−{-# INLINE fork #-}
363−fork :: MonadAsync m => Context a -> AsyncT m a -> m ()
364−fork ctx m = do
352+{-# INLINE forkWorker #-}
353+forkWorker :: MonadAsync m => Context m a -> m ()
354+forkWorker ctx = do
365355 let chan = childChannel ctx
366− tid <- doFork (push ctx m) (handleChildException chan)
356+ tid <- doFork (push ctx) (handleChildException chan)
367357 liftIO $ atomically $ writeTBQueue chan (ChildCreate tid)
368358
359+{-# INLINE queueWork #-}
360+queueWork :: MonadAsync m => Context m a -> AsyncT m a -> m ()
361+queueWork ctx m = do
362+ -- To guarantee deadlock avoidance we need to dispatch a worker when the
363+ -- workQueue goes full. Otherwise all the worker threads might be waiting
364+ -- on the queue and never wakeup.
365+ --
366+ -- TBD If we run out of threads we can also evaluate the action completely
367+ -- right here, disallowing any further child workers and turning the
368+ -- parallel composition into interleaved serial composition.
369+ workQFull <- liftIO $ atomically $ isFullTBQueue (pendingWork ctx)
370+ when (workQFull) $ forkWorker ctx
371+ liftIO $ atomically $ writeTBQueue (pendingWork ctx) m
372+
373+{-# INLINE dequeueLoop #-}
374+dequeueLoop :: MonadAsync m => Context m a -> AsyncT m a
375+dequeueLoop ctx = AsyncT $ \_ stp yld -> do
376+ work <- liftIO $ atomically $ tryReadTBQueue (pendingWork ctx)
377+ case work of
378+ Nothing -> stp
379+ Just m -> do
380+ let loop = (runAsyncT (dequeueLoop ctx)) Nothing stp yld
381+ yielder a c Nothing = yld a c (Just (dequeueLoop ctx))
382+ yielder a c r = yld a c r
383+ (runAsyncT m) (Just ctx) loop yielder
384+
369385instance MonadAsync m => Alternative (AsyncT m) where
370386 empty = mempty
371387
@@ -374,8 +390,13 @@ instance MonadAsync m => Alternative (AsyncT m) where
374390 {-# INLINE (<|>) #-}
375391 m1 <|> m2 = AsyncT $ \ctx stp yld -> do
376392 case ctx of
377− Nothing -> (runAsyncT (pullFork m1 m2)) ctx stp yld
378− Just c -> fork c m2 >> (runAsyncT m1) ctx stp yld
393+ Nothing -> (runAsyncT (pullFork m1 m2)) Nothing stp yld
394+ Just c -> do
395+ -- we open up both the branches fairly but on a given level we
396+ -- go left to right. If the left one keeps producing results we
397+ -- may or may not run the right one.
398+ queueWork c m1 >> queueWork c m2
399+ (runAsyncT (dequeueLoop c)) Nothing stp yld
379400
380401instance MonadAsync m => MonadPlus (AsyncT m) where
381402 mzero = empty