{-# LANGUAGE DeriveDataTypeable #-}

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 qualified UnliftIO.Exception as E
import Control.Reaper
import Data.Typeable (Typeable)
import Data.IORef (IORef)
import qualified Data.IORef as I

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

-- | 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 !(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.
           | Canceled  -- Manager removes it without timeout action.

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

-- | 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 = forall workload item.
ReaperSettings workload item -> IO (Reaper workload item)
mkReaper forall item. ReaperSettings [item] item
defaultReaperSettings
        { reaperAction :: [Handle] -> IO ([Handle] -> [Handle])
reaperAction = forall item item'.
(item -> IO (Maybe item')) -> [item] -> IO ([item'] -> [item'])
mkListAction Handle -> IO (Maybe Handle)
prune
        , reaperDelay :: Int
reaperDelay = Int
timeout
        }
  where
    prune :: Handle -> IO (Maybe Handle)
prune m :: Handle
m@(Handle IORef TimeoutAction
actionRef IORef State
stateRef) = do
        State
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 <- forall a. IORef a -> IO a
I.readIORef IORef TimeoutAction
actionRef
                TimeoutAction
onTimeout forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> (e -> m a) -> m a
`E.catch` SomeException -> TimeoutAction
ignoreAll
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
            State
Canceled -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
            State
_        -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ 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 = forall (m :: * -> *) a. MonadUnliftIO m => m a -> m a
E.mask_ (forall workload item. Reaper workload item -> IO workload
reaperStop Manager
mgr forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 IORef TimeoutAction
actionRef IORef State
_) = do
        TimeoutAction
onTimeout <- forall a. IORef a -> IO a
I.readIORef IORef TimeoutAction
actionRef
        TimeoutAction
onTimeout 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
_ = forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Killing timeout manager immediately without firing onTimeout.
killManager :: Manager -> IO ()
killManager :: Manager -> TimeoutAction
killManager = 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 <- forall a. a -> IO (IORef a)
I.newIORef TimeoutAction
onTimeout
    IORef State
stateRef  <- forall a. a -> IO (IORef a)
I.newIORef State
Active
    let h :: Handle
h = IORef TimeoutAction -> IORef State -> Handle
Handle IORef TimeoutAction
actionRef IORef State
stateRef
    forall workload item. Reaper workload item -> item -> TimeoutAction
reaperAdd Manager
mgr Handle
h
    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 forall a b. (a -> b) -> a -> b
$ TimeoutAction
onTimeout forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`E.finally` 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 = forall e. Exception e => e -> SomeException
E.asyncExceptionToException
    fromException :: SomeException -> Maybe TimeoutThread
fromException = 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 IORef TimeoutAction
_ IORef State
stateRef) = forall a. IORef a -> a -> TimeoutAction
I.writeIORef IORef State
stateRef State
Active

-- | Setting the state to canceled.
--   'Manager' eventually removes this without timeout action.
cancel :: Handle -> IO ()
cancel :: Handle -> TimeoutAction
cancel (Handle IORef TimeoutAction
actionRef IORef State
stateRef) = do
    forall a. IORef a -> a -> TimeoutAction
I.writeIORef IORef TimeoutAction
actionRef (forall (m :: * -> *) a. Monad m => a -> m a
return ()) -- ensuring to release ThreadId
    forall a. IORef a -> a -> TimeoutAction
I.writeIORef IORef State
stateRef State
Canceled

-- | Setting the state to paused.
--   'Manager' does not change the value.
pause :: Handle -> IO ()
pause :: Handle -> TimeoutAction
pause (Handle IORef TimeoutAction
_ IORef State
stateRef) = 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 = 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 = 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