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

    -- | A thread manager including a time manager.
    --   The manager has responsibility to kill managed threads.
    module System.ThreadManager (
        ThreadManager,
        newThreadManager,
        stopAfter,
        KilledByThreadManager (..),

        -- * Fork
        forkManaged,
        forkManagedFinally,
        forkManagedUnmask,
        forkManagedTimeout,
        forkManagedTimeoutFinally,

        -- * Synchronization
        waitUntilAllGone,
        isAllGone,

        -- * Re-exports
        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
    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 manage the thread and the timer.
    data ThreadManager = ThreadManager T.Manager (TVar ManagedThreads)

    type Key = Word64
    type ManagedThreads = Map Key ManagedThread

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

    -- 'IORef' prevents race between WAI TimeManager (TimeoutThread)
    -- and stopAfter (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)

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

    -- | Starting a thread manager.
    --   Its action is initially set to 'return ()' and should be set
    --   by 'setAction'. This allows that the action can include
    --   the manager itself.
    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

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

    -- | 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.
    --
    -- 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.
    stopAfter :: ThreadManager -> IO a -> (Maybe SomeException -> IO ()) -> IO a
    stopAfter :: forall a.
    ThreadManager -> IO a -> (Maybe SomeException -> IO ()) -> IO a
    stopAfter (ThreadManager Manager
    _timmgr TVar ManagedThreads
    var) IO a
    action Maybe SomeException -> IO ()
    cleanup = do
        ((forall a. IO a -> IO a) -> IO a) -> IO a
    forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
    E.mask (((forall a. IO a -> IO a) -> IO a) -> IO a)
    -> ((forall a. IO a -> IO a) -> IO a) -> IO a
    forall a b. (a -> b) -> a -> b
    $ \forall a. IO a -> IO a
    unmask -> do
            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
            m <- atomically $ do
                m0 <- readTVar var
                writeTVar var Map.empty
                return m0
            let ths = ManagedThreads -> [ManagedThread]
    forall k a. Map k a -> [a]
    Map.elems ManagedThreads
    m
                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 = Maybe SomeException -> KilledByThreadManager
    KilledByThreadManager Maybe SomeException
    er
            forM_ ths $ \(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
            case ma of
                Left SomeException
    err -> Maybe SomeException -> IO ()
    cleanup (SomeException -> Maybe SomeException
    forall a. a -> Maybe a
    Just SomeException
    err) 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
    E.throwIO SomeException
    err
                Right a
    a -> Maybe SomeException -> IO ()
    cleanup Maybe SomeException
    forall a. Maybe a
    Nothing 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
    >> a -> IO a
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return a
    a

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

    -- | Fork a managed thread.
    --
    -- This guarantees that the thread ID is added to the manager's queue before
    -- the thread starts, and is removed again when the thread 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 -> ((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

    -- | Like 'forkManaged', but run action with exceptions masked
    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

    -- | Fork a managed thread with a handle created by a timeout manager.
    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
    $ (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
    _n, Weak ThreadId
    wtid, IORef Bool
    ref) ->
                -- 'TimeoutThread' is ignored by 'withHandle'.
                IO (Maybe ()) -> IO ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
    void (IO (Maybe ()) -> IO ()) -> IO (Maybe ()) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ Manager -> IO () -> (Handle -> IO ()) -> IO (Maybe ())
    forall a. Manager -> IO () -> (Handle -> IO a) -> IO (Maybe a)
    T.withHandle Manager
    timmgr (Weak ThreadId -> IORef Bool -> TimeoutThread -> IO ()
    forall e. Exception e => Weak ThreadId -> IORef Bool -> e -> IO ()
    lockAndKill Weak ThreadId
    wtid IORef Bool
    ref TimeoutThread
    T.TimeoutThread) Handle -> IO ()
    io

    -- | Fork a managed thread with a cleanup function.
    forkManagedFinally :: ThreadManager -> String -> IO () -> IO () -> IO ()
    forkManagedFinally :: ThreadManager -> String -> IO () -> IO () -> IO ()
    forkManagedFinally ThreadManager
    mgr String
    label 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 -> IO () -> IO ()
    forkManaged
            ThreadManager
    mgr
            String
    label
            (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)

    -- | Fork a managed thread with a handle created by a timeout manager
    -- and with a cleanup function.
    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
    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
        (wtid, n) <- IO (Weak ThreadId, Key)
    myWeakThradId
        ref <- newIORef False
        let ent = Weak ThreadId -> IORef Bool -> ManagedThread
    ManagedThread Weak ThreadId
    wtid IORef Bool
    ref
        -- asking to throw KilledByThreadManager to me
        atomically $ modifyTVar' var $ Map.insert n ent
        return (n, wtid, 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
        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
        unless alreadyLocked $ do
            mtid <- deRefWeak wtid
            case 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 ()

    -- | Wait until all managed thread are finished.
    waitUntilAllGone :: ThreadManager -> IO ()
    waitUntilAllGone :: ThreadManager -> IO ()
    waitUntilAllGone (ThreadManager Manager
    _timmgr TVar ManagedThreads
    var) = STM () -> IO ()
    forall a. STM a -> IO a
    atomically (STM () -> IO ()) -> STM () -> IO ()
    forall a b. (a -> b) -> a -> b
    $ do
        m <- TVar ManagedThreads -> STM ManagedThreads
    forall a. TVar a -> STM a
    readTVar TVar ManagedThreads
    var
        check (Map.size m == 0)

    isAllGone :: ThreadManager -> STM Bool
    isAllGone :: ThreadManager -> STM Bool
    isAllGone (ThreadManager Manager
    _timmgr TVar ManagedThreads
    var) = do
        m <- TVar ManagedThreads -> STM ManagedThreads
    forall a. TVar a -> STM a
    readTVar TVar ManagedThreads
    var
        return (Map.size m == 0)

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

    myWeakThradId :: IO (Weak ThreadId, Key)
    myWeakThradId :: IO (Weak ThreadId, Key)
    myWeakThradId = do
        tid <- IO ThreadId
    myThreadId
        wtid <- mkWeakThreadId tid
        let n = ThreadId -> Key
    fromThreadId ThreadId
    tid
        return (wtid, n)

    labelMe :: String -> IO ()
    labelMe :: String -> IO ()
    labelMe String
    l = do
        tid <- IO ThreadId
    myThreadId
        labelThread tid l

    withHandle
        :: ThreadManager -> T.TimeoutAction -> (T.Handle -> IO a) -> IO (Maybe a)
    withHandle :: forall a.
    ThreadManager -> IO () -> (Handle -> IO a) -> IO (Maybe a)
    withHandle (ThreadManager Manager
    timmgr TVar ManagedThreads
    _) = Manager -> IO () -> (Handle -> IO a) -> IO (Maybe a)
    forall a. Manager -> IO () -> (Handle -> IO a) -> IO (Maybe a)
    T.withHandle Manager
    timmgr

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