Fix supervisor and proper rate limiting

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
 

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.

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)
-      }
+      }

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

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)

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