module Network.Wai.Handler.Warp.Timeout (
Manager
, TimeoutAction
, Handle
, initialize
, stopManager
, withManager
, register
, registerKillThread
, tickle
, cancel
, pause
, resume
, TimeoutThread (..)
) where
#ifndef MIN_VERSION_base
#define MIN_VERSION_base(x,y,z) 1
#endif
#if MIN_VERSION_base(4,6,0)
import Control.Concurrent (mkWeakThreadId, ThreadId)
#else
import GHC.Conc (ThreadId(..))
import GHC.Exts (mkWeak#)
import GHC.IO (IO (IO))
#endif
import Control.Concurrent (myThreadId)
import qualified Control.Exception as E
import GHC.Weak (Weak (..))
import System.Mem.Weak (deRefWeak)
import Data.Typeable (Typeable)
import Control.Reaper
#if USE_ATOMIC_PRIMOPS
import qualified Data.Atomics.Counter.Unboxed as C
import Data.Atomics.Counter.Unboxed (casCounter)
#else
import Network.Wai.Handler.Warp.IORef
#endif
type Manager = Reaper [Handle] Handle
type TimeoutAction = IO ()
data Handle = Handle TimeoutAction
#if USE_ATOMIC_PRIMOPS
!C.AtomicCounter
#else
!(IORef Int)
#endif
initialize :: Int -> IO Manager
initialize timeout = mkReaper defaultReaperSettings
{ reaperAction = mkListAction prune
, reaperDelay = timeout
}
where
prune m@(Handle onTimeout iactive) = do
(wasActive, newState) <- casCounter iactive 1 2
case newState of
2 | not wasActive -> do
onTimeout `E.catch` ignoreAll
return Nothing
4 -> return Nothing
_ -> return $ Just m
#if !USE_ATOMIC_PRIMOPS
casCounter :: IORef Int -> Int -> Int -> IO (Bool, Int)
casCounter ref old new = atomicModifyIORef' ref $ \curr ->
if old == curr
then (new, (True, new))
else (old, (False, old))
#endif
stopManager :: Manager -> IO ()
stopManager mgr = E.mask_ (reaperStop mgr >>= mapM_ fire)
where
fire (Handle onTimeout _) = onTimeout `E.catch` ignoreAll
ignoreAll :: E.SomeException -> IO ()
ignoreAll _ = return ()
register :: Manager -> TimeoutAction -> IO Handle
register mgr onTimeout = do
#if USE_ATOMIC_PRIMOPS
iactive <- C.newCounter 1
#else
iactive <- newIORef 1
#endif
let h = Handle onTimeout iactive
reaperAdd mgr h
return h
registerKillThread :: Manager -> IO Handle
registerKillThread m = do
wtid <- myThreadId >>= mkWeakThreadId
register m $ killIfExist wtid
killIfExist :: Weak ThreadId -> TimeoutAction
killIfExist wtid = deRefWeak wtid >>= maybe (return ()) (`E.throwTo` TimeoutThread)
data TimeoutThread = TimeoutThread
deriving Typeable
instance E.Exception TimeoutThread
instance Show TimeoutThread where
show TimeoutThread = "Thread killed by Warp's timeout reaper"
#if !MIN_VERSION_base(4,6,0)
mkWeakThreadId :: ThreadId -> IO (Weak ThreadId)
mkWeakThreadId t@(ThreadId t#) = IO $ \s ->
case mkWeak# t# t Nothing s of
(# s1, w #) -> (# s1, Weak w #)
#endif
writeCounter :: Int -> Handle -> IO ()
#if USE_ATOMIC_PRIMOPS
writeCounter i (Handle _ iactive) = C.writeCounter iactive i
#else
writeCounter i (Handle _ iactive) = writeIORef iactive i
#endif
tickle :: Handle -> IO ()
tickle = writeCounter 1
cancel :: Handle -> IO ()
cancel = writeCounter 4
pause :: Handle -> IO ()
pause = writeCounter 3
resume :: Handle -> IO ()
resume = tickle
withManager :: Int
-> (Manager -> IO a)
-> IO a
withManager timeout f = do
man <- initialize timeout
f man