Fix supervisor and proper rate limiting

This commit is contained in:
2026-09-16 12:28:48 +03:00
parent 64561718a5
commit 3d7f6cb3c6
12 changed files with 164 additions and 247 deletions
+4 -4
View File
@@ -1,8 +1,8 @@
{ mkDerivation, aeson, annotated-exception, async, base, bytestring { mkDerivation, aeson, annotated-exception, async, base, bytestring
, cereal, cereal-conduit, conduit, containers, directory, ekg-core , cereal, cereal-conduit, conduit, containers, directory, ekg-core
, filepath, hedgehog, hspec, hspec-hedgehog, http-types, katip , filepath, hedgehog, hspec, hspec-hedgehog, http-types, katip
, lens, lens-aeson, lib, network, process, scotty, stm, text, time , lens, lens-aeson, lib, network, process, retry, scotty, stm, text
, unordered-containers, uuid, wai, warp, websockets , time, unordered-containers, uuid, wai, warp, websockets
}: }:
mkDerivation { mkDerivation {
pname = "home-assistant-controller"; pname = "home-assistant-controller";
@@ -13,8 +13,8 @@ mkDerivation {
libraryHaskellDepends = [ libraryHaskellDepends = [
aeson annotated-exception async base bytestring cereal aeson annotated-exception async base bytestring cereal
cereal-conduit conduit containers directory ekg-core filepath cereal-conduit conduit containers directory ekg-core filepath
http-types katip lens lens-aeson network process scotty stm text http-types katip lens lens-aeson network process retry scotty stm
time unordered-containers uuid wai warp websockets text time unordered-containers uuid wai warp websockets
]; ];
executableHaskellDepends = [ base ]; executableHaskellDepends = [ base ];
testHaskellDepends = [ testHaskellDepends = [
+4 -2
View File
@@ -72,6 +72,7 @@ library
, HomeAssistant.Runtime.Graphing , HomeAssistant.Runtime.Graphing
, HomeAssistant.Runtime.Metrics , HomeAssistant.Runtime.Metrics
, HomeAssistant.Runtime.Supervisor , HomeAssistant.Runtime.Supervisor
, HomeAssistant.Runtime.RateLimit
-- Modules included in this library but not exported. -- Modules included in this library but not exported.
-- other-modules: -- other-modules:
@@ -85,7 +86,9 @@ library
, aeson , aeson
, lens-aeson , lens-aeson
, lens , lens
, exceptions
, text , text
, unliftio
, network , network
, bytestring , bytestring
, time , time
@@ -108,6 +111,7 @@ library
, http-types , http-types
, warp , warp
, wai , wai
, retry
-- Directories containing source files. -- Directories containing source files.
hs-source-dirs: src hs-source-dirs: src
@@ -150,14 +154,12 @@ test-suite home-assistant-controller-test
-- Modules included in this executable, other than Main. -- Modules included in this executable, other than Main.
other-modules: AFRPLawsSpec other-modules: AFRPLawsSpec
, AFRPSpec , AFRPSpec
, BackoffProp
, BedroomSpec , BedroomSpec
, BusSpec , BusSpec
, ConnectionSpec , ConnectionSpec
, GraphingSpec , GraphingSpec
, MetricsSpec , MetricsSpec
, RuntimeSpec , RuntimeSpec
, SupervisorSpec
, Support , Support
-- LANGUAGE extensions used by modules in this package. -- LANGUAGE extensions used by modules in this package.
+14 -11
View File
@@ -14,7 +14,6 @@ module HomeAssistant.Runtime
) where ) where
import AFRP (Event (..), Mealy (..), Request (..), Auto, stepAutoSerializing, load, DecodedAuto (..)) import AFRP (Event (..), Mealy (..), Request (..), Auto, stepAutoSerializing, load, DecodedAuto (..))
import Control.Concurrent.Async (async, waitAny)
import Control.Concurrent.STM (atomically, dupTChan, readTChan) import Control.Concurrent.STM (atomically, dupTChan, readTChan)
import Data.Aeson (Value) import Data.Aeson (Value)
import qualified Data.Text as T import qualified Data.Text as T
@@ -22,8 +21,7 @@ import Data.Time (getCurrentTime, getCurrentTimeZone)
import Data.Void (Void, absurd) import Data.Void (Void, absurd)
import HomeAssistant.Controller (HASS, HASSEff (..)) import HomeAssistant.Controller (HASS, HASSEff (..))
import HomeAssistant.Runtime.Bus import HomeAssistant.Runtime.Bus
import HomeAssistant.Runtime.Connection (readerAction, writerAction) import HomeAssistant.Runtime.Connection (writerAction, readerAction)
import HomeAssistant.Runtime.Supervisor (defaultBackoff, supervised)
import Network.Socket (withSocketsDo) import Network.Socket (withSocketsDo)
import System.Environment (getEnv, lookupEnv) import System.Environment (getEnv, lookupEnv)
import HomeAssistant.Controller.Bedroom (bedroomPresenceController, bedroomButtonController, bedroomDrawerController, humidifierController) import HomeAssistant.Controller.Bedroom (bedroomPresenceController, bedroomButtonController, bedroomDrawerController, humidifierController)
@@ -36,11 +34,14 @@ import HomeAssistant.Controller.Children (schoolLightController)
import Data.Maybe (fromMaybe) import Data.Maybe (fromMaybe)
import qualified System.Metrics import qualified System.Metrics
import qualified HomeAssistant.Runtime.Metrics import qualified HomeAssistant.Runtime.Metrics
import qualified HomeAssistant.Runtime.Graphing
import System.FilePath ((</>)) import System.FilePath ((</>))
import Text.Read (readMaybe) import Text.Read (readMaybe)
import HomeAssistant.Controller.Kitchen (kitchenMotionController) import HomeAssistant.Controller.Kitchen (kitchenMotionController)
import HomeAssistant.Controller.Livingroom (livingroomPresenceController) 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 :: (MonadIO m) => FilePath -> UUID -> Auto m a b -> a -> m (b, Auto m a b)
step path trace st a = do step path trace st a = do
@@ -66,8 +67,8 @@ controllers =
-- | Steps the machine for every inbound message; service calls go to the -- | 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 -- bus. A restart re-dups the inbound channel and starts from the machine's
-- initial state; messages broadcast during the restart window are lost. -- initial state; messages broadcast during the restart window are lost.
runController :: FilePath -> Bus -> Controller -> IO Void runController :: MonadIO m => FilePath -> Bus -> Controller -> m Void
runController rootDir bus (Controller name machine _enabled) = do runController rootDir bus (Controller name machine _enabled) = liftIO $ do
inbound <- atomically (dupTChan (busInbound bus)) inbound <- atomically (dupTChan (busInbound bus))
let ns = Namespace [name] let ns = Namespace [name]
let workerDefinition = runMealy machine (runKatipContextT (busLogEnv bus) () ns . channelHassEval bus) let workerDefinition = runMealy machine (runKatipContextT (busLogEnv bus) () ns . channelHassEval bus)
@@ -98,6 +99,7 @@ defaultMain = withSocketsDo $ do
store <- System.Metrics.newStore store <- System.Metrics.newStore
System.Metrics.registerGcMetrics store System.Metrics.registerGcMetrics store
appMetrics <- HomeAssistant.Runtime.Metrics.registerAppMetrics store appMetrics <- HomeAssistant.Runtime.Metrics.registerAppMetrics store
rateLimitMetrics <- registerRateLimitMetrics store
withBus severity appMetrics $ \bus -> do withBus severity appMetrics $ \bus -> do
token <- getEnv "HA_TOKEN" token <- getEnv "HA_TOKEN"
host <- getEnv "HA_HOST" host <- getEnv "HA_HOST"
@@ -105,16 +107,17 @@ defaultMain = withSocketsDo $ do
rrdPath <- fromMaybe "hass-controller.rrd" <$> lookupEnv "HA_RRD_PATH" rrdPath <- fromMaybe "hass-controller.rrd" <$> lookupEnv "HA_RRD_PATH"
rrdtool <- fromMaybe "rrdtool" <$> lookupEnv "HA_RRDTOOL" rrdtool <- fromMaybe "rrdtool" <$> lookupEnv "HA_RRDTOOL"
metricsPort <- lookupPort metricsPort <- lookupPort
writerLimiter <- slidingWindowLimiter rateLimitMetrics 10 20
let active = [c | c@(Controller _ _ True) <- controllers] let active = [c | c@(Controller _ _ True) <- controllers]
ents = foldMap (\(Controller _ m _) -> entities m) active ents = foldMap (\(Controller _ m _) -> entities m) active
workers = workers =
[ ("reader", readerAction host 8123 token ents bus) [ ("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 ] ] ++ [ (name, runController rootPath bus c) | c@(Controller name _ True) <- controllers ]
++ [ ("metrics", HomeAssistant.Runtime.Metrics.metricsAction store rrdPath rrdtool) as <- runKatipContextT (busLogEnv bus) () mempty $
, ("metrics-http", HomeAssistant.Runtime.Graphing.graphAction rrdPath rrdtool metricsPort) mapM (\(name, act) -> async (supervised name act)) workers
]
as <- mapM (\(name, act) -> async (supervised name defaultBackoff act)) workers
(_, v) <- waitAny as (_, v) <- waitAny as
absurd v absurd v
+21 -34
View File
@@ -5,17 +5,14 @@ module HomeAssistant.Runtime.Connection
( readerAction ( readerAction
, writerAction , writerAction
, encodeService , encodeService
, dedupeBatch
, isTriggerEvent , isTriggerEvent
) where ) where
import Control.Concurrent.STM import Control.Concurrent.STM
( TChan ( atomically
, atomically
, readTChan , readTChan
, readTVar , readTVar
, retry , retry
, tryReadTChan
, writeTChan , writeTChan
, writeTVar , writeTVar
) )
@@ -27,8 +24,6 @@ import Control.Lens ((^?))
import Control.Monad (forever, forM_, when) import Control.Monad (forever, forM_, when)
import Data.Aeson (Value, eitherDecode, encode, object, (.=)) import Data.Aeson (Value, eitherDecode, encode, object, (.=))
import Data.Aeson.Lens (key, _String) 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.Set as S
import qualified Data.Text as T import qualified Data.Text as T
import Data.Void (Void) import Data.Void (Void)
@@ -36,16 +31,18 @@ import HomeAssistant.Controller (Service (..), Target (..))
import HomeAssistant.Runtime.Bus import HomeAssistant.Runtime.Bus
import HomeAssistant.Runtime.Supervisor (Fatal (..)) import HomeAssistant.Runtime.Supervisor (Fatal (..))
import qualified Network.WebSockets as WS 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 Data.UUID (toText)
import AFRP (Request(..), Event(..)) 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. -- | Connect, authenticate, subscribe, then receive and broadcast forever.
-- Restarting this action reconnects. All setup sends happen before the -- Restarting this action reconnects. All setup sends happen before the
-- connection is published in the bus, so only the writer sends afterwards. -- connection is published in the bus, so only the writer sends afterwards.
readerAction :: String -> Int -> String -> S.Set T.Text -> Bus -> IO Void readerAction :: MonadIO m => String -> Int -> String -> S.Set T.Text -> Bus -> m Void
readerAction host port token ents bus = readerAction host port token ents bus = liftIO $
WS.runClient host port "/api/websocket" $ \conn -> do WS.runClient host port "/api/websocket" $ \conn -> do
handshake conn token handshake conn token
subscribe bus conn ents subscribe bus conn ents
@@ -114,20 +111,21 @@ receiveJSON conn = do
sendWithId :: Bus -> WS.Connection -> Request -> Service -> IO () sendWithId :: (MonadIO m, KatipContext m) => Bus -> WS.Connection -> Service -> RateLimited m ()
sendWithId bus conn request svc = sendWithId bus conn svc = katipAddNamespace "connection" $ do
runKatipContextT (busLogEnv bus) (sl "traceId" (toText (requestTraceId request))) "connection" $ do callId <- liftIO $ generateCallId (busGen bus)
callId <- liftIO $ generateCallId (busGen bus) let textData = encode $ encodeService callId svc
let textData = encode $ encodeService callId svc logFM DebugS (ls textData)
logFM DebugS (ls textData) liftIO $ recordOutbound bus
liftIO $ recordOutbound bus liftIO $ WS.sendTextData conn textData
liftIO $ WS.sendTextData conn textData
writerAction :: Bus -> IO Void writerAction :: (MonadIO m, KatipContext m, MonadCatch m) => RateLimiter -> Bus -> m Void
writerAction bus = forever $ do writerAction rateLimiter bus = forever $ do
(request, svc) <- atomically $ readTChan (busOutbound bus) (request, svc) <- liftIO $ atomically $ readTChan (busOutbound bus)
conn <- atomically $ readTVar (busConn bus) >>= maybe retry pure katipAddContext (sl "traceId" (toText (requestTraceId request))) $ do
sendWithId bus conn request svc 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 :: Int -> Service -> Value
encodeService callId Service{..} = object $ encodeService callId Service{..} = object $
@@ -138,17 +136,6 @@ encodeService callId Service{..} = object $
, "target" .= targetObject serviceTarget , "target" .= targetObject serviceTarget
] <> maybe [] (\d -> ["service_data" .= d]) serviceData ] <> 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 -- | A single target encodes as a scalar; multiple encode as a list. Empty
-- lists are omitted so Home Assistant receives only populated keys. -- lists are omitted so Home Assistant receives only populated keys.
+4 -4
View File
@@ -16,7 +16,7 @@ module HomeAssistant.Runtime.Graphing
import Control.Exception (evaluate) import Control.Exception (evaluate)
import Control.Monad (forever) 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 as BS
import qualified Data.ByteString.Lazy as BL import qualified Data.ByteString.Lazy as BL
import Data.Text (Text) import Data.Text (Text)
@@ -155,9 +155,9 @@ handler rrdPath rrdtool = do
app :: FilePath -> FilePath -> ScottyM () app :: FilePath -> FilePath -> ScottyM ()
app rrdPath rrdtool = get "/metrics/:name" (handler rrdPath rrdtool) app rrdPath rrdtool = get "/metrics/:name" (handler rrdPath rrdtool)
graphAction :: FilePath -> FilePath -> Int -> IO Void graphAction :: MonadIO m => FilePath -> FilePath -> Int -> m Void
graphAction rrdPath rrdtool port = forever $ scottyOpts opts (app rrdPath rrdtool) graphAction rrdPath rrdtool port = liftIO $ forever $ scottyOpts opts (app rrdPath rrdtool)
where where
opts = defaultOptions opts = defaultOptions
{ settings = setHost "127.0.0.1" (setPort port defaultSettings) { settings = setHost "127.0.0.1" (setPort port defaultSettings)
} }
+23 -18
View File
@@ -35,6 +35,9 @@ import System.Metrics.Counter (Counter)
import System.Directory (doesFileExist, renameFile) import System.Directory (doesFileExist, renameFile)
import System.Exit (ExitCode (..)) import System.Exit (ExitCode (..))
import System.Process (callProcess, readProcessWithExitCode) 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 data DsType = Derive | Gauge
deriving (Eq, Ord, Show) deriving (Eq, Ord, Show)
@@ -81,6 +84,8 @@ registerAppMetrics store =
<$> M.createCounter "hass.trigger.in" store <$> M.createCounter "hass.trigger.in" store
<*> M.createCounter "hass.service.out" store <*> M.createCounter "hass.service.out" store
buildCreateArgs :: FilePath -> Int -> [(String, DsType)] -> [String] buildCreateArgs :: FilePath -> Int -> [(String, DsType)] -> [String]
buildCreateArgs path step specs = buildCreateArgs path step specs =
["create", path, "--step", show step] ["create", path, "--step", show step]
@@ -134,44 +139,44 @@ lookupValue sample name = case HM.lookup name sample of
Just (M.Gauge n) -> Just n Just (M.Gauge n) -> Just n
_ -> Nothing _ -> Nothing
readExistingDs :: FilePath -> FilePath -> IO [(String, DsType)] readExistingDs :: (MonadCatch m, MonadIO m) => FilePath -> FilePath -> m [(String, DsType)]
readExistingDs rrdtool rrdPath = do readExistingDs rrdtool rrdPath = checkpointCallStack $ do
(code, out, _) <- readProcessWithExitCode rrdtool ["info", rrdPath] "" (code, out, _) <- liftIO $ readProcessWithExitCode rrdtool ["info", rrdPath] ""
pure $ case code of pure $ case code of
ExitSuccess -> parseInfoDs out ExitSuccess -> parseInfoDs out
_ -> [] _ -> []
backupPath :: FilePath -> IO FilePath backupPath :: MonadIO m => FilePath -> m FilePath
backupPath path = do backupPath path = liftIO $ do
now <- getCurrentTime now <- getCurrentTime
pure $ path ++ ".bak-" ++ formatTime defaultTimeLocale "%Y%m%dT%H%M%S" now pure $ path ++ ".bak-" ++ formatTime defaultTimeLocale "%Y%m%dT%H%M%S" now
ensureRrd :: FilePath -> FilePath -> [DsSpec] -> IO () ensureRrd :: (MonadCatch m, MonadIO m) => FilePath -> FilePath -> [DsSpec] -> m ()
ensureRrd rrdtool rrdPath schema = do ensureRrd rrdtool rrdPath schema = checkpointCallStack $ do
exists <- doesFileExist rrdPath exists <- liftIO $ doesFileExist rrdPath
if not exists if not exists
then create then create
else do else do
current <- readExistingDs rrdtool rrdPath current <- readExistingDs rrdtool rrdPath
unless (schemaMatches schema current) $ do unless (schemaMatches schema current) $ do
backup <- backupPath rrdPath backup <- backupPath rrdPath
renameFile rrdPath backup liftIO $ renameFile rrdPath backup
create create
where 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) toPair s = (dsName s, dsType s)
sampleAndUpdate :: M.Store -> FilePath -> FilePath -> [DsSpec] -> IO () sampleAndUpdate :: (MonadCatch m, MonadIO m) => M.Store -> FilePath -> FilePath -> [DsSpec] -> m ()
sampleAndUpdate store rrdtool rrdPath schema = do sampleAndUpdate store rrdtool rrdPath schema = checkpointCallStack $ do
sample <- M.sampleAll store sample <- liftIO $ M.sampleAll store
let names = map dsName schema let names = map dsName schema
values = map (lookupValue sample . dsEkgName) 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 :: (MonadCatch m, MonadIO m) => M.Store -> FilePath -> FilePath -> m Void
metricsAction store rrdPath rrdtool = do metricsAction store rrdPath rrdtool = checkpointCallStack $ do
schema <- buildSchema <$> M.sampleAll store schema <- buildSchema <$> liftIO (M.sampleAll store)
ensureRrd rrdtool rrdPath schema ensureRrd rrdtool rrdPath schema
forever $ do forever $ do
sampleAndUpdate store rrdtool rrdPath schema sampleAndUpdate store rrdtool rrdPath schema
threadDelay 10000000 liftIO $ threadDelay 10000000
+74
View File
@@ -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)
+19 -47
View File
@@ -2,26 +2,23 @@
module HomeAssistant.Runtime.Supervisor module HomeAssistant.Runtime.Supervisor
( supervised ( supervised
, Backoff(..)
, defaultBackoff
, backoffDelay
, nextAttempt
, Fatal(..) , Fatal(..)
) where ) where
import Control.Concurrent (threadDelay)
import Control.Exception.Annotated import Control.Exception.Annotated
( Exception ( Exception
, Handler (..) , Handler (..)
, SomeException , SomeException
, catches , checkpoint, Annotation (..)
, displayException
, throw
) )
import Data.Text (Text) import Data.Text (Text)
import qualified Data.Text as T
import Data.Time (NominalDiffTime, diffUTCTime, getCurrentTime)
import Data.Void (Void) 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 -- | A failure that cannot be fixed by restarting; propagates out of
-- 'supervised' and terminates the process. -- 'supervised' and terminates the process.
@@ -30,28 +27,11 @@ newtype Fatal = Fatal Text
instance Exception Fatal 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 data SupervisorException = RetriesExhausted
defaultBackoff = Backoff 0.1 30 30 deriving Show
-- | Delay before the @attempt@-th restart: doubles from base, clamped at cap. instance Exception SupervisorException
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
-- | Runs the action forever, restarting it with backoff after synchronous -- | Runs the action forever, restarting it with backoff after synchronous
-- exceptions; async exceptions propagate. 'Fatal' is rethrown (crashing the -- 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 -- must not be wrapped in checkpoints around 'Fatal'-throwing code: a
-- doubly-wrapped 'Fatal' is indistinguishable from a crash and would be -- doubly-wrapped 'Fatal' is indistinguishable from a crash and would be
-- restarted instead of escalated. -- restarted instead of escalated.
supervised :: Text -> Backoff -> IO Void -> IO Void supervised :: (MonadIO m, MonadCatch m, KatipContext m, MonadMask m) => Text -> m Void -> m Void
supervised name backoff action = go 1 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 where
go attempt = do handlers =
start <- getCurrentTime [ \_retryStatus -> Handler $ \(_ :: Fatal) -> pure False
outcome <- (action >> pure (Nothing :: Maybe (Either Fatal SomeException))) `catches` , \_retryStatus -> Handler $ \(_ :: SomeException) -> pure True
[ Handler $ \(f :: Fatal) -> pure (Just (Left f))
, Handler $ \(e :: SomeException) -> pure (Just (Right e))
] ]
end <- getCurrentTime retryPolicy = fullJitterBackoff 50000 <> limitRetries 5
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)
-21
View File
@@ -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))
+1 -29
View File
@@ -9,7 +9,7 @@ import Data.Text (Text)
import Data.Time (UTCTime (..), utc) import Data.Time (UTCTime (..), utc)
import Data.UUID (fromString) import Data.UUID (fromString)
import HomeAssistant.Controller (Service (..), Target(..)) import HomeAssistant.Controller (Service (..), Target(..))
import HomeAssistant.Runtime.Connection (encodeService, dedupeBatch, isTriggerEvent) import HomeAssistant.Runtime.Connection (encodeService, isTriggerEvent)
import Test.Hspec import Test.Hspec
spec :: Spec spec :: Spec
@@ -44,34 +44,6 @@ spec = do
it "is false when there is no type" $ it "is false when there is no type" $
isTriggerEvent (object ["id" .= (1 :: Int)]) `shouldBe` False 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 :: Int -> Request
req n = Request (UTCTime (toEnum 0) (fromIntegral (0 :: Int))) utc req n = Request (UTCTime (toEnum 0) (fromIntegral (0 :: Int))) utc
-4
View File
@@ -3,14 +3,12 @@ module Main (main) where
import Test.Hspec (hspec) import Test.Hspec (hspec)
import qualified AFRPLawsSpec import qualified AFRPLawsSpec
import qualified AFRPSpec import qualified AFRPSpec
import qualified BackoffProp
import qualified BedroomSpec import qualified BedroomSpec
import qualified BusSpec import qualified BusSpec
import qualified ConnectionSpec import qualified ConnectionSpec
import qualified GraphingSpec import qualified GraphingSpec
import qualified MetricsSpec import qualified MetricsSpec
import qualified RuntimeSpec import qualified RuntimeSpec
import qualified SupervisorSpec
main :: IO () main :: IO ()
main = hspec $ do main = hspec $ do
@@ -22,5 +20,3 @@ main = hspec $ do
GraphingSpec.spec GraphingSpec.spec
MetricsSpec.spec MetricsSpec.spec
RuntimeSpec.spec RuntimeSpec.spec
SupervisorSpec.spec
BackoffProp.spec
-73
View File
@@ -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