@@ -51,6 +51,8 @@ import Control.Monad.Reader.Class (MonadReader(..))
5151import Control.Monad.State.Class (MonadState(..))
5252import Control.Monad.Trans.Class (MonadTrans (lift))
5353import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
54+import Data.Concurrent.Queue.MichaelScott (LinkedQueue, newQ, pushL,
55+ tryPopR, nullQ)
5456import Data.Functor (void)
5557import Data.IORef (IORef, modifyIORef, newIORef,
5658 writeIORef, readIORef)
@@ -98,10 +100,11 @@ data ChildEvent a =
98100data Context m a =
99101 Context { outputQueue :: IORef [ChildEvent a]
100102 , 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
102106 , runningThreads :: IORef (Set ThreadId)
103107 , doneThreads :: IORef (Set ThreadId)
104− , fairSched :: Bool
105108 }
106109
107110-- The 'Maybe (AsyncT m a)' is redundant as we can use 'stop' value for the
@@ -222,43 +225,73 @@ send ctx msg = liftIO $ do
222225 void $ tryPutMVar (synchOutQ ctx) True
223226
224227-- Note: Left associated operations can grow this queue to a large size
225−{-# INLINE queueWork #-}
226−queueWork :: MonadAsync m => IORef [AsyncT m a] -> AsyncT m a -> m ()
227−queueWork q m = liftIO $ atomicModifyIORefCAS_ q $ \ ms -> m : ms
228−
229−-- Note: This is performance sensitive code.
230−{-# INLINE dequeueLoop #-}
231−dequeueLoop :: MonadAsync m
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 ->
228+{-# INLINE enqueueLIFO #-}
229+enqueueLIFO :: IORef [AsyncT m a] -> AsyncT m a -> IO ()
230+enqueueLIFO q m = atomicModifyIORefCAS_ q $ \ ms -> m : ms
231+
232+{-# INLINE dequeueLIFO #-}
233+dequeueLIFO :: IORef [AsyncT m a] -> IO (Maybe (AsyncT m a))
234+dequeueLIFO q =
235+ atomicModifyIORefCAS q $ \ ms ->
235236 case ms of
236237 [] -> ([], Nothing)
237238 x : xs -> (xs, Just x)
239+
240+{-# INLINE isEmptyLIFO #-}
241+isEmptyLIFO :: IORef [AsyncT m a] -> IO Bool
242+isEmptyLIFO q = do
243+ work <- readIORef q
244+ return $ case work of
245+ [] -> True
246+ _ -> False
247+
248+-- Note: This is performance sensitive code.
249+{-# INLINE dequeueLoopLIFO #-}
250+dequeueLoopLIFO :: MonadAsync m
251+ => Maybe (Context m a)
252+ -> IORef [AsyncT m a]
253+ -> AsyncT m a
254+dequeueLoopLIFO ctx q = AsyncT $ \_ stp yld -> do
255+ work <- liftIO $ dequeueLIFO q
238256 case work of
239257 Nothing -> stp
240258 Just m -> do
241− let stop = (runAsyncT (dequeueLoop ctx q fair)) Nothing stp yld
242− continue a c = yld a c (Just (dequeueLoop ctx q fair))
243− yield a c Nothing = continue a c
244− yield a c r | not fair = yld a c r
245− yield a c (Just r) = queueWork q r >> continue a c
259+ let continue = dequeueLoopLIFO ctx q
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
264+
265+{-# INLINE dequeueLoopFIFO #-}
266+dequeueLoopFIFO :: MonadAsync m
267+ => Maybe (Context m a)
268+ -> LinkedQueue (AsyncT m a)
269+ -> AsyncT m a
270+dequeueLoopFIFO ctx q = AsyncT $ \_ stp yld -> do
271+ work <- liftIO $ tryPopR q
272+ case work of
273+ Nothing -> stp
274+ Just m -> do
275+ let continue = dequeueLoopFIFO ctx q
276+ stop = (runAsyncT continue) Nothing stp yld
277+ yield a c Nothing = yld a c (Just continue)
278+ yield a c (Just r) = do
279+ liftIO $ pushL q r
280+ yld a c (Just continue)
246281 (runAsyncT m) ctx stop yield
247282
248283{-# NOINLINE push #-}
249284push :: MonadAsync m => Context m a -> m ()
250−push ctx = run (Just ctx) (dequeueLoop (Just ctx) (workQueue ctx)
251− (fairSched ctx))
285+push ctx = run (Just ctx) (dequeueLoop ctx)
252286
253287 where
254288
255289 stop = do
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
290+ done <- liftIO $ hasNoWork ctx
291+ if done then do
292+ tid <- liftIO $ myThreadId
293+ send ctx (ChildStop tid Nothing)
294+ else push ctx
262295 yield a _ Nothing = do
263296 tid <- liftIO myThreadId
264297 send ctx (ChildDone tid a)
@@ -376,25 +409,43 @@ pullWorker ctx = AsyncT $ \pctx stp yld -> do
376409pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a
377410pullFork m1 m2 fair = AsyncT $ \_ stp yld -> do
378411 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
382413
383414 where
384415
385416 newContext = do
386417 outQ <- newIORef []
387418 outQMv <- newEmptyMVar
388− workQ <- newIORef []
389419 running <- newIORef S.empty
390420 done <- newIORef S.empty
391− return $ Context { outputQueue = outQ
392− , synchOutQ = outQMv
393− , runningThreads = running
394− , doneThreads = done
395− , workQueue = workQ
396− , fairSched = fair
397− }
421+
422+ case fair of
423+ True -> do
424+ q <- newQ
425+ pushL q m1 >> pushL q m2
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
434+ }
435+ in return ctx
436+ False -> do
437+ q <- newIORef []
438+ enqueueLIFO q m2 >> enqueueLIFO q m1
439+ let ctx =
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
447+ }
448+ in return ctx
398449
399450-- Concurrency rate control. Our objective is to create more threads on
400451-- 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
435486 case ctx of
436487 Nothing -> (runAsyncT (pullFork m1 m2 fair)) Nothing stp yld
437488 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.
441− let q = workQueue c
442− continue = (runAsyncT (dequeueLoop ctx q fair)) Nothing stp yld
443− queueWork q m2
444− (runAsyncT m1) ctx continue yld
489+ let dqloop = (runAsyncT (dequeueLoop c)) Nothing stp yld
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
493+ (runAsyncT m1) ctx dqloop yld
445494
446495instance MonadAsync m => Alternative (AsyncT m) where
447496 empty = mempty