@@ -12,6 +12,11 @@ module Strands.Threads
1212
1313 , wait
1414 , gather
15+
16+ , logged
17+ , suspend
18+ , waitLogged
19+ , gatherLogged
1520 )
1621where
1722
@@ -27,7 +32,7 @@ import qualified Control.Exception.Lifted as EL
2732import Control.Monad.Catch (MonadCatch, MonadThrow, throwM,
2833 try)
2934import Control.Monad.IO.Class (MonadIO (..))
30−import Control.Monad.State (get, gets, modify, mzero,
35+import Control.Monad.State (get, gets, modify, mzero, put,
3136 runStateT, when, StateT)
3237import Control.Monad.Trans.Class (MonadTrans (lift))
3338import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
@@ -496,14 +501,14 @@ noTrans x = AsyncT $ x >>= return . Just
496501-- | Run an 'AsyncT m' computation and collect the results generated by each
497502-- thread of the computation in a list.
498503waitAsync :: forall m a. (MonadIO m, MonadCatch m)
499− => (a -> AsyncT m a) -> AsyncT m a -> m ()
500−waitAsync finalizer m = do
504+ => (a -> AsyncT m a) -> Maybe (IORef [Log]) -> AsyncT m a -> m ()
505+waitAsync finalizer lref m = do
501506 childChan <- liftIO $ atomically newTChan
502507 pendingRef <- liftIO $ newIORef []
503508 credit <- liftIO $ newIORef maxBound
504509
505510 let ctx = initContext (empty :: AsyncT m a) childChan pendingRef credit
506− finalizer
511+ finalizer lref
507512
508513 r <- try $ runStateT (runAsyncT $ m >>= finalizer) ctx
509514
@@ -532,11 +537,101 @@ gather :: forall m a. (MonadIO m, MonadCatch m)
532537 => AsyncT m a -> m [a]
533538gather m = do
534539 resultsRef <- liftIO $ newIORef []
535− waitAsync (gatherResult resultsRef) m
540+ waitAsync (gatherResult resultsRef) Nothing m
536541 liftIO $ readIORef resultsRef
537542
538543-- | Run an 'AsyncT m' computation, wait for it to finish and discard the
539544-- results.
540545wait :: forall m a. (MonadIO m, MonadCatch m)
541546 => AsyncT m a -> m ()
542−wait m = waitAsync (const mzero) m
547+wait m = waitAsync (const mzero) Nothing m
548+
549+------------------------------------------------------------------------------
550+-- Logging
551+------------------------------------------------------------------------------
552+
553+logged :: (Loggable a, MonadIO m) => AsyncT m a -> AsyncT m a
554+logged m = AsyncT $ do
555+ ctx <- get
556+ case logsRef ctx of
557+ Nothing -> runAsyncT m
558+ Just _ -> withLog ctx m
559+ where
560+
561+ withLog ctx action =
562+ case journal ctx of
563+ -- no replay
564+ j@(CtxLog _ []) -> do
565+ preLog ctx j
566+ x <- runAsyncT action
567+ postLog ctx j x
568+ return x
569+
570+ -- replaying the log
571+ CtxLog ls (r:rs) -> do
572+ let j = CtxLog ls rs
573+ preLog ctx j
574+ case r of
575+ Executing -> do
576+ x <- runAsyncT action
577+ postLog ctx (CtxLog ls []) x
578+ return x
579+ Result val -> do
580+ let x = fmap read val
581+ postLog ctx j x
582+ return x
583+
584+ -- add an Executing entry at the head of the log
585+ preLog = appendLog Executing
586+ postLog ctx j x = appendLog (Result (fmap show x)) ctx j
587+ appendLog l ctx (CtxLog ls rs) = put $ ctx {journal = CtxLog (l : ls) rs}
588+
589+-- XXX enable or disable the suspension points by sending an exception to the
590+-- computation.
591+--
592+-- | Suspend works as an exit point with current thread state (partially done)
593+-- saved. The computation exits only if all the threads in the computation hit
594+-- the suspend point or exit normally. The threads which hit the suspension and
595+-- were suspended can be re-started by replaying the log. If there are threads
596+-- which did not suspend or did not exit, the computation will not exit. For
597+-- completed threads output is returned and for suspended threads the logs to
598+-- resume them are returned.
599+
600+suspend :: MonadIO m => AsyncT m ()
601+suspend = logged $ AsyncT $ do
602+ ctx <- get
603+ case logsRef ctx of
604+ Nothing -> return $ Just ()
605+ Just ref ->
606+ case journal ctx of
607+ CtxLog ls [] -> do
608+ -- replace the "Executing" entry at the head of the log
609+ -- with a "()" so that we do not enter suspend on replay
610+ liftIO $ atomicModifyIORef ref $ \logs ->
611+ (Log (logResult (Just ()) : tail ls) : logs, ())
612+ return Nothing
613+ _ -> error "Bug: replay inside suspend"
614+ where logResult x = Result (fmap show x)
615+
616+-- | Run an 'AsyncT m' computation with logging enabled and collect the results
617+-- or logs generated by each thread of the computation.
618+gatherLogged :: forall m a. (MonadIO m, MonadCatch m)
619+ => AsyncT m a -> m ([a], [Log])
620+gatherLogged m = do
621+ resultsRef <- liftIO $ newIORef []
622+ lref <- liftIO $ newIORef []
623+ waitAsync (gatherResult resultsRef) (Just lref) m
624+ res <- liftIO $ readIORef resultsRef
625+ logs <- liftIO $ readIORef lref
626+ return (res, logs)
627+
628+-- | Run an 'AsyncT m' computation with logging enabled, wait for it to finish
629+-- and discard the results. If the computation suspends collect the logs to
630+-- replay later.
631+waitLogged :: forall m a. (MonadIO m, MonadCatch m)
632+ => AsyncT m a -> m [Log]
633+waitLogged m = do
634+ lref <- liftIO $ newIORef []
635+ waitAsync (const mzero) (Just lref) m
636+ logs <- liftIO $ readIORef lref
637+ return logs