TimeoutManager.hs

{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Snap.Internal.Http.Server.TimeoutManager
  ( TimeoutManager
  , TimeoutThread
  , initialize
  , stop
  , register
  , tickle
  , set
  , modify
  , cancel
  ) where

------------------------------------------------------------------------------
import           Control.Exception                (evaluate, finally)
import qualified Control.Exception                as E
import           Control.Monad                    (Monad (return, (>>=)), mapM_, void, when)
import qualified Data.ByteString.Char8            as S
import           Data.IORef                       (IORef, newIORef, readIORef, writeIORef)
import           Prelude                          (Bool, Double, IO, Int, Show (..), const, fromIntegral, max, null, otherwise, round, ($), ($!), (+), (++), (-), (.), (<=), (==))
------------------------------------------------------------------------------
import           Control.Concurrent               (MVar, newEmptyMVar, putMVar, readMVar, takeMVar, tryPutMVar)
------------------------------------------------------------------------------
import           Snap.Internal.Http.Server.Clock  (ClockTime)
import qualified Snap.Internal.Http.Server.Clock  as Clock
import           Snap.Internal.Http.Server.Common (atomicModifyIORef', eatException)
import qualified Snap.Internal.Http.Server.Thread as T


------------------------------------------------------------------------------
type State = ClockTime

canceled :: State
canceled = 0

isCanceled :: State -> Bool
isCanceled = (== 0)


------------------------------------------------------------------------------
data TimeoutThread = TimeoutThread {
      _thread     :: !T.SnapThread
    , _state      :: !(IORef State)
    , _hGetTime   :: !(IO ClockTime)
    }

instance Show TimeoutThread where
    show = show . _thread


------------------------------------------------------------------------------
-- | Given a 'State' value and the current time, apply the given modification
-- function to the amount of time remaining.
--
smap :: ClockTime -> (ClockTime -> ClockTime) -> State -> State
smap now f deadline | isCanceled deadline = deadline
                    | otherwise = t'
  where
    remaining    = max 0 (deadline - now)
    newremaining = f remaining
    t'           = now + newremaining


------------------------------------------------------------------------------
data TimeoutManager = TimeoutManager {
      _defaultTimeout :: !ClockTime
    , _pollInterval   :: !ClockTime
    , _getTime        :: !(IO ClockTime)
    , _threads        :: !(IORef [TimeoutThread])
    , _morePlease     :: !(MVar ())
    , _managerThread  :: !(MVar T.SnapThread)
    }


------------------------------------------------------------------------------
-- | Create a new TimeoutManager.
initialize :: Double            -- ^ default timeout
           -> Double            -- ^ poll interval
           -> IO ClockTime      -- ^ function to get current time
           -> IO TimeoutManager
initialize defaultTimeout interval getTime = E.uninterruptibleMask_ $ do
    conns <- newIORef []
    mp    <- newEmptyMVar
    mthr  <- newEmptyMVar

    let tm = TimeoutManager (Clock.fromSecs defaultTimeout)
                            (Clock.fromSecs interval)
                            getTime
                            conns
                            mp
                            mthr

    thr <- T.fork "snap-server: timeout manager" $ managerThread tm
    putMVar mthr thr
    return tm


------------------------------------------------------------------------------
-- | Stop a TimeoutManager.
stop :: TimeoutManager -> IO ()
stop tm = readMVar (_managerThread tm) >>= T.cancelAndWait


------------------------------------------------------------------------------
wakeup :: TimeoutManager -> IO ()
wakeup tm = void $ tryPutMVar (_morePlease tm) $! ()


------------------------------------------------------------------------------
-- | Register a new thread with the TimeoutManager.
register :: TimeoutManager                        -- ^ manager to register
                                                  --   with
         -> S.ByteString                          -- ^ thread label
         -> ((forall a . IO a -> IO a) -> IO ())  -- ^ thread action to run
         -> IO TimeoutThread
register tm label action = do
    now <- getTime
    let !state = now + defaultTimeout
    stateRef <- newIORef state
    th <- E.uninterruptibleMask_ $ do
        t <- T.fork label action
        let h = TimeoutThread t stateRef getTime
        atomicModifyIORef' threads (\x -> (h:x, ())) >>= evaluate
        return $! h
    wakeup tm
    return th

  where
    getTime        = _getTime tm
    threads        = _threads tm
    defaultTimeout = _defaultTimeout tm


------------------------------------------------------------------------------
-- | Tickle the timeout on a connection to be at least N seconds into the
-- future. If the existing timeout is set for M seconds from now, where M > N,
-- then the timeout is unaffected.
tickle :: TimeoutThread -> Int -> IO ()
tickle th = modify th . max
{-# INLINE tickle #-}


------------------------------------------------------------------------------
-- | Set the timeout on a connection to be N seconds into the future.
set :: TimeoutThread -> Int -> IO ()
set th = modify th . const
{-# INLINE set #-}


------------------------------------------------------------------------------
-- | Modify the timeout with the given function.
modify :: TimeoutThread -> (Int -> Int) -> IO ()
modify th f = do
    now   <- getTime
    state <- readIORef stateRef
    let !state' = smap now f' state
    writeIORef stateRef state'

  where
    f' !x    = Clock.fromSecs $! fromIntegral $ f $ round $ Clock.toSecs x
    getTime  = _hGetTime th
    stateRef = _state th
{-# INLINE modify #-}


------------------------------------------------------------------------------
-- | Cancel a timeout.
cancel :: TimeoutThread -> IO ()
cancel h = E.uninterruptibleMask_ $ do
    writeIORef (_state h) canceled
    T.cancel $ _thread h
{-# INLINE cancel #-}


------------------------------------------------------------------------------
managerThread :: TimeoutManager -> (forall a. IO a -> IO a) -> IO ()
managerThread tm restore = restore loop `finally` cleanup
  where
    cleanup = E.uninterruptibleMask_ $
              eatException (readIORef threads >>= destroyAll)

    --------------------------------------------------------------------------
    getTime      = _getTime tm
    morePlease   = _morePlease tm
    pollInterval = _pollInterval tm
    threads      = _threads tm

    --------------------------------------------------------------------------
    loop = do
        now <- getTime
        E.uninterruptibleMask $ \restore' -> do
            handles <- atomicModifyIORef' threads (\x -> ([], x))
            if null handles
              then do restore' $ takeMVar morePlease
              else do
                handles' <- processHandles now handles
                atomicModifyIORef' threads (\x -> (handles' ++ x, ()))
                    >>= evaluate
        Clock.sleepFor pollInterval
        loop

    --------------------------------------------------------------------------
    processHandles now handles = go handles []
      where
        go [] !kept = return $! kept

        go (x:xs) !kept = do
            !state <- readIORef $ _state x
            !kept' <-
                if isCanceled state
                  then do b <- T.isFinished (_thread x)
                          return $! if b
                                      then kept
                                      else (x:kept)
                  else do when (state <= now) $ do
                            T.cancel (_thread x)
                            writeIORef (_state x) canceled
                          return (x:kept)
            go xs kept'

    --------------------------------------------------------------------------
    destroyAll xs = do
        mapM_ (T.cancel . _thread) xs
        mapM_ (T.wait . _thread) xs