{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | A thread manager.
--   The manager has responsibility to kill managed threads.
--
-- Because this is based on the accompanying "System.TimeManager" module,
-- the same caveats apply:
--
--   * Only works for GHC.
--   * Only works with a threaded runtime.
--   * Users of older versions should check the current semantics.
--   * Using 32-bit systems means the max timeout is @'maxBound' :: Int@
--     (2147483647) microseconds, which is less than 36 minutes.
--   * Using the same 'Handle' in different threads might cause issues in some
--     edge cases. (i.e. using cancel/pause in one thread, and resume in another)
module System.ThreadManager (
    ThreadManager,
    newThreadManager,
    stopAfter,
    stopAfterWithResult,
    KilledByThreadManager (..),

    -- * Fork
    Reservation,
    reserveManaged,
    takeReservation,
    forkManaged,
    forkManagedFinally,
    forkManagedUnmask,
    forkManagedUnmaskFinally,
    forkManagedBy,
    forkManagedByFinally,
    forkManagedTimeout,
    forkManagedTimeoutFinally,

    -- * Synchronization
    waitUntilAllGone,
    isAllGone,
    countManaged,
) where

import Control.Concurrent
import Control.Concurrent.STM
import Control.Exception (
    AssertionFailed (..),
    Exception (..),
    SomeException (..),
 )
import qualified Control.Exception as E
import Control.Monad (unless, void, when)
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

----------------------------------------------------------------

-- | Manager to own threads.
data ThreadManager = ThreadManager
    { ThreadManager -> TVar ManagedThreads
managedThreads :: TVar ManagedThreads
    , ThreadManager -> TVar Int
reservedPlaces :: TVar Int
    -- ^ Places kept by 'reserveManaged' for threads not started yet.
    , ThreadManager -> TVar Bool
managerStopped :: TVar Bool
    -- ^ Stopped by 'stopAfter'.  A thread which registers after that
    --   is killed on arrival: it was forked before the stop and would
    --   otherwise run on with nobody owning it.
    }

type Key = Word64
type ManagedThreads = Map Key ManagedThread

----------------------------------------------------------------

-- The 'IORef' prevents a race between the timeout of
-- 'forkManagedTimeout' and 'forkManagedTimeoutFinally', which the time
-- manager runs, and 'stopAfter'.  Both throw 'KilledByThreadManager'.
-- It is initialized with 'False' and turned into 'True' when locked.
-- The winner can throw an asynchronous exception.
data ManagedThread = ManagedThread (Weak ThreadId) (IORef Bool)

----------------------------------------------------------------

-- | Create a thread manager.
--
-- A manager owns threads; it does not time them out.  A caller which wants
-- both passes its own 'T.Manager' to 'forkManagedTimeout'.
newThreadManager :: IO ThreadManager
newThreadManager :: IO ThreadManager
newThreadManager = do
    TVar ManagedThreads
managedThreads <- ManagedThreads -> IO (TVar ManagedThreads)
forall a. a -> IO (TVar a)
newTVarIO ManagedThreads
forall k a. Map k a
Map.empty
    TVar Int
reservedPlaces <- Int -> IO (TVar Int)
forall a. a -> IO (TVar a)
newTVarIO Int
0
    TVar Bool
managerStopped <- Bool -> IO (TVar Bool)
forall a. a -> IO (TVar a)
newTVarIO Bool
False
    ThreadManager -> IO ThreadManager
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ThreadManager{TVar Bool
TVar Int
TVar ManagedThreads
managedThreads :: TVar ManagedThreads
reservedPlaces :: TVar Int
managerStopped :: TVar Bool
managedThreads :: TVar ManagedThreads
reservedPlaces :: TVar Int
managerStopped :: TVar Bool
..}

----------------------------------------------------------------

-- | A place kept for a thread that is not running yet.
--
--   Registering a thread needs its 'ThreadId', and it is the thread
--   itself which provides it, by asking for its own.  A caller which
--   forks through a function that hands no identity back cannot do it on
--   the thread's behalf, so until the thread arrives and registers
--   itself, what is counted is the place kept for it.
newtype Reservation = Reservation (IORef Bool)

-- | Keeping a place for a thread about to be forked.  It counts towards
--   'countManaged' from here, so that a manager stopped in between does
--   not see a connection that has been accepted as one that is gone.
--   A place that is never taken is never given back, so a caller which
--   reserves and then fails to fork leaves the manager counting it.
reserveManaged :: ThreadManager -> IO Reservation
reserveManaged :: ThreadManager -> IO Reservation
reserveManaged ThreadManager{TVar Bool
TVar Int
TVar ManagedThreads
managedThreads :: ThreadManager -> TVar ManagedThreads
reservedPlaces :: ThreadManager -> TVar Int
managerStopped :: ThreadManager -> TVar Bool
managedThreads :: TVar ManagedThreads
reservedPlaces :: TVar Int
managerStopped :: TVar Bool
..} = do
    STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
reservedPlaces (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    IORef Bool -> Reservation
Reservation (IORef Bool -> Reservation) -> IO (IORef Bool) -> IO Reservation
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False

-- | Running the action of a thread which was reserved for, taking the
--   place kept for it: the thread records itself, so that 'stopAfter' can
--   reach it, and gives the reservation back.
--
--   The reservation is given back only once the thread is recorded, so
--   that it is counted throughout.
--   A thread which gets here after 'stopAfter' has stopped the manager is
--   killed rather than registered.
--
--   The reservation must be one kept by this manager.  Taking one kept by
--   another lowers this manager's count instead of that one's: the other
--   manager never reaches zero, and once this count goes negative,
--   'countManaged', 'isAllGone' and 'waitUntilAllGone' throw 'AssertionFailed'.
takeReservation :: ThreadManager -> Reservation -> IO () -> IO ()
takeReservation :: ThreadManager -> Reservation -> IO () -> IO ()
takeReservation ThreadManager
mgr Reservation
rsv IO ()
io = ThreadManager
-> Reservation -> IO () -> ((Key, ManagedThread) -> IO ()) -> IO ()
taking ThreadManager
mgr Reservation
rsv (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) (((Key, ManagedThread) -> IO ()) -> IO ())
-> ((Key, ManagedThread) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Key, ManagedThread)
_ -> IO ()
io

-- | Giving a place back, at most once however often this is called: the
--   thread which takes it gives it back, and so does a caller whose fork
--   throws, and only one of them can have it.
release :: ThreadManager -> Reservation -> IO ()
release :: ThreadManager -> Reservation -> IO ()
release ThreadManager{TVar Bool
TVar Int
TVar ManagedThreads
managedThreads :: ThreadManager -> TVar ManagedThreads
reservedPlaces :: ThreadManager -> TVar Int
managerStopped :: ThreadManager -> TVar Bool
managedThreads :: TVar ManagedThreads
reservedPlaces :: TVar Int
managerStopped :: TVar Bool
..} (Reservation IORef Bool
ref) = do
    Bool
taken <- 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
taken (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
reservedPlaces (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)

----------------------------------------------------------------

-- | An exception used internally to kill a managed thread.
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

-- | Stopping the manager.
--
-- @
-- stopAfter threadManager action cleanup
-- @
--
-- The action is run in the scope of an exception handler that catches all
-- exceptions (including asynchronous ones); this allows the cleanup handler
-- to cleanup in all circumstances. If an exception is caught, it is rethrown
-- after the cleanup is complete.
--
-- This tells the managed threads to stop; it does not wait for them to be
-- gone, and neither does the cleanup handler.  They are the manager\'s
-- until they end, so 'countManaged' goes on counting them -- which is what
-- a caller waits on with 'waitUntilAllGone' before it frees what those
-- threads were using.
--
-- A thread whose place is still being kept by 'reserveManaged' is not told
-- but refused when it arrives, and never runs.
--
-- A thread which swallows the exception and never ends is therefore
-- counted for ever, and 'waitUntilAllGone' on such a manager never
-- returns.  Nothing here can make it end; a caller which cannot rule that
-- out should put its own timeout around the wait.
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

-- | Generalization of 'stopAfter' where the cleanup handler is allowed to construct the final result
--
-- Unlike in 'stopAfter', if an exception is thrown, it is not re-thrown after the cleanup
-- handler completes; the cleanup handler itself can decide to rethrow it or compute a result.
--
-- @since 0.3.2
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{TVar Bool
TVar Int
TVar ManagedThreads
managedThreads :: ThreadManager -> TVar ManagedThreads
reservedPlaces :: ThreadManager -> TVar Int
managerStopped :: ThreadManager -> TVar Bool
managedThreads :: TVar ManagedThreads
reservedPlaces :: TVar Int
managerStopped :: TVar Bool
..} 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
        -- The threads are left in the map rather than taken out of it:
        -- 'throwTo' waits for the exception to be taken, not for the
        -- thread to be gone, and a manager which reports that it owns
        -- nothing while its threads are still unwinding is no use to a
        -- caller which frees what they were using.  They should take
        -- themselves out as they end, and 'managerStopped' is what keeps
        -- new ones from arriving.
        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
            TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Bool
managerStopped Bool
True
            TVar ManagedThreads -> STM ManagedThreads
forall a. TVar a -> STM a
readTVar TVar ManagedThreads
managedThreads
        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

----------------------------------------------------------------

-- | Fork a managed thread.
--
-- A place is kept for the thread before it is forked, so that it counts
-- towards 'countManaged' from this call rather than from the moment the
-- run-time system gives the thread its first step.  The thread adds
-- itself to the manager's queue as the first thing it does, takes the
-- place over, and is removed again when it terminates (normally or
-- abnormally).
forkManaged
    :: ThreadManager
    -> String
    -- ^ Thread name
    -> IO ()
    -- ^ Action
    -> IO ()
forkManaged :: ThreadManager -> String -> IO () -> IO ()
forkManaged ThreadManager
mgr String
label IO ()
io = ThreadManager -> String -> IO () -> IO () -> IO ()
forkManagedFinally ThreadManager
mgr String
label IO ()
io (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Like 'forkManaged', but with a cleanup function, run regardless of how
--   the action ends.  An exception from the action is not caught: it ends
--   the thread as it would under 'forkManaged'.
forkManagedFinally
    :: ThreadManager
    -> String
    -- ^ Thread name
    -> IO ()
    -- ^ Action
    -> IO ()
    -- ^ Cleanup function
    -> 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 () -> IO ()
forkManagedUnmaskFinally ThreadManager
mgr String
label (\forall a. IO a -> IO a
restore -> IO () -> IO ()
forall a. IO a -> IO a
restore IO ()
io) IO ()
final

-- | Like 'forkManaged', but run the action with exceptions masked
--   until the provided `unmask` is used. (i.e. @(forall x. IO x -> IO x)@)
forkManagedUnmask
    :: ThreadManager
    -> String
    -- ^ Thread name
    -> ((forall x. IO x -> IO x) -> IO ())
    -- ^ Action with unmask argument
    -> IO ()
forkManagedUnmask :: ThreadManager
-> String -> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forkManagedUnmask ThreadManager
mgr String
label (forall a. IO a -> IO a) -> IO ()
io = ThreadManager
-> String -> ((forall a. IO a -> IO a) -> IO ()) -> IO () -> IO ()
forkManagedUnmaskFinally ThreadManager
mgr String
label (forall a. IO a -> IO a) -> IO ()
io (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Like 'forkManagedUnmask', but with a cleanup function, run regardless of
--   how the action ends.  An exception from the action is not caught: it
--   ends the thread as it would under 'forkManagedUnmask'.
forkManagedUnmaskFinally
    :: ThreadManager
    -> String
    -- ^ Thread name
    -> ((forall x. IO x -> IO x) -> IO ())
    -- ^ Action with unmask argument
    -> IO ()
    -- ^ Cleanup: will always run, regardless of how the thread ends
    -> IO ()
forkManagedUnmaskFinally :: ThreadManager
-> String -> ((forall a. IO a -> IO a) -> IO ()) -> IO () -> IO ()
forkManagedUnmaskFinally ThreadManager
mgr String
label (forall a. IO a -> IO a) -> IO ()
io IO ()
final =
    IO () -> IO ()
forall a. IO a -> IO a
E.mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        ThreadManager
-> (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> String
-> ((forall a. IO a -> IO a) -> IO ())
-> IO ()
-> IO ()
forall a.
ThreadManager
-> (((forall a. IO a -> IO a) -> IO ()) -> IO a)
-> String
-> ((forall a. IO a -> IO a) -> IO ())
-> IO ()
-> IO a
forkManagedByFinally ThreadManager
mgr (\(forall a. IO a -> IO a) -> IO ()
body -> 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
$ ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask (forall a. IO a -> IO a) -> IO ()
body) String
label (forall a. IO a -> IO a) -> IO ()
io IO ()
final

-- | Fork a managed thread with a fork of the caller's own, for a caller
--   which does not fork with 'forkIO'.  (The @warp@ library, for example,
--   forks the thread serving a connection through a function of its
--   settings, so that the user can replace it.)
--
--   A place is kept for the thread before the fork, so that it counts
--   towards 'countManaged' from here rather than from the moment the
--   thread starts, and the place is given back whether the fork makes the
--   thread or throws.  A caller which forks by itself can do the same
--   with 'reserveManaged' and 'takeReservation', which this is made of.
forkManagedBy
    :: ThreadManager
    -> (((forall x. IO x -> IO x) -> IO ()) -> IO a)
    -- ^ The fork to make the thread with
    -> String
    -- ^ Thread name
    -> ((forall x. IO x -> IO x) -> IO ())
    -- ^ Action with unmask argument
    -> IO a
forkManagedBy :: forall a.
ThreadManager
-> (((forall a. IO a -> IO a) -> IO ()) -> IO a)
-> String
-> ((forall a. IO a -> IO a) -> IO ())
-> IO a
forkManagedBy ThreadManager
mgr ((forall a. IO a -> IO a) -> IO ()) -> IO a
fork String
label (forall a. IO a -> IO a) -> IO ()
io = ThreadManager
-> (((forall a. IO a -> IO a) -> IO ()) -> IO a)
-> String
-> ((forall a. IO a -> IO a) -> IO ())
-> IO ()
-> IO a
forall a.
ThreadManager
-> (((forall a. IO a -> IO a) -> IO ()) -> IO a)
-> String
-> ((forall a. IO a -> IO a) -> IO ())
-> IO ()
-> IO a
forkManagedByFinally ThreadManager
mgr ((forall a. IO a -> IO a) -> IO ()) -> IO a
fork String
label (forall a. IO a -> IO a) -> IO ()
io (IO () -> IO a) -> IO () -> IO a
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Like 'forkManagedBy', but with a cleanup function.
forkManagedByFinally
    :: ThreadManager
    -> (((forall x. IO x -> IO x) -> IO ()) -> IO a)
    -- ^ The fork to make the thread with
    -> String
    -- ^ Thread name
    -> ((forall x. IO x -> IO x) -> IO ())
    -- ^ Action with unmask argument
    -> IO ()
    -- ^ Cleanup: will always run, regardless of how the thread ends
    -> IO a
forkManagedByFinally :: forall a.
ThreadManager
-> (((forall a. IO a -> IO a) -> IO ()) -> IO a)
-> String
-> ((forall a. IO a -> IO a) -> IO ())
-> IO ()
-> IO a
forkManagedByFinally ThreadManager
mgr ((forall a. IO a -> IO a) -> IO ()) -> IO a
fork String
label (forall a. IO a -> IO a) -> IO ()
io IO ()
final =
    IO Reservation
-> (Reservation -> IO ()) -> (Reservation -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracketOnError (ThreadManager -> IO Reservation
reserveManaged ThreadManager
mgr) (ThreadManager -> Reservation -> IO ()
release ThreadManager
mgr) ((Reservation -> IO a) -> IO a) -> (Reservation -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Reservation
rsv ->
        ((forall a. IO a -> IO a) -> IO ()) -> IO a
fork (((forall a. IO a -> IO a) -> IO ()) -> IO a)
-> ((forall a. IO a -> IO a) -> IO ()) -> IO a
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> ThreadManager
-> Reservation
-> String
-> IO ()
-> ((Key, ManagedThread) -> IO ())
-> IO ()
managed ThreadManager
mgr Reservation
rsv String
label IO ()
final (((Key, ManagedThread) -> IO ()) -> IO ())
-> ((Key, ManagedThread) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Key, ManagedThread)
_ -> (forall a. IO a -> IO a) -> IO ()
io IO x -> IO x
forall a. IO a -> IO a
unmask

-- | Fork a managed thread with a handle created by a timeout manager.
--
--   The time manager is the caller's, since this manager has none: owning a
--   thread and timing it out are separate things, and only this pair of
--   functions needs both.  What needs them together is the lock that keeps
--   the timeout and 'stopAfter' from both throwing at the thread.
forkManagedTimeout
    :: ThreadManager
    -> T.Manager
    -- ^ Time manager
    -> String
    -- ^ Thread name
    -> (T.Handle -> IO ())
    -- ^ Action with timeout handle
    -> IO ()
forkManagedTimeout :: ThreadManager -> Manager -> String -> (Handle -> IO ()) -> IO ()
forkManagedTimeout ThreadManager
mgr Manager
timmgr String
label Handle -> IO ()
io =
    ThreadManager
-> Manager -> String -> (Handle -> IO ()) -> IO () -> IO ()
forkManagedTimeoutFinally ThreadManager
mgr Manager
timmgr String
label Handle -> IO ()
io (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Like 'forkManagedTimeout', but with a cleanup function, run regardless of
--   how the action ends.  An exception from the action is not caught: it
--   ends the thread as it would under 'forkManagedTimeout'.
forkManagedTimeoutFinally
    :: ThreadManager
    -> T.Manager
    -- ^ Time manager
    -> String
    -- ^ Thread name
    -> (T.Handle -> IO ())
    -- ^ Action with timeout handle
    -> IO ()
    -- ^ Cleanup: will always run, regardless of how the thread ends
    -> IO ()
forkManagedTimeoutFinally :: ThreadManager
-> Manager -> String -> (Handle -> IO ()) -> IO () -> IO ()
forkManagedTimeoutFinally ThreadManager
mgr Manager
timmgr String
label Handle -> IO ()
io IO ()
final =
    IO Reservation
-> (Reservation -> IO ()) -> (Reservation -> IO ()) -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracketOnError (ThreadManager -> IO Reservation
reserveManaged ThreadManager
mgr) (ThreadManager -> Reservation -> IO ()
release ThreadManager
mgr) ((Reservation -> IO ()) -> IO ())
-> (Reservation -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Reservation
rsv ->
        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
$
            ThreadManager
-> Reservation
-> String
-> IO ()
-> ((Key, ManagedThread) -> IO ())
-> IO ()
managed ThreadManager
mgr Reservation
rsv String
label IO ()
final (((Key, ManagedThread) -> IO ()) -> IO ())
-> ((Key, ManagedThread) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Key
_n, ManagedThread Weak ThreadId
wtid IORef Bool
ref) ->
                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

-- | The body of a managed thread, under its own label, taking the place
--   kept for it.
managed
    :: ThreadManager
    -> Reservation
    -> String
    -> IO ()
    -> ((Key, ManagedThread) -> IO ())
    -> IO ()
managed :: ThreadManager
-> Reservation
-> String
-> IO ()
-> ((Key, ManagedThread) -> IO ())
-> IO ()
managed ThreadManager
mgr Reservation
rsv String
label IO ()
final (Key, ManagedThread) -> IO ()
body = do
    String -> IO ()
labelMe String
label
    ThreadManager
-> Reservation -> IO () -> ((Key, ManagedThread) -> IO ()) -> IO ()
taking ThreadManager
mgr Reservation
rsv IO ()
final (Key, ManagedThread) -> IO ()
body

-- | Running the body of a thread which had a place kept for it.  The
--   place is given back only once the thread is recorded, so that it is
--   counted throughout.
taking
    :: ThreadManager
    -> Reservation
    -> IO ()
    -> ((Key, ManagedThread) -> IO ())
    -> IO ()
taking :: ThreadManager
-> Reservation -> IO () -> ((Key, ManagedThread) -> IO ()) -> IO ()
taking ThreadManager
mgr Reservation
rsv IO ()
final (Key, ManagedThread) -> IO ()
body =
    ThreadManager -> IO () -> ((Key, ManagedThread) -> IO ()) -> IO ()
recorded ThreadManager
mgr (IO ()
give IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
final) (((Key, ManagedThread) -> IO ()) -> IO ())
-> ((Key, ManagedThread) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Key, ManagedThread)
key -> IO ()
give IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Key, ManagedThread) -> IO ()
body (Key, ManagedThread)
key
  where
    give :: IO ()
give = ThreadManager -> Reservation -> IO ()
release ThreadManager
mgr Reservation
rsv

-- | Running the body of a managed thread with the thread recorded.
--
--   The cleanup runs regardless of how the body ends, and while the thread is still
--   recorded, so that a manager which has none left has none whose cleanup
--   is still running either.
--
--   A thread forked before the manager was stopped and scheduled after it
--   is not recorded and does not run its body: there is nobody left to
--   kill it, and running it is what the stop was for.
recorded
    :: ThreadManager
    -> IO ()
    -> ((Key, ManagedThread) -> IO ())
    -> IO ()
recorded :: ThreadManager -> IO () -> ((Key, ManagedThread) -> IO ()) -> IO ()
recorded ThreadManager
mgr IO ()
final (Key, ManagedThread) -> IO ()
body =
    (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
$
        IO (Maybe (Key, ManagedThread))
-> (Maybe (Key, ManagedThread) -> IO ())
-> (Maybe (Key, ManagedThread) -> IO ())
-> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracket (ThreadManager -> IO (Maybe (Key, ManagedThread))
setup ThreadManager
mgr) (((Key, ManagedThread) -> IO ())
-> Maybe (Key, ManagedThread) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (((Key, ManagedThread) -> IO ())
 -> Maybe (Key, ManagedThread) -> IO ())
-> ((Key, ManagedThread) -> IO ())
-> Maybe (Key, ManagedThread)
-> IO ()
forall a b. (a -> b) -> a -> b
$ ThreadManager -> (Key, ManagedThread) -> IO ()
clear ThreadManager
mgr) ((Maybe (Key, ManagedThread) -> IO ()) -> IO ())
-> (Maybe (Key, ManagedThread) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Maybe (Key, ManagedThread)
mkey ->
            ((Key, ManagedThread) -> IO ())
-> Maybe (Key, ManagedThread) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Key, ManagedThread) -> IO ()
body Maybe (Key, ManagedThread)
mkey IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`E.finally` IO ()
final

-- | Recording this thread, or 'Nothing' once the manager has stopped.
setup
    :: ThreadManager
    -> IO (Maybe (Key, ManagedThread))
setup :: ThreadManager -> IO (Maybe (Key, ManagedThread))
setup ThreadManager{TVar Bool
TVar Int
TVar ManagedThreads
managedThreads :: ThreadManager -> TVar ManagedThreads
reservedPlaces :: ThreadManager -> TVar Int
managerStopped :: ThreadManager -> TVar Bool
managedThreads :: TVar ManagedThreads
reservedPlaces :: TVar Int
managerStopped :: TVar Bool
..} = 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
    -- asking to throw KilledByThreadManager to me
    STM (Maybe (Key, ManagedThread)) -> IO (Maybe (Key, ManagedThread))
forall a. STM a -> IO a
atomically (STM (Maybe (Key, ManagedThread))
 -> IO (Maybe (Key, ManagedThread)))
-> STM (Maybe (Key, ManagedThread))
-> IO (Maybe (Key, ManagedThread))
forall a b. (a -> b) -> a -> b
$ do
        Bool
stopped <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
managerStopped
        if Bool
stopped
            then Maybe (Key, ManagedThread) -> STM (Maybe (Key, ManagedThread))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Key, ManagedThread)
forall a. Maybe a
Nothing
            else do
                TVar ManagedThreads -> (ManagedThreads -> ManagedThreads) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar ManagedThreads
managedThreads ((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
                Maybe (Key, ManagedThread) -> STM (Maybe (Key, ManagedThread))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Key, ManagedThread) -> STM (Maybe (Key, ManagedThread)))
-> Maybe (Key, ManagedThread) -> STM (Maybe (Key, ManagedThread))
forall a b. (a -> b) -> a -> b
$ (Key, ManagedThread) -> Maybe (Key, ManagedThread)
forall a. a -> Maybe a
Just (Key
n, ManagedThread
ent)

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)) -- try to lock
    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
    :: ThreadManager
    -> (Key, ManagedThread)
    -> IO ()
clear :: ThreadManager -> (Key, ManagedThread) -> IO ()
clear ThreadManager{TVar Bool
TVar Int
TVar ManagedThreads
managedThreads :: ThreadManager -> TVar ManagedThreads
reservedPlaces :: ThreadManager -> TVar Int
managerStopped :: ThreadManager -> TVar Bool
managedThreads :: TVar ManagedThreads
reservedPlaces :: TVar Int
managerStopped :: TVar Bool
..} (Key
n, ManagedThread
_) =
    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
managedThreads ((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 ()

-- | Wait until all managed threads are finished.
--
--   After 'stopAfter' this waits for the threads it stopped to end, not
--   merely for them to have been told.  See there for the thread that
--   never ends.
--
--   A place kept by 'reserveManaged' and never taken is never given back,
--   so this never returns for a caller which reserves and then fails to
--   fork.
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

-- | STM action that checks if all managed threads are finished.
isAllGone :: ThreadManager -> STM Bool
isAllGone :: ThreadManager -> STM Bool
isAllGone ThreadManager
tm = (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (Int -> Bool) -> STM Int -> STM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ThreadManager -> STM Int
countManaged ThreadManager
tm

-- | STM action for how many threads are managed, counting the places
--   kept by 'reserveManaged' for threads which have not started, and the
--   threads 'stopAfter' has stopped which have not ended yet.
--
--   Throws 'AssertionFailed' if the count of places has gone negative,
--   which happens when a reservation is taken by a manager other than the
--   one that kept it.  A mistake of that kind which does not drive the
--   count negative goes unnoticed.
countManaged :: ThreadManager -> STM Int
countManaged :: ThreadManager -> STM Int
countManaged ThreadManager{TVar Bool
TVar Int
TVar ManagedThreads
managedThreads :: ThreadManager -> TVar ManagedThreads
reservedPlaces :: ThreadManager -> TVar Int
managerStopped :: ThreadManager -> TVar Bool
managedThreads :: TVar ManagedThreads
reservedPlaces :: TVar Int
managerStopped :: TVar Bool
..} = do
    Int
places <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
reservedPlaces
    Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
places Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (STM () -> STM ()) -> STM () -> STM ()
forall a b. (a -> b) -> a -> b
$
        AssertionFailed -> STM ()
forall e a. Exception e => e -> STM a
throwSTM (AssertionFailed -> STM ()) -> AssertionFailed -> STM ()
forall a b. (a -> b) -> a -> b
$
            String -> AssertionFailed
AssertionFailed (String -> AssertionFailed) -> String -> AssertionFailed
forall a b. (a -> b) -> a -> b
$
                String
"System.ThreadManager: "
                    String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
places
                    String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" reserved places; a reservation was taken by a"
                    String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" manager other than its own"
    (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
places) (Int -> Int) -> (ManagedThreads -> Int) -> ManagedThreads -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ManagedThreads -> Int
forall k a. Map k a -> Int
Map.size (ManagedThreads -> Int) -> STM ManagedThreads -> STM Int
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
managedThreads

----------------------------------------------------------------

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

#if __GLASGOW_HASKELL__ < 908
fromThreadId :: ThreadId -> Word64
fromThreadId tid = read (drop 9 $ show tid)
#endif