@@ -25,9 +25,11 @@ module Asyncly.AsyncT
2525where
2626
2727import Control.Applicative (Alternative (..))
28−import Control.Concurrent (ThreadId, forkIO, myThreadId)
28+import Control.Concurrent (ThreadId, forkIO,
29+ myThreadId, threadDelay)
2930import Control.Concurrent.STM (TChan, atomically, newTChan,
30− readTChan, writeTChan)
31+ readTChan, tryReadTChan,
32+ writeTChan)
3133import Control.Exception (fromException, try,
3234 SomeException (..),
3335 BlockedIndefinitelyOnSTM(..))
@@ -44,6 +46,8 @@ import Control.Monad.Trans.Control (ComposeSt, MonadBaseControl (..),
4446 defaultLiftBaseWith,
4547 defaultRestoreM, liftBaseWith)
4648 -}
49+import Data.IORef (IORef, writeIORef, newIORef,
50+ readIORef)
4751import Data.Maybe (maybe)
4852import Data.Monoid ((<>))
4953
@@ -78,7 +82,11 @@ data ChildEvent a =
7882-- State threaded around the monad for thread management
7983------------------------------------------------------------------------------
8084
81−newtype Context a = Context { childChannel :: TChan (ChildEvent a) }
85+data Context m a =
86+ Context { childChannel :: TChan (ChildEvent a)
87+ , pullSide :: Bool
88+ , pendingWork :: IORef (AsyncT m a)
89+ }
8290
8391-- The 'Maybe (AsyncT m a)' is redundant as we can use 'stop' value for the
8492-- Nothing case, but it makes the fold using '<|>' 25% faster. Need to try
@@ -91,9 +99,9 @@ newtype Context a = Context { childChannel :: TChan (ChildEvent a) }
9199newtype AsyncT m a =
92100 AsyncT {
93101 runAsyncT :: forall r.
94− Maybe (Context a) -- state
102+ Maybe (Context m a) -- state
95103 -> m r -- stop
96− -> (a -> Maybe (Context a) -> Maybe (AsyncT m a) -> m r) -- yield
104+ -> (a -> Maybe (Context m a) -> Maybe (AsyncT m a) -> m r) -- yield
97105 -> m r
98106 }
99107
@@ -170,7 +178,7 @@ doFork action preExit =
170178 return ()
171179
172180{-# NOINLINE push #-}
173−push :: MonadIO m => Context a -> AsyncT m a -> m ()
181+push :: MonadIO m => Context m a -> AsyncT m a -> m ()
174182push context action = run (Just context) action
175183
176184 where
@@ -203,52 +211,107 @@ handlePushException pchan e = do
203211
204212-- | run m1 in a new thread, pushing its results to a pull channel and then run
205213-- m2 in the parent thread. Any exceptions are also pushed to the channel.
206−{-# INLINE pushFork #-}
207−pushFork :: (MonadIO m, MonadBaseControl IO m)
208− => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a
209−pushFork c m1 m2 = AsyncT $ \ctx stp yld -> do
210− _ <- doFork (push c m1) (handlePushException (childChannel c))
211− -- XXX run m2 only when the channel becomes empty
212− (runAsyncT m2) ctx stp yld
213−
214+{-# INLINE pullSideDispatch #-}
215+pullSideDispatch :: MonadAsync m
216+ => Context m a -> AsyncT m a -> AsyncT m a -> AsyncT m a
217+pullSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
218+ let chan = childChannel ctx
219+ _ <- doFork (push (ctx {pullSide = False}) m1) (handlePushException chan)
220+ liftIO $ threadDelay 0
221+ liftIO $ writeIORef (pendingWork ctx) m2
222+ (runAsyncT (pullDispatch ctx False)) Nothing stp yld
223+
224+{-# NOINLINE pushSideDispatch #-}
225+pushSideDispatch :: MonadAsync m
226+ => Context m a -> AsyncT m a -> AsyncT m a -> AsyncT m a
227+pushSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do
228+ let chan = childChannel ctx
229+ _ <- doFork (push ctx m1) (handlePushException chan)
230+ (runAsyncT m2) (Just ctx) stp yld
231+
232+-- XXX the TBQueue size should be proportional to the pending threads.
214233-- We re-raise any exceptions received from the child threads, that way
215234-- exceptions get propagated to the top level computation and can be handled
216235-- there.
217−{-# NOINLINE pull #-}
218−pull :: (MonadIO m, MonadThrow m) => Context a -> AsyncT m a
219−pull ctx = AsyncT $ \_ stp yld -> do
220− res <- liftIO $ try (atomically $ readTChan (childChannel ctx))
221− case res of
222− Left e ->
223− case fromException e of
224− Just BlockedIndefinitelyOnSTM -> stp
225− Nothing -> throwM e
226− Right ev ->
227− let yielder a = yld a Nothing (Just (pull ctx))
228− in case ev of
229− ChildYield a -> yielder a
230− ChildDone _ a -> yielder a
231− ChildStop _ e -> do
232− let continue = (runAsyncT (pull ctx)) (Just ctx) stp yld
233− in maybe continue throwM e
236+{-# NOINLINE pullDispatch #-}
237+pullDispatch :: (MonadIO m, MonadThrow m)
238+ => Context m a -> Bool -> AsyncT m a
239+pullDispatch ctx dispatch = AsyncT $ \_ stp yld -> do
240+ let chan = childChannel ctx
241+ eHandler e = handleException e stp
242+ evHandler ev = handleEvent ev stp yld
243+ maybeEvHandler ev = maybe (continue stp yld) evHandler ev
244+
245+ if dispatch then do
246+ m <- liftIO $ readIORef (pendingWork ctx)
247+ (runAsyncT m) (Just ctx) stp yld
248+ else do
249+ res <- liftIO $ try $ atomically $ tryReadTChan chan
250+ either eHandler maybeEvHandler res
251+
252+ where
253+
254+ handleException e stp =
255+ case fromException e of
256+ Just BlockedIndefinitelyOnSTM -> stp
257+ Nothing -> throwM e
258+
259+ {-# INLINE continue #-}
260+ continue stp yld = (runAsyncT (pullDispatch ctx True)) (Just ctx) stp yld
261+
262+ {-# INLINE handleEvent #-}
263+ handleEvent ev stp yld =
264+ case ev of
265+ ChildYield a -> yld a Nothing (Just (pullDispatch ctx False))
266+ ChildDone _ a -> yld a Nothing (Just (pullDispatch ctx True))
267+ -- XXX kill threads
268+ ChildStop _ e -> maybe (continue stp yld) throwM e
269+
270+{-# NOINLINE pullDrain #-}
271+pullDrain :: (MonadIO m, MonadThrow m) => TChan (ChildEvent a) -> AsyncT m a
272+pullDrain chan = AsyncT $ \_ stp yld -> do
273+ let eHandler e = handleException e stp
274+ evHandler ev = handleEvent ev stp yld
275+
276+ res <- liftIO $ try (atomically $ readTChan chan)
277+ either eHandler evHandler res
278+
279+ where
280+
281+ handleException e stp =
282+ case fromException e of
283+ Just BlockedIndefinitelyOnSTM -> stp
284+ Nothing -> throwM e
285+
286+ {-# INLINE continue #-}
287+ continue stp yld = (runAsyncT (pullDrain chan)) Nothing stp yld
288+
289+ {-# INLINE handleEvent #-}
290+ handleEvent ev stp yld =
291+ let yielder a = yld a Nothing (Just (pullDrain chan))
292+ in case ev of
293+ ChildYield a -> yielder a
294+ ChildDone _ a -> yielder a
295+ -- XXX kill threads
296+ ChildStop _ e -> maybe (continue stp yld) throwM e
234297
235298-- | Split the original computation in a pull-push pair. The original
236299-- computation pulls from a Channel while m1 and m2 push to the channel.
237300{-# NOINLINE pullFork #-}
238301pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
239302pullFork m1 m2 = AsyncT $ \_ stp yld -> do
240− c <- liftIO $ newContext
241− _ <- doFork (push c (m1 <|> m2)) (handlePushException (childChannel c))
242− (runAsyncT (pull c)) (Just c) stp yld
303+ ctx <- liftIO $ newContext (m1 <|> m2)
304+ (runAsyncT (pullDispatch ctx True)) Nothing stp yld
243305
244306 where
245307
246− newContext = do
247− childChan <- atomically newTChan
248− return $ Context { childChannel = childChan }
249−
250−canFork :: IO Bool
251−canFork = return False
308+ newContext m = do
309+ channel <- atomically newTChan
310+ work <- liftIO $ newIORef (m <|> pullDrain channel)
311+ return $ Context { childChannel = channel
312+ , pullSide = True
313+ , pendingWork = work
314+ }
252315
253316-- Concurrency rate control. Our objective is to create more threads on
254317-- demand if the consumer is running faster than us. As soon as we
@@ -260,25 +323,42 @@ canFork = return False
260323-- problematic if the initial production latency is high, we may end up
261324-- creating too many threads. So we need some way to monitor and use the
262325-- latency as well.
326+--
263327-- TBD For quick response we may have to increase the rate in the middle of
264328-- a serially running computation. For that we can use a state flag to fork
265329-- the rest of the computation at any point of time inside the Monad bind
266330-- operation if the consumer is running at a faster speed.
267331--
332+-- TBD the alternative composition allows us to dispatch a chunkSize of only 1.
333+-- If we have to dispatch in arbitrary chunksizes we will need to compose the
334+-- parallel actions using a data constructor instead so that we can divide it
335+-- in chunks of arbitrary size before dispatch. When batching we can convert
336+-- the structure into Alternative batches of Monoid composition. That will also
337+-- allow us to dispatch more work to existing threads rather than creating new
338+-- threads always.
339+--
340+-- TBD for pure work (when we are not in the IO monad) we can divide it into
341+-- just the number of CPUs.
342+--
343+-- XXX to rate control left folded structrues we will have to return the
344+-- residual work back to the dispatcher. It will also consume a lot of
345+-- memory due to queueing of all the work before execution starts.
346+--
268347instance MonadAsync m => Alternative (AsyncT m) where
269348 empty = mempty
270349
271− -- XXX pullFork can have more than 2x penalty for smallish tasks. Can we
272− -- avoid this automatically if the task is small?
350+ -- Note: This is designed to scale for right associated compositions,
351+ -- therefore always use a right fold for folding bigger structures.
273352 {-# INLINE (<|>) #-}
274353 m1 <|> m2 = AsyncT $ \ctx stp yld -> do
275354 case ctx of
276355 Nothing -> (runAsyncT (pullFork m1 m2)) ctx stp yld
277356 Just c -> do
278− doAsync <- liftIO $ canFork
279− if doAsync then
280− (runAsyncT (pushFork c m1 m2)) ctx stp yld
281− else (runAsyncT (m1 <> m2)) ctx stp yld
357+ if pullSide c then
358+ (runAsyncT (pullSideDispatch c m1 m2)) ctx stp yld
359+ else
360+ -- for left associated compositions
361+ (runAsyncT (pushSideDispatch c m1 m2)) ctx stp yld
282362
283363instance MonadAsync m => MonadPlus (AsyncT m) where
284364 mzero = empty