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
data Graceful = Graceful
{ activeConnections :: MVar Int
, shutdownTrigger :: MVar Bool
}
initGraceful :: IO Graceful
initGraceful = Graceful <$> newMVar 0
<*> newMVar False
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..."
gracefulShutdown :: Int
-> 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
waitForTermination :: IO ()
waitForTermination = do
mv <- newEmptyMVar
installHandler softwareTermination (CatchOnce (putMVar mv ())) Nothing
takeMVar mv