graceful-shutdown.hs

#!/usr/bin/env nix-shell
#!nix-shell --pure -i runhaskell -p "haskellPackages.ghcWithPackages (p: [p.wai p.warp p.http-types p.stm p.async])"

{-# LANGUAGE OverloadedStrings #-}

import qualified Control.Concurrent.Async as Async
import qualified Control.Concurrent.STM as STM
import qualified Network.HTTP.Types as HTTP
import qualified Network.Wai as Wai
import qualified Network.Wai.Handler.Warp as Warp

main :: IO ()
main = do
    shutdownSignal <- STM.newTVarIO False
    -- NOTE: If you want to shut down based on eg. OS signals, then register
    -- those here and have the handler flip this bool to `True`.
    let settings
            = Warp.setPort 8080
            . Warp.setInstallShutdownHandler (installShutdownHandler shutdownSignal)
            . Warp.setGracefulShutdownTimeout (Just 30) -- seconds
            . Warp.setBeforeMainLoop beforeMainLoop
            $ Warp.defaultSettings
    print "warp is starting"
    Warp.runSettings settings $ app shutdownSignal
    print "warp is done"

-- | Hook which runs after warp acquires the socket.
--
-- NOTE: This is a good place to drop privileges.
--
-- NOTE: If you're a Type=notify systemd service, tell systemd that you've
-- started up here so that your dependencies can make connections immediately.
beforeMainLoop :: IO ()
beforeMainLoop = do
    print "socket is listening"

-- | Spawn a thread to wait for the shutdown signal and initiate shutdown.
installShutdownHandler :: STM.TVar Bool -> (IO ()) -> IO ()
installShutdownHandler shutdownSignal closeSocket = do
    _ <- Async.async $ do
        STM.atomically $ STM.check =<< STM.readTVar shutdownSignal
        closeSocket
    return ()

app :: STM.TVar Bool -> Wai.Application
app shutdownSignal req respond = do
    print ("Request from", Wai.remoteHost req)
    case Wai.rawPathInfo req of

        "/shutdown" -> do
            STM.atomically $ STM.writeTVar shutdownSignal True
            respond $ Wai.responseLBS HTTP.ok200 [] "shutting down"

        _ -> do
            respond $ Wai.responseLBS HTTP.ok200 [] "hello"