{-# LANGUAGE DeriveDataTypeable #-}

module System.TimeManager (
  -- ** Types
    Manager
  , TimeoutAction
  , Handle
  -- ** Manager
  , initialize
  , stopManager
  , killManager
  , withManager
  -- ** Registration
  , register
  , registerKillThread
  -- ** Control
  , tickle
  , cancel
  , pause
  , resume
  -- ** Exceptions
  , TimeoutThread (..)
  ) where

import Control.Concurrent (myThreadId)
import qualified Control.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 = 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 :: [Handle] -> IO ([Handle] -> [Handle])
reaperAction = (Handle -> IO (Maybe Handle))
-> [Handle] -> IO ([Handle] -> [Handle])
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 <- 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 e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` SomeException -> TimeoutAction
ignoreAll
                Maybe Handle -> IO (Maybe Handle)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Handle
forall a. Maybe a
Nothing
            State
Canceled -> Maybe Handle -> IO (Maybe Handle)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Handle
forall a. Maybe a
Nothing
            State
_        -> Maybe Handle -> IO (Maybe Handle)
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 a. IO a -> IO a
E.mask_ (Manager -> IO [Handle]
forall workload item. Reaper workload item -> IO workload
reaperStop Manager
mgr IO [Handle] -> ([Handle] -> TimeoutAction) -> TimeoutAction
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 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 e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` SomeException -> TimeoutAction
ignoreAll

ignoreAll :: E.SomeException -> IO ()
ignoreAll :: SomeException -> TimeoutAction
ignoreAll SomeException
_ = () -> TimeoutAction
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 = IORef TimeoutAction -> IORef State -> Handle
Handle 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 (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 a b. IO a -> IO b -> IO a
`E.finally` ThreadId -> TimeoutThread -> TimeoutAction
forall e. Exception e => ThreadId -> e -> TimeoutAction
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 IORef TimeoutAction
_ IORef State
stateRef) = IORef State -> State -> TimeoutAction
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
    IORef TimeoutAction -> TimeoutAction -> TimeoutAction
forall a. IORef a -> a -> TimeoutAction
I.writeIORef IORef TimeoutAction
actionRef (() -> TimeoutAction
forall (m :: * -> *) a. Monad m => a -> m a
return ()) -- ensuring to release ThreadId
    IORef State -> State -> TimeoutAction
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) = 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.
withManager :: Int -- ^ timeout in microseconds
            -> (Manager -> IO a)
            -> IO a
withManager :: Int -> (Manager -> IO a) -> IO a
withManager Int
timeout Manager -> IO a
f = do
    -- FIXME when stopManager is available, use it
    Manager
man <- Int -> IO Manager
initialize Int
timeout
    Manager -> IO a
f Manager
man