-- | A thread pool manager.
--   The manager has responsibility to spawn and kill
--   worker threads.
module Network.HTTP2.Arch.Manager (
    Manager
  , Action
  , start
  , setAction
  , stop
  , spawnAction
  , forkManaged
  , deleteMyId
  , timeoutKillThread
  , timeoutClose
  ) where

import Control.Exception
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 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 = forall (m :: * -> *) a. Monad m => a -> m a
return ()

data Command = Stop | Spawn | Add ThreadId | Delete ThreadId

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

-- | Starting a thread pool 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 <- forall (m :: * -> *) a. MonadIO m => m (TQueue a)
newTQueueIO
    IORef Action
ref <- forall a. a -> IO (IORef a)
newIORef Action
noAction
    forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadUnliftIO m => m () -> m ThreadId
forkIO forall a b. (a -> b) -> a -> b
$ TQueue Command -> Set ThreadId -> IORef Action -> Action
go TQueue Command
q forall a. Set a
Set.empty IORef Action
ref
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ TQueue Command -> IORef Action -> Manager -> Manager
Manager TQueue Command
q IORef Action
ref Manager
timmgr
  where
    go :: TQueue Command -> Set ThreadId -> IORef Action -> Action
go TQueue Command
q Set ThreadId
tset0 IORef Action
ref = do
        Command
x <- forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TQueue a -> STM a
readTQueue TQueue Command
q
        case Command
x of
            Command
Stop          -> Set ThreadId -> Action
kill Set ThreadId
tset0
            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 Action -> Action
go TQueue Command
q Set ThreadId
tset IORef Action
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 Action -> Action
go TQueue Command
q Set ThreadId
tset IORef Action
ref
      where
        next :: Set ThreadId -> Action
next Set ThreadId
tset = do
            Action
action <- forall a. IORef a -> IO a
readIORef IORef Action
ref
            ThreadId
newtid <- forall (m :: * -> *). MonadUnliftIO m => m () -> m ThreadId
forkIO Action
action
            let tset' :: Set ThreadId
tset' = ThreadId -> Set ThreadId -> Set ThreadId
add ThreadId
newtid Set ThreadId
tset
            TQueue Command -> Set ThreadId -> IORef Action -> Action
go TQueue Command
q Set ThreadId
tset' IORef Action
ref

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

-- | Stopping the manager.
stop :: Manager -> IO ()
stop :: Manager -> Action
stop (Manager TQueue Command
q IORef Action
_ Manager
_) = forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Command
q Command
Stop

-- | Spawning the action.
spawnAction :: Manager -> IO ()
spawnAction :: Manager -> Action
spawnAction (Manager TQueue Command
q IORef Action
_ Manager
_) = forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically forall a b. (a -> b) -> a -> b
$ 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 =
    forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadUnliftIO m =>
((forall a. m a -> m a) -> m ()) -> m ThreadId
forkIOWithUnmask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
      Manager -> Action
addMyId Manager
mgr
      ()
r <- forall a. IO a -> IO a
unmask Action
io forall a b. IO a -> IO b -> IO a
`onException` Manager -> Action
deleteMyId Manager
mgr
      Manager -> Action
deleteMyId Manager
mgr
      forall (m :: * -> *) a. Monad m => a -> m a
return ()
r

-- | 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
_ Manager
_) = do
    ThreadId
tid <- forall (m :: * -> *). MonadIO m => m ThreadId
myThreadId
    forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Command
q 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
_ Manager
_) = do
    ThreadId
tid <- forall (m :: * -> *). MonadIO m => m ThreadId
myThreadId
    forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Command
q 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' = 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' = forall a. Ord a => a -> Set a -> Set a
Set.delete ThreadId
tid Set ThreadId
set

kill :: Set ThreadId -> IO ()
kill :: Set ThreadId -> Action
kill Set ThreadId
set = forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ forall (m :: * -> *). MonadIO m => ThreadId -> m ()
killThread Set ThreadId
set

-- | Killing the IO action of the second argument on timeout.
timeoutKillThread :: Manager -> (T.Handle -> IO ()) -> IO ()
timeoutKillThread :: Manager -> (Handle -> Action) -> Action
timeoutKillThread (Manager TQueue Command
_ IORef Action
_ Manager
tmgr) Handle -> Action
action = 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 -> Action
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
_ Manager
tmgr) Action
closer = do
    Handle
th <- Manager -> Action -> IO Handle
T.register Manager
tmgr Action
closer
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Handle -> Action
T.tickle Handle
th