[api] Pre-reserve Home Assistant state subscriptions vector (#20986)

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

esphome/components/api/api_server.cpp

@@ -65,6 +65,9 @@ void APIServer::socket_failed_(const LogString *msg) {
 }
 
 void APIServer::setup() {
+#if defined(USE_API_HOMEASSISTANT_STATES) && defined(HOMEASSISTANT_STATE_SUBSCRIPTION_COUNT)
+  this->state_subs_.reserve(HOMEASSISTANT_STATE_SUBSCRIPTION_COUNT);
+#endif
 #ifdef USE_API_NOISE
   // Always reserve the slot: flash preferences are positional on esp8266, so
   // a yaml key build must keep the layout of a runtime key build

esphome/components/homeassistant/__init__.py

@@ -1,8 +1,11 @@
 from collections.abc import Callable, Iterable
+from dataclasses import dataclass
 
 import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome.const import CONF_ATTRIBUTE, CONF_ENTITY_ID, CONF_ID, CONF_INTERNAL
+from esphome.core import CORE, coroutine_with_priority
+from esphome.coroutine import CoroPriority
 from esphome.cpp_generator import MockObj
 import esphome.final_validate as fv
 from esphome.types import ConfigType
@@ -59,7 +62,28 @@ def final_validate_entity_id(config: ConfigType) -> ConfigType:
     return config
 
 
-def setup_home_assistant_entity(var: MockObj, config: ConfigType) -> None:
+@dataclass
+class HomeAssistantData:
+    state_subscription_count: int = 0
+
+
+def _get_data() -> HomeAssistantData:
+    if DOMAIN not in CORE.data:
+        CORE.data[DOMAIN] = HomeAssistantData()
+    return CORE.data[DOMAIN]
+
+
+@coroutine_with_priority(CoroPriority.FINAL)
+async def _emit_state_subscription_count_define() -> None:
+    """Emit the HOMEASSISTANT_STATE_SUBSCRIPTION_COUNT define with the total count."""
+    count = _get_data().state_subscription_count
+    if count > 0:
+        cg.add_define("HOMEASSISTANT_STATE_SUBSCRIPTION_COUNT", count)
+
+
+def setup_home_assistant_entity(
+    var: MockObj, config: ConfigType, state_subscription_count: int = 1
+) -> None:
     from esphome.components.api import wizard
 
     if (buffer := wizard.wizard_input_buffer(config[CONF_ID])) is not None:
@@ -70,3 +94,8 @@ def setup_home_assistant_entity(var: MockObj, config: ConfigType) -> None:
     if CONF_ATTRIBUTE in config:
         cg.add(var.set_attribute(config[CONF_ATTRIBUTE]))
     cg.add_define("USE_API_HOMEASSISTANT_STATES")
+    if state_subscription_count > 0:
+        data = _get_data()
+        if data.state_subscription_count == 0:
+            CORE.add_job(_emit_state_subscription_count_define)
+        data.state_subscription_count += state_subscription_count

esphome/components/homeassistant/button/__init__.py

@@ -37,4 +37,5 @@ async def to_code(config: ConfigType) -> None:
     cg.add_define("USE_API_HOMEASSISTANT_SERVICES")
     var = await button.new_button(config)
     await cg.register_component(var, config)
-    setup_home_assistant_entity(var, config)
+    # Button only calls services, does not subscribe to state
+    setup_home_assistant_entity(var, config, state_subscription_count=0)

esphome/components/homeassistant/number/__init__.py

@@ -39,4 +39,5 @@ async def to_code(config: ConfigType) -> None:
         step=0,
     )
     await cg.register_component(var, config)
-    setup_home_assistant_entity(var, config)
+    # 1 state + min/max/step fetches in HomeassistantNumber::setup()
+    setup_home_assistant_entity(var, config, state_subscription_count=4)

esphome/components/homeassistant/select/__init__.py

@@ -55,4 +55,5 @@ async def to_code(config: ConfigType) -> None:
         options=[],
     )
     await cg.register_component(var, config)
-    setup_home_assistant_entity(var, config)
+    # 1 state + options fetch in HomeassistantSelect::setup()
+    setup_home_assistant_entity(var, config, state_subscription_count=2)

esphome/components/homeassistant/text/__init__.py

@@ -37,4 +37,5 @@ async def to_code(config: ConfigType) -> None:
     cg.add_define("USE_API_HOMEASSISTANT_SERVICES")
     var = await text.new_text(config)
     await cg.register_component(var, config)
-    setup_home_assistant_entity(var, config)
+    # 1 state + min/max/mode fetches in HomeassistantText::setup()
+    setup_home_assistant_entity(var, config, state_subscription_count=4)

esphome/core/defines.h

@@ -319,6 +319,7 @@
 #define USE_API_HOMEASSISTANT_ACTION_RESPONSES_JSON
 #define USE_API_HOMEASSISTANT_SERVICES
 #define USE_API_HOMEASSISTANT_STATES
+#define HOMEASSISTANT_STATE_SUBSCRIPTION_COUNT 1
 #define USE_API_NOISE
 #define USE_API_OUTGOING_CONNECTION
 #define API_OUTGOING_CONNECTION_PORT 6054

tests/component_tests/homeassistant/mixed_entities.yaml

@@ -0,0 +1,65 @@
+esphome:
+  name: test-mixed-ha
+
+esp32:
+  variant: esp32
+
+wifi:
+  ssid: SomeNetwork
+  password: SomePassword
+
+logger:
+
+api:
+
+sensor:
+  - platform: homeassistant
+    id: ha_sensor_1
+    name: HA Sensor 1
+    entity_id: sensor.temp_1
+  - platform: homeassistant
+    id: ha_sensor_2
+    name: HA Sensor 2
+    entity_id: sensor.temp_2
+
+binary_sensor:
+  - platform: homeassistant
+    id: ha_binary_sensor
+    name: HA Binary Sensor
+    entity_id: binary_sensor.motion
+
+text_sensor:
+  - platform: homeassistant
+    id: ha_text_sensor
+    name: HA Text Sensor
+    entity_id: sensor.status_text
+
+switch:
+  - platform: homeassistant
+    id: ha_switch
+    name: HA Switch
+    entity_id: switch.relay
+
+select:
+  - platform: homeassistant
+    id: ha_select
+    name: HA Select
+    entity_id: select.mode
+
+number:
+  - platform: homeassistant
+    id: ha_number
+    name: HA Number
+    entity_id: number.target_temp
+
+text:
+  - platform: homeassistant
+    id: ha_text
+    name: HA Text
+    entity_id: text.custom_note
+
+button:
+  - platform: homeassistant
+    id: ha_button
+    name: HA Button
+    entity_id: button.restart

tests/component_tests/homeassistant/test_state_subscription_count.py

@@ -0,0 +1,51 @@
+"""Tests for Home Assistant state subscription count calculation and pre-reservation."""
+
+from collections.abc import Callable
+from pathlib import Path
+
+from esphome.core import CORE
+
+DIR = Path(__file__).parent
+
+
+def test_mixed_entities_state_subscription_count(
+    generate_main: Callable[[str | Path], str],
+) -> None:
+    """A configuration with mixed entities computes the exact weighted subscription count."""
+    generate_main(DIR / "mixed_entities.yaml")
+
+    # sensor (2*1) + binary_sensor (1) + text_sensor (1) + switch (1)
+    # + select (2) + number (4) + text (4) + button (0) = 15
+    define_match = [
+        d for d in CORE.defines if d.name == "HOMEASSISTANT_STATE_SUBSCRIPTION_COUNT"
+    ]
+    assert len(define_match) == 1
+    assert str(define_match[0].value) == "15"
+    assert any(d.name == "USE_API_HOMEASSISTANT_STATES" for d in CORE.defines)
+
+
+def test_button_only_no_state_subscription_count(
+    generate_main: Callable[[str | Path], str],
+) -> None:
+    """A configuration with only button entities has no state subscriptions."""
+    generate_main(DIR / "button_valid.yaml")
+
+    assert not any(
+        d.name == "HOMEASSISTANT_STATE_SUBSCRIPTION_COUNT" for d in CORE.defines
+    )
+    assert any(d.name == "USE_API_HOMEASSISTANT_STATES" for d in CORE.defines)
+    assert any(d.name == "USE_API_HOMEASSISTANT_SERVICES" for d in CORE.defines)
+
+
+def test_select_only_state_subscription_count(
+    generate_main: Callable[[str | Path], str],
+) -> None:
+    """A configuration with 2 select entities calculates 2 * 2 = 4 subscriptions."""
+    generate_main(DIR / "select_valid.yaml")
+
+    define_match = [
+        d for d in CORE.defines if d.name == "HOMEASSISTANT_STATE_SUBSCRIPTION_COUNT"
+    ]
+    assert len(define_match) == 1
+    assert str(define_match[0].value) == "4"
+    assert any(d.name == "USE_API_HOMEASSISTANT_STATES" for d in CORE.defines)