@@ -51,6 +51,8 @@ import Control.Monad.Reader.Class (MonadReader(..))
51import Control.Monad.State.Class (MonadState(..))
51import Control.Monad.State.Class (MonadState(..))
52import Control.Monad.Trans.Class (MonadTrans (lift))
52import Control.Monad.Trans.Class (MonadTrans (lift))
53import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
53import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
54+import Data.Concurrent.Queue.MichaelScott (LinkedQueue, newQ, pushL,
54import Data.Functor (void)
56import Data.Functor (void)
55import Data.IORef (IORef, modifyIORef, newIORef,
57import Data.IORef (IORef, modifyIORef, newIORef,
56 writeIORef, readIORef)
58 writeIORef, readIORef)
@@ -98,10 +100,11 @@ data ChildEvent a =
98data Context m a =
100data Context m a =
99 Context { outputQueue :: IORef [ChildEvent a]
101 Context { outputQueue :: IORef [ChildEvent a]
100 , synchOutQ :: MVar Bool -- wakeup mechanism for outQ
102 , synchOutQ :: MVar Bool -- wakeup mechanism for outQ
101− , workQueue :: IORef [AsyncT m a]
103+ , dequeueLoop :: AsyncT m a
104+ , enqueueWork :: AsyncT m a -> IO ()
105+ , hasNoWork :: IO Bool
102 , runningThreads :: IORef (Set ThreadId)
106 , runningThreads :: IORef (Set ThreadId)
103 , doneThreads :: IORef (Set ThreadId)
107 , doneThreads :: IORef (Set ThreadId)
107-- The 'Maybe (AsyncT m a)' is redundant as we can use 'stop' value for the
110-- The 'Maybe (AsyncT m a)' is redundant as we can use 'stop' value for the
@@ -222,43 +225,73 @@ send ctx msg = liftIO $ do
222 void $ tryPutMVar (synchOutQ ctx) True
225 void $ tryPutMVar (synchOutQ ctx) True
224-- Note: Left associated operations can grow this queue to a large size
227-- Note: Left associated operations can grow this queue to a large size
225−{-# INLINE queueWork #-}
228+{-# INLINE enqueueLIFO #-}
226−queueWork :: MonadAsync m => IORef [AsyncT m a] -> AsyncT m a -> m ()
229+enqueueLIFO :: IORef [AsyncT m a] -> AsyncT m a -> IO ()
227−queueWork q m = liftIO $ atomicModifyIORefCAS_ q $ \ ms -> m : ms
230+enqueueLIFO q m = atomicModifyIORefCAS_ q $ \ ms -> m : ms
229−-- Note: This is performance sensitive code.
232+{-# INLINE dequeueLIFO #-}
230−{-# INLINE dequeueLoop #-}
233+dequeueLIFO :: IORef [AsyncT m a] -> IO (Maybe (AsyncT m a))
231−dequeueLoop :: MonadAsync m
234+dequeueLIFO q =
232− => Maybe (Context m a) -> IORef [AsyncT m a] -> Bool -> AsyncT m a
233−dequeueLoop ctx q fair = AsyncT $ \_ stp yld -> do
234− work <- liftIO $ atomicModifyIORefCAS q $ \ ms ->
235+ atomicModifyIORefCAS q $ \ ms ->
235 case ms of
236 case ms of
236 [] -> ([], Nothing)
237 [] -> ([], Nothing)
237 x : xs -> (xs, Just x)
238 x : xs -> (xs, Just x)
240+{-# INLINE isEmptyLIFO #-}
241+isEmptyLIFO :: IORef [AsyncT m a] -> IO Bool
244+ return $ case work of
248+-- Note: This is performance sensitive code.
249+{-# INLINE dequeueLoopLIFO #-}
250+dequeueLoopLIFO :: MonadAsync m
251+ => Maybe (Context m a)
252+ -> IORef [AsyncT m a]
254+dequeueLoopLIFO ctx q = AsyncT $ \_ stp yld -> do
255+ work <- liftIO $ dequeueLIFO q
238 case work of
256 case work of
239 Nothing -> stp
257 Nothing -> stp
240 Just m -> do
258 Just m -> do
241− let stop = (runAsyncT (dequeueLoop ctx q fair)) Nothing stp yld
259+ let continue = dequeueLoopLIFO ctx q
242− continue a c = yld a c (Just (dequeueLoop ctx q fair))
260+ stop = (runAsyncT continue) Nothing stp yld
261+ yield a c Nothing = yld a c (Just continue)
262+ yield a c r = yld a c r
263+ (runAsyncT m) ctx stop yield
265+{-# INLINE dequeueLoopFIFO #-}
266+dequeueLoopFIFO :: MonadAsync m
267+ => Maybe (Context m a)
268+ -> LinkedQueue (AsyncT m a)
270+dequeueLoopFIFO ctx q = AsyncT $ \_ stp yld -> do
271+ work <- liftIO $ tryPopR q
243− yield a c Nothing = continue a c
273+ Nothing -> stp
275+ let continue = dequeueLoopFIFO ctx q
276+ stop = (runAsyncT continue) Nothing stp yld
244− yield a c r | not fair = yld a c r
277+ yield a c Nothing = yld a c (Just continue)
245− yield a c (Just r) = queueWork q r >> continue a c
278+ yield a c (Just r) = do
280+ yld a c (Just continue)
246 (runAsyncT m) ctx stop yield
281 (runAsyncT m) ctx stop yield
248{-# NOINLINE push #-}
283{-# NOINLINE push #-}
249push :: MonadAsync m => Context m a -> m ()
284push :: MonadAsync m => Context m a -> m ()
250−push ctx = run (Just ctx) (dequeueLoop (Just ctx) (workQueue ctx)
285+push ctx = run (Just ctx) (dequeueLoop ctx)
255 stop = do
289 stop = do
256− work <- liftIO $ readIORef (workQueue ctx)
290+ done <- liftIO $ hasNoWork ctx
257− case work of
291+ if done then do
259− tid <- liftIO $ myThreadId
292+ tid <- liftIO $ myThreadId
260− send ctx (ChildStop tid Nothing)
293+ send ctx (ChildStop tid Nothing)
261− _ -> push ctx
294+ else push ctx
262 yield a _ Nothing = do
295 yield a _ Nothing = do
263 tid <- liftIO myThreadId
296 tid <- liftIO myThreadId
264 send ctx (ChildDone tid a)
297 send ctx (ChildDone tid a)
@@ -376,25 +409,43 @@ pullWorker ctx = AsyncT $ \pctx stp yld -> do
376pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a
409pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a
377pullFork m1 m2 fair = AsyncT $ \_ stp yld -> do
410pullFork m1 m2 fair = AsyncT $ \_ stp yld -> do
378 ctx <- liftIO $ newContext
411 ctx <- liftIO $ newContext
379− let q = workQueue ctx
380− queueWork q m1 >> queueWork q m2 >> pushWorker ctx
381− (runAsyncT (pullWorker ctx)) Nothing stp yld
412+ pushWorker ctx >> (runAsyncT (pullWorker ctx)) Nothing stp yld
385 newContext = do
416 newContext = do
386 outQ <- newIORef []
417 outQ <- newIORef []
387 outQMv <- newEmptyMVar
418 outQMv <- newEmptyMVar
388− workQ <- newIORef []
389 running <- newIORef S.empty
419 running <- newIORef S.empty
390 done <- newIORef S.empty
420 done <- newIORef S.empty
391− return $ Context { outputQueue = outQ
421+
392− , synchOutQ = outQMv
422+ case fair of
393− , runningThreads = running
423+ True -> do
394− , doneThreads = done
424+ q <- newQ
395− , workQueue = workQ
425+ pushL q m1 >> pushL q m2
396− , fairSched = fair
426+ let ctx =
427+ Context { outputQueue = outQ
428+ , synchOutQ = outQMv
429+ , runningThreads = running
430+ , doneThreads = done
431+ , dequeueLoop = dequeueLoopFIFO (Just ctx) q
432+ , enqueueWork = pushL q
433+ , hasNoWork = nullQ q
438+ enqueueLIFO q m2 >> enqueueLIFO q m1
440+ Context { outputQueue = outQ
441+ , synchOutQ = outQMv
442+ , runningThreads = running
443+ , doneThreads = done
444+ , dequeueLoop = dequeueLoopLIFO (Just ctx) q
445+ , enqueueWork = enqueueLIFO q
446+ , hasNoWork = isEmptyLIFO q
399-- Concurrency rate control. Our objective is to create more threads on
450-- Concurrency rate control. Our objective is to create more threads on
400-- demand if the consumer is running faster than us. As soon as we
451-- demand if the consumer is running faster than us. As soon as we
@@ -435,13 +486,11 @@ parallel m1 m2 fair = AsyncT $ \ctx stp yld -> do
435 case ctx of
486 case ctx of
436 Nothing -> (runAsyncT (pullFork m1 m2 fair)) Nothing stp yld
487 Nothing -> (runAsyncT (pullFork m1 m2 fair)) Nothing stp yld
437 Just c -> do
488 Just c -> do
438− -- we open up both the branches fairly but on a given level we
439− -- go left to right. If the left one keeps producing results we
440− -- may or may not run the right one.
442− continue = (runAsyncT (dequeueLoop ctx q fair)) Nothing stp yld
489+ let dqloop = (runAsyncT (dequeueLoop c)) Nothing stp yld
443− queueWork q m2
490+ -- XXX We can statically decide which function to call here, that
491+ -- will inline the call and improve perf a little bit.
492+ liftIO $ (enqueueWork c) m2
444− (runAsyncT m1) ctx continue yld
493+ (runAsyncT m1) ctx dqloop yld
446instance MonadAsync m => Alternative (AsyncT m) where
495instance MonadAsync m => Alternative (AsyncT m) where
447 empty = mempty
496 empty = mempty