1 {-# LANGUAGE BangPatterns #-}
    2 
    3 module Snap.Internal.Http.Server.TimeoutManager
    4   ( TimeoutManager
    5   , TimeoutHandle
    6   , initialize
    7   , stop
    8   , register
    9   , tickle
   10   , cancel
   11   ) where
   12 
   13 ------------------------------------------------------------------------------
   14 import           Control.Concurrent
   15 import           Control.Exception
   16 import           Control.Monad
   17 import           Data.IORef
   18 import           Foreign.C.Types
   19 
   20 ------------------------------------------------------------------------------
   21 data State = Deadline !CTime
   22            | Canceled
   23 
   24 
   25 ------------------------------------------------------------------------------
   26 data TimeoutHandle = TimeoutHandle {
   27       _killAction :: !(IO ())
   28     , _state      :: !(IORef State)
   29     , _hGetTime   :: !(IO CTime)
   30     }
   31 
   32 
   33 ------------------------------------------------------------------------------
   34 data TimeoutManager = TimeoutManager {
   35       _defaultTimeout :: !Int
   36     , _getTime        :: !(IO CTime)
   37     , _connections    :: !(IORef [TimeoutHandle])
   38     , _inactivity     :: !(IORef Bool)
   39     , _morePlease     :: !(MVar ())
   40     , _managerThread  :: !(MVar ThreadId)
   41     }
   42 
   43 
   44 ------------------------------------------------------------------------------
   45 -- | Create a new TimeoutManager.
   46 initialize :: Int               -- ^ default timeout
   47            -> IO CTime          -- ^ function to get current time
   48            -> IO TimeoutManager
   49 initialize defaultTimeout getTime = do
   50     conns <- newIORef []
   51     inact <- newIORef False
   52     mp    <- newEmptyMVar
   53     mthr  <- newEmptyMVar
   54 
   55     let tm = TimeoutManager defaultTimeout getTime conns inact mp mthr
   56 
   57     thr <- forkIO $ managerThread tm
   58     putMVar mthr thr
   59     return tm
   60 
   61 
   62 ------------------------------------------------------------------------------
   63 -- | Stop a TimeoutManager.
   64 stop :: TimeoutManager -> IO ()
   65 stop tm = readMVar (_managerThread tm) >>= killThread
   66 
   67 
   68 ------------------------------------------------------------------------------
   69 -- | Register a new connection with the TimeoutManager.
   70 register :: IO ()               -- ^ action to run when the timeout deadline is
   71                                 -- exceeded.
   72          -> TimeoutManager      -- ^ manager to register with.
   73          -> IO TimeoutHandle
   74 register killAction tm = do
   75     now <- getTime
   76     let !state = Deadline $ now + toEnum defaultTimeout
   77     stateRef <- newIORef state
   78 
   79     let !h = TimeoutHandle killAction stateRef getTime
   80     atomicModifyIORef connections $ \x -> (h:x, ())
   81 
   82     inact <- readIORef inactivity
   83     when inact $ do
   84         -- wake up manager thread
   85         writeIORef inactivity False
   86         _ <- tryPutMVar morePlease ()
   87         return ()
   88     return h
   89 
   90   where
   91     getTime        = _getTime tm
   92     inactivity     = _inactivity tm
   93     morePlease     = _morePlease tm
   94     connections    = _connections tm
   95     defaultTimeout = _defaultTimeout tm
   96 
   97 
   98 ------------------------------------------------------------------------------
   99 -- | Tickle the timeout on a connection to be N seconds into the future.
  100 tickle :: TimeoutHandle -> Int -> IO ()
  101 tickle th n = do
  102     now <- getTime
  103 
  104     let state = Deadline $ now + toEnum n
  105     writeIORef stateRef state
  106 
  107   where
  108     getTime  = _hGetTime th
  109     stateRef = _state th
  110 
  111 
  112 ------------------------------------------------------------------------------
  113 -- | Cancel a timeout.
  114 cancel :: TimeoutHandle -> IO ()
  115 cancel h = writeIORef (_state h) Canceled
  116 
  117 
  118 ------------------------------------------------------------------------------
  119 managerThread :: TimeoutManager -> IO ()
  120 managerThread tm = loop `finally` (readIORef connections >>= destroyAll)
  121   where
  122     --------------------------------------------------------------------------
  123     connections = _connections tm
  124     getTime     = _getTime tm
  125     inactivity  = _inactivity tm
  126     morePlease  = _morePlease tm
  127     waitABit    = threadDelay 5000000
  128 
  129     --------------------------------------------------------------------------
  130     loop = do
  131         waitABit
  132         handles <- atomicModifyIORef connections (\x -> ([],x))
  133 
  134         if null handles
  135           then do
  136             -- we're inactive, go to sleep until we get new threads
  137             writeIORef inactivity True
  138             takeMVar morePlease
  139           else do
  140             now   <- getTime
  141             dlist <- processHandles now handles id
  142             atomicModifyIORef connections (\x -> (dlist x, ()))
  143 
  144         loop
  145 
  146     --------------------------------------------------------------------------
  147     processHandles !now handles initDlist = go handles initDlist
  148       where
  149         go [] !dlist = return dlist
  150 
  151         go (x:xs) !dlist = do
  152             state   <- readIORef $ _state x
  153             !dlist' <- case state of
  154                          Canceled   -> return dlist
  155                          Deadline t -> if t <= now
  156                                          then do
  157                                            _killAction x
  158                                            return dlist
  159                                          else return (dlist . (x:))
  160             go xs dlist'
  161 
  162     --------------------------------------------------------------------------
  163     destroyAll = mapM_ diediedie
  164 
  165     --------------------------------------------------------------------------
  166     diediedie x = do
  167         state <- readIORef $ _state x
  168         case state of
  169           Canceled -> return ()
  170           _        -> _killAction x