Add supervisor with backoff restarts

default.nix

@@ -1,5 +1,6 @@
-{ mkDerivation, aeson, async, base, bytestring, hspec, lens
+{ mkDerivation, aeson, annotated-exception, async, base, bytestring
-, lens-aeson, lib, network, stm, text, time, websockets
+, hedgehog, hspec, hspec-hedgehog, lens, lens-aeson, lib, network
+, stm, text, time, websockets
 }:
 mkDerivation {
   pname = "home-assistant-controller";
@@ -8,11 +9,14 @@ mkDerivation {
   isLibrary = true;
   isExecutable = true;
   libraryHaskellDepends = [
-    aeson async base bytestring lens lens-aeson network stm text time
+    aeson annotated-exception async base bytestring lens lens-aeson
-    websockets
+    network stm text time websockets
   ];
   executableHaskellDepends = [ base ];
-  testHaskellDepends = [ aeson async base hspec stm text ];
+  testHaskellDepends = [
+    aeson annotated-exception async base hedgehog hspec hspec-hedgehog
+    stm text time
+  ];
   license = lib.meta.getLicenseFromSpdxId "BSD-3-Clause";
   mainProgram = "home-assistant-controller";
 }

home-assistant-controller.cabal

@@ -64,6 +64,7 @@ library
                     , HomeAssistant.Runtime
                     , HomeAssistant.Runtime.Bus
                     , HomeAssistant.Runtime.Connection
+                    , HomeAssistant.Runtime.Supervisor
 
     -- Modules included in this library but not exported.
     -- other-modules:
@@ -83,6 +84,7 @@ library
                     , time
                     , stm
                     , async
+                    , annotated-exception
 
     -- Directories containing source files.
     hs-source-dirs:   src
@@ -125,6 +127,8 @@ test-suite home-assistant-controller-test
     other-modules:    BusSpec
                     , ConnectionSpec
                     , RuntimeSpec
+                    , SupervisorSpec
+                    , BackoffProp
 
     -- LANGUAGE extensions used by modules in this package.
     -- other-extensions:
@@ -146,4 +150,8 @@ test-suite home-assistant-controller-test
         stm,
         aeson,
         text,
-        async
+        async,
+        hedgehog,
+        hspec-hedgehog,
+        annotated-exception,
+        time

src/HomeAssistant/Runtime/Supervisor.hs

@@ -0,0 +1,80 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module HomeAssistant.Runtime.Supervisor
+  ( supervised
+  , Backoff(..)
+  , defaultBackoff
+  , backoffDelay
+  , nextAttempt
+  , Fatal(..)
+  ) where
+
+import Control.Concurrent (threadDelay)
+import Control.Exception.Annotated
+  ( Exception
+  , Handler (..)
+  , SomeException
+  , catches
+  , displayException
+  , throw
+  )
+import Data.Text (Text)
+import qualified Data.Text as T
+import Data.Time (NominalDiffTime, diffUTCTime, getCurrentTime)
+import Data.Void (Void)
+
+-- | A failure that cannot be fixed by restarting; propagates out of
+-- 'supervised' and terminates the process.
+newtype Fatal = Fatal Text
+  deriving (Show)
+
+instance Exception Fatal
+
+data Backoff = Backoff
+  { backoffBase  :: NominalDiffTime -- ^ Delay before the first restart
+  , backoffCap   :: NominalDiffTime -- ^ Maximum delay between restarts
+  , backoffQuiet :: NominalDiffTime -- ^ Uptime after which the delay resets
+  }
+  deriving (Eq, Show)
+
+defaultBackoff :: Backoff
+defaultBackoff = Backoff 0.1 30 30
+
+-- | Delay before the @attempt@-th restart: doubles from base, clamped at cap.
+backoffDelay :: Backoff -> Int -> NominalDiffTime
+backoffDelay (Backoff base cap _) attempt = go (attempt - 1) base
+  where
+    go 0 d = d
+    go n d = go (n - 1) (min cap (d * 2))
+
+-- | Attempt number to use after a crash that ran for the given uptime.
+nextAttempt :: Backoff -> NominalDiffTime -> Int -> Int
+nextAttempt (Backoff _ _ quiet) uptime attempt
+  | uptime >= quiet = 1
+  | otherwise = attempt + 1
+
+-- | Runs the action forever, restarting it with backoff after synchronous
+-- exceptions; async exceptions propagate. 'Fatal' is rethrown (crashing the
+-- caller) rather than restarted. The action must never return normally and
+-- must not be wrapped in checkpoints around 'Fatal'-throwing code: a
+-- doubly-wrapped 'Fatal' is indistinguishable from a crash and would be
+-- restarted instead of escalated.
+supervised :: Text -> Backoff -> IO Void -> IO Void
+supervised name backoff action = go 1
+  where
+    go attempt = do
+      start <- getCurrentTime
+      outcome <- (action >> pure (Nothing :: Maybe (Either Fatal SomeException))) `catches`
+        [ Handler $ \(f :: Fatal) -> pure (Just (Left f))
+        , Handler $ \(e :: SomeException) -> pure (Just (Right e))
+        ]
+      case outcome of
+        Nothing -> error "unreachable: supervised action returned"
+        Just (Left f) -> throw f
+        Just (Right e) -> do
+          putStrLn $ "[" <> T.unpack name <> "] attempt " <> show attempt <> " crashed: " <> displayException e
+          let delay = backoffDelay backoff attempt
+          putStrLn $ "[" <> T.unpack name <> "] restarting in " <> show delay <> "s"
+          threadDelay (round (realToFrac delay * 1000000 :: Double))
+          end <- getCurrentTime
+          go (nextAttempt backoff (diffUTCTime end start) attempt)

test/BackoffProp.hs

@@ -0,0 +1,21 @@
+module BackoffProp (spec) where
+
+import Data.Maybe (listToMaybe)
+import Data.Time (NominalDiffTime)
+import qualified Hedgehog.Gen as Gen
+import qualified Hedgehog.Range as Range
+import HomeAssistant.Runtime.Supervisor (Backoff (..), backoffDelay)
+import Test.Hspec (Spec, describe, it)
+import Test.Hspec.Hedgehog (hedgehog, forAll, (===))
+
+spec :: Spec
+spec = describe "backoffDelay" $
+  it "doubles from base, clamped at cap" $ hedgehog $ do
+    baseD <- forAll $ Gen.double (Range.constant 0.0001 10)
+    ratio <- forAll $ Gen.double (Range.constant 1 100)
+    let base = realToFrac baseD :: NominalDiffTime
+        cap = realToFrac (baseD * ratio) :: NominalDiffTime
+        backoff = Backoff base cap 1
+        delays = map (backoffDelay backoff) [1 .. 100 :: Int]
+    listToMaybe delays === Just (min cap base)
+    mapM_ (\(a, b) -> b === min cap (a * 2)) (zip delays (drop 1 delays))

test/Main.hs

@@ -1,12 +1,16 @@
 module Main (main) where
 
 import Test.Hspec (hspec)
+import qualified BackoffProp
 import qualified BusSpec
 import qualified ConnectionSpec
 import qualified RuntimeSpec
+import qualified SupervisorSpec
 
 main :: IO ()
 main = hspec $ do
   BusSpec.spec
   ConnectionSpec.spec
   RuntimeSpec.spec
+  SupervisorSpec.spec
+  BackoffProp.spec

test/SupervisorSpec.hs

@@ -0,0 +1,73 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module SupervisorSpec (spec) where
+
+import Control.Concurrent (newEmptyMVar, putMVar, readMVar, threadDelay)
+import Control.Concurrent.Async (async, cancel, poll, waitCatch)
+import Control.Exception (fromException)
+import Control.Exception.Annotated (AnnotatedException (..), throw)
+import Control.Monad (forever)
+import Data.IORef (atomicModifyIORef', newIORef, readIORef)
+import Data.Maybe (isJust, isNothing)
+import HomeAssistant.Runtime.Supervisor
+import Test.Hspec
+
+tinyBackoff :: Backoff
+tinyBackoff = Backoff 0.001 0.002 0.001
+
+spec :: Spec
+spec = describe "supervised" $ do
+  it "restarts a crashing action until it stays up" $ do
+    counter <- newIORef (0 :: Int)
+    up <- newEmptyMVar
+    let action = do
+          n <- atomicModifyIORef' counter (\c -> (c + 1, c + 1))
+          if n < 3
+            then ioError (userError "boom")
+            else do putMVar up (); forever (threadDelay 1000000)
+    sup <- async (supervised "test" tinyBackoff action)
+    readMVar up
+    threadDelay 50000
+    status <- poll sup
+    isNothing status `shouldBe` True
+    readIORef counter `shouldReturn` 3
+    cancel sup
+
+  it "rethrows Fatal instead of restarting" $ do
+    counter <- newIORef (0 :: Int)
+    let action = do
+          _ <- atomicModifyIORef' counter (\c -> (c + 1, c + 1))
+          throw (Fatal "auth_invalid")
+    sup <- async (supervised "test" tinyBackoff action)
+    res <- waitCatch sup
+    case res of
+      Left se -> case fromException se :: Maybe (AnnotatedException Fatal) of
+        Just _  -> pure ()
+        Nothing -> expectationFailure "expected Fatal to propagate"
+      Right _ -> expectationFailure "supervised returned"
+    threadDelay 50000
+    readIORef counter `shouldReturn` 1
+
+  it "does not restart on async exceptions" $ do
+    counter <- newIORef (0 :: Int)
+    let action = do
+          _ <- atomicModifyIORef' counter (\c -> (c + 1, c + 1))
+          forever (threadDelay 1000000)
+    sup <- async (supervised "test" tinyBackoff action)
+    threadDelay 100000
+    cancel sup
+    threadDelay 100000
+    readIORef counter `shouldReturn` 1
+    status <- poll sup
+    isJust status `shouldBe` True
+
+  describe "nextAttempt" $ do
+    it "resets after a quiet period" $
+      nextAttempt tinyBackoff 0.001 5 `shouldBe` 1
+    it "increments otherwise" $
+      nextAttempt tinyBackoff 0.0005 5 `shouldBe` 6
+
+  describe "backoffDelay" $ do
+    it "starts at base" $ backoffDelay tinyBackoff 1 `shouldBe` 0.001
+    it "doubles" $ backoffDelay tinyBackoff 2 `shouldBe` 0.002
+    it "clamps at cap" $ backoffDelay tinyBackoff 3 `shouldBe` 0.002