module ConIO.Core
(
ConIO,
runConIO,
runConIOCancel,
Task (..),
launch,
wait,
cancel,
cancelAll,
linkTasks,
ConScope,
withConScope,
useConScope,
UnsafeConScope,
toUnsafeConScope,
fromUnsafeConScope,
ConIOException (..),
ConIOKillThread (..),
)
where
import ConIO.MonadSTM
import Control.Concurrent
import Control.Concurrent.STM
import Control.Exception
import Control.Monad (void)
import Control.Monad.Fix
import Control.Monad.IO.Class
import Control.Monad.Trans.Reader
import Data.Foldable (forM_, traverse_)
import Data.IORef
import Data.Map qualified as M
import Data.Maybe (fromMaybe)
import Data.Set qualified as S
data ConEnv = ConEnv
{ ConEnv -> MVar Bool
enabled :: MVar Bool,
ConEnv -> IORef (Map ThreadId (IO ()))
children :: IORef (M.Map ThreadId (IO ()))
}
makeConEnv :: IO ConEnv
makeConEnv :: IO ConEnv
makeConEnv = MVar Bool -> IORef (Map ThreadId (IO ())) -> ConEnv
ConEnv (MVar Bool -> IORef (Map ThreadId (IO ())) -> ConEnv)
-> IO (MVar Bool) -> IO (IORef (Map ThreadId (IO ())) -> ConEnv)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> IO (MVar Bool)
forall a. a -> IO (MVar a)
newMVar Bool
True IO (IORef (Map ThreadId (IO ())) -> ConEnv)
-> IO (IORef (Map ThreadId (IO ()))) -> IO ConEnv
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map ThreadId (IO ()) -> IO (IORef (Map ThreadId (IO ())))
forall a. a -> IO (IORef a)
newIORef Map ThreadId (IO ())
forall k a. Map k a
M.empty
newtype ConIO a = ConIO (ReaderT ConEnv IO a) deriving newtype ((forall a b. (a -> b) -> ConIO a -> ConIO b)
-> (forall a b. a -> ConIO b -> ConIO a) -> Functor ConIO
forall a b. a -> ConIO b -> ConIO a
forall a b. (a -> b) -> ConIO a -> ConIO b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> ConIO a -> ConIO b
fmap :: forall a b. (a -> b) -> ConIO a -> ConIO b
$c<$ :: forall a b. a -> ConIO b -> ConIO a
<$ :: forall a b. a -> ConIO b -> ConIO a
Functor, Functor ConIO
Functor ConIO =>
(forall a. a -> ConIO a)
-> (forall a b. ConIO (a -> b) -> ConIO a -> ConIO b)
-> (forall a b c. (a -> b -> c) -> ConIO a -> ConIO b -> ConIO c)
-> (forall a b. ConIO a -> ConIO b -> ConIO b)
-> (forall a b. ConIO a -> ConIO b -> ConIO a)
-> Applicative ConIO
forall a. a -> ConIO a
forall a b. ConIO a -> ConIO b -> ConIO a
forall a b. ConIO a -> ConIO b -> ConIO b
forall a b. ConIO (a -> b) -> ConIO a -> ConIO b
forall a b c. (a -> b -> c) -> ConIO a -> ConIO b -> ConIO c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> ConIO a
pure :: forall a. a -> ConIO a
$c<*> :: forall a b. ConIO (a -> b) -> ConIO a -> ConIO b
<*> :: forall a b. ConIO (a -> b) -> ConIO a -> ConIO b
$cliftA2 :: forall a b c. (a -> b -> c) -> ConIO a -> ConIO b -> ConIO c
liftA2 :: forall a b c. (a -> b -> c) -> ConIO a -> ConIO b -> ConIO c
$c*> :: forall a b. ConIO a -> ConIO b -> ConIO b
*> :: forall a b. ConIO a -> ConIO b -> ConIO b
$c<* :: forall a b. ConIO a -> ConIO b -> ConIO a
<* :: forall a b. ConIO a -> ConIO b -> ConIO a
Applicative, Applicative ConIO
Applicative ConIO =>
(forall a b. ConIO a -> (a -> ConIO b) -> ConIO b)
-> (forall a b. ConIO a -> ConIO b -> ConIO b)
-> (forall a. a -> ConIO a)
-> Monad ConIO
forall a. a -> ConIO a
forall a b. ConIO a -> ConIO b -> ConIO b
forall a b. ConIO a -> (a -> ConIO b) -> ConIO b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. ConIO a -> (a -> ConIO b) -> ConIO b
>>= :: forall a b. ConIO a -> (a -> ConIO b) -> ConIO b
$c>> :: forall a b. ConIO a -> ConIO b -> ConIO b
>> :: forall a b. ConIO a -> ConIO b -> ConIO b
$creturn :: forall a. a -> ConIO a
return :: forall a. a -> ConIO a
Monad, Monad ConIO
Monad ConIO => (forall a. IO a -> ConIO a) -> MonadIO ConIO
forall a. IO a -> ConIO a
forall (m :: * -> *).
Monad m =>
(forall a. IO a -> m a) -> MonadIO m
$cliftIO :: forall a. IO a -> ConIO a
liftIO :: forall a. IO a -> ConIO a
MonadIO, Monad ConIO
Monad ConIO => (forall a. String -> ConIO a) -> MonadFail ConIO
forall a. String -> ConIO a
forall (m :: * -> *).
Monad m =>
(forall a. String -> m a) -> MonadFail m
$cfail :: forall a. String -> ConIO a
fail :: forall a. String -> ConIO a
MonadFail, Monad ConIO
Monad ConIO =>
(forall a. (a -> ConIO a) -> ConIO a) -> MonadFix ConIO
forall a. (a -> ConIO a) -> ConIO a
forall (m :: * -> *).
Monad m =>
(forall a. (a -> m a) -> m a) -> MonadFix m
$cmfix :: forall a. (a -> ConIO a) -> ConIO a
mfix :: forall a. (a -> ConIO a) -> ConIO a
MonadFix, Monad ConIO
Monad ConIO =>
(forall a. STM a -> ConIO a)
-> (forall a. STM a -> IO a -> ConIO a) -> MonadSTM ConIO
forall a. STM a -> ConIO a
forall a. STM a -> IO a -> ConIO a
forall (m :: * -> *).
Monad m =>
(forall a. STM a -> m a)
-> (forall a. STM a -> IO a -> m a) -> MonadSTM m
$cliftSTM :: forall a. STM a -> ConIO a
liftSTM :: forall a. STM a -> ConIO a
$cliftSTM_IO :: forall a. STM a -> IO a -> ConIO a
liftSTM_IO :: forall a. STM a -> IO a -> ConIO a
MonadSTM)
getEnv :: ConIO (ConEnv)
getEnv :: ConIO ConEnv
getEnv = ReaderT ConEnv IO ConEnv -> ConIO ConEnv
forall a. ReaderT ConEnv IO a -> ConIO a
ConIO ReaderT ConEnv IO ConEnv
forall (m :: * -> *) r. Monad m => ReaderT r m r
ask
runConIO :: (MonadIO m) => ConIO a -> m a
runConIO :: forall (m :: * -> *) a. MonadIO m => ConIO a -> m a
runConIO (ConIO ReaderT ConEnv IO a
r) = IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ do
env <- IO ConEnv
makeConEnv
catch
( do
a <- runReaderT r env
readIORef env.children >>= sequence_
pure a
)
(\(SomeException
e :: SomeException) -> ConEnv -> IO ()
killConIO ConEnv
env IO () -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SomeException -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e)
runConIOCancel :: (MonadIO m) => ConIO a -> m a
runConIOCancel :: forall (m :: * -> *) a. MonadIO m => ConIO a -> m a
runConIOCancel (ConIO ReaderT ConEnv IO a
r) = IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ do
env <- IO ConEnv
makeConEnv
catch
( do
a <- runReaderT r env
killConIO env
pure a
)
(\(SomeException
e :: SomeException) -> ConEnv -> IO ()
killConIO ConEnv
env IO () -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SomeException -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e)
launch :: IO a -> ConIO (Task a)
launch :: forall a. IO a -> ConIO (Task a)
launch IO a
action = do
env <- ConIO ConEnv
getEnv
myId <- liftIO $ myThreadId
maybeA <- liftIO $ withLock env $ do
tvar <- newTVarIO StillRunning
tId <- mask $ \forall a. IO a -> IO a
restore -> IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ do
couldBeA <- IO a -> IO (Either SomeException a)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO a -> IO (Either SomeException a))
-> IO a -> IO (Either SomeException a)
forall a b. (a -> b) -> a -> b
$ IO a -> IO a
forall a. IO a -> IO a
restore IO a
action
case couldBeA of
Right a
a -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Result a) -> Result a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Result a)
tvar (a -> Result a
forall a. a -> Result a
Success a
a)
Left SomeException
err -> do
case forall e. Exception e => SomeException -> Maybe e
fromException @ConIOKillThread SomeException
err of
Just ConIOKillThread
_conIOKillThread -> do
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Result a) -> Result a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Result a)
tvar (SomeException -> Result a
forall a. SomeException -> Result a
Failure (ConIOKillThread -> SomeException
forall e. Exception e => e -> SomeException
toException ConIOKillThread
ConIOKillThread))
Maybe ConIOKillThread
Nothing -> do
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Result a) -> Result a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Result a)
tvar (SomeException -> Result a
forall a. SomeException -> Result a
Failure (SomeException -> SomeException
forall e. Exception e => e -> SomeException
toException SomeException
err))
ThreadId -> ConIOException -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
myId (SomeException -> ConIOException
ConIOTaskException SomeException
err)
let readValue = do
value <- TVar (Result a) -> STM (Result a)
forall a. TVar a -> STM a
readTVar TVar (Result a)
tvar
case value of
Result a
StillRunning -> STM a
forall a. STM a
retry
Success a
a -> a -> STM a
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
Failure SomeException
err -> SomeException -> STM a
forall e a. Exception e => e -> STM a
throwSTM SomeException
err
let isDone = do
value <- TVar (Result a) -> STM (Result a)
forall a. TVar a -> STM a
readTVar TVar (Result a)
tvar
case value of
Result a
StillRunning -> STM ()
forall a. STM a
retry
Success a
_a -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
Failure SomeException
_err -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
liftIO $ atomicModifyIORef' env.children $ \Map ThreadId (IO ())
children -> (ThreadId -> IO () -> Map ThreadId (IO ()) -> Map ThreadId (IO ())
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert ThreadId
tId (STM () -> IO ()
forall a. STM a -> IO a
atomically STM ()
isDone) Map ThreadId (IO ())
children, ())
pure $
Task
{ payload = readValue,
threadIds = S.singleton tId
}
pure $ fromMaybe (Task {payload = throwSTM ConIODisabled, threadIds = S.empty}) maybeA
data Task a = Task
{ forall a. Task a -> STM a
payload :: STM a,
forall a. Task a -> Set ThreadId
threadIds :: S.Set ThreadId
}
wait :: (MonadSTM m) => Task a -> m a
wait :: forall (m :: * -> *) a. MonadSTM m => Task a -> m a
wait (Task STM a
payload Set ThreadId
_) = STM a -> m a
forall a. STM a -> m a
forall (m :: * -> *) a. MonadSTM m => STM a -> m a
liftSTM STM a
payload
instance Functor Task where
fmap :: forall a b. (a -> b) -> Task a -> Task b
fmap a -> b
f Task a
task = Task a
task {payload = f <$> task.payload}
instance Applicative Task where
pure :: forall a. a -> Task a
pure a
a =
Task
{ payload :: STM a
payload = a -> STM a
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a,
threadIds :: Set ThreadId
threadIds = Set ThreadId
forall a. Set a
S.empty
}
(Task STM (a -> b)
payloadF Set ThreadId
threadIdsF) <*> :: forall a b. Task (a -> b) -> Task a -> Task b
<*> (Task STM a
payloadA Set ThreadId
threadIdsA) =
Task
{ payload :: STM b
payload = STM (a -> b)
payloadF STM (a -> b) -> STM a -> STM b
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> STM a
payloadA,
threadIds :: Set ThreadId
threadIds = Set ThreadId
threadIdsF Set ThreadId -> Set ThreadId -> Set ThreadId
forall a. Semigroup a => a -> a -> a
<> Set ThreadId
threadIdsA
}
linkTasks :: Task a -> Task b -> Task b
linkTasks :: forall a b. Task a -> Task b -> Task b
linkTasks Task a
taskA Task b
taskB = Task b
taskB {threadIds = taskA.threadIds <> taskB.threadIds}
conIOkillThread :: ThreadId -> IO ()
conIOkillThread :: ThreadId -> IO ()
conIOkillThread ThreadId
tId = ThreadId -> ConIOKillThread -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
tId ConIOKillThread
ConIOKillThread
cancel :: (MonadIO m) => Task a -> m ()
cancel :: forall (m :: * -> *) a. MonadIO m => Task a -> m ()
cancel (Task STM a
_payload Set ThreadId
tIds) = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
IO () -> IO ()
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ (ThreadId -> IO ()) -> Set ThreadId -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ ThreadId -> IO ()
conIOkillThread Set ThreadId
tIds
() -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
withLock :: ConEnv -> IO a -> IO (Maybe a)
withLock :: forall a. ConEnv -> IO a -> IO (Maybe a)
withLock ConEnv
env IO a
action = do
mEnabled <- MVar Bool -> IO (Maybe Bool)
forall a. MVar a -> IO (Maybe a)
tryReadMVar ConEnv
env.enabled
if mEnabled == Just False
then pure Nothing
else
bracket
(takeMVar env.enabled)
(putMVar env.enabled)
( \Bool
enabled ->
if Bool
enabled
then a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> IO a -> IO (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO a
action
else Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
)
killConIO :: ConEnv -> IO ()
killConIO :: ConEnv -> IO ()
killConIO ConEnv
env = IO (Either SomeException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either SomeException ()) -> IO ())
-> IO (Either SomeException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException (IO () -> IO (Either SomeException ()))
-> IO () -> IO (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ do
MVar Bool -> (Bool -> IO Bool) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVarMasked_ ConEnv
env.enabled ((Bool -> IO Bool) -> IO ()) -> (Bool -> IO Bool) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Bool
enabled ->
if Bool
enabled
then do
IO (Either ConIOException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either ConIOException ()) -> IO ())
-> IO (Either ConIOException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @ConIOException (IO () -> IO (Either ConIOException ()))
-> IO () -> IO (Either ConIOException ())
forall a b. (a -> b) -> a -> b
$ do
children <- IORef (Map ThreadId (IO ())) -> IO (Map ThreadId (IO ()))
forall a. IORef a -> IO a
readIORef ConEnv
env.children
liftIO $ traverse_ conIOkillThread (M.keys children)
liftIO $ sequence_ children
Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
else Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
cancelAll :: ConIO ()
cancelAll :: ConIO ()
cancelAll = do
env <- ConIO ConEnv
getEnv
_ <- liftIO $ void $ killConIO env
pure ()
newtype ConScope s = ConScope ConEnv
withConScope :: (forall s. ConScope s -> IO a) -> ConIO a
withConScope :: forall {k} a. (forall (s :: k). ConScope s -> IO a) -> ConIO a
withConScope forall (s :: k). ConScope s -> IO a
f = do
env <- ReaderT ConEnv IO ConEnv -> ConIO ConEnv
forall a. ReaderT ConEnv IO a -> ConIO a
ConIO ReaderT ConEnv IO ConEnv
forall (m :: * -> *) r. Monad m => ReaderT r m r
ask
a <- liftIO $ f (ConScope env)
pure a
useConScope :: ConScope s -> ConIO a -> IO a
useConScope :: forall {k} (s :: k) a. ConScope s -> ConIO a -> IO a
useConScope (ConScope ConEnv
env) (ConIO ReaderT ConEnv IO a
r) = ReaderT ConEnv IO a -> ConEnv -> IO a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT ConEnv IO a
r ConEnv
env
newtype UnsafeConScope = UnsafeConScope ConEnv
toUnsafeConScope :: ConScope s -> UnsafeConScope
toUnsafeConScope :: forall {k} (s :: k). ConScope s -> UnsafeConScope
toUnsafeConScope (ConScope ConEnv
env) = ConEnv -> UnsafeConScope
UnsafeConScope ConEnv
env
fromUnsafeConScope :: UnsafeConScope -> (forall s. ConScope s -> a) -> a
fromUnsafeConScope :: forall {k} a.
UnsafeConScope -> (forall (s :: k). ConScope s -> a) -> a
fromUnsafeConScope (UnsafeConScope ConEnv
env) forall (s :: k). ConScope s -> a
f = ConScope Any -> a
forall (s :: k). ConScope s -> a
f (ConEnv -> ConScope Any
forall {k} (s :: k). ConEnv -> ConScope s
ConScope ConEnv
env)
data ConIOException
=
ConIODisabled
|
ConIOTaskException SomeException
deriving (Int -> ConIOException -> ShowS
[ConIOException] -> ShowS
ConIOException -> String
(Int -> ConIOException -> ShowS)
-> (ConIOException -> String)
-> ([ConIOException] -> ShowS)
-> Show ConIOException
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ConIOException -> ShowS
showsPrec :: Int -> ConIOException -> ShowS
$cshow :: ConIOException -> String
show :: ConIOException -> String
$cshowList :: [ConIOException] -> ShowS
showList :: [ConIOException] -> ShowS
Show)
instance Exception ConIOException
data ConIOKillThread = ConIOKillThread deriving (Int -> ConIOKillThread -> ShowS
[ConIOKillThread] -> ShowS
ConIOKillThread -> String
(Int -> ConIOKillThread -> ShowS)
-> (ConIOKillThread -> String)
-> ([ConIOKillThread] -> ShowS)
-> Show ConIOKillThread
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ConIOKillThread -> ShowS
showsPrec :: Int -> ConIOKillThread -> ShowS
$cshow :: ConIOKillThread -> String
show :: ConIOKillThread -> String
$cshowList :: [ConIOKillThread] -> ShowS
showList :: [ConIOKillThread] -> ShowS
Show)
instance Exception ConIOKillThread
data Result a = StillRunning | Success a | Failure SomeException