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

-- | A thread manager.
--   The manager has responsibility to spawn and kill
--   worker threads.
module Network.HTTP2.H2.Manager (
    Manager,
    Action,
    start,
    setAction,
    stopAfter,
    spawnAction,
    forkManaged,
    forkManagedUnmask,
    timeoutKillThread,
    timeoutClose,
    KilledByHttp2ThreadManager (..),
    incCounter,
    decCounter,
    waitCounter0,
) where

import Data.Foldable
import Data.IORef
import Data.Set (Set)
import qualified Data.Set as Set
import qualified System.TimeManager as T
import UnliftIO.Concurrent
import UnliftIO.Exception
import qualified UnliftIO.Exception as E
import UnliftIO.STM

import Imports

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

-- | Action to be spawned by the manager.
type Action = IO ()

noAction :: Action
noAction :: Action
noAction = () -> Action
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

data Command = Stop (Maybe SomeException) | Spawn | Add ThreadId | Delete ThreadId

-- | Manager to manage the thread and the timer.
data Manager = Manager (TQueue Command) (IORef Action) (TVar Int) T.Manager

-- | 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.
start :: T.Manager -> IO Manager
start :: Manager -> IO Manager
start Manager
timmgr = do
    TQueue Command
q <- IO (TQueue Command)
forall (m :: * -> *) a. MonadIO m => m (TQueue a)
newTQueueIO
    IORef Action
ref <- Action -> IO (IORef Action)
forall a. a -> IO (IORef a)
newIORef Action
noAction
    TVar Int
cnt <- Int -> IO (TVar Int)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Int
0
    IO ThreadId -> Action
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> Action) -> IO ThreadId -> Action
forall a b. (a -> b) -> a -> b
$ Action -> IO ThreadId
forall (m :: * -> *). MonadUnliftIO m => m () -> m ThreadId
forkIO (Action -> IO ThreadId) -> Action -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ TQueue Command -> Set ThreadId -> IORef Action -> Action
forall {a}.
TQueue Command -> Set ThreadId -> IORef (IO a) -> Action
go TQueue Command
q Set ThreadId
forall a. Set a
Set.empty IORef Action
ref
    Manager -> IO Manager
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Manager -> IO Manager) -> Manager -> IO Manager
forall a b. (a -> b) -> a -> b
$ TQueue Command -> IORef Action -> TVar Int -> Manager -> Manager
Manager TQueue Command
q IORef Action
ref TVar Int
cnt Manager
timmgr
  where
    go :: TQueue Command -> Set ThreadId -> IORef (IO a) -> Action
go TQueue Command
q Set ThreadId
tset0 IORef (IO a)
ref = do
        Command
x <- STM Command -> IO Command
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM Command -> IO Command) -> STM Command -> IO Command
forall a b. (a -> b) -> a -> b
$ TQueue Command -> STM Command
forall a. TQueue a -> STM a
readTQueue TQueue Command
q
        case Command
x of
            Stop Maybe SomeException
err -> Set ThreadId -> Maybe SomeException -> Action
kill Set ThreadId
tset0 Maybe SomeException
err
            Command
Spawn -> Set ThreadId -> Action
next Set ThreadId
tset0
            Add ThreadId
newtid ->
                let tset :: Set ThreadId
tset = ThreadId -> Set ThreadId -> Set ThreadId
add ThreadId
newtid Set ThreadId
tset0
                 in TQueue Command -> Set ThreadId -> IORef (IO a) -> Action
go TQueue Command
q Set ThreadId
tset IORef (IO a)
ref
            Delete ThreadId
oldtid ->
                let tset :: Set ThreadId
tset = ThreadId -> Set ThreadId -> Set ThreadId
del ThreadId
oldtid Set ThreadId
tset0
                 in TQueue Command -> Set ThreadId -> IORef (IO a) -> Action
go TQueue Command
q Set ThreadId
tset IORef (IO a)
ref
      where
        next :: Set ThreadId -> Action
next Set ThreadId
tset = do
            IO a
action <- IORef (IO a) -> IO (IO a)
forall a. IORef a -> IO a
readIORef IORef (IO a)
ref
            ThreadId
newtid <- IO a -> (Either SomeException a -> Action) -> IO ThreadId
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> (Either SomeException a -> m ()) -> m ThreadId
forkFinally IO a
action ((Either SomeException a -> Action) -> IO ThreadId)
-> (Either SomeException a -> Action) -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \Either SomeException a
_ -> do
                ThreadId
mytid <- IO ThreadId
forall (m :: * -> *). MonadIO m => m ThreadId
myThreadId
                STM () -> Action
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Action) -> STM () -> Action
forall a b. (a -> b) -> a -> b
$ TQueue Command -> Command -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Command
q (Command -> STM ()) -> Command -> STM ()
forall a b. (a -> b) -> a -> b
$ ThreadId -> Command
Delete ThreadId
mytid
            let tset' :: Set ThreadId
tset' = ThreadId -> Set ThreadId -> Set ThreadId
add ThreadId
newtid Set ThreadId
tset
            TQueue Command -> Set ThreadId -> IORef (IO a) -> Action
go TQueue Command
q Set ThreadId
tset' IORef (IO a)
ref

-- | Setting the action to be spawned.
setAction :: Manager -> Action -> IO ()
setAction :: Manager -> Action -> Action
setAction (Manager TQueue Command
_ IORef Action
ref TVar Int
_ Manager
_) Action
action = IORef Action -> Action -> Action
forall a. IORef a -> a -> Action
writeIORef IORef Action
ref Action
action

-- | Stopping the manager.
stopAfter :: Manager -> IO a -> (Either SomeException a -> IO b) -> IO b
stopAfter :: forall a b.
Manager -> IO a -> (Either SomeException a -> IO b) -> IO b
stopAfter (Manager TQueue Command
q IORef Action
_ TVar Int
_ Manager
_) IO a
action Either SomeException a -> IO b
cleanup = do
    ((forall a. IO a -> IO a) -> IO b) -> IO b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> m a) -> m b) -> m b
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 (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> m (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
unmask IO a
action
        STM () -> Action
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Action) -> STM () -> Action
forall a b. (a -> b) -> a -> b
$ TQueue Command -> Command -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Command
q (Command -> STM ()) -> Command -> STM ()
forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> Command
Stop ((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)
        Either SomeException a -> IO b
cleanup Either SomeException a
ma

-- | Spawning the action.
spawnAction :: Manager -> IO ()
spawnAction :: Manager -> Action
spawnAction (Manager TQueue Command
q IORef Action
_ TVar Int
_ Manager
_) = STM () -> Action
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Action) -> STM () -> Action
forall a b. (a -> b) -> a -> b
$ TQueue Command -> Command -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Command
q Command
Spawn

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

-- | Fork 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 :: Manager -> IO () -> IO ()
forkManaged :: Manager -> Action -> Action
forkManaged Manager
mgr Action
io =
    Manager -> ((forall a. IO a -> IO a) -> Action) -> Action
forkManagedUnmask Manager
mgr (((forall a. IO a -> IO a) -> Action) -> Action)
-> ((forall a. IO a -> IO a) -> Action) -> Action
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> Action -> Action
forall a. IO a -> IO a
unmask Action
io

-- | Like 'forkManaged', but run action with exceptions masked
forkManagedUnmask :: Manager -> ((forall x. IO x -> IO x) -> IO ()) -> IO ()
forkManagedUnmask :: Manager -> ((forall a. IO a -> IO a) -> Action) -> Action
forkManagedUnmask Manager
mgr (forall a. IO a -> IO a) -> Action
io =
    IO ThreadId -> Action
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> Action) -> IO ThreadId -> Action
forall a b. (a -> b) -> a -> b
$ IO ThreadId -> IO ThreadId
forall (m :: * -> *) a. MonadUnliftIO m => m a -> m a
mask_ (IO ThreadId -> IO ThreadId) -> IO ThreadId -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ ((forall a. IO a -> IO a) -> Action) -> IO ThreadId
forall (m :: * -> *).
MonadUnliftIO m =>
((forall a. m a -> m a) -> m ()) -> m ThreadId
forkIOWithUnmask (((forall a. IO a -> IO a) -> Action) -> IO ThreadId)
-> ((forall a. IO a -> IO a) -> Action) -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
        Manager -> Action
addMyId Manager
mgr
        -- We catch the exception and do not rethrow it: we don't want the
        -- exception printed to stderr.
        (forall a. IO a -> IO a) -> Action
io IO x -> IO x
forall a. IO a -> IO a
unmask Action -> (SomeException -> Action) -> Action
forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> (e -> m a) -> m a
`catch` \(SomeException
_e :: SomeException) -> () -> Action
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Manager -> Action
deleteMyId Manager
mgr

-- | Adding my thread id to the kill-thread list on stopping.
--
-- This is not part of the public API; see 'forkManaged' instead.
addMyId :: Manager -> IO ()
addMyId :: Manager -> Action
addMyId (Manager TQueue Command
q IORef Action
_ TVar Int
_ Manager
_) = do
    ThreadId
tid <- IO ThreadId
forall (m :: * -> *). MonadIO m => m ThreadId
myThreadId
    STM () -> Action
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Action) -> STM () -> Action
forall a b. (a -> b) -> a -> b
$ TQueue Command -> Command -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Command
q (Command -> STM ()) -> Command -> STM ()
forall a b. (a -> b) -> a -> b
$ ThreadId -> Command
Add ThreadId
tid

-- | Deleting my thread id from the kill-thread list on stopping.
--
-- This is /only/ necessary when you want to remove the thread's ID from
-- the manager /before/ the thread terminates (thereby assuming responsibility
-- for thread cleanup yourself).
deleteMyId :: Manager -> IO ()
deleteMyId :: Manager -> Action
deleteMyId (Manager TQueue Command
q IORef Action
_ TVar Int
_ Manager
_) = do
    ThreadId
tid <- IO ThreadId
forall (m :: * -> *). MonadIO m => m ThreadId
myThreadId
    STM () -> Action
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Action) -> STM () -> Action
forall a b. (a -> b) -> a -> b
$ TQueue Command -> Command -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Command
q (Command -> STM ()) -> Command -> STM ()
forall a b. (a -> b) -> a -> b
$ ThreadId -> Command
Delete ThreadId
tid

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

add :: ThreadId -> Set ThreadId -> Set ThreadId
add :: ThreadId -> Set ThreadId -> Set ThreadId
add ThreadId
tid Set ThreadId
set = Set ThreadId
set'
  where
    set' :: Set ThreadId
set' = ThreadId -> Set ThreadId -> Set ThreadId
forall a. Ord a => a -> Set a -> Set a
Set.insert ThreadId
tid Set ThreadId
set

del :: ThreadId -> Set ThreadId -> Set ThreadId
del :: ThreadId -> Set ThreadId -> Set ThreadId
del ThreadId
tid Set ThreadId
set = Set ThreadId
set'
  where
    set' :: Set ThreadId
set' = ThreadId -> Set ThreadId -> Set ThreadId
forall a. Ord a => a -> Set a -> Set a
Set.delete ThreadId
tid Set ThreadId
set

kill :: Set ThreadId -> Maybe SomeException -> IO ()
kill :: Set ThreadId -> Maybe SomeException -> Action
kill Set ThreadId
set Maybe SomeException
err = (ThreadId -> Action) -> Set ThreadId -> Action
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (\ThreadId
tid -> ThreadId -> KilledByHttp2ThreadManager -> Action
forall e (m :: * -> *).
(Exception e, MonadIO m) =>
ThreadId -> e -> m ()
E.throwTo ThreadId
tid (KilledByHttp2ThreadManager -> Action)
-> KilledByHttp2ThreadManager -> Action
forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> KilledByHttp2ThreadManager
KilledByHttp2ThreadManager Maybe SomeException
err) Set ThreadId
set

-- | Killing the IO action of the second argument on timeout.
timeoutKillThread :: Manager -> (T.Handle -> IO a) -> IO a
timeoutKillThread :: forall a. Manager -> (Handle -> IO a) -> IO a
timeoutKillThread (Manager TQueue Command
_ IORef Action
_ TVar Int
_ Manager
tmgr) Handle -> IO a
action = IO Handle -> (Handle -> Action) -> (Handle -> IO a) -> IO a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
E.bracket IO Handle
register Handle -> Action
T.cancel Handle -> IO a
action
  where
    register :: IO Handle
register = Manager -> Action -> IO Handle
T.registerKillThread Manager
tmgr Action
noAction

-- | Registering closer for a resource and
--   returning a timer refresher.
timeoutClose :: Manager -> IO () -> IO (IO ())
timeoutClose :: Manager -> Action -> IO Action
timeoutClose (Manager TQueue Command
_ IORef Action
_ TVar Int
_ Manager
tmgr) Action
closer = do
    Handle
th <- Manager -> Action -> IO Handle
T.register Manager
tmgr Action
closer
    Action -> IO Action
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Action -> IO Action) -> Action -> IO Action
forall a b. (a -> b) -> a -> b
$ Handle -> Action
T.tickle Handle
th

data KilledByHttp2ThreadManager = KilledByHttp2ThreadManager (Maybe SomeException)
    deriving (Int -> KilledByHttp2ThreadManager -> ShowS
[KilledByHttp2ThreadManager] -> ShowS
KilledByHttp2ThreadManager -> String
(Int -> KilledByHttp2ThreadManager -> ShowS)
-> (KilledByHttp2ThreadManager -> String)
-> ([KilledByHttp2ThreadManager] -> ShowS)
-> Show KilledByHttp2ThreadManager
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> KilledByHttp2ThreadManager -> ShowS
showsPrec :: Int -> KilledByHttp2ThreadManager -> ShowS
$cshow :: KilledByHttp2ThreadManager -> String
show :: KilledByHttp2ThreadManager -> String
$cshowList :: [KilledByHttp2ThreadManager] -> ShowS
showList :: [KilledByHttp2ThreadManager] -> ShowS
Show)

instance Exception KilledByHttp2ThreadManager where
    toException :: KilledByHttp2ThreadManager -> SomeException
toException = KilledByHttp2ThreadManager -> SomeException
forall e. Exception e => e -> SomeException
asyncExceptionToException
    fromException :: SomeException -> Maybe KilledByHttp2ThreadManager
fromException = SomeException -> Maybe KilledByHttp2ThreadManager
forall e. Exception e => SomeException -> Maybe e
asyncExceptionFromException

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

incCounter :: Manager -> IO ()
incCounter :: Manager -> Action
incCounter (Manager TQueue Command
_ IORef Action
_ TVar Int
cnt Manager
_) = STM () -> Action
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Action) -> STM () -> Action
forall a b. (a -> b) -> a -> b
$ TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
cnt (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

decCounter :: Manager -> IO ()
decCounter :: Manager -> Action
decCounter (Manager TQueue Command
_ IORef Action
_ TVar Int
cnt Manager
_) = STM () -> Action
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Action) -> STM () -> Action
forall a b. (a -> b) -> a -> b
$ TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
cnt (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)

waitCounter0 :: Manager -> IO ()
waitCounter0 :: Manager -> Action
waitCounter0 (Manager TQueue Command
_ IORef Action
_ TVar Int
cnt Manager
_) = STM () -> Action
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> Action) -> STM () -> Action
forall a b. (a -> b) -> a -> b
$ do
    Int
n <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
cnt
    Bool -> STM ()
checkSTM (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1)