Merge feat/outbound-rate-limiting: drop excess outbound calls

src/AFRP.hs

@@ -52,7 +52,7 @@ data Request = Request
   { requestTime :: !UTCTime
   , requestTimeZone :: !TimeZone
   , requestTraceId :: !UUID
-  } deriving Show
+  } deriving (Show, Eq)
 
 -- | The set of entity ids an arrow subscribes to. Static: it does not
 -- change as the machine steps, so the runtime can read it once to build

src/HomeAssistant/Controller.hs

@@ -39,7 +39,7 @@ import qualified Data.Text.Lens as TL
 import Data.Bool (bool)
 
 data Target = EntityId !T.Text | AreaId !T.Text
-  deriving (Show,Eq)
+  deriving (Show,Eq,Ord)
 
 data Service = Service
   { serviceDomain :: T.Text

src/HomeAssistant/Runtime/Connection.hs

@@ -5,13 +5,16 @@ module HomeAssistant.Runtime.Connection
   ( readerAction
   , writerAction
   , encodeService
+  , dedupeBatch
   ) where
 
 import Control.Concurrent.STM
-  ( atomically
+  ( TChan
+  , atomically
   , readTChan
   , readTVar
   , retry
+  , tryReadTChan
   , writeTChan
   , writeTVar
   )
@@ -23,6 +26,8 @@ import Control.Lens ((^?))
 import Control.Monad (forever, forM_)
 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)
@@ -100,16 +105,41 @@ receiveJSON conn = do
     Left err -> throw (Fatal $ "Invalid JSON from Home Assistant: " <> T.pack err)
     Right x  -> pure x
 
-writerAction :: Bus -> IO Void
-writerAction bus = forever $ do
-  (request, svc) <- atomically $ readTChan (busOutbound bus)
-  conn <- atomically $ readTVar (busConn bus) >>= maybe retry pure
+-- | Floor between sends within a batch: 100ms, so a many-distinct-target
+-- flood still caps at ~10 sends/sec even after dedupe.
+minInterval :: Int
+minInterval = 100000
+
+-- | Non-blocking drain of everything queued on a channel. Returns items
+-- oldest-first (FIFO from the channel), so prepending the blocking
+-- `readTChan` item keeps the whole batch oldest-first for `dedupeBatch`.
+drainTry :: TChan a -> IO [a]
+drainTry chan = go []
+  where
+    go acc = do
+      m <- atomically $ tryReadTChan chan
+      case m of
+        Nothing -> pure (reverse acc)
+        Just x  -> go (x : acc)
+
+sendWithId :: Bus -> WS.Connection -> Request -> Service -> IO ()
+sendWithId bus conn request svc = do
   callId <- generateCallId (busGen bus)
   let textData = encode $ encodeService callId svc
   runKatipContextT (busLogEnv bus) (sl "traceId" (toText (requestTraceId request))) "connection" $
     logFM DebugS (ls textData)
   WS.sendTextData conn textData
 
+writerAction :: Bus -> IO Void
+writerAction bus = forever $ do
+  first <- atomically $ readTChan (busOutbound bus)
+  rest  <- drainTry (busOutbound bus)
+  let deduped = dedupeBatch (first : rest)
+  conn <- atomically $ readTVar (busConn bus) >>= maybe retry pure
+  forM_ deduped $ \(request, svc) -> do
+    sendWithId bus conn request svc
+    threadDelay minInterval
+
 encodeService :: Int -> Service -> Value
 encodeService callId Service{..} = object $
   [ "id"      .= callId
@@ -119,6 +149,18 @@ 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.
 targetObject :: [Target] -> Value

test/ConnectionSpec.hs

@@ -2,31 +2,78 @@
 
 module ConnectionSpec (spec) where
 
+import AFRP (Request(..))
 import Data.Aeson (object, (.=))
+import Data.Maybe (fromJust)
 import Data.Text (Text)
+import Data.Time (UTCTime (..), utc)
+import Data.UUID (fromString)
 import HomeAssistant.Controller (Service (..), Target(..))
-import HomeAssistant.Runtime.Connection (encodeService)
+import HomeAssistant.Runtime.Connection (encodeService, dedupeBatch)
 import Test.Hspec
 
 spec :: Spec
-spec = describe "encodeService" $ do
-  it "encodes a call_service message" $
-    encodeService 7 (Service "light" "turn_on" Nothing [EntityId "light.bedroom_masse"])
-      `shouldBe` object
-        [ "id" .= (7 :: Int)
-        , "type" .= ("call_service" :: Text)
-        , "domain" .= ("light" :: Text)
-        , "service" .= ("turn_on" :: Text)
-        , "target" .= object ["entity_id" .= ("light.bedroom_masse" :: Text)]
-        ]
+spec = do
+  describe "encodeService" $ do
+    it "encodes a call_service message" $
+      encodeService 7 (Service "light" "turn_on" Nothing [EntityId "light.bedroom_masse"])
+        `shouldBe` object
+          [ "id" .= (7 :: Int)
+          , "type" .= ("call_service" :: Text)
+          , "domain" .= ("light" :: Text)
+          , "service" .= ("turn_on" :: Text)
+          , "target" .= object ["entity_id" .= ("light.bedroom_masse" :: Text)]
+          ]
 
-  it "includes service_data when present" $
-    encodeService 8 (Service "light" "turn_on" (Just (object ["brightness" .= (200 :: Int)])) [EntityId "light.bedroom_masse"])
-      `shouldBe` object
-        [ "id" .= (8 :: Int)
-        , "type" .= ("call_service" :: Text)
-        , "domain" .= ("light" :: Text)
-        , "service" .= ("turn_on" :: Text)
-        , "target" .= object ["entity_id" .= ("light.bedroom_masse" :: Text)]
-        , "service_data" .= object ["brightness" .= (200 :: Int)]
-        ]
+    it "includes service_data when present" $
+      encodeService 8 (Service "light" "turn_on" (Just (object ["brightness" .= (200 :: Int)])) [EntityId "light.bedroom_masse"])
+        `shouldBe` object
+          [ "id" .= (8 :: Int)
+          , "type" .= ("call_service" :: Text)
+          , "domain" .= ("light" :: Text)
+          , "service" .= ("turn_on" :: Text)
+          , "target" .= object ["entity_id" .= ("light.bedroom_masse" :: Text)]
+          , "service_data" .= object ["brightness" .= (200 :: Int)]
+          ]
+
+  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
+          (fromJust (fromString uuid))
+  where
+    pad i = replicate (12 - length (show i)) '0' <> show i
+    uuid = "00000000-0000-0000-0000-" <> pad n
+
+lightOn :: [Target] -> Service
+lightOn targets = Service "light" "turn_on" Nothing targets
+
+lightOff :: [Target] -> Service
+lightOff targets = Service "light" "turn_off" Nothing targets