From 3d7f6cb3c6d4c2d3b41df2125dc9e2dfc07599a6 Mon Sep 17 00:00:00 2001 From: Mats Rauhala Date: Wed, 16 Sep 2026 12:28:48 +0300 Subject: [PATCH] Fix supervisor and proper rate limiting --- default.nix | 8 +-- home-assistant-controller.cabal | 6 +- src/HomeAssistant/Runtime.hs | 25 +++++---- src/HomeAssistant/Runtime/Connection.hs | 55 +++++++----------- src/HomeAssistant/Runtime/Graphing.hs | 8 +-- src/HomeAssistant/Runtime/Metrics.hs | 41 ++++++++------ src/HomeAssistant/Runtime/RateLimit.hs | 74 +++++++++++++++++++++++++ src/HomeAssistant/Runtime/Supervisor.hs | 66 +++++++--------------- test/BackoffProp.hs | 21 ------- test/ConnectionSpec.hs | 30 +--------- test/Main.hs | 4 -- test/SupervisorSpec.hs | 73 ------------------------ 12 files changed, 164 insertions(+), 247 deletions(-) create mode 100644 src/HomeAssistant/Runtime/RateLimit.hs delete mode 100644 test/BackoffProp.hs delete mode 100644 test/SupervisorSpec.hs diff --git a/default.nix b/default.nix index 9025a71..51ccfba 100644 --- a/default.nix +++ b/default.nix @@ -1,8 +1,8 @@ { mkDerivation, aeson, annotated-exception, async, base, bytestring , cereal, cereal-conduit, conduit, containers, directory, ekg-core , filepath, hedgehog, hspec, hspec-hedgehog, http-types, katip -, lens, lens-aeson, lib, network, process, scotty, stm, text, time -, unordered-containers, uuid, wai, warp, websockets +, lens, lens-aeson, lib, network, process, retry, scotty, stm, text +, time, unordered-containers, uuid, wai, warp, websockets }: mkDerivation { pname = "home-assistant-controller"; @@ -13,8 +13,8 @@ mkDerivation { libraryHaskellDepends = [ aeson annotated-exception async base bytestring cereal cereal-conduit conduit containers directory ekg-core filepath - http-types katip lens lens-aeson network process scotty stm text - time unordered-containers uuid wai warp websockets + http-types katip lens lens-aeson network process retry scotty stm + text time unordered-containers uuid wai warp websockets ]; executableHaskellDepends = [ base ]; testHaskellDepends = [ diff --git a/home-assistant-controller.cabal b/home-assistant-controller.cabal index 9cc6521..922bf4a 100644 --- a/home-assistant-controller.cabal +++ b/home-assistant-controller.cabal @@ -72,6 +72,7 @@ library , HomeAssistant.Runtime.Graphing , HomeAssistant.Runtime.Metrics , HomeAssistant.Runtime.Supervisor + , HomeAssistant.Runtime.RateLimit -- Modules included in this library but not exported. -- other-modules: @@ -85,7 +86,9 @@ library , aeson , lens-aeson , lens + , exceptions , text + , unliftio , network , bytestring , time @@ -108,6 +111,7 @@ library , http-types , warp , wai + , retry -- Directories containing source files. hs-source-dirs: src @@ -150,14 +154,12 @@ test-suite home-assistant-controller-test -- Modules included in this executable, other than Main. other-modules: AFRPLawsSpec , AFRPSpec - , BackoffProp , BedroomSpec , BusSpec , ConnectionSpec , GraphingSpec , MetricsSpec , RuntimeSpec - , SupervisorSpec , Support -- LANGUAGE extensions used by modules in this package. diff --git a/src/HomeAssistant/Runtime.hs b/src/HomeAssistant/Runtime.hs index f1e14d3..541edaa 100644 --- a/src/HomeAssistant/Runtime.hs +++ b/src/HomeAssistant/Runtime.hs @@ -14,7 +14,6 @@ module HomeAssistant.Runtime ) where import AFRP (Event (..), Mealy (..), Request (..), Auto, stepAutoSerializing, load, DecodedAuto (..)) -import Control.Concurrent.Async (async, waitAny) import Control.Concurrent.STM (atomically, dupTChan, readTChan) import Data.Aeson (Value) import qualified Data.Text as T @@ -22,8 +21,7 @@ import Data.Time (getCurrentTime, getCurrentTimeZone) import Data.Void (Void, absurd) import HomeAssistant.Controller (HASS, HASSEff (..)) import HomeAssistant.Runtime.Bus -import HomeAssistant.Runtime.Connection (readerAction, writerAction) -import HomeAssistant.Runtime.Supervisor (defaultBackoff, supervised) +import HomeAssistant.Runtime.Connection (writerAction, readerAction) import Network.Socket (withSocketsDo) import System.Environment (getEnv, lookupEnv) import HomeAssistant.Controller.Bedroom (bedroomPresenceController, bedroomButtonController, bedroomDrawerController, humidifierController) @@ -36,11 +34,14 @@ import HomeAssistant.Controller.Children (schoolLightController) import Data.Maybe (fromMaybe) import qualified System.Metrics import qualified HomeAssistant.Runtime.Metrics -import qualified HomeAssistant.Runtime.Graphing import System.FilePath (()) import Text.Read (readMaybe) import HomeAssistant.Controller.Kitchen (kitchenMotionController) import HomeAssistant.Controller.Livingroom (livingroomPresenceController) +import HomeAssistant.Runtime.RateLimit (slidingWindowLimiter, registerRateLimitMetrics) +import UnliftIO.Async +import HomeAssistant.Runtime.Supervisor (supervised) +import qualified HomeAssistant.Runtime.Graphing step :: (MonadIO m) => FilePath -> UUID -> Auto m a b -> a -> m (b, Auto m a b) step path trace st a = do @@ -66,8 +67,8 @@ controllers = -- | Steps the machine for every inbound message; service calls go to the -- bus. A restart re-dups the inbound channel and starts from the machine's -- initial state; messages broadcast during the restart window are lost. -runController :: FilePath -> Bus -> Controller -> IO Void -runController rootDir bus (Controller name machine _enabled) = do +runController :: MonadIO m => FilePath -> Bus -> Controller -> m Void +runController rootDir bus (Controller name machine _enabled) = liftIO $ do inbound <- atomically (dupTChan (busInbound bus)) let ns = Namespace [name] let workerDefinition = runMealy machine (runKatipContextT (busLogEnv bus) () ns . channelHassEval bus) @@ -98,6 +99,7 @@ defaultMain = withSocketsDo $ do store <- System.Metrics.newStore System.Metrics.registerGcMetrics store appMetrics <- HomeAssistant.Runtime.Metrics.registerAppMetrics store + rateLimitMetrics <- registerRateLimitMetrics store withBus severity appMetrics $ \bus -> do token <- getEnv "HA_TOKEN" host <- getEnv "HA_HOST" @@ -105,16 +107,17 @@ defaultMain = withSocketsDo $ do rrdPath <- fromMaybe "hass-controller.rrd" <$> lookupEnv "HA_RRD_PATH" rrdtool <- fromMaybe "rrdtool" <$> lookupEnv "HA_RRDTOOL" metricsPort <- lookupPort + writerLimiter <- slidingWindowLimiter rateLimitMetrics 10 20 let active = [c | c@(Controller _ _ True) <- controllers] ents = foldMap (\(Controller _ m _) -> entities m) active workers = [ ("reader", readerAction host 8123 token ents bus) - , ("writer", writerAction bus) + , ("writer", writerAction writerLimiter bus) + , ("metrics", HomeAssistant.Runtime.Metrics.metricsAction store rrdPath rrdtool) + , ("metrics-http", HomeAssistant.Runtime.Graphing.graphAction rrdPath rrdtool metricsPort) ] ++ [ (name, runController rootPath bus c) | c@(Controller name _ True) <- controllers ] - ++ [ ("metrics", HomeAssistant.Runtime.Metrics.metricsAction store rrdPath rrdtool) - , ("metrics-http", HomeAssistant.Runtime.Graphing.graphAction rrdPath rrdtool metricsPort) - ] - as <- mapM (\(name, act) -> async (supervised name defaultBackoff act)) workers + as <- runKatipContextT (busLogEnv bus) () mempty $ + mapM (\(name, act) -> async (supervised name act)) workers (_, v) <- waitAny as absurd v diff --git a/src/HomeAssistant/Runtime/Connection.hs b/src/HomeAssistant/Runtime/Connection.hs index ecad52e..616b30d 100644 --- a/src/HomeAssistant/Runtime/Connection.hs +++ b/src/HomeAssistant/Runtime/Connection.hs @@ -5,17 +5,14 @@ module HomeAssistant.Runtime.Connection ( readerAction , writerAction , encodeService - , dedupeBatch , isTriggerEvent ) where import Control.Concurrent.STM - ( TChan - , atomically + ( atomically , readTChan , readTVar , retry - , tryReadTChan , writeTChan , writeTVar ) @@ -27,8 +24,6 @@ import Control.Lens ((^?)) import Control.Monad (forever, forM_, when) import Data.Aeson (Value, eitherDecode, encode, object, (.=)) import Data.Aeson.Lens (key, _String) -import Data.List (sort) -import qualified Data.Map.Strict as M import qualified Data.Set as S import qualified Data.Text as T import Data.Void (Void) @@ -36,16 +31,18 @@ import HomeAssistant.Controller (Service (..), Target (..)) import HomeAssistant.Runtime.Bus import HomeAssistant.Runtime.Supervisor (Fatal (..)) import qualified Network.WebSockets as WS -import Katip (runKatipContextT, sl, logFM, Severity (..), ls) +import Katip (sl, logFM, Severity (..), ls, KatipContext, katipAddNamespace, katipAddContext) import Data.UUID (toText) import AFRP (Request(..), Event(..)) -import Control.Monad.IO.Class (liftIO) +import Control.Monad.IO.Class (liftIO, MonadIO) +import HomeAssistant.Runtime.RateLimit (RateLimited, RateLimiter, runRateLimited) +import Control.Monad.Catch (MonadCatch) -- | Connect, authenticate, subscribe, then receive and broadcast forever. -- Restarting this action reconnects. All setup sends happen before the -- connection is published in the bus, so only the writer sends afterwards. -readerAction :: String -> Int -> String -> S.Set T.Text -> Bus -> IO Void -readerAction host port token ents bus = +readerAction :: MonadIO m => String -> Int -> String -> S.Set T.Text -> Bus -> m Void +readerAction host port token ents bus = liftIO $ WS.runClient host port "/api/websocket" $ \conn -> do handshake conn token subscribe bus conn ents @@ -114,20 +111,21 @@ receiveJSON conn = do -sendWithId :: Bus -> WS.Connection -> Request -> Service -> IO () -sendWithId bus conn request svc = - runKatipContextT (busLogEnv bus) (sl "traceId" (toText (requestTraceId request))) "connection" $ do - callId <- liftIO $ generateCallId (busGen bus) - let textData = encode $ encodeService callId svc - logFM DebugS (ls textData) - liftIO $ recordOutbound bus - liftIO $ WS.sendTextData conn textData +sendWithId :: (MonadIO m, KatipContext m) => Bus -> WS.Connection -> Service -> RateLimited m () +sendWithId bus conn svc = katipAddNamespace "connection" $ do + callId <- liftIO $ generateCallId (busGen bus) + let textData = encode $ encodeService callId svc + logFM DebugS (ls textData) + liftIO $ recordOutbound bus + liftIO $ WS.sendTextData conn textData -writerAction :: Bus -> IO Void -writerAction bus = forever $ do - (request, svc) <- atomically $ readTChan (busOutbound bus) - conn <- atomically $ readTVar (busConn bus) >>= maybe retry pure - sendWithId bus conn request svc +writerAction :: (MonadIO m, KatipContext m, MonadCatch m) => RateLimiter -> Bus -> m Void +writerAction rateLimiter bus = forever $ do + (request, svc) <- liftIO $ atomically $ readTChan (busOutbound bus) + katipAddContext (sl "traceId" (toText (requestTraceId request))) $ do + conn <- liftIO $ atomically $ readTVar (busConn bus) >>= maybe retry pure + -- runRateLimited rateLimiter $ sendWithId bus conn request svc + runRateLimited rateLimiter $ sendWithId bus conn svc encodeService :: Int -> Service -> Value encodeService callId Service{..} = object $ @@ -138,17 +136,6 @@ encodeService callId Service{..} = object $ , "target" .= targetObject serviceTarget ] <> maybe [] (\d -> ["service_data" .= d]) serviceData --- | Collapse a drained batch of outbound calls: the newest call per --- `(domain, service, sorted-targets)` survives; older duplicates are --- dropped. `serviceData` is not part of the key, so a newer `turn_on` --- with different brightness supersedes an older one to the same target. -dedupeBatch :: [(Request, Service)] -> [(Request, Service)] -dedupeBatch = M.elems . foldl' ins M.empty - where - ins m (req, svc) = M.insert (dedupeKey svc) (req, svc) m - -dedupeKey :: Service -> (T.Text, T.Text, [Target]) -dedupeKey Service{..} = (serviceDomain, serviceName, sort serviceTarget) -- | A single target encodes as a scalar; multiple encode as a list. Empty -- lists are omitted so Home Assistant receives only populated keys. diff --git a/src/HomeAssistant/Runtime/Graphing.hs b/src/HomeAssistant/Runtime/Graphing.hs index a2dc3d6..afb93d7 100644 --- a/src/HomeAssistant/Runtime/Graphing.hs +++ b/src/HomeAssistant/Runtime/Graphing.hs @@ -16,7 +16,7 @@ module HomeAssistant.Runtime.Graphing import Control.Exception (evaluate) import Control.Monad (forever) -import Control.Monad.IO.Class (liftIO) +import Control.Monad.IO.Class (liftIO, MonadIO) import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as BL import Data.Text (Text) @@ -155,9 +155,9 @@ handler rrdPath rrdtool = do app :: FilePath -> FilePath -> ScottyM () app rrdPath rrdtool = get "/metrics/:name" (handler rrdPath rrdtool) -graphAction :: FilePath -> FilePath -> Int -> IO Void -graphAction rrdPath rrdtool port = forever $ scottyOpts opts (app rrdPath rrdtool) +graphAction :: MonadIO m => FilePath -> FilePath -> Int -> m Void +graphAction rrdPath rrdtool port = liftIO $ forever $ scottyOpts opts (app rrdPath rrdtool) where opts = defaultOptions { settings = setHost "127.0.0.1" (setPort port defaultSettings) - } \ No newline at end of file + } diff --git a/src/HomeAssistant/Runtime/Metrics.hs b/src/HomeAssistant/Runtime/Metrics.hs index 5d9d92f..02ca9fd 100644 --- a/src/HomeAssistant/Runtime/Metrics.hs +++ b/src/HomeAssistant/Runtime/Metrics.hs @@ -35,6 +35,9 @@ import System.Metrics.Counter (Counter) import System.Directory (doesFileExist, renameFile) import System.Exit (ExitCode (..)) import System.Process (callProcess, readProcessWithExitCode) +import Control.Monad.IO.Class (MonadIO, liftIO) +import Control.Exception.Annotated (checkpointCallStack) +import Control.Monad.Catch (MonadCatch) data DsType = Derive | Gauge deriving (Eq, Ord, Show) @@ -81,6 +84,8 @@ registerAppMetrics store = <$> M.createCounter "hass.trigger.in" store <*> M.createCounter "hass.service.out" store + + buildCreateArgs :: FilePath -> Int -> [(String, DsType)] -> [String] buildCreateArgs path step specs = ["create", path, "--step", show step] @@ -134,44 +139,44 @@ lookupValue sample name = case HM.lookup name sample of Just (M.Gauge n) -> Just n _ -> Nothing -readExistingDs :: FilePath -> FilePath -> IO [(String, DsType)] -readExistingDs rrdtool rrdPath = do - (code, out, _) <- readProcessWithExitCode rrdtool ["info", rrdPath] "" +readExistingDs :: (MonadCatch m, MonadIO m) => FilePath -> FilePath -> m [(String, DsType)] +readExistingDs rrdtool rrdPath = checkpointCallStack $ do + (code, out, _) <- liftIO $ readProcessWithExitCode rrdtool ["info", rrdPath] "" pure $ case code of ExitSuccess -> parseInfoDs out _ -> [] -backupPath :: FilePath -> IO FilePath -backupPath path = do +backupPath :: MonadIO m => FilePath -> m FilePath +backupPath path = liftIO $ do now <- getCurrentTime pure $ path ++ ".bak-" ++ formatTime defaultTimeLocale "%Y%m%dT%H%M%S" now -ensureRrd :: FilePath -> FilePath -> [DsSpec] -> IO () -ensureRrd rrdtool rrdPath schema = do - exists <- doesFileExist rrdPath +ensureRrd :: (MonadCatch m, MonadIO m) => FilePath -> FilePath -> [DsSpec] -> m () +ensureRrd rrdtool rrdPath schema = checkpointCallStack $ do + exists <- liftIO $ doesFileExist rrdPath if not exists then create else do current <- readExistingDs rrdtool rrdPath unless (schemaMatches schema current) $ do backup <- backupPath rrdPath - renameFile rrdPath backup + liftIO $ renameFile rrdPath backup create where - create = callProcess rrdtool (buildCreateArgs rrdPath 10 (map toPair schema)) + create = liftIO $ callProcess rrdtool (buildCreateArgs rrdPath 10 (map toPair schema)) toPair s = (dsName s, dsType s) -sampleAndUpdate :: M.Store -> FilePath -> FilePath -> [DsSpec] -> IO () -sampleAndUpdate store rrdtool rrdPath schema = do - sample <- M.sampleAll store +sampleAndUpdate :: (MonadCatch m, MonadIO m) => M.Store -> FilePath -> FilePath -> [DsSpec] -> m () +sampleAndUpdate store rrdtool rrdPath schema = checkpointCallStack $ do + sample <- liftIO $ M.sampleAll store let names = map dsName schema values = map (lookupValue sample . dsEkgName) schema - callProcess rrdtool (buildUpdateArgs rrdPath names values) + liftIO $ callProcess rrdtool (buildUpdateArgs rrdPath names values) -metricsAction :: M.Store -> FilePath -> FilePath -> IO Void -metricsAction store rrdPath rrdtool = do - schema <- buildSchema <$> M.sampleAll store +metricsAction :: (MonadCatch m, MonadIO m) => M.Store -> FilePath -> FilePath -> m Void +metricsAction store rrdPath rrdtool = checkpointCallStack $ do + schema <- buildSchema <$> liftIO (M.sampleAll store) ensureRrd rrdtool rrdPath schema forever $ do sampleAndUpdate store rrdtool rrdPath schema - threadDelay 10000000 + liftIO $ threadDelay 10000000 diff --git a/src/HomeAssistant/Runtime/RateLimit.hs b/src/HomeAssistant/Runtime/RateLimit.hs new file mode 100644 index 0000000..26ca078 --- /dev/null +++ b/src/HomeAssistant/Runtime/RateLimit.hs @@ -0,0 +1,74 @@ +{-# Language DerivingVia #-} +{-# Language OverloadedStrings #-} +module HomeAssistant.Runtime.RateLimit where +import Control.Monad.IO.Class (MonadIO, liftIO) +import Data.Time (NominalDiffTime) +import Control.Concurrent.STM (newTVarIO, atomically, readTVar, writeTVar) +import Data.Time.Clock.POSIX (getPOSIXTime) +import Katip.Monadic (KatipContext(..)) +import Control.Retry (fullJitterBackoff, limitRetries, retrying, RetryStatus (..)) +import Data.Maybe (isNothing) +import Control.Exception.Annotated (Exception, checkpointCallStack, throw) +import Control.Monad.Catch (MonadCatch) +import Katip (Katip, Severity (WarningS), logFM, ls) +import Control.Monad (when) +import System.Metrics.Counter (Counter) +import qualified System.Metrics.Counter as Counter +import System.Metrics (Store, createCounter) + +data RateLimitMetrics = RateLimitMetrics + { rlRetries :: Counter + , rlRejects :: Counter + } + + +registerRateLimitMetrics :: MonadIO m => Store -> m RateLimitMetrics +registerRateLimitMetrics store = liftIO $ RateLimitMetrics + <$> createCounter "ratelimit.retry" store + <*> createCounter "ratelimit.exhausted" store + +newtype RateLimited m a = RateLimited (m a) + deriving (Functor, Applicative, Monad, MonadIO, KatipContext, Katip) via m + + +data RateLimiter = RateLimiter + { metrics :: RateLimitMetrics + , limiter :: forall m a. (KatipContext m, MonadIO m) => RateLimited m a -> m (Maybe a) + } + + +slidingWindowLimiter + :: RateLimitMetrics + -> NominalDiffTime -- ^ Within how many seconds + -> Int -- ^ How many requests we can make + -> IO RateLimiter +slidingWindowLimiter metrics timeWindow size = do + window <- newTVarIO [] + return $ RateLimiter metrics $ \(RateLimited action) -> do + now <- liftIO getPOSIXTime + allowed <- liftIO $ atomically $ do + elements <- takeWhile (\x -> x > now - timeWindow) <$> readTVar window + if length elements >= size + then pure False + else True <$ writeTVar window (now : elements) + if allowed + then Just <$> action + else pure Nothing + +data RateLimitException = RateLimitExhausted + deriving Show + +instance Exception RateLimitException + +runRateLimited :: (MonadIO m, KatipContext m, MonadCatch m) => RateLimiter -> RateLimited m a -> m a +runRateLimited (RateLimiter metrics run) action = checkpointCallStack $ do + x <- retrying retryPolicy (const $ pure . isNothing) $ \RetryStatus{rsIterNumber} -> do + when (rsIterNumber > 0) $ do + logFM WarningS $ ls $ "Rate limited, attempt no. " <> show rsIterNumber + incRetry + run action + maybe (incReject >> throw RateLimitExhausted) pure x + where + retryPolicy = fullJitterBackoff 50000 <> limitRetries 5 + incRetry = liftIO $ Counter.inc (rlRetries metrics) + incReject = liftIO $ Counter.inc (rlRejects metrics) diff --git a/src/HomeAssistant/Runtime/Supervisor.hs b/src/HomeAssistant/Runtime/Supervisor.hs index a8e7454..b9b3db3 100644 --- a/src/HomeAssistant/Runtime/Supervisor.hs +++ b/src/HomeAssistant/Runtime/Supervisor.hs @@ -2,26 +2,23 @@ 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 + , checkpoint, Annotation (..) ) import Data.Text (Text) -import qualified Data.Text as T -import Data.Time (NominalDiffTime, diffUTCTime, getCurrentTime) import Data.Void (Void) +import Control.Monad.IO.Class (MonadIO) +import Control.Monad.Catch (MonadCatch, MonadMask) +import Katip (KatipContext, logFM, ls) +import Control.Retry (RetryStatus(..), recovering, limitRetries, fullJitterBackoff) +import Control.Monad (when) +import Katip.Core (Severity(..)) -- | A failure that cannot be fixed by restarting; propagates out of -- 'supervised' and terminates the process. @@ -30,28 +27,11 @@ newtype Fatal = Fatal Text 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 +data SupervisorException = RetriesExhausted + deriving Show --- | Delay before the @attempt@-th restart: doubles from base, clamped at cap. -backoffDelay :: Backoff -> Int -> NominalDiffTime -backoffDelay (Backoff base cap _) attempt = go (max 0 (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 +instance Exception SupervisorException -- | Runs the action forever, restarting it with backoff after synchronous -- exceptions; async exceptions propagate. 'Fatal' is rethrown (crashing the @@ -59,22 +39,14 @@ nextAttempt (Backoff _ _ quiet) uptime attempt -- 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 +supervised :: (MonadIO m, MonadCatch m, KatipContext m, MonadMask m) => Text -> m Void -> m Void +supervised name action = checkpoint (Annotation name) $ do + recovering retryPolicy handlers $ \RetryStatus{rsIterNumber} -> do + when (rsIterNumber > 0) $ logFM WarningS $ ls $ "Child crashed, retry no " <> show rsIterNumber + action 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)) + handlers = + [ \_retryStatus -> Handler $ \(_ :: Fatal) -> pure False + , \_retryStatus -> Handler $ \(_ :: SomeException) -> pure True ] - end <- getCurrentTime - 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 - threadDelay (round (realToFrac delay * 1000000 :: Double)) - go (nextAttempt backoff (diffUTCTime end start) attempt) + retryPolicy = fullJitterBackoff 50000 <> limitRetries 5 diff --git a/test/BackoffProp.hs b/test/BackoffProp.hs deleted file mode 100644 index 11c64cf..0000000 --- a/test/BackoffProp.hs +++ /dev/null @@ -1,21 +0,0 @@ -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)) diff --git a/test/ConnectionSpec.hs b/test/ConnectionSpec.hs index b279380..733fb36 100644 --- a/test/ConnectionSpec.hs +++ b/test/ConnectionSpec.hs @@ -9,7 +9,7 @@ import Data.Text (Text) import Data.Time (UTCTime (..), utc) import Data.UUID (fromString) import HomeAssistant.Controller (Service (..), Target(..)) -import HomeAssistant.Runtime.Connection (encodeService, dedupeBatch, isTriggerEvent) +import HomeAssistant.Runtime.Connection (encodeService, isTriggerEvent) import Test.Hspec spec :: Spec @@ -44,34 +44,6 @@ spec = do it "is false when there is no type" $ isTriggerEvent (object ["id" .= (1 :: Int)]) `shouldBe` False - describe "dedupeBatch" $ do - it "collapses identical calls to one" $ - let batch = [ (req 1, lightOn [AreaId "x"]) - , (req 2, lightOn [AreaId "x"]) - , (req 3, lightOn [AreaId "x"]) - ] - in dedupeBatch batch `shouldBe` [(req 3, lightOn [AreaId "x"])] - - it "keeps same-target different-service calls separate" $ - let batch = [ (req 1, lightOn [AreaId "x"]) - , (req 2, lightOff [AreaId "x"]) - ] - result = dedupeBatch batch - in length result `shouldBe` 2 - - it "newest call wins for the same key" $ - let batch = [ (req 1, lightOn [AreaId "x"]) - , (req 2, lightOn [AreaId "x"]) - , (req 3, lightOn [AreaId "x"]) - ] - in map requestTraceId (map fst (dedupeBatch batch)) `shouldBe` - [fromJust (fromString "00000000-0000-0000-0000-000000000003")] - - it "treats target lists in different order as the same key" $ - let batch = [ (req 1, lightOn [EntityId "a", EntityId "b"]) - , (req 2, lightOn [EntityId "b", EntityId "a"]) - ] - in length (dedupeBatch batch) `shouldBe` 1 req :: Int -> Request req n = Request (UTCTime (toEnum 0) (fromIntegral (0 :: Int))) utc diff --git a/test/Main.hs b/test/Main.hs index d27dce3..fdd0264 100644 --- a/test/Main.hs +++ b/test/Main.hs @@ -3,14 +3,12 @@ module Main (main) where import Test.Hspec (hspec) import qualified AFRPLawsSpec import qualified AFRPSpec -import qualified BackoffProp import qualified BedroomSpec import qualified BusSpec import qualified ConnectionSpec import qualified GraphingSpec import qualified MetricsSpec import qualified RuntimeSpec -import qualified SupervisorSpec main :: IO () main = hspec $ do @@ -22,5 +20,3 @@ main = hspec $ do GraphingSpec.spec MetricsSpec.spec RuntimeSpec.spec - SupervisorSpec.spec - BackoffProp.spec diff --git a/test/SupervisorSpec.hs b/test/SupervisorSpec.hs deleted file mode 100644 index 077981d..0000000 --- a/test/SupervisorSpec.hs +++ /dev/null @@ -1,73 +0,0 @@ -{-# 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