Simple Warp server that can be gracefully shutdown over HTTP.

gistfile1.hs

{-# LANGUAGE OverloadedStrings #-}

module Main (main) where

import Control.Concurrent (forkIO)
import Control.Concurrent.STM
import Control.Monad (when)
import Control.Monad.Trans (liftIO)
import Network.HTTP.Types
import Network.Wai as Wai
import Network.Wai.Handler.Warp as Warp

app :: TMVar () -> Application
app shutdown req = liftIO . atomically $ do
  let textResponse = return . responseLBS ok200 [(hContentType, "text/plain")]
  shouldRun <- isEmptyTMVar shutdown
  case (shouldRun, pathInfo req) of
    -- if the TMVar is full we're shutting down
    (False, _)                 -> return $ responseLBS serviceUnavailable503 [] ""
    -- if it's not full and a shutdown was requested, signal the shutdown TMVar
    (True, ["shutdownByMVar"]) -> putTMVar shutdown () >> textResponse "shutting down!"
    -- otherwise handle the request as normal...
    _                          -> textResponse "ok"

main :: IO ()
main = do
  -- The pair of (shutdown, activeConnections) will be used to signal our exit criteria
  shutdown <- newEmptyTMVarIO
  activeConnections <- newTVarIO (0 :: Int)

  -- Hook the OnOpen/OnClose events to manage the activeConnection count
  let settings = defaultSettings
        { settingsOnOpen  = atomically $ modifyTVar' activeConnections (+1)
        , settingsOnClose = atomically $ modifyTVar' activeConnections (subtract 1)
        }

  _ <- forkIO $ Warp.runSettings settings (app shutdown)

  -- Once the shutdown has been triggered, wait until the active connection count
  -- drops to zero before exiting.
  atomically $ do
    takeTMVar shutdown
    conns <- readTVar activeConnections
    when (conns /= 0)
      retry