git.iko.soy/streamly
on 154a610 help sign in
streamly/src/Asyncly/AsyncT.hs
631 lines 24 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 : BSD3
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 , take
25 , drop
26 , interleave
27 , (<=>)
28 , parAhead
29 , (<>|)
30 , parLeft
31 , (<|)
32 , foldWith
33 , foldMapWith
34 , forEachWith
35 , fromCallback
36 )
37where
39import Control.Applicative (Alternative (..), liftA2)
40import Control.Concurrent (ThreadId, forkIO,
41 myThreadId, threadDelay)
42import Control.Concurrent.MVar (MVar, newEmptyMVar, tryTakeMVar,
43 tryPutMVar, takeMVar)
44import Control.Exception (SomeException (..))
45import qualified Control.Exception.Lifted as EL
46import Control.Monad (ap, liftM, MonadPlus(..), mzero,
47 when)
48import Control.Monad.Base (MonadBase (..), liftBaseDefault)
49import Control.Monad.Catch (MonadThrow, throwM)
50import Control.Monad.Error.Class (MonadError(..))
51import Control.Monad.IO.Class (MonadIO(..))
52import Control.Monad.Reader.Class (MonadReader(..))
53import Control.Monad.State.Class (MonadState(..))
54import Control.Monad.Trans.Class (MonadTrans (lift))
55import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
56import Data.Concurrent.Queue.MichaelScott (LinkedQueue, newQ, pushL,
57 tryPopR, nullQ)
58import Data.Functor (void)
59import Data.IORef (IORef, modifyIORef, newIORef,
60 readIORef)
61import Data.Atomics (atomicModifyIORefCAS,
62 atomicModifyIORefCAS_)
63import Data.Maybe (isNothing)
64import Data.Monoid ((<>))
65import Data.Set (Set)
66import qualified Data.Set as S
67import Prelude hiding (take, drop)
69------------------------------------------------------------------------------
70-- Concurrency Semantics
71------------------------------------------------------------------------------
72--
73-- Asyncly is essentially a concurrent list transformer. To concatenate the
74-- lists it provides three different ways of composing them. The Monoid
75-- instance concatenates lists in a non-concurrent in-order fashion. The
76-- Alternative instance concatenates list in a concurrent manner potentially
77-- running each action in parallel and therefore the order of the resulting
78-- items is not deterministic. Computations in Alternative composition may or
79-- may not run in parallel. Thirdly, the <||> operator provides a way to
80-- compose computations that are guaranteed to run in parallel. This provides a
81-- way to run infinite loops or long running computations without blocking
82-- or starving other computations.
84------------------------------------------------------------------------------
85-- Parent child thread communication types
86------------------------------------------------------------------------------
88data ChildEvent a =
89 ChildYield a
90 | ChildStop ThreadId (Maybe SomeException)
92------------------------------------------------------------------------------
93-- State threaded around the monad for thread management
94------------------------------------------------------------------------------
96data Context m a =
97 Context { outputQueue :: IORef [ChildEvent a]
98 , doorBell :: MVar Bool -- wakeup mechanism for outQ
99 , enqueue :: AsyncT m a -> IO ()
100 , runqueue :: m ()
101 , runningThreads :: IORef (Set ThreadId)
102 , queueEmpty :: m Bool
103 }
105newtype AsyncT m a =
106 AsyncT {
107 runAsyncT :: forall r.
108 Maybe (Context m a) -- state
109 -> m r -- stop
110 -> (a -> Maybe (Context m a) -> Maybe (AsyncT m a) -> m r) -- yield
111 -> m r
112 }
114type MonadAsync m = (MonadIO m, MonadBaseControl IO m, MonadThrow m)
116------------------------------------------------------------------------------
117-- Monad
118------------------------------------------------------------------------------
120-- | Appends the results of two AsyncT computations in order.
121instance Monad m => Monoid (AsyncT m a) where
122 mempty = AsyncT $ \_ stp _ -> stp
123 mappend (AsyncT m1) m2 = AsyncT $ \ctx stp yld ->
124 let stop = (runAsyncT m2) ctx stp yld
125 yield a c Nothing = yld a c (Just m2)
126 yield a c (Just r) = yld a c (Just (mappend r m2))
127 in m1 ctx stop yield
129-- | Interleaves the results of two 'AsyncT' streams
130interleave :: AsyncT m a -> AsyncT m a -> AsyncT m a
131interleave m1 m2 = AsyncT $ \ctx stp yld -> do
132 let stop = (runAsyncT m2) ctx stp yld
133 yield a c Nothing = yld a c (Just m2)
134 yield a c (Just r) = yld a c (Just (interleave m2 r))
135 (runAsyncT m1) ctx stop yield
137infixr 5 <=>
139(<=>) :: AsyncT m a -> AsyncT m a -> AsyncT m a
140(<=>) = interleave
142-- We do not use bind for parallelism. That is, we do not start each iteration
143-- of the list in parallel. That will introduce too much uncontrolled
144-- parallelism. Instead we control parallelism using <|>, it is used for
145-- actions that 'may' be run in parallel. The library decides whether they will
146-- be started in parallel. This puts the parallel semantics in the hands of the
147-- user and operational semantics in the hands of the library. Ideally we would
148-- want parallelism to be completely hidden from the user, for example we can
149-- automatically decide where to start an action in parallel so that it is most
150-- efficient. We can decide whether a particular bind should fork a parallel
151-- thread or not based on the level of the bind in the tree. Maybe at some
152-- point we shall be able to do that?
154instance Monad m => Monad (AsyncT m) where
155 return a = AsyncT $ \ctx _ yld -> yld a ctx Nothing
157 -- | A thread context is valid only until the next bind. Upon a bind we
158 -- reset the context to Nothing.
159 AsyncT m >>= f = AsyncT $ \_ stp yld ->
160 let run x = (runAsyncT x) Nothing stp yld
161 yield a _ Nothing = run $ f a
162 yield a _ (Just r) = run $ f a <> (r >>= f)
163 in m Nothing stp yield
165------------------------------------------------------------------------------
166-- Functor
167------------------------------------------------------------------------------
169instance Monad m => Functor (AsyncT m) where
170 fmap = liftM
172------------------------------------------------------------------------------
173-- Applicative
174------------------------------------------------------------------------------
176-- Note: We cannot have the applicative instance run the two actions in
177-- parallel. This is because the first action can always return empty and in
178-- case it returns empty we should not be starting the second action. Therefore
179-- we can only start the second action after getting a result from the first.
180-- For running applicative actions in parallel we can use the async package and
181-- lift the result into AsyncT.
183instance Monad m => Applicative (AsyncT m) where
184 pure = return
185 (<*>) = ap
187------------------------------------------------------------------------------
188-- Alternative
189------------------------------------------------------------------------------
191{-# INLINE doFork #-}
192doFork :: MonadBaseControl IO m
193 => m ()
194 -> (SomeException -> m ())
195 -> m ThreadId
196doFork action exHandler =
197 EL.mask $ \restore ->
198 liftBaseWith $ \runInIO -> forkIO $ do
199 -- XXX test the exception handling
200 _ <- runInIO $ EL.catch (restore action) exHandler
201 -- XXX restore state here?
202 return ()
204-- XXX exception safety of all atomic/MVar operations
206{-# INLINE send #-}
207send :: MonadIO m => Context m a -> ChildEvent a -> m ()
208send ctx msg = liftIO $ do
209 atomicModifyIORefCAS_ (outputQueue ctx) $ \es -> msg : es
210 -- XXX need a memory barrier? The wake up must happen only after the
211 -- store has finished otherwise we can have lost wakeup problems.
212 void $ tryPutMVar (doorBell ctx) True
214{-# INLINE sendStop #-}
215sendStop :: MonadIO m => Context m a -> m ()
216sendStop ctx = liftIO myThreadId >>= \tid -> send ctx (ChildStop tid Nothing)
218-- Note: Left associated operations can grow this queue to a large size
219{-# INLINE enqueueLIFO #-}
220enqueueLIFO :: IORef [AsyncT m a] -> AsyncT m a -> IO ()
221enqueueLIFO q m = atomicModifyIORefCAS_ q $ \ ms -> m : ms
223runqueueLIFO :: MonadIO m => Context m a -> IORef [AsyncT m a] -> m ()
224runqueueLIFO ctx q = run
226 where
228 run = do
229 work <- dequeue
230 case work of
231 Nothing -> sendStop ctx
232 Just m -> (runAsyncT m) (Just ctx) run yield
234 sendit a = send ctx (ChildYield a)
235 yield a _ Nothing = sendit a >> run
236 yield a c (Just r) = sendit a >> (runAsyncT r) c run yield
238 dequeue = liftIO $ atomicModifyIORefCAS q $ \ ms ->
239 case ms of
240 [] -> ([], Nothing)
241 x : xs -> (xs, Just x)
243{-# INLINE enqueueFIFO #-}
244enqueueFIFO :: LinkedQueue (AsyncT m a) -> AsyncT m a -> IO ()
245enqueueFIFO = pushL
247runqueueFIFO :: MonadIO m => Context m a -> LinkedQueue (AsyncT m a) -> m ()
248runqueueFIFO ctx q = run
250 where
252 run = do
253 work <- dequeue
254 case work of
255 Nothing -> sendStop ctx
256 Just m -> (runAsyncT m) (Just ctx) run yield
258 dequeue = liftIO $ tryPopR q
259 sendit a = send ctx (ChildYield a)
260 yield a _ Nothing = sendit a >> run
261 -- XXX do we need to save the context as well?
262 -- we can enqueue "(runAsyncT m) ctx run yield" instead
263 yield a _ (Just r) = sendit a >> liftIO (enqueueFIFO q r) >> run
265-- Thread tracking is needed for two reasons:
266--
267-- 1) Killing threads on exceptions. Threads may not be allowed to go away by
268-- themselves because they may run for significant times before going away or
269-- worse they may be stuck in IO and never go away.
270--
271-- 2) To know when all threads are done.
273{-# NOINLINE addThread #-}
274addThread :: MonadIO m => Context m a -> ThreadId -> m ()
275addThread ctx tid =
276 liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)
278{-# INLINE delThread #-}
279delThread :: MonadIO m => Context m a -> ThreadId -> m ()
280delThread ctx tid =
281 liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.delete tid s)
283{-# INLINE allThreadsDone #-}
284allThreadsDone :: MonadIO m => Context m a -> m Bool
285allThreadsDone ctx = liftIO $ do
286 readIORef (runningThreads ctx) >>= return . S.null
288{-# NOINLINE handleChildException #-}
289handleChildException :: MonadIO m => Context m a -> SomeException -> m ()
290handleChildException ctx e = do
291 tid <- liftIO myThreadId
292 send ctx (ChildStop tid (Just e))
294{-# NOINLINE pushWorker #-}
295pushWorker :: MonadAsync m => Context m a -> m ()
296pushWorker ctx =
297 doFork (runqueue ctx) (handleChildException ctx) >>= addThread ctx
299{-# INLINE sendWorkerWait #-}
300sendWorkerWait :: MonadAsync m => Context m a -> m ()
301sendWorkerWait ctx = dispatch >> void (liftIO $ takeMVar (doorBell ctx))
303 where
305 dispatch = do
306 liftIO $ threadDelay 200
307 output <- liftIO $ readIORef (outputQueue ctx)
308 when (null output) $ do
309 done <- queueEmpty ctx
310 when (not done) $ (pushWorker ctx) >> dispatch
313-- Note: This is performance sensitive code.
314{-# NOINLINE pullWorker #-}
315pullWorker :: MonadAsync m => Context m a -> AsyncT m a
316pullWorker ctx = AsyncT $ \pctx stp yld -> do
317 res <- liftIO $ tryTakeMVar (doorBell ctx)
318 when (isNothing res) $ sendWorkerWait ctx
319 list <- liftIO $ atomicModifyIORefCAS (outputQueue ctx) $ \x -> ([], x)
320 (runAsyncT $ processEvents list) pctx stp yld
322 where
324 handleException e tid = do
325 delThread ctx tid
326 -- XXX implement kill async exception handling
327 -- liftIO $ readIORef (runningThreads ctx) >>= mapM_ killThread
328 throwM e
330 {-# INLINE processEvents #-}
331 processEvents [] = AsyncT $ \pctx stp yld -> do
332 done <- allThreadsDone ctx
333 if not done
334 then (runAsyncT (pullWorker ctx)) pctx stp yld
335 else stp
337 processEvents (ev : es) = AsyncT $ \pctx stp yld -> do
338 let continue = (runAsyncT (processEvents es)) pctx stp yld
339 yield a = yld a pctx (Just (processEvents es))
341 case ev of
342 ChildYield a -> yield a
343 ChildStop tid e ->
344 case e of
345 Nothing -> delThread ctx tid >> continue
346 Just ex -> handleException ex tid
348-- | Split the original computation in a pull-push pair. The original
349-- computation pulls from a Channel while m1 and m2 push to the channel.
350{-# NOINLINE pullFork #-}
351pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a
352pullFork m1 m2 fifo = AsyncT $ \_ stp yld -> do
353 ctx <- liftIO $ newContext
354 pushWorker ctx >> (runAsyncT (pullWorker ctx)) Nothing stp yld
356 where
358 newContext = do
359 outQ <- newIORef []
360 outQMv <- newEmptyMVar
361 running <- newIORef S.empty
363 case fifo of
364 True -> do
365 q <- newQ
366 pushL q m1 >> pushL q m2
367 let ctx =
368 Context { outputQueue = outQ
369 , doorBell = outQMv
370 , runningThreads = running
371 , runqueue = runqueueFIFO ctx q
372 , enqueue = pushL q
373 , queueEmpty = liftIO $ nullQ q
374 }
375 in return ctx
376 False -> do
377 q <- newIORef []
378 enqueueLIFO q m2 >> enqueueLIFO q m1
379 let checkEmpty = liftIO (readIORef q) >>= return . null
380 let ctx =
381 Context { outputQueue = outQ
382 , doorBell = outQMv
383 , runningThreads = running
384 , runqueue = runqueueLIFO ctx q
385 , enqueue = enqueueLIFO q
386 , queueEmpty = checkEmpty
387 }
388 in return ctx
390-- Concurrency rate control. Our objective is to create more threads on
391-- demand if the consumer is running faster than us. As soon as we
392-- encounter an Alternative composition we create a push pull pair of
393-- threads. We use a channel for communication between the consumer that
394-- pulls from the channel and the producer that pushes to the channel. The
395-- producer creates more threads if the channel becomes empty at times,
396-- that is the consumer is running faster. However this mechanism can be
397-- problematic if the initial production latency is high, we may end up
398-- creating too many threads. So we need some way to monitor and use the
399-- latency as well.
400--
401-- TBD For quick response we may have to increase the rate in the middle of
402-- a serially running computation. For that we can use a state flag to fork
403-- the rest of the computation at any point of time inside the Monad bind
404-- operation if the consumer is running at a faster speed.
405--
406-- TBD the alternative composition allows us to dispatch a chunkSize of only 1.
407-- If we have to dispatch in arbitrary chunksizes we will need to compose the
408-- parallel actions using a data constructor instead so that we can divide it
409-- in chunks of arbitrary size before dispatch. When batching we can convert
410-- the structure into Alternative batches of Monoid composition. That will also
411-- allow us to dispatch more work to existing threads rather than creating new
412-- threads always.
413--
414-- TBD for pure work (when we are not in the IO monad) we can divide it into
415-- just the number of CPUs.
416--
417-- XXX to rate control left folded structrues we will have to return the
418-- residual work back to the dispatcher. It will also consume a lot of
419-- memory due to queueing of all the work before execution starts.
421-- Note: This is designed to scale for right associated compositions,
422-- therefore always use a right fold for folding bigger structures.
423{-# INLINE parallel #-}
424parallel :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a
425parallel m1 m2 fifo = AsyncT $ \ctx stp yld -> do
426 case ctx of
427 Nothing -> (runAsyncT (pullFork m1 m2 fifo)) Nothing stp yld
428 Just c -> liftIO ((enqueue c) m2) >> (runAsyncT m1) ctx stp yld
430instance MonadAsync m => Alternative (AsyncT m) where
431 empty = mempty
433 {-# INLINE (<|>) #-}
434 m1 <|> m2 = parallel m1 m2 True
436-- | Just like '<>' except that it can execute the action on the right in
437-- parallel ahead of time. Returns the results in serial order like '<>' from
438-- left to right.
439parAhead :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
440parAhead = undefined
442(<>|) :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
443(<>|) = parAhead
445-- | Left biased parallel execution. Actions may run in parallel but not
446-- necessarily, the action on the left is executed before the one on the right
447-- when parallelism is not needed. This combinator is useful when fairness is
448-- not required.
449{-# INLINE parLeft #-}
450parLeft :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
451parLeft m1 m2 = parallel m1 m2 False
453-- | Same as 'parLeft'.
454{-# INLINE (<|) #-}
455(<|) :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a
456(<|) = parLeft
458instance MonadAsync m => MonadPlus (AsyncT m) where
459 mzero = empty
460 mplus = (<|>)
462------------------------------------------------------------------------------
463-- Num
464------------------------------------------------------------------------------
466instance (Monad m, Num a) => Num (AsyncT m a) where
467 fromInteger n = pure (fromInteger n)
469 negate = fmap negate
470 abs = fmap abs
471 signum = fmap signum
473 (+) = liftA2 (+)
474 (*) = liftA2 (*)
475 (-) = liftA2 (-)
477instance (Monad m, Fractional a) => Fractional (AsyncT m a) where
478 fromRational n = pure (fromRational n)
480 recip = fmap recip
482 (/) = liftA2 (/)
484instance (Monad m, Floating a) => Floating (AsyncT m a) where
485 pi = pure pi
487 exp = fmap exp
488 sqrt = fmap sqrt
489 log = fmap log
490 sin = fmap sin
491 tan = fmap tan
492 cos = fmap cos
493 asin = fmap asin
494 atan = fmap atan
495 acos = fmap acos
496 sinh = fmap sinh
497 tanh = fmap tanh
498 cosh = fmap cosh
499 asinh = fmap asinh
500 atanh = fmap atanh
501 acosh = fmap acosh
503 (**) = liftA2 (**)
504 logBase = liftA2 logBase
506-------------------------------------------------------------------------------
507-- AsyncT transformer
508-------------------------------------------------------------------------------
510instance MonadTrans AsyncT where
511 lift mx = AsyncT $ \ctx _ yld -> mx >>= (\a -> (yld a ctx Nothing))
513instance (MonadBase b m, Monad m) => MonadBase b (AsyncT m) where
514 liftBase = liftBaseDefault
516------------------------------------------------------------------------------
517-- Standard transformer instances
518------------------------------------------------------------------------------
520instance MonadIO m => MonadIO (AsyncT m) where
521 liftIO = lift . liftIO
523instance MonadThrow m => MonadThrow (AsyncT m) where
524 throwM = lift . throwM
526-- XXX handle and test cross thread state transfer
527instance MonadError e m => MonadError e (AsyncT m) where
528 throwError = lift . throwError
529 catchError m h = AsyncT $ \ctx stp yld ->
530 let handle r = r `catchError` \e -> (runAsyncT (h e)) ctx stp yld
531 yield a c Nothing = yld a c Nothing
532 yield a c (Just r) = yld a c (Just (catchError r h))
533 in handle $ (runAsyncT m) ctx stp yield
535instance MonadReader r m => MonadReader r (AsyncT m) where
536 ask = lift ask
537 local f m = AsyncT $ \ctx stp yld ->
538 let yield a c Nothing = local f $ yld a c Nothing
539 yield a c (Just r) = local f $ yld a c (Just (local f r))
540 in (runAsyncT m) ctx (local f stp) yield
542instance MonadState s m => MonadState s (AsyncT m) where
543 get = lift get
544 put x = lift (put x)
545 state k = lift (state k)
547------------------------------------------------------------------------------
548-- Running the monad
549------------------------------------------------------------------------------
551-- | Run an 'AsyncT m' computation, wait for it to finish and discard the
552-- results.
553runAsyncly :: MonadAsync m => AsyncT m a -> m ()
554runAsyncly m = run Nothing m
556 where
558 stop = return ()
560 {-# INLINE yield #-}
561 yield _ _ Nothing = stop
562 yield _ c (Just x) = run c x
564 run ct mx = (runAsyncT mx) ct stop yield
566-- | Run an 'AsyncT m' computation and collect the results generated by each
567-- thread of the computation in a list.
568{-# INLINABLE toList #-}
569toList :: MonadAsync m => AsyncT m a -> m [a]
570toList m = run Nothing m
572 where
574 stop = return []
576 {-# INLINE yield #-}
577 yield a _ Nothing = return [a]
578 yield a c (Just x) = liftM (a :) (run c x)
580 {-# INLINE run #-}
581 run ctx mx = (runAsyncT mx) ctx stop yield
583------------------------------------------------------------------------------
584-- Transformation
585------------------------------------------------------------------------------
587take :: MonadAsync m => Int -> AsyncT m a -> AsyncT m a
588take n m = AsyncT $ \ctx stp yld -> do
589 let yield a c Nothing = yld a c Nothing
590 yield a c (Just x) = yld a c (Just (take (n - 1) x))
591 if (n <= 0)
592 then stp
593 else (runAsyncT m) ctx stp yield
595drop :: MonadAsync m => Int -> AsyncT m a -> AsyncT m a
596drop n m = AsyncT $ \ctx stp yld -> do
597 let yield _ _ Nothing = stp
598 yield _ c (Just x) = (runAsyncT $ drop (n - 1) x) c stp yld
599 if (n <= 0)
600 then (runAsyncT m) ctx stp yld
601 else (runAsyncT m) ctx stp yield
603------------------------------------------------------------------------------
604-- Utilities
605------------------------------------------------------------------------------
607-- | Fold a 'Foldable' container using the given function.
608{-# INLINABLE foldWith #-}
609foldWith :: (Monoid b, Foldable t) => (a -> b -> b) -> t a -> b
610foldWith f = foldr f mempty
612-- | Fold a 'Foldable' container using a function that is a composition of the
613-- two arguments.
614{-# INLINABLE foldMapWith #-}
615foldMapWith :: (Monoid b, Foldable t) =>
616 (b1 -> b -> b) -> (a -> b1) -> t a -> b
617foldMapWith f g = foldr (f . g) mempty
619-- | Fold a 'Foldable' container using a function that is a compostioin of the
620-- first and the third argument.
621{-# INLINABLE forEachWith #-}
622forEachWith :: (Monoid b, Foldable t) =>
623 (b1 -> b -> b) -> t a -> (a -> b1) -> b
624forEachWith f xs g = foldr (f . g) mempty xs
626------------------------------------------------------------------------------
627-- Convert a callback into an 'AsyncT' computation
628------------------------------------------------------------------------------
630fromCallback :: (forall r. (a -> m r) -> m r) -> AsyncT m a
631fromCallback k = AsyncT $ \ctx _ yld -> k (\a -> yld a ctx Nothing)