git.iko.soy/streamly
on 57e6232 help sign in
streamly/src/Asyncly/AsyncT.hs
521 lines 22 kB history raw
1{-# LANGUAGE ConstraintKinds #-}
2{-# LANGUAGE EmptyCase #-}
3{-# LANGUAGE FlexibleContexts #-}
4{-# LANGUAGE FlexibleInstances #-}
5{-# LANGUAGE MultiParamTypeClasses #-}
6{-# LANGUAGE RankNTypes #-}
7{-# LANGUAGE UndecidableInstances #-} -- XXX
9-- |
10-- Module : Asyncly.AsyncT
11-- Copyright : (c) 2017 Harendra Kumar
12--
13-- License : MIT-style
14-- Maintainer : harendra.kumar@gmail.com
15-- Stability : experimental
16-- Portability : GHC
17--
18--
19module Asyncly.AsyncT
20 ( AsyncT (..)
21 , MonadAsync
22 , runAsyncly
23 , toList
24 , each
25 , for
26-- , async
27-- , makeAsync
28 )
29where
31import Control.Applicative (Alternative (..))
32import Control.Concurrent (ThreadId, forkIO, killThread,
33 myThreadId, threadDelay)
34import Control.Concurrent.STM (TBQueue, atomically, newTBQueue,
35 tryReadTBQueue, writeTBQueue,
36 isEmptyTBQueue, isFullTBQueue,
37 peekTBQueue)
38import Control.Exception (SomeException (..))
39import qualified Control.Exception.Lifted as EL
40import Control.Monad (ap, liftM, MonadPlus(..), mzero,
41 when)
42import Control.Monad.Base (MonadBase (..), liftBaseDefault)
43import Control.Monad.Catch (MonadThrow, throwM)
44import Control.Monad.IO.Class (MonadIO(..))
45import Control.Monad.Trans.Class (MonadTrans (lift))
46import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
47import Data.Functor (void)
48import Data.IORef (IORef, modifyIORef, newIORef,
49 writeIORef, readIORef)
50import Data.Monoid ((<>))
51import Data.Set (Set)
52import qualified Data.Set as S
54import Control.Monad.Trans.Recorder (MonadRecorder(..))
57------------------------------------------------------------------------------
58-- Concurrency Semantics
59------------------------------------------------------------------------------
60--
61-- Asyncly is essentially a concurrent list transformer. To concatenate the
62-- lists it provides three different ways of composing them. The Monoid
63-- instance concatenates lists in a non-concurrent in-order fashion. The
64-- Alternative instance concatenates list in a concurrent manner potentially
65-- running each action in parallel and therefore the order of the resulting
66-- items is not deterministic. Computations in Alternative composition may or
67-- may not run in parallel. Thirdly, the <||> operator provides a way to
68-- compose computations that are guaranteed to run in parallel. This provides a
69-- way to run infinite loops or long running computations without blocking
70-- or starving other computations.
72------------------------------------------------------------------------------
73-- Parent child thread communication types
74------------------------------------------------------------------------------
76data ChildEvent a =
77 ChildYield a
78 | ChildDone ThreadId a
79 | ChildStop ThreadId (Maybe SomeException)
80 | ChildCreate ThreadId
82------------------------------------------------------------------------------
83-- State threaded around the monad for thread management
84------------------------------------------------------------------------------
86data Context m a =
87 Context { outputQueue :: TBQueue (ChildEvent a)
88 , workQueue :: TBQueue (AsyncT m a)
89 , runningThreads :: IORef (Set ThreadId)
90 , doneThreads :: IORef (Set ThreadId)
91 }
93-- The 'Maybe (AsyncT m a)' is redundant as we can use 'stop' value for the
94-- Nothing case, but it makes the fold using '<|>' 25% faster. Need to try
95-- again as the code has gone through many changes after this was tested.
96-- With the Maybe, 'stop' is required only to represent 'empty' in an
97-- Alternative composition.
98--
99-- Currently the only state we need is the thread context, For generality we
100-- can parameterize the type with a state type 's'.
101newtype AsyncT m a =
102 AsyncT {
103 runAsyncT :: forall r.
104 Maybe (Context m a) -- state
105 -> m r -- stop
106 -> (a -> Maybe (Context m a) -> Maybe (AsyncT m a) -> m r) -- yield
107 -> m r
108 }
110type MonadAsync m = (MonadIO m, MonadBaseControl IO m, MonadThrow m)
112------------------------------------------------------------------------------
113-- Monad
114------------------------------------------------------------------------------
116-- | Appends the results of two AsyncT computations in order.
117instance Monad m => Monoid (AsyncT m a) where
118 mempty = AsyncT $ \_ stp _ -> stp
119 mappend (AsyncT m1) m2 = AsyncT $ \ctx stp yld ->
120 let stop = (runAsyncT m2) ctx stp yld
121 yield a c Nothing = yld a c (Just m2)
122 yield a c (Just r) = yld a c (Just (mappend r m2))
123 in m1 ctx stop yield
125-- We do not use bind for parallelism. That is, we do not start each iteration
126-- of the list in parallel. That will introduce too much uncontrolled
127-- parallelism. Instead we control parallelism using <|>, it is used for
128-- actions that 'may' be run in parallel. The library decides whether they will
129-- be started in parallel. This puts the parallel semantics in the hands of the
130-- user and operational semantics in the hands of the library. Ideally we would
131-- want parallelism to be completely hidden from the user, for example we can
132-- automatically decide where to start an action in parallel so that it is most
133-- efficient. We can decide whether a particular bind should fork a parallel
134-- thread or not based on the level of the bind in the tree. Maybe at some
135-- point we shall be able to do that?
137instance Monad m => Monad (AsyncT m) where
138 return a = AsyncT $ \ctx _ yld -> yld a ctx Nothing
140 -- | A thread context is valid only until the next bind. Upon a bind we
141 -- reset the context to Nothing.
142 AsyncT m >>= f = AsyncT $ \_ stp yld ->
143 let run x = (runAsyncT x) Nothing stp yld
144 yield a _ Nothing = run $ f a
145 yield a _ (Just r) = run $ f a <> (r >>= f)
146 in m Nothing stp yield
148------------------------------------------------------------------------------
149-- Functor
150------------------------------------------------------------------------------
152instance Monad m => Functor (AsyncT m) where
153 fmap = liftM
155------------------------------------------------------------------------------
156-- Applicative
157------------------------------------------------------------------------------
159-- Note: We cannot have the applicative instance run the two actions in
160-- parallel. This is because the first action can always return empty and in
161-- case it returns empty we should not be starting the second action. Therefore
162-- we can only start the second action after getting a result from the first.
163-- For running applicative actions in parallel we can use the async package and
164-- lift the result into AsyncT.
166instance Monad m => Applicative (AsyncT m) where
167 pure = return
168 (<*>) = ap
170------------------------------------------------------------------------------
171-- Alternative
172------------------------------------------------------------------------------
174{-# INLINE doFork #-}
175doFork :: (MonadIO m, MonadBaseControl IO m)
176 => m ()
177 -> (SomeException -> IO ())
178 -> m ThreadId
179doFork action exHandler =
180 EL.mask $ \restore ->
181 liftBaseWith $ \runInIO -> forkIO $ do
182 -- XXX test the exception handling
183 _ <- runInIO $ EL.catch (restore action) (liftIO . exHandler)
184 -- XXX restore state here?
185 return ()
187{-# NOINLINE push #-}
188push :: MonadAsync m => Context m a -> m ()
189push ctx = run (Just ctx) (dequeueLoop ctx)
191 where
193 send msg = atomically $ writeTBQueue (outputQueue ctx) msg
194 stop = do
195 workQEmpty <- liftIO $ atomically $ isEmptyTBQueue (workQueue ctx)
196 if (not workQEmpty) then push ctx
197 else liftIO $ myThreadId >>= \tid -> send (ChildStop tid Nothing)
198 yield a _ Nothing = liftIO $ myThreadId >>= \tid -> send (ChildDone tid a)
199 yield a c (Just r) = liftIO (send (ChildYield a)) >> run c r
200 run c m = (runAsyncT m) c stop yield
202-- Thread tracking has a significant performance overhead (~20% on empty
203-- threads, it will be lower for heavy threads). It is needed for two reasons:
204--
205-- 1) Killing threads on exceptions. Threads may not be allowed to go away by
206-- themselves because they may run for significant times before going away or
207-- worse they may be stuck in IO and never go away.
208--
209-- 2) To know when all threads are done. This can be acheived by detecting a
210-- BlockedIndefinitelyOnSTM exception too. But we will have to trigger a GC to
211-- make sure that we detect it promptly.
212--
213-- This is a bit messy because ChildCreate and ChildDone events can arrive out
214-- of order in case of pushSideDispatch. Returns whether we are done draining
215-- threads.
216{-# INLINE accountThread #-}
217accountThread :: MonadIO m
218 => ThreadId -> IORef (Set ThreadId) -> IORef (Set ThreadId) -> m Bool
219accountThread tid ref1 ref2 = liftIO $ do
220 s1 <- readIORef ref1
221 s2 <- readIORef ref2
223 if (S.member tid s1) then do
224 let r = S.delete tid s1
225 writeIORef ref1 r
226 return $ S.null r && S.null s2
227 else do
228 liftIO $ writeIORef ref2 (S.insert tid s2)
229 return False
231{-# NOINLINE addThread #-}
232addThread :: MonadIO m => Context m a -> ThreadId -> m Bool
233addThread ctx tid = accountThread tid (doneThreads ctx) (runningThreads ctx)
235{-# INLINE delThread #-}
236delThread :: MonadIO m => Context m a -> ThreadId -> m Bool
237delThread ctx tid = accountThread tid (runningThreads ctx) (doneThreads ctx)
239handleException :: (MonadIO m, MonadThrow m)
240 => SomeException -> Context m a -> ThreadId -> m r
241handleException e ctx tid = do
242 void (delThread ctx tid)
243 liftIO $ readIORef (runningThreads ctx) >>= mapM_ killThread
244 throwM e
246-- We re-raise any exceptions received from the child threads, that way
247-- exceptions get propagated to the top level computation and can be handled
248-- there.
249{-# NOINLINE pullWorker #-}
250pullWorker :: MonadAsync m => Context m a -> AsyncT m a
251pullWorker ctx = AsyncT $ \pctx stp yld -> do
252 let continue = (runAsyncT (pullWorker ctx)) pctx stp yld
253 yield a = yld a pctx (Just (pullWorker ctx))
254 threadOp tid f finish cont = do
255 done <- f ctx tid
256 if done then finish else cont
258 res <- liftIO $ atomically $ tryReadTBQueue (outputQueue ctx)
259 case res of
260 Nothing -> do
261 liftIO $ threadDelay 4
262 let workQ = workQueue ctx
263 outQ = outputQueue ctx
264 workQEmpty <- liftIO $ atomically $ isEmptyTBQueue workQ
265 outQEmpty <- liftIO $ atomically $ isEmptyTBQueue outQ
266 when (not workQEmpty && outQEmpty) $ pushWorker ctx
267 void $ liftIO $ atomically $ peekTBQueue (outputQueue ctx)
268 continue
269 Just ev ->
270 case ev of
271 ChildYield a -> yield a
272 ChildDone tid a ->
273 threadOp tid delThread (yld a pctx Nothing) (yield a)
274 ChildStop tid e ->
275 case e of
276 Nothing -> threadOp tid delThread stp continue
277 Just ex -> handleException ex ctx tid
278 ChildCreate tid -> threadOp tid addThread stp continue
280-- If an exception occurs we push it to the channel so that it can handled by
281-- the parent. 'Paused' exceptions are to be collected at the top level.
282-- XXX Paused exceptions should only bubble up to the runRecorder computation
283{-# NOINLINE handleChildException #-}
284handleChildException :: TBQueue (ChildEvent a) -> SomeException -> IO ()
285handleChildException pchan e = do
286 tid <- myThreadId
287 atomically $ writeTBQueue pchan (ChildStop tid (Just e))
289-- This function is different than "forkWorker" because we have to directly
290-- insert the threadIds here and cannot use the channel to send ChildCreate
291-- unlike on the push side. If we do that, the first thread's done message
292-- may arrive even before the second thread is forked, in that case
293-- pullWorker will falsely detect that all threads are over.
294{-# INLINE pushWorker #-}
295pushWorker :: MonadAsync m => Context m a -> m ()
296pushWorker ctx = do
297 tid <- doFork (push ctx) (handleChildException (outputQueue ctx))
298 liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)
300-- | Split the original computation in a pull-push pair. The original
301-- computation pulls from a Channel while m1 and m2 push to the channel.
302{-# NOINLINE pullFork #-}
303pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
304pullFork m1 m2 = AsyncT $ \_ stp yld -> do
305 ctx <- liftIO $ newContext
306 queueWork ctx m1 >> queueWork ctx m2 >> pushWorker ctx
307 (runAsyncT (pullWorker ctx)) Nothing stp yld
309 where
311 newContext = do
312 outQ <- atomically $ newTBQueue 32
313 workQ <- atomically $ newTBQueue 32
314 running <- newIORef S.empty
315 done <- newIORef S.empty
316 return $ Context { outputQueue = outQ
317 , runningThreads = running
318 , doneThreads = done
319 , workQueue = workQ
320 }
322-- Concurrency rate control. Our objective is to create more threads on
323-- demand if the consumer is running faster than us. As soon as we
324-- encounter an Alternative composition we create a push pull pair of
325-- threads. We use a channel for communication between the consumer that
326-- pulls from the channel and the producer that pushes to the channel. The
327-- producer creates more threads if the channel becomes empty at times,
328-- that is the consumer is running faster. However this mechanism can be
329-- problematic if the initial production latency is high, we may end up
330-- creating too many threads. So we need some way to monitor and use the
331-- latency as well.
332--
333-- TBD For quick response we may have to increase the rate in the middle of
334-- a serially running computation. For that we can use a state flag to fork
335-- the rest of the computation at any point of time inside the Monad bind
336-- operation if the consumer is running at a faster speed.
337--
338-- TBD the alternative composition allows us to dispatch a chunkSize of only 1.
339-- If we have to dispatch in arbitrary chunksizes we will need to compose the
340-- parallel actions using a data constructor instead so that we can divide it
341-- in chunks of arbitrary size before dispatch. When batching we can convert
342-- the structure into Alternative batches of Monoid composition. That will also
343-- allow us to dispatch more work to existing threads rather than creating new
344-- threads always.
345--
346-- TBD for pure work (when we are not in the IO monad) we can divide it into
347-- just the number of CPUs.
348--
349-- XXX to rate control left folded structrues we will have to return the
350-- residual work back to the dispatcher. It will also consume a lot of
351-- memory due to queueing of all the work before execution starts.
353{-# INLINE forkWorker #-}
354forkWorker :: MonadAsync m => Context m a -> m ()
355forkWorker ctx = do
356 let q = outputQueue ctx
357 tid <- doFork (push ctx) (handleChildException q)
358 liftIO $ atomically $ writeTBQueue q (ChildCreate tid)
360{-# INLINE queueWork #-}
361queueWork :: MonadAsync m => Context m a -> AsyncT m a -> m ()
362queueWork ctx m = do
363 -- To guarantee deadlock avoidance we need to dispatch a worker when the
364 -- workQueue goes full. Otherwise all the worker threads might be waiting
365 -- on the queue and never wakeup.
366 --
367 -- TBD If we run out of threads we can also evaluate the action completely
368 -- right here, disallowing any further child workers and turning the
369 -- parallel composition into interleaved serial composition.
370 workQFull <- liftIO $ atomically $ isFullTBQueue (workQueue ctx)
371 when (workQFull) $ forkWorker ctx
372 liftIO $ atomically $ writeTBQueue (workQueue ctx) m
374{-# INLINE dequeueLoop #-}
375dequeueLoop :: MonadAsync m => Context m a -> AsyncT m a
376dequeueLoop ctx = AsyncT $ \_ stp yld -> do
377 work <- liftIO $ atomically $ tryReadTBQueue (workQueue ctx)
378 case work of
379 Nothing -> stp
380 Just m -> do
381 let stop = (runAsyncT (dequeueLoop ctx)) Nothing stp yld
382 yield a c Nothing = yld a c (Just (dequeueLoop ctx))
383 yield a c r = yld a c r
384 (runAsyncT m) (Just ctx) stop yield
386instance MonadAsync m => Alternative (AsyncT m) where
387 empty = mempty
389 -- Note: This is designed to scale for right associated compositions,
390 -- therefore always use a right fold for folding bigger structures.
391 {-# INLINE (<|>) #-}
392 m1 <|> m2 = AsyncT $ \ctx stp yld -> do
393 case ctx of
394 Nothing -> (runAsyncT (pullFork m1 m2)) Nothing stp yld
395 Just c -> do
396 -- we open up both the branches fairly but on a given level we
397 -- go left to right. If the left one keeps producing results we
398 -- may or may not run the right one.
399 queueWork c m1 >> queueWork c m2
400 (runAsyncT (dequeueLoop c)) Nothing stp yld
402instance MonadAsync m => MonadPlus (AsyncT m) where
403 mzero = empty
404 mplus = (<|>)
406------------------------------------------------------------------------------
407-- Num
408------------------------------------------------------------------------------
410instance (Num a, Monad (AsyncT m)) => Num (AsyncT m a) where
411 fromInteger = return . fromInteger
412 mf + mg = (+) <$> mf <*> mg
413 mf * mg = (*) <$> mf <*> mg
414 negate f = f >>= return . negate
415 abs f = f >>= return . abs
416 signum f = f >>= return . signum
418-------------------------------------------------------------------------------
419-- AsyncT transformer
420-------------------------------------------------------------------------------
422instance MonadTrans AsyncT where
423 lift mx = AsyncT $ \ctx _ yld -> mx >>= (\a -> (yld a ctx Nothing))
425instance (MonadBase b m, Monad m) => MonadBase b (AsyncT m) where
426 liftBase = liftBaseDefault
428------------------------------------------------------------------------------
429-- Standard transformer instances
430------------------------------------------------------------------------------
432instance MonadIO m => MonadIO (AsyncT m) where
433 liftIO = lift . liftIO
435instance MonadThrow m => MonadThrow (AsyncT m) where
436 throwM = lift . throwM
438------------------------------------------------------------------------------
439-- MonadRecorder
440------------------------------------------------------------------------------
442instance MonadRecorder m => MonadRecorder (AsyncT m) where
443 getJournal = lift getJournal
444 putJournal = lift . putJournal
445 play = lift . play
447------------------------------------------------------------------------------
448-- Async primitives
449------------------------------------------------------------------------------
450--
451-- Only those actions that are marked with 'async' are guaranteed to be
452-- asynchronous. Asyncly is free to run other actions synchronously or
453-- asynchronously and it should not matter to the semantics of the program, if
454-- it does then use async to force.
455--
456-- Why not make async as default and ask the programmer to use a 'sync'
457-- primitive to force an action to run synchronously? But then we would not
458-- have the freedom to convert the async actions to sync dynamically. Note that
459-- it is safe to convert a sync action to async but vice-versa is not true.
460-- Converting an async to sync can cause change in semantics if the async
461-- action was an infinite loop for example.
462--
463-- | In an 'Alternative' composition, force the action to run asynchronously.
464-- The @\<|\>@ operator implies "can be parallel", whereas 'async' implies
465-- "must be parallel". Note that outside an 'Alternative' composition 'async'
466-- is not useful and should not be used. Even in an 'Alternative' composition
467-- 'async' is not useful as the last action as the last action always runs in
468-- the current thread.
469{-
470async :: Monad m => AsyncT m a -> AsyncT m a
471async action = AsyncT $ runAsyncTask True (runAsyncT action)
473makeAsync :: Monad m => ((a -> m ()) -> m ()) -> AsyncT m a
474makeAsync = AsyncT . makeCont
475-}
477------------------------------------------------------------------------------
478-- Running the monad
479------------------------------------------------------------------------------
481-- | Run an 'AsyncT m' computation, wait for it to finish and discard the
482-- results.
483runAsyncly :: MonadAsync m => AsyncT m a -> m ()
484runAsyncly m = run Nothing m
486 where
488 stop = return ()
490 {-# INLINE yield #-}
491 yield _ _ Nothing = stop
492 yield _ c (Just x) = run c x
494 run ct mx = (runAsyncT mx) ct stop yield
496-- | Run an 'AsyncT m' computation and collect the results generated by each
497-- thread of the computation in a list.
498{-# INLINABLE toList #-}
499toList :: MonadAsync m => AsyncT m a -> m [a]
500toList m = run Nothing m
502 where
504 stop = return []
506 {-# INLINE yield #-}
507 yield a _ Nothing = return [a]
508 yield a c (Just x) = liftM (a :) (run c x)
510 {-# INLINE run #-}
511 run ctx mx = (runAsyncT mx) ctx stop yield
513-- | Run a given function concurrently on the list and collect the results.
514{-# INLINABLE for #-}
515for :: MonadAsync m => [a] -> (a -> AsyncT m b) -> AsyncT m b
516for xs f = foldr (<|>) empty $ map f xs
518-- XXX rename to fromList?
519{-# INLINABLE each #-}
520each :: MonadAsync m => [a] -> AsyncT m a
521each xs = foldr (<>) mempty $ map return xs