    {-# LANGUAGE CPP #-}
    {-# LANGUAGE BangPatterns #-}
    {-# LANGUAGE DeriveDataTypeable #-}
    {-# LANGUAGE UnboxedTuples, MagicHash #-}
    module Network.Wai.Handler.Warp.Timeout (
        Manager
      , Handle
      , initialize
      , stopManager
      , register
      , registerKillThread
      , tickle
      , pause
      , resume
      , cancel
      , withManager
      , dummyHandle
      ) where

    import System.Mem.Weak (deRefWeak)
    #if MIN_VERSION_base(4,6,0)
    import Control.Concurrent (mkWeakThreadId)
    #else
    import GHC.Weak (Weak (..))
    import GHC.Conc.Sync (ThreadId (..))
    import GHC.IO (IO (IO))
    import GHC.Exts (mkWeak#)
    #endif
    import Control.Concurrent (forkIO, threadDelay, myThreadId, killThread)
    import qualified Control.Exception as E
    import Control.Monad (forever, void)
    import qualified Data.IORef as I
    import System.IO.Unsafe (unsafePerformIO)
    import Data.Typeable (Typeable)

    -- | A timeout manager
    newtype Manager = Manager (I.IORef [Handle])

    -- | A handle used by 'Manager'
    --
    -- First field is action to be performed on timeout.
    data Handle = Handle (IO ()) (I.IORef State)

    -- | A dummy @Handle@.
    dummyHandle :: Handle
    dummyHandle = Handle (return ()) (unsafePerformIO $ I.newIORef Active)

    data State = Active | Inactive | Paused | Canceled

    initialize :: Int -> IO Manager
    initialize timeout = do
        ref <- I.newIORef []
        void . forkIO $ E.handle ignoreStop $ forever $ do
            threadDelay timeout
            ms <- I.atomicModifyIORef ref (\x -> ([], x))
            ms' <- go ms id
            I.atomicModifyIORef ref (\x -> (ms' x, ()))
        return $ Manager ref
      where
        ignoreStop TimeoutManagerStopped = return ()

        go [] front = return front
        go (m@(Handle onTimeout iactive):rest) front = do
            state <- I.atomicModifyIORef iactive (\x -> (go' x, x))
            case state of
                Inactive -> do
                    onTimeout `E.catch` ignoreAll
                    go rest front
                Canceled -> go rest front
                _ -> go rest (front . (:) m)
        go' Active = Inactive
        go' x = x

    data TimeoutManagerStopped = TimeoutManagerStopped
        deriving (Show, Typeable)
    instance E.Exception TimeoutManagerStopped

    stopManager :: Manager -> IO ()
    stopManager (Manager ihandles) = E.mask_ $ do
        -- Put an undefined value in the IORef to kill the worker thread (yes, it's
        -- a bit of a hack)
        !handles <- I.atomicModifyIORef ihandles $ \h -> (E.throw TimeoutManagerStopped, h)
        mapM_ go handles
      where
        go (Handle onTimeout _) = onTimeout `E.catch` ignoreAll

    ignoreAll :: E.SomeException -> IO ()
    ignoreAll _ = return ()

    register :: Manager -> IO () -> IO Handle
    register (Manager ref) onTimeout = do
        iactive <- I.newIORef Active
        let h = Handle onTimeout iactive
        I.atomicModifyIORef ref (\x -> (h : x, ()))
        return h

    registerKillThread :: Manager -> IO Handle
    registerKillThread m = do
        wtid <- myThreadId >>= mkWeakThreadId
        register m $ deRefWeak wtid >>= maybe (return ()) killThread

    #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

    tickle, pause, resume, cancel :: Handle -> IO ()
    tickle (Handle _ iactive) = I.writeIORef iactive Active
    pause (Handle _ iactive) = I.writeIORef iactive Paused
    resume = tickle
    cancel (Handle _ iactive) = I.writeIORef iactive Canceled

    -- | Call the inner function with a timeout manager.
    withManager :: Int -- ^ timeout in microseconds
                -> (Manager -> IO a)
                -> IO a
    withManager timeout f = do
        -- FIXME when stopManager is available, use it
        man <- initialize timeout
        f man