{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE BangPatterns #-}

module System.TimeManager (
    -- ** Types
    Manager,
    TimeoutAction,
    Handle,

    -- ** Manager
    initialize,
    stopManager,
    killManager,
    withManager,
    withManager',

    -- ** Registration
    register,
    registerKillThread,

    -- ** Control
    tickle,
    cancel,
    pause,
    resume,

    -- ** Exceptions
    TimeoutThread (..),
) where

import Control.Concurrent (myThreadId)
import Control.Reaper
import Data.IORef (IORef)
import qualified Data.IORef as I
import Data.Typeable (Typeable)
import qualified UnliftIO.Exception as E

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

-- | A timeout manager
type Manager = Reaper [Handle] Handle

-- | An action to be performed on timeout.
type TimeoutAction = IO ()

-- | A handle used by 'Manager'
data Handle = Handle Manager !(IORef TimeoutAction) !(IORef State)

data State
    = Active -- Manager turns it to Inactive.
    | Inactive -- Manager removes it with timeout action.
    | Paused -- Manager does not change it.

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

-- | Creating timeout manager which works every N micro seconds
--   where N is the first argument.
initialize :: Int -> IO Manager
initialize :: Int -> IO Manager
initialize Int
timeout =
    ReaperSettings [Handle] Handle -> IO Manager
forall workload item.
ReaperSettings workload item -> IO (Reaper workload item)
mkReaper
        ReaperSettings [Handle] Handle
forall item. ReaperSettings [item] item
defaultReaperSettings
            { reaperAction = mkListAction prune
            , reaperDelay = timeout
            }
  where
    prune :: Handle -> IO (Maybe Handle)
prune m :: Handle
m@(Handle Manager
_ IORef TimeoutAction
actionRef IORef State
stateRef) = do
        State
state <- IORef State -> (State -> (State, State)) -> IO State
forall a b. IORef a -> (a -> (a, b)) -> IO b
I.atomicModifyIORef' IORef State
stateRef (\State
x -> (State -> State
inactivate State
x, State
x))
        case State
state of
            State
Inactive -> do
                TimeoutAction
onTimeout <- IORef TimeoutAction -> IO TimeoutAction
forall a. IORef a -> IO a
I.readIORef IORef TimeoutAction
actionRef
                TimeoutAction
onTimeout TimeoutAction -> (SomeException -> TimeoutAction) -> TimeoutAction
forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> (e -> m a) -> m a
`E.catch` SomeException -> TimeoutAction
ignoreAll
                Maybe Handle -> IO (Maybe Handle)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Handle
forall a. Maybe a
Nothing
            State
_ -> Maybe Handle -> IO (Maybe Handle)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Handle -> IO (Maybe Handle))
-> Maybe Handle -> IO (Maybe Handle)
forall a b. (a -> b) -> a -> b
$ Handle -> Maybe Handle
forall a. a -> Maybe a
Just Handle
m

    inactivate :: State -> State
inactivate State
Active = State
Inactive
    inactivate State
x = State
x

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

-- | Stopping timeout manager with onTimeout fired.
stopManager :: Manager -> IO ()
stopManager :: Manager -> TimeoutAction
stopManager Manager
mgr = TimeoutAction -> TimeoutAction
forall (m :: * -> *) a. MonadUnliftIO m => m a -> m a
E.mask_ (Manager -> IO [Handle]
forall workload item. Reaper workload item -> IO workload
reaperStop Manager
mgr IO [Handle] -> ([Handle] -> TimeoutAction) -> TimeoutAction
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Handle -> TimeoutAction) -> [Handle] -> TimeoutAction
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Handle -> TimeoutAction
fire)
  where
    fire :: Handle -> TimeoutAction
fire (Handle Manager
_ IORef TimeoutAction
actionRef IORef State
_) = do
        TimeoutAction
onTimeout <- IORef TimeoutAction -> IO TimeoutAction
forall a. IORef a -> IO a
I.readIORef IORef TimeoutAction
actionRef
        TimeoutAction
onTimeout TimeoutAction -> (SomeException -> TimeoutAction) -> TimeoutAction
forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> (e -> m a) -> m a
`E.catch` SomeException -> TimeoutAction
ignoreAll

ignoreAll :: E.SomeException -> IO ()
ignoreAll :: SomeException -> TimeoutAction
ignoreAll SomeException
_ = () -> TimeoutAction
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Killing timeout manager immediately without firing onTimeout.
killManager :: Manager -> IO ()
killManager :: Manager -> TimeoutAction
killManager = Manager -> TimeoutAction
forall workload item. Reaper workload item -> TimeoutAction
reaperKill

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

-- | Registering a timeout action.
register :: Manager -> TimeoutAction -> IO Handle
register :: Manager -> TimeoutAction -> IO Handle
register Manager
mgr !TimeoutAction
onTimeout = do
    IORef TimeoutAction
actionRef <- TimeoutAction -> IO (IORef TimeoutAction)
forall a. a -> IO (IORef a)
I.newIORef TimeoutAction
onTimeout
    IORef State
stateRef <- State -> IO (IORef State)
forall a. a -> IO (IORef a)
I.newIORef State
Active
    let h :: Handle
h = Manager -> IORef TimeoutAction -> IORef State -> Handle
Handle Manager
mgr IORef TimeoutAction
actionRef IORef State
stateRef
    Manager -> Handle -> TimeoutAction
forall workload item. Reaper workload item -> item -> TimeoutAction
reaperAdd Manager
mgr Handle
h
    Handle -> IO Handle
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Handle
h

-- | Registering a timeout action of killing this thread.
registerKillThread :: Manager -> TimeoutAction -> IO Handle
registerKillThread :: Manager -> TimeoutAction -> IO Handle
registerKillThread Manager
m TimeoutAction
onTimeout = do
    -- If we hold ThreadId, the stack and data of the thread is leaked.
    -- If we hold Weak ThreadId, the stack is released. However, its
    -- data is still leaked probably because of a bug of GHC.
    -- So, let's just use ThreadId and release ThreadId by
    -- overriding the timeout action by "cancel".
    ThreadId
tid <- IO ThreadId
myThreadId
    -- First run the timeout action in case the child thread is masked.
    Manager -> TimeoutAction -> IO Handle
register Manager
m (TimeoutAction -> IO Handle) -> TimeoutAction -> IO Handle
forall a b. (a -> b) -> a -> b
$ TimeoutAction
onTimeout TimeoutAction -> TimeoutAction -> TimeoutAction
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`E.finally` ThreadId -> TimeoutThread -> TimeoutAction
forall e (m :: * -> *).
(Exception e, MonadIO m) =>
ThreadId -> e -> m ()
E.throwTo ThreadId
tid TimeoutThread
TimeoutThread

data TimeoutThread = TimeoutThread
    deriving (Typeable)
instance E.Exception TimeoutThread where
    toException :: TimeoutThread -> SomeException
toException = TimeoutThread -> SomeException
forall e. Exception e => e -> SomeException
E.asyncExceptionToException
    fromException :: SomeException -> Maybe TimeoutThread
fromException = SomeException -> Maybe TimeoutThread
forall e. Exception e => SomeException -> Maybe e
E.asyncExceptionFromException
instance Show TimeoutThread where
    show :: TimeoutThread -> String
show TimeoutThread
TimeoutThread = String
"Thread killed by timeout manager"

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

-- | Setting the state to active.
--   'Manager' turns active to inactive repeatedly.
tickle :: Handle -> IO ()
tickle :: Handle -> TimeoutAction
tickle (Handle Manager
_ IORef TimeoutAction
_ IORef State
stateRef) = IORef State -> State -> TimeoutAction
forall a. IORef a -> a -> TimeoutAction
I.writeIORef IORef State
stateRef State
Active

-- | Removing the 'Handle' from the 'Manager' immediately.
cancel :: Handle -> IO ()
cancel :: Handle -> TimeoutAction
cancel (Handle Manager
mgr IORef TimeoutAction
_ IORef State
stateRef) = do
    [Handle]
_ <- Manager -> ([Handle] -> [Handle]) -> IO [Handle]
forall workload item.
Reaper workload item -> (workload -> workload) -> IO workload
reaperModify Manager
mgr [Handle] -> [Handle]
filt
    () -> TimeoutAction
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  where
    -- It's very important that this function forces the whole workload so we
    -- don't retain old handles, otherwise disasterous leaks occur.
    filt :: [Handle] -> [Handle]
filt [] = []
    filt (h :: Handle
h@(Handle Manager
_ IORef TimeoutAction
_ IORef State
stateRef') : [Handle]
hs)
        | IORef State
stateRef IORef State -> IORef State -> Bool
forall a. Eq a => a -> a -> Bool
== IORef State
stateRef' =
            [Handle]
hs
        | Bool
otherwise =
            let !hs' :: [Handle]
hs'= [Handle] -> [Handle]
filt [Handle]
hs
             in Handle
h Handle -> [Handle] -> [Handle]
forall a. a -> [a] -> [a]
: [Handle]
hs'

-- | Setting the state to paused.
--   'Manager' does not change the value.
pause :: Handle -> IO ()
pause :: Handle -> TimeoutAction
pause (Handle Manager
_ IORef TimeoutAction
_ IORef State
stateRef) = IORef State -> State -> TimeoutAction
forall a. IORef a -> a -> TimeoutAction
I.writeIORef IORef State
stateRef State
Paused

-- | Setting the paused state to active.
--   This is an alias to 'tickle'.
resume :: Handle -> IO ()
resume :: Handle -> TimeoutAction
resume = Handle -> TimeoutAction
tickle

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

-- | Call the inner function with a timeout manager.
--   'stopManager' is used after that.
withManager
    :: Int
    -- ^ timeout in microseconds
    -> (Manager -> IO a)
    -> IO a
withManager :: forall a. Int -> (Manager -> IO a) -> IO a
withManager Int
timeout Manager -> IO a
f =
    IO Manager
-> (Manager -> TimeoutAction) -> (Manager -> IO a) -> IO a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
E.bracket
        (Int -> IO Manager
initialize Int
timeout)
        Manager -> TimeoutAction
stopManager
        Manager -> IO a
f

-- | Call the inner function with a timeout manager.
--   'killManager' is used after that.
withManager'
    :: Int
    -- ^ timeout in microseconds
    -> (Manager -> IO a)
    -> IO a
withManager' :: forall a. Int -> (Manager -> IO a) -> IO a
withManager' Int
timeout Manager -> IO a
f =
    IO Manager
-> (Manager -> TimeoutAction) -> (Manager -> IO a) -> IO a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
E.bracket
        (Int -> IO Manager
initialize Int
timeout)
        Manager -> TimeoutAction
killManager
        Manager -> IO a
f