on 1be7567 help sign in
streamly/src/Asyncly/Threads.hs
121 lines 4.5 kB history raw
1{-# LANGUAGE FlexibleContexts #-}
2{-# LANGUAGE ScopedTypeVariables #-}
4-- |
5-- Module : Asyncly.Threads
6-- Copyright : (c) 2017 Harendra Kumar
7--
8-- License : MIT-style
9-- Maintainer : harendra.kumar@gmail.com
10-- Stability : experimental
11-- Portability : GHC
12--
13module Asyncly.Threads
14 ( wait
15 , wait_
16 , gather
18 , waitRecord_
19 , waitRecord
20 , playRecordings
21 )
22where
24import Control.Concurrent.STM (atomically, newTChan)
25import Control.Monad.Catch (MonadCatch, throwM, try)
26import Control.Monad.IO.Class (MonadIO (..))
27import Control.Monad.State (mzero, runStateT)
28import Control.Monad.Trans.Class (MonadTrans (lift))
29import Data.IORef (IORef, atomicModifyIORef,
30 newIORef, readIORef)
32import Control.Monad.Trans.Recorder (MonadRecorder(..), RecorderT,
33 Recording, blank, runRecorderT)
34import Asyncly.AsyncT
36------------------------------------------------------------------------------
37-- Running the monad
38------------------------------------------------------------------------------
40-- | Run an 'AsyncT m' computation and collect the results generated by each
41-- thread of the computation in a list.
42waitAsync :: (MonadAsync m, MonadCatch m)
43 => (a -> AsyncT m a) -> Maybe (IORef [Recording]) -> AsyncT m a -> m ()
44waitAsync finalizer lref m = do
45 childChan <- liftIO $ atomically newTChan
46 pendingRef <- liftIO $ newIORef []
47 credit <- liftIO $ newIORef maxBound
49 let ctx = initContext childChan pendingRef credit finalizer lref
51 r <- try $ runStateT (runAsyncT $ m >>= finalizer) ctx
52 e <- handleResult ctx r
53 maybe (return ()) throwM e
55-- TBD throttling of producer based on conumption rate.
57-- | Invoked to store the result of the computation in the context and finish
58-- the computation when the computation is done
59gatherResult :: MonadAsync m => IORef [a] -> a -> AsyncT m a
60gatherResult ref r = do
61 liftIO $ atomicModifyIORef ref $ \rs -> (r : rs, ())
62 mzero
64gather :: (MonadAsync m, MonadCatch m) => AsyncT m a -> AsyncT m [a]
65gather m = AsyncT $ do
66 resultsRef <- liftIO $ newIORef []
67 lift $ waitAsync (gatherResult resultsRef) Nothing m
68 r <- liftIO $ readIORef resultsRef
69 return $ Just r
71-- | Run an 'AsyncT m' computation and collect the results generated by each
72-- thread of the computation in a list.
73wait :: (MonadAsync m, MonadCatch m) => AsyncT m a -> m [a]
74wait m = do
75 resultsRef <- liftIO $ newIORef []
76 waitAsync (gatherResult resultsRef) Nothing m
77 liftIO $ readIORef resultsRef
79-- | Run an 'AsyncT m' computation, wait for it to finish and discard the
80-- results.
81wait_ :: (MonadAsync m, MonadCatch m) => AsyncT m a -> m ()
82wait_ m = waitAsync (const mzero) Nothing m
84------------------------------------------------------------------------------
85-- Logging
86------------------------------------------------------------------------------
88-- | Compose a computation using previously captured logs
89playRecording :: (MonadAsync m, MonadRecorder m)
90 => AsyncT m a -> Recording -> AsyncT m a
91playRecording m recording = play recording >> m
93-- | Resume an 'AsyncT' computation using previously recorded logs. The
94-- recording consists of a list of journals one for each thread in the
95-- computation.
96playRecordings :: (MonadAsync m, MonadRecorder m)
97 => AsyncT m a -> [Recording] -> AsyncT m a
98playRecordings m logs = each logs >>= playRecording m
100-- | Run an 'AsyncT' computation with recording enabled, wait for it to finish
101-- returning results for completed threads and recordings for paused threads.
102waitRecord :: (MonadAsync m, MonadCatch m)
103 => AsyncT m a -> m ([a], [Recording])
104waitRecord m = do
105 resultsRef <- liftIO $ newIORef []
106 lref <- liftIO $ newIORef []
107 waitAsync (gatherResult resultsRef) (Just lref) m
108 res <- liftIO $ readIORef resultsRef
109 logs <- liftIO $ readIORef lref
110 return (res, logs)
112-- | Run an 'AsyncT' computation with recording enabled, wait for it to finish
113-- and discard the results and return the recordings for paused threads, if
114-- any.
115waitRecord_ :: (MonadAsync m, MonadCatch m)
116 => AsyncT (RecorderT m) a -> m [Recording]
117waitRecord_ m = do
118 lref <- liftIO $ newIORef []
119 runRecorderT blank (waitAsync (const mzero) (Just lref) m)
120 logs <- liftIO $ readIORef lref
121 return logs