Supervisor.hs

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Rank2Types #-}

-- |
-- Module      : Control.Concurrent.NQE.Supervisor
-- Copyright   : No rights reserved
-- License     : UNLICENSE
-- Maintainer  : xenog@protonmail.com
-- Stability   : experimental
-- Portability : POSIX
--
-- Supervisors run and monitor processes, including other supervisors. A supervisor
-- has a corresponding 'Strategy' that controls its behaviour if a child stops.
-- Supervisors deal with exceptions in concurrent processes so that their code does
-- not need to be written in an overly-defensive style. They help prevent problems
-- caused by processes dying quietly in the background, potentially locking an
-- entire application.
module Control.Concurrent.NQE.Supervisor
  ( ChildAction,
    Child,
    SupervisorMessage,
    Supervisor,
    Strategy (..),
    withSupervisor,
    supervisor,
    supervisorProcess,
    addChild,
    removeChild,
  )
where

import Control.Applicative
import Control.Concurrent.NQE.Process
import Control.Concurrent.STM (retry)
import Control.Monad
import Data.List
import UnliftIO

-- | Alias for child action to be executed asynchronously by supervisor.
type ChildAction = IO ()

-- | Thread handler for child.
type Child = Async ()

-- | Send this message to a supervisor to add or remove a child.
data SupervisorMessage
  = AddChild
      !ChildAction
      !(Listen Child)
  | RemoveChild
      !Child
      !(Listen ())

-- | Alias for supervisor process.
type Supervisor = Process SupervisorMessage

-- | Supervisor strategies to decide what to do when a child stops.
data Strategy
  = -- | send a 'SupervisorNotif' to 'Mailbox' when child dies
    Notify (Listen (Child, Maybe SomeException))
  | -- | kill all processes and propagate exception upstream
    KillAll
  | -- | ignore processes that stop without raising an exception
    IgnoreGraceful
  | -- | keep running if a child dies and ignore it
    IgnoreAll

-- | Run a supervisor asynchronously and pass its mailbox to a function.
-- Supervisor will be stopped along with all its children when the function
-- ends.
withSupervisor ::
  (MonadUnliftIO m) =>
  Strategy ->
  (Supervisor -> m a) ->
  m a
withSupervisor = withProcess . supervisorProcess

-- | Run a supervisor as an asynchronous process.
supervisor :: (MonadUnliftIO m) => Strategy -> m Supervisor
supervisor strat = process (supervisorProcess strat)

-- | Run a supervisor in the current thread.
supervisorProcess ::
  (MonadUnliftIO m) =>
  Strategy ->
  Inbox SupervisorMessage ->
  m ()
supervisorProcess strat i = do
  state <- newTVarIO []
  finally (loop state) (stopAll state)
  where
    loop state = do
      e <- atomically $ Right <$> receiveSTM i <|> Left <$> waitForChild state
      again <-
        case e of
          Right m -> processMessage state m
          Left x -> processDead state strat x
      when again $ loop state

-- | Add a new 'ChildAction' to the supervisor. Will return the 'Child' that was
-- just started. This function will not block or raise an exception if the child
-- dies.
addChild :: (MonadIO m) => Supervisor -> ChildAction -> m Child
addChild sup action = AddChild action `query` sup

-- | Stop a 'Child' controlled by this supervisor. Will block until the child
-- dies.
removeChild :: (MonadIO m) => Supervisor -> Child -> m ()
removeChild sup c = RemoveChild c `query` sup

-- | Internal function to stop all children.
stopAll :: (MonadUnliftIO m) => TVar [Child] -> m ()
stopAll state = mask_ $ do
  as <- readTVarIO state
  mapM_ cancel as

-- | Internal function to wait for a child process to finish running.
waitForChild :: TVar [Child] -> STM (Child, Either SomeException ())
waitForChild state = do
  as <- readTVar state
  when (null as) retry
  waitAnyCatchSTM as

-- | Internal function to process incoming supervisor message.
processMessage ::
  (MonadUnliftIO m) => TVar [Child] -> SupervisorMessage -> m Bool
processMessage state (AddChild ch r) = do
  a <- startChild state ch
  atomically $ r a
  return True
processMessage state (RemoveChild a r) = do
  stopChild state a
  atomically $ r ()
  return True

-- | Internal function to run when a child process dies.
processDead ::
  (MonadUnliftIO m) =>
  TVar [Child] ->
  Strategy ->
  (Child, Either SomeException ()) ->
  m Bool
processDead state IgnoreAll (a, _) = do
  atomically . modifyTVar' state $ filter (/= a)
  return True
processDead state KillAll (a, e) = do
  atomically $ modifyTVar' state . filter $ (/= a)
  stopAll state
  case e of
    Left x -> throwIO x
    Right () -> return False
processDead state IgnoreGraceful (a, Right ()) = do
  atomically (modifyTVar' state (filter (/= a)))
  return True
processDead state IgnoreGraceful (a, Left e) = do
  atomically $ modifyTVar' state (filter (/= a))
  stopAll state
  throwIO e
processDead state (Notify notif) (a, ee) = do
  atomically $ do
    as <- readTVar state
    case find (== a) as of
      Just p -> notif (p, me)
      Nothing -> return ()
    modifyTVar state (filter (/= a))
  return True
  where
    me =
      case ee of
        Left e -> Just e
        Right () -> Nothing

-- | Internal function to start a child process.
startChild :: (MonadUnliftIO m) => TVar [Child] -> ChildAction -> m Child
startChild state ch = mask_ $ do
  a <- liftIO $ async ch
  atomically $ modifyTVar' state (a :)
  return a

-- | Internal fuction to stop a child process.
stopChild :: (MonadUnliftIO m) => TVar [Child] -> Child -> m ()
stopChild state a = mask_ $ do
  isChild <-
    atomically $ do
      cur <- readTVar state
      let new = filter (/= a) cur
      writeTVar state new
      return (cur /= new)
  when isChild $ cancel a