Graceful handling of SIGTERM for Haskell programs

GracefulTermination.hs

module System.GracefulTermination where

import Control.Concurrent.Async (race)
import Control.Concurrent.MVar
import Data.Functor (void)
import System.Posix.Signals

-- | Ensure that @SIGTERM@ is handled gracefully, because it's how containers are stopped.
--
-- @action@ will receive an 'AsyncCancelled' exception if @SIGTERM@ is received by the process.
--
-- Typical use:
--
-- > main :: IO ()
-- > main = withGracefulTermination_ $ do
--
-- Note that although the Haskell runtime handles @SIGINT@ it doesn't do anything with @SIGTERM@.
-- Therefore, the program will be killed immediately and no cleanup will be performed. In particular,
-- exception handlers in `bracket`, `finally`, `onException`, etc. won't be run. However, if the
-- program is running as PID 1 in a container, @SIGTERM@ will be ignored and the program will keep
-- running. This will likely result in a @SIGKILL@ being sent a short while later and cleanup still
-- won't be performed.
withGracefulTermination :: IO a -> IO (Maybe a)
withGracefulTermination action = do
  var <- newEmptyMVar
  let terminate = void $ tryPutMVar var ()
      waitForTermination = takeMVar var
  void $ installHandler sigTERM (CatchOnce terminate) Nothing
  either (const Nothing) Just <$> race waitForTermination action

-- | Like 'withGracefulTermination' but ignoring the return value
withGracefulTermination_ :: IO a -> IO ()
withGracefulTermination_ = void . withGracefulTermination