{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
module System.ThreadManager (
ThreadManager,
newThreadManager,
stopAfter,
stopAfterWithResult,
KilledByThreadManager (..),
forkManaged,
forkManagedFinally,
forkManagedUnmask,
forkManagedTimeout,
forkManagedTimeoutFinally,
waitUntilAllGone,
isAllGone,
T.Manager,
withHandle,
T.Handle,
T.tickle,
T.pause,
T.resume,
) where
import Control.Concurrent
import Control.Concurrent.STM
import Control.Exception (Exception (..), SomeException (..))
import qualified Control.Exception as E
import Control.Monad (unless, void)
import Data.Foldable (forM_)
import Data.IORef (IORef, atomicModifyIORef', newIORef)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Word (Word64)
import GHC.Conc.Sync (labelThread)
#if __GLASGOW_HASKELL__ >= 908
import GHC.Conc.Sync (fromThreadId)
#endif
import System.Mem.Weak (Weak, deRefWeak)
import qualified System.TimeManager as T
data ThreadManager = ThreadManager T.Manager (TVar ManagedThreads)
type Key = Word64
type ManagedThreads = Map Key ManagedThread
data ManagedThread = ManagedThread (Weak ThreadId) (IORef Bool)
newThreadManager :: T.Manager -> IO ThreadManager
newThreadManager :: Manager -> IO ThreadManager
newThreadManager Manager
timmgr = Manager -> TVar ManagedThreads -> ThreadManager
ThreadManager Manager
timmgr (TVar ManagedThreads -> ThreadManager)
-> IO (TVar ManagedThreads) -> IO ThreadManager
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ManagedThreads -> IO (TVar ManagedThreads)
forall a. a -> IO (TVar a)
newTVarIO ManagedThreads
forall k a. Map k a
Map.empty
data KilledByThreadManager = KilledByThreadManager (Maybe SomeException)
deriving (Int -> KilledByThreadManager -> ShowS
[KilledByThreadManager] -> ShowS
KilledByThreadManager -> String
(Int -> KilledByThreadManager -> ShowS)
-> (KilledByThreadManager -> String)
-> ([KilledByThreadManager] -> ShowS)
-> Show KilledByThreadManager
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> KilledByThreadManager -> ShowS
showsPrec :: Int -> KilledByThreadManager -> ShowS
$cshow :: KilledByThreadManager -> String
show :: KilledByThreadManager -> String
$cshowList :: [KilledByThreadManager] -> ShowS
showList :: [KilledByThreadManager] -> ShowS
Show)
instance Exception KilledByThreadManager where
toException :: KilledByThreadManager -> SomeException
toException = KilledByThreadManager -> SomeException
forall e. Exception e => e -> SomeException
E.asyncExceptionToException
fromException :: SomeException -> Maybe KilledByThreadManager
fromException = SomeException -> Maybe KilledByThreadManager
forall e. Exception e => SomeException -> Maybe e
E.asyncExceptionFromException
stopAfter :: ThreadManager -> IO a -> (Maybe SomeException -> IO ()) -> IO a
stopAfter :: forall a.
ThreadManager -> IO a -> (Maybe SomeException -> IO ()) -> IO a
stopAfter ThreadManager
mgr IO a
action Maybe SomeException -> IO ()
cleanup =
ThreadManager -> IO a -> (Either SomeException a -> IO a) -> IO a
forall a b.
ThreadManager -> IO a -> (Either SomeException a -> IO b) -> IO b
stopAfterWithResult ThreadManager
mgr IO a
action ((Either SomeException a -> IO a) -> IO a)
-> (Either SomeException a -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Either SomeException a
mResult ->
case Either SomeException a
mResult of
Left SomeException
err -> do
Maybe SomeException -> IO ()
cleanup (SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just SomeException
err)
SomeException -> IO a
forall e a. Exception e => e -> IO a
E.throwIO SomeException
err
Right a
result -> do
Maybe SomeException -> IO ()
cleanup Maybe SomeException
forall a. Maybe a
Nothing
a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
result
stopAfterWithResult
:: ThreadManager -> IO a -> (Either SomeException a -> IO b) -> IO b
stopAfterWithResult :: forall a b.
ThreadManager -> IO a -> (Either SomeException a -> IO b) -> IO b
stopAfterWithResult (ThreadManager Manager
_timmgr TVar ManagedThreads
var) IO a
action Either SomeException a -> IO b
cleanup = do
((forall a. IO a -> IO a) -> IO b) -> IO b
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
E.mask (((forall a. IO a -> IO a) -> IO b) -> IO b)
-> ((forall a. IO a -> IO a) -> IO b) -> IO b
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
Either SomeException a
ma <- IO a -> IO (Either SomeException a)
forall e a. Exception e => IO a -> IO (Either e a)
E.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
unmask IO a
action
ManagedThreads
m <- STM ManagedThreads -> IO ManagedThreads
forall a. STM a -> IO a
atomically (STM ManagedThreads -> IO ManagedThreads)
-> STM ManagedThreads -> IO ManagedThreads
forall a b. (a -> b) -> a -> b
$ do
ManagedThreads
m0 <- TVar ManagedThreads -> STM ManagedThreads
forall a. TVar a -> STM a
readTVar TVar ManagedThreads
var
TVar ManagedThreads -> ManagedThreads -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ManagedThreads
var ManagedThreads
forall k a. Map k a
Map.empty
ManagedThreads -> STM ManagedThreads
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ManagedThreads
m0
let ths :: [ManagedThread]
ths = ManagedThreads -> [ManagedThread]
forall k a. Map k a -> [a]
Map.elems ManagedThreads
m
er :: Maybe SomeException
er = (SomeException -> Maybe SomeException)
-> (a -> Maybe SomeException)
-> Either SomeException a
-> Maybe SomeException
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just (Maybe SomeException -> a -> Maybe SomeException
forall a b. a -> b -> a
const Maybe SomeException
forall a. Maybe a
Nothing) Either SomeException a
ma
ex :: KilledByThreadManager
ex = Maybe SomeException -> KilledByThreadManager
KilledByThreadManager Maybe SomeException
er
[ManagedThread] -> (ManagedThread -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [ManagedThread]
ths ((ManagedThread -> IO ()) -> IO ())
-> (ManagedThread -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(ManagedThread Weak ThreadId
wtid IORef Bool
ref) -> Weak ThreadId -> IORef Bool -> KilledByThreadManager -> IO ()
forall e. Exception e => Weak ThreadId -> IORef Bool -> e -> IO ()
lockAndKill Weak ThreadId
wtid IORef Bool
ref KilledByThreadManager
ex
Either SomeException a -> IO b
cleanup Either SomeException a
ma
forkManaged
:: ThreadManager
-> String
-> IO ()
-> IO ()
forkManaged :: ThreadManager -> String -> IO () -> IO ()
forkManaged ThreadManager
mgr String
label IO ()
io =
ThreadManager
-> String -> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forkManagedUnmask ThreadManager
mgr String
label (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> IO () -> IO ()
forall a. IO a -> IO a
unmask IO ()
io
forkManagedUnmask
:: ThreadManager
-> String
-> ((forall x. IO x -> IO x) -> IO ())
-> IO ()
forkManagedUnmask :: ThreadManager
-> String -> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forkManagedUnmask (ThreadManager Manager
_timmgr TVar ManagedThreads
var) String
label (forall a. IO a -> IO a) -> IO ()
io =
IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO ThreadId -> IO ThreadId
forall a. IO a -> IO a
E.mask_ (IO ThreadId -> IO ThreadId) -> IO ThreadId -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask (((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId)
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> (KilledByThreadManager -> IO ()) -> IO () -> IO ()
forall e a. Exception e => (e -> IO a) -> IO a -> IO a
E.handle KilledByThreadManager -> IO ()
ignore (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
String -> IO ()
labelMe String
label
IO (Key, Weak ThreadId, IORef Bool)
-> ((Key, Weak ThreadId, IORef Bool) -> IO ())
-> ((Key, Weak ThreadId, IORef Bool) -> IO ())
-> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracket (TVar ManagedThreads -> IO (Key, Weak ThreadId, IORef Bool)
setup TVar ManagedThreads
var) (TVar ManagedThreads -> (Key, Weak ThreadId, IORef Bool) -> IO ()
clear TVar ManagedThreads
var) (((Key, Weak ThreadId, IORef Bool) -> IO ()) -> IO ())
-> ((Key, Weak ThreadId, IORef Bool) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Key, Weak ThreadId, IORef Bool)
_ -> (forall a. IO a -> IO a) -> IO ()
io IO x -> IO x
forall a. IO a -> IO a
unmask
forkManagedTimeout
:: ThreadManager
-> String
-> (T.Handle -> IO ())
-> IO ()
forkManagedTimeout :: ThreadManager -> String -> (Handle -> IO ()) -> IO ()
forkManagedTimeout (ThreadManager Manager
timmgr TVar ManagedThreads
var) String
label Handle -> IO ()
io =
IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ do
String -> IO ()
labelMe String
label
IO (Key, Weak ThreadId, IORef Bool)
-> ((Key, Weak ThreadId, IORef Bool) -> IO ())
-> ((Key, Weak ThreadId, IORef Bool) -> IO ())
-> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracket (TVar ManagedThreads -> IO (Key, Weak ThreadId, IORef Bool)
setup TVar ManagedThreads
var) (TVar ManagedThreads -> (Key, Weak ThreadId, IORef Bool) -> IO ()
clear TVar ManagedThreads
var) (((Key, Weak ThreadId, IORef Bool) -> IO ()) -> IO ())
-> ((Key, Weak ThreadId, IORef Bool) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Key
_n, Weak ThreadId
wtid, IORef Bool
ref) ->
(KilledByThreadManager -> IO ()) -> IO () -> IO ()
forall e a. Exception e => (e -> IO a) -> IO a -> IO a
E.handle KilledByThreadManager -> IO ()
ignore (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Manager -> IO () -> (Handle -> IO ()) -> IO ()
forall a. Manager -> IO () -> (Handle -> IO a) -> IO a
T.withHandle Manager
timmgr (Weak ThreadId -> IORef Bool -> KilledByThreadManager -> IO ()
forall e. Exception e => Weak ThreadId -> IORef Bool -> e -> IO ()
lockAndKill Weak ThreadId
wtid IORef Bool
ref KilledByThreadManager
ex) Handle -> IO ()
io
where
ex :: KilledByThreadManager
ex = Maybe SomeException -> KilledByThreadManager
KilledByThreadManager Maybe SomeException
forall a. Maybe a
Nothing
forkManagedFinally
:: ThreadManager
-> String
-> IO ()
-> IO ()
-> IO ()
forkManagedFinally :: ThreadManager -> String -> IO () -> IO () -> IO ()
forkManagedFinally ThreadManager
mgr String
label IO ()
io IO ()
final =
ThreadManager
-> String -> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forkManagedUnmask ThreadManager
mgr String
label (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore ->
IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
E.try (IO () -> IO ()
forall a. IO a -> IO a
restore IO ()
io) IO (Either SomeException ())
-> (Either SomeException () -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(Either SomeException ()
_ :: Either E.SomeException ()) -> IO ()
final
forkManagedTimeoutFinally
:: ThreadManager
-> String
-> (T.Handle -> IO ())
-> IO ()
-> IO ()
forkManagedTimeoutFinally :: ThreadManager -> String -> (Handle -> IO ()) -> IO () -> IO ()
forkManagedTimeoutFinally ThreadManager
mgr String
label Handle -> IO ()
io IO ()
final = ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
E.mask (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore ->
ThreadManager -> String -> (Handle -> IO ()) -> IO ()
forkManagedTimeout ThreadManager
mgr String
label ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle
th ->
IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
E.try (IO () -> IO ()
forall a. IO a -> IO a
restore (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Handle -> IO ()
io Handle
th) IO (Either SomeException ())
-> (Either SomeException () -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(Either SomeException ()
_ :: Either E.SomeException ()) -> IO ()
final
setup :: TVar (Map Key ManagedThread) -> IO (Key, Weak ThreadId, IORef Bool)
setup :: TVar ManagedThreads -> IO (Key, Weak ThreadId, IORef Bool)
setup TVar ManagedThreads
var = do
(Weak ThreadId
wtid, Key
n) <- IO (Weak ThreadId, Key)
myWeakThreadId
IORef Bool
ref <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
let ent :: ManagedThread
ent = Weak ThreadId -> IORef Bool -> ManagedThread
ManagedThread Weak ThreadId
wtid IORef Bool
ref
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar ManagedThreads -> (ManagedThreads -> ManagedThreads) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar ManagedThreads
var ((ManagedThreads -> ManagedThreads) -> STM ())
-> (ManagedThreads -> ManagedThreads) -> STM ()
forall a b. (a -> b) -> a -> b
$ Key -> ManagedThread -> ManagedThreads -> ManagedThreads
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Key
n ManagedThread
ent
(Key, Weak ThreadId, IORef Bool)
-> IO (Key, Weak ThreadId, IORef Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Key
n, Weak ThreadId
wtid, IORef Bool
ref)
lockAndKill :: Exception e => Weak ThreadId -> IORef Bool -> e -> IO ()
lockAndKill :: forall e. Exception e => Weak ThreadId -> IORef Bool -> e -> IO ()
lockAndKill Weak ThreadId
wtid IORef Bool
ref e
e = do
Bool
alreadyLocked <- IORef Bool -> (Bool -> (Bool, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Bool
ref (\Bool
b -> (Bool
True, Bool
b))
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
alreadyLocked (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Maybe ThreadId
mtid <- Weak ThreadId -> IO (Maybe ThreadId)
forall v. Weak v -> IO (Maybe v)
deRefWeak Weak ThreadId
wtid
case Maybe ThreadId
mtid of
Maybe ThreadId
Nothing -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
Just ThreadId
tid -> ThreadId -> e -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
E.throwTo ThreadId
tid e
e
clear
:: TVar (Map Key ManagedThread)
-> (Key, Weak ThreadId, IORef Bool)
-> IO ()
clear :: TVar ManagedThreads -> (Key, Weak ThreadId, IORef Bool) -> IO ()
clear TVar ManagedThreads
var (Key
n, Weak ThreadId
_, IORef Bool
_) = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar ManagedThreads -> (ManagedThreads -> ManagedThreads) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar ManagedThreads
var ((ManagedThreads -> ManagedThreads) -> STM ())
-> (ManagedThreads -> ManagedThreads) -> STM ()
forall a b. (a -> b) -> a -> b
$ Key -> ManagedThreads -> ManagedThreads
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Key
n
ignore :: KilledByThreadManager -> IO ()
ignore :: KilledByThreadManager -> IO ()
ignore (KilledByThreadManager Maybe SomeException
_) = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
waitUntilAllGone :: ThreadManager -> IO ()
waitUntilAllGone :: ThreadManager -> IO ()
waitUntilAllGone ThreadManager
tm =
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$
ThreadManager -> STM Bool
isAllGone ThreadManager
tm STM Bool -> (Bool -> STM ()) -> STM ()
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
check
isAllGone :: ThreadManager -> STM Bool
isAllGone :: ThreadManager -> STM Bool
isAllGone (ThreadManager Manager
_timmgr TVar ManagedThreads
var) =
ManagedThreads -> Bool
forall k a. Map k a -> Bool
Map.null (ManagedThreads -> Bool) -> STM ManagedThreads -> STM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar ManagedThreads -> STM ManagedThreads
forall a. TVar a -> STM a
readTVar TVar ManagedThreads
var
myWeakThreadId :: IO (Weak ThreadId, Key)
myWeakThreadId :: IO (Weak ThreadId, Key)
myWeakThreadId = do
ThreadId
tid <- IO ThreadId
myThreadId
Weak ThreadId
wtid <- ThreadId -> IO (Weak ThreadId)
mkWeakThreadId ThreadId
tid
let n :: Key
n = ThreadId -> Key
fromThreadId ThreadId
tid
(Weak ThreadId, Key) -> IO (Weak ThreadId, Key)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Weak ThreadId
wtid, Key
n)
labelMe :: String -> IO ()
labelMe :: String -> IO ()
labelMe String
l = do
ThreadId
tid <- IO ThreadId
myThreadId
ThreadId -> String -> IO ()
labelThread ThreadId
tid String
l
withHandle
:: ThreadManager -> T.TimeoutAction -> (T.Handle -> IO a) -> IO a
withHandle :: forall a. ThreadManager -> IO () -> (Handle -> IO a) -> IO a
withHandle (ThreadManager Manager
timmgr TVar ManagedThreads
_) = Manager -> IO () -> (Handle -> IO a) -> IO a
forall a. Manager -> IO () -> (Handle -> IO a) -> IO a
T.withHandle Manager
timmgr
#if __GLASGOW_HASKELL__ < 908
fromThreadId :: ThreadId -> Word64
fromThreadId tid = read (drop 9 $ show tid)
#endif