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