{-# LANGUAGE OverloadedStrings #-}

-- | 
--
-- > app req = do
-- >     liftIO $ somethingImportant -- like threadDelay 100000000
-- >     LBSResponse status200 [("Content-Type", "text/plain")] "Finished!"
-- >
-- > main = do
-- >     gs <- initGraceful
-- >
-- >     forkIO $ run 8000 (graceful gs app)
-- >
-- >     waitForTermination >> gracefulShutdown gs
-- >     putStrLn "Bye!"

module Network.Wai.Graceful
    ( initGraceful, Graceful(..)
    , graceful
    , gracefulShutdown
    , waitForTermination
    ) where

import Control.Applicative
import Control.Monad.Trans (liftIO)
import Control.Monad.Trans.Resource (register)
import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.MVar
import Network.Wai (Application, Response, responseLBS)
import Network.HTTP.Types (status503)

import System.Posix.Signals -- (installHandler, sigTERM)

data Graceful = Graceful
    { activeConnections :: MVar Int
    , shutdownTrigger :: MVar Bool
    }

initGraceful :: IO Graceful
initGraceful = Graceful <$> newMVar 0
                        <*> newMVar False

-- | WAI Middleware that keeps track of active connections
--   and blocks further requests with HTTP503.
graceful :: Graceful -> Application -> Application
graceful (Graceful active shutdown) app req = do
    halt <- liftIO $ readMVar shutdown
    if halt
        then return shuttingDown
        else do
            liftIO $ modifyMVar_ active (\c -> return $! c + 1)
            register $ liftIO $ modifyMVar_ active (\c -> return $! c - 1)
            app req

    where
        shuttingDown =
            responseLBS
                status503
                [("Content-Type", "text/plain")]
                "Shutdown in progress..."

-- | Trigger shutdown and monitor progress.
gracefulShutdown :: Int      -- ^ Timeout in seconds.
                 -> Graceful
                 -> IO ()
gracefulShutdown seconds (Graceful active shutdown) = do
    modifyMVar_ shutdown (const $ return True)

    done <- newEmptyMVar

    let waitActive = do count <- readMVar active
                        if count > 0
                            then threadDelay 100000 >> waitActive
                            else putMVar done ()

    let timeout = threadDelay (1000000 * seconds) >> putMVar done ()

    forkIO waitActive
    forkIO timeout
    takeMVar done

-- | Wait for TERM signal.
waitForTermination :: IO ()
waitForTermination = do
    mv <- newEmptyMVar
    installHandler softwareTermination (CatchOnce (putMVar mv ())) Nothing
    takeMVar mv