[core] Shared slot count factory for codegen sized listener storage (#18057)

esphome/codegen.py

@@ -49,11 +49,13 @@ from esphome.cpp_helpers import (  # noqa: F401
     build_registry_entry,
     build_registry_list,
     extract_registry_entry_config,
+    get_slot_count,
     gpio_pin_expression,
     past_safe_mode,
     register_component,
     register_parented,
     set_setup_priority,
+    slot_counter,
 )
 from esphome.cpp_types import (  # noqa: F401
     NAN,

esphome/components/bk72xx_ble/__init__.py

@@ -24,7 +24,6 @@ from esphome.components import libretiny
 from esphome.components.libretiny.const import FAMILY_BK7231N, FAMILY_BK7238
 import esphome.config_validation as cv
 from esphome.const import CONF_ENABLE_ON_BOOT, CONF_ID
-from esphome.core import CORE, CoroPriority, coroutine_with_priority
 from esphome.types import ConfigType
 
 DEPENDENCIES = ["bk72xx"]
@@ -46,21 +45,9 @@ CONFIG_SCHEMA = cv.Schema(
 ).extend(cv.COMPONENT_SCHEMA)
 
 
-KEY_SCAN_LISTENER_COUNT = "bk72xx_ble_scan_listener_count"
-
-
-def request_scan_listener_slot() -> None:
-    """Called from a consumer's codegen once per registered scan listener; sizes
-    the controller's StaticVector listener storage (heap-free, mirrors the
-    tracker's ble_device_base listener storage)."""
-    CORE.data[KEY_SCAN_LISTENER_COUNT] = CORE.data.get(KEY_SCAN_LISTENER_COUNT, 0) + 1
-
-
-@coroutine_with_priority(CoroPriority.FINAL)
-async def _add_listener_count() -> None:
-    # FINAL: every consumer's to_code has requested its slot by now.
-    if count := CORE.data.get(KEY_SCAN_LISTENER_COUNT, 0):
-        cg.add_define("BK72XX_BLE_SCAN_LISTENER_COUNT", count)
+# Once per registered scan listener; sizes the controller's StaticVector
+# listener storage.
+request_scan_listener_slot = cg.slot_counter("BK72XX_BLE_SCAN_LISTENER_COUNT")
 
 
 async def to_code(config: ConfigType) -> None:
@@ -110,5 +97,3 @@ async def to_code(config: ConfigType) -> None:
         )
 
     cg.add_define("USE_BK72XX_BLE")
-
-    CORE.add_job(_add_listener_count)

esphome/components/bk72xx_ble_tracker/__init__.py

@@ -24,7 +24,6 @@ from esphome.components import bk72xx_ble, ble_device_base, ota
 from esphome.components.const import CONF_SCAN_PARAMETERS, CONF_WINDOW
 import esphome.config_validation as cv
 from esphome.const import CONF_CONTINUOUS, CONF_DURATION, CONF_ID, CONF_INTERVAL
-from esphome.core import CORE, CoroPriority, coroutine_with_priority
 from esphome.types import ConfigType
 
 CONF_BK72XX_BLE_ID = "bk72xx_ble_id"
@@ -53,16 +52,6 @@ CONFIG_SCHEMA = cv.Schema(
 ).extend(cv.COMPONENT_SCHEMA)
 
 
-# Runs at FINAL priority so every BLE sensor has registered through
-# ble_device_base (and any tracker-owned listeners have been counted) before
-# the StaticVector size is emitted. Same pattern as esp32_ble_tracker.
-@coroutine_with_priority(CoroPriority.FINAL)
-async def _emit_listener_count() -> None:
-    count = ble_device_base.get_listener_count()
-    if count > 0:
-        cg.add_define("ESPHOME_BLE_DEVICE_BASE_LISTENER_COUNT", count)
-
-
 async def to_code(config: ConfigType) -> None:
     var = cg.new_Pvariable(config[CONF_ID])
     await cg.register_component(var, config)
@@ -81,5 +70,3 @@ async def to_code(config: ConfigType) -> None:
     cg.add(var.set_scan_window(ble_device_base.to_ble_units(scan[CONF_WINDOW])))
     cg.add(var.set_scan_duration(scan[CONF_DURATION].total_milliseconds))
     cg.add(var.set_scan_continuous(scan[CONF_CONTINUOUS]))
-
-    CORE.add_job(_emit_listener_count)

esphome/components/ble_device_base/__init__.py

@@ -23,17 +23,16 @@ import esphome.codegen as cg
 from esphome.components.const import CONF_WINDOW
 import esphome.config_validation as cv
 from esphome.const import CONF_ACTIVE, CONF_CONTINUOUS, CONF_DURATION, CONF_INTERVAL
-from esphome.core import CORE
 from esphome.types import ConfigType
 
 CODEOWNERS = ["@Bl00d-B0b"]
 
 CONF_BLE_HUB_ID = "ble_hub_id"
 
-# CORE.data key: number of parsed-advertisement listeners registered in this
-# build. Trackers whose codegen sizes storage at compile time (esp32's
-# StaticVector count define) read it in their final coroutine.
-KEY_BLE_LISTENER_COUNT = "ble_device_base_listener_count"
+# Number of parsed-advertisement listeners registered in this build; read via
+# cg.get_slot_count() by esp32_ble_tracker's feature coupling.
+LISTENER_COUNT_DEFINE = "ESPHOME_BLE_DEVICE_BASE_LISTENER_COUNT"
+
 
 ble_device_base_ns = cg.esphome_ns.namespace("ble_device_base")
 
@@ -64,16 +63,14 @@ def request_irk_support() -> None:
     cg.add_define("USE_BLE_DEVICE_IRK")
 
 
-def get_listener_count() -> int:
-    """Number of parsed listeners registered so far (for tracker codegen)."""
-    return CORE.data.get(KEY_BLE_LISTENER_COUNT, 0)
+_request_listener_slot = cg.slot_counter(LISTENER_COUNT_DEFINE)
 
 
 async def register_ble_device(var: cg.MockObj, config: ConfigType) -> cg.MockObj:
     """Register `var` as a parsed-advertisement listener on the configured hub."""
     hub = await cg.get_variable(config[CONF_BLE_HUB_ID])
     cg.add(hub.register_listener(var))
-    CORE.data[KEY_BLE_LISTENER_COUNT] = CORE.data.get(KEY_BLE_LISTENER_COUNT, 0) + 1
+    _request_listener_slot()
     return var
 
 

esphome/components/esp32_ble_tracker/__init__.py

@@ -308,12 +308,10 @@ async def _add_ble_features():
     required_features = _get_required_features()
     # Sensors registered through the neutral ble_device_base path (BLEHub) need
     # the parsed-device pipeline compiled in, exactly like esp32-path listeners.
-    neutral_listener_count = ble_device_base.get_listener_count()
-    if neutral_listener_count > 0:
+    if cg.get_slot_count(ble_device_base.LISTENER_COUNT_DEFINE):
+        # The neutral (BLEHub) listener count define itself is emitted by
+        # ble_device_base's own job; only the feature coupling lives here.
         required_features.add(BLEFeatures.ESP_BT_DEVICE)
-        # StaticVector sizing for the neutral (BLEHub) listener list — same
-        # pattern as the esp32-path registration counts below.
-        cg.add_define("ESPHOME_BLE_DEVICE_BASE_LISTENER_COUNT", neutral_listener_count)
     if BLEFeatures.ESP_BT_DEVICE in required_features:
         cg.add_define("USE_ESP32_BLE_DEVICE")
         cg.add_define("USE_ESP32_BLE_UUID")

esphome/components/ln882h_ble/__init__.py

@@ -12,7 +12,6 @@ component does (LibreTiny v1.13.0+).
 import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome.const import CONF_ENABLE_ON_BOOT, CONF_ID
-from esphome.core import CORE, CoroPriority, coroutine_with_priority
 from esphome.types import ConfigType
 
 DEPENDENCIES = ["ln882x"]
@@ -32,21 +31,9 @@ CONFIG_SCHEMA = cv.Schema(
 ).extend(cv.COMPONENT_SCHEMA)
 
 
-KEY_SCAN_LISTENER_COUNT = "ln882h_ble_scan_listener_count"
-
-
-def request_scan_listener_slot() -> None:
-    """Called from a consumer's codegen once per registered scan listener; sizes
-    the controller's StaticVector listener storage (heap-free, mirrors the
-    tracker's ble_device_base listener storage)."""
-    CORE.data[KEY_SCAN_LISTENER_COUNT] = CORE.data.get(KEY_SCAN_LISTENER_COUNT, 0) + 1
-
-
-@coroutine_with_priority(CoroPriority.FINAL)
-async def _add_listener_count() -> None:
-    # FINAL: every consumer's to_code has requested its slot by now.
-    if count := CORE.data.get(KEY_SCAN_LISTENER_COUNT, 0):
-        cg.add_define("LN882H_BLE_SCAN_LISTENER_COUNT", count)
+# Once per registered scan listener; sizes the controller's StaticVector
+# listener storage.
+request_scan_listener_slot = cg.slot_counter("LN882H_BLE_SCAN_LISTENER_COUNT")
 
 
 async def to_code(config: ConfigType) -> None:
@@ -64,5 +51,3 @@ async def to_code(config: ConfigType) -> None:
     cg.add_platformio_option("custom_options.proj_config#h", ["CFG_SUPPORT_BLE=1"])
 
     cg.add_define("USE_LN882H_BLE")
-
-    CORE.add_job(_add_listener_count)

esphome/components/rp2_ble_tracker/__init__.py

@@ -20,7 +20,6 @@ from esphome.const import (
     CONF_ID,
     CONF_INTERVAL,
 )
-from esphome.core import CORE, CoroPriority, coroutine_with_priority
 from esphome.types import ConfigType
 
 CONF_RP2040_BLE_ID = "rp2040_ble_id"
@@ -54,16 +53,6 @@ CONFIG_SCHEMA = cv.Schema(
 ).extend(cv.COMPONENT_SCHEMA)
 
 
-# Runs at FINAL priority so every BLE sensor has registered through
-# ble_device_base (and any tracker-owned listeners have been counted) before
-# the StaticVector size is emitted. Same pattern as esp32_ble_tracker.
-@coroutine_with_priority(CoroPriority.FINAL)
-async def _emit_listener_count() -> None:
-    count = ble_device_base.get_listener_count()
-    if count > 0:
-        cg.add_define("ESPHOME_BLE_DEVICE_BASE_LISTENER_COUNT", count)
-
-
 async def to_code(config: ConfigType) -> None:
     var = cg.new_Pvariable(config[CONF_ID])
     await cg.register_component(var, config)
@@ -80,5 +69,3 @@ async def to_code(config: ConfigType) -> None:
     cg.add(var.set_scan_duration(scan[CONF_DURATION].total_milliseconds))
     cg.add(var.set_scan_active(scan[CONF_ACTIVE]))
     cg.add(var.set_scan_continuous(scan[CONF_CONTINUOUS]))
-
-    CORE.add_job(_emit_listener_count)

esphome/cpp_helpers.py

@@ -1,3 +1,4 @@
+from collections.abc import Callable
 from dataclasses import dataclass, field
 import logging
 
@@ -136,6 +137,67 @@ async def _generate_component_source_table() -> None:
         )
 
 
+_SLOT_COUNTER_DOMAIN = "slot_counter"
+
+
+@dataclass
+class _SlotCounterState:
+    """Per-run slot counter state: requested counts and already-emitted defines."""
+
+    counts: dict[str, int] = field(default_factory=dict)
+    emitted: set[str] = field(default_factory=set)
+
+
+def _get_slot_counter_state() -> _SlotCounterState:
+    """Get or create the slot counter state from CORE.data."""
+    if _SLOT_COUNTER_DOMAIN not in CORE.data:
+        CORE.data[_SLOT_COUNTER_DOMAIN] = _SlotCounterState()
+    return CORE.data[_SLOT_COUNTER_DOMAIN]
+
+
+def get_slot_count(define: str) -> int:
+    """Number of slots requested so far for `define`."""
+    return _get_slot_counter_state().counts.get(define, 0)
+
+
+def slot_counter(define: str) -> Callable[[], None]:
+    """Create a request_slot function for codegen-sized storage.
+
+    The pattern behind a StaticVector listener array: a consumer's to_code
+    calls the returned function once per slot it will occupy at runtime, and
+    at FINAL priority — after every consumer's to_code has run — `define` is
+    emitted with the requested count. No requests, no define: the guarded
+    storage and its registration method compile out entirely.
+
+    The counts live in a table under CORE.data, which clears between runs.
+    A request arriving after the define was already emitted raises instead of
+    silently undercounting: the define would keep the stale smaller value and
+    StaticVector::push_back would drop the extra listener at runtime.
+    """
+
+    @coroutine_with_priority(CoroPriority.FINAL)
+    async def emit_job() -> None:
+        state = _get_slot_counter_state()
+        state.emitted.add(define)
+        # Scheduled only by the first request, so the count is always >= 1 here.
+        add_define(define, state.counts[define])
+
+    def request_slot() -> None:
+        state = _get_slot_counter_state()
+        if define in state.emitted:
+            raise ValueError(
+                f"slot_counter('{define}'): slot requested after the count "
+                f"define was emitted; request slots from to_code, not from a "
+                f"job running after FINAL emission"
+            )
+        counts = state.counts
+        counts[define] = (count := counts.get(define, 0) + 1)
+        if count == 1:
+            CORE.add_job(emit_job)
+
+    return request_slot
+
+
 async def gpio_pin_expression(conf):
     """Generate an expression for the given pin option.
 

tests/component_tests/ble_device_base/config/bk72xx_controller_only.yaml

@@ -0,0 +1,7 @@
+esphome:
+  name: slotcount-controller
+
+bk72xx:
+  board: generic-bk7252
+
+bk72xx_ble:

tests/component_tests/ble_device_base/config/bk72xx_tracker.yaml

@@ -0,0 +1,7 @@
+esphome:
+  name: slotcount-tracker
+
+bk72xx:
+  board: generic-bk7252
+
+bk72xx_ble_tracker:

tests/component_tests/ble_device_base/test_slot_counter.py

@@ -0,0 +1,63 @@
+"""Tests for the shared slot_counter codegen factory.
+
+The factory is exercised end to end through the real controllers: a tracker
+config must emit the platform's scan listener count define, and a
+controller-only config must emit nothing so the guarded StaticVector storage
+compiles out.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from pathlib import Path
+
+from esphome.core import CORE
+
+from ..helpers import get_define_value
+
+
+def test_tracker_requests_one_slot(
+    generate_main: Callable[[str | Path], str],
+    component_config_path: Callable[[str], Path],
+) -> None:
+    """The tracker's to_code requests a slot; the FINAL job emits the count.
+
+    The neutral listener count must stay absent from the same build: no BLE
+    consumer registered through register_ble_device().
+    """
+    generate_main(component_config_path("bk72xx_tracker.yaml"))
+    assert get_define_value("BK72XX_BLE_SCAN_LISTENER_COUNT") == "1"
+    assert get_define_value("ESPHOME_BLE_DEVICE_BASE_LISTENER_COUNT") is None
+
+
+def test_controller_only_emits_no_count(
+    generate_main: Callable[[str | Path], str],
+    component_config_path: Callable[[str], Path],
+) -> None:
+    """No consumer, no define — the guarded listener storage compiles out."""
+    generate_main(component_config_path("bk72xx_controller_only.yaml"))
+    assert get_define_value("BK72XX_BLE_SCAN_LISTENER_COUNT") is None
+
+
+def test_neutral_listener_count_emitted_when_requested() -> None:
+    """Registering through register_ble_device() emits the neutral count.
+
+    No in-tree sensor registers through ble_device_base.register_ble_device()
+    yet (consumer migration is a follow-up), so the coroutine is driven with a
+    mock hub instead of a config; every tracker's #ifdef-guarded listener
+    storage keys on this define, and a broken emit path would compile the
+    storage out silently.
+    """
+    import esphome.codegen as cg
+    from esphome.components import ble_device_base
+    from esphome.core import ID
+
+    hub_id = ID("hub", type=ble_device_base.BLEHub)
+    CORE.register_variable(hub_id, cg.MockObj("hub"))
+    CORE.add_job(
+        ble_device_base.register_ble_device,
+        cg.MockObj("listener"),
+        {ble_device_base.CONF_BLE_HUB_ID: hub_id},
+    )
+    CORE.flush_tasks()
+    assert get_define_value(ble_device_base.LISTENER_COUNT_DEFINE) == "1"

tests/component_tests/helpers.py

@@ -27,3 +27,19 @@ def extract_packed_value(main_cpp: str, var_name: str) -> int:
     match = re.search(combined_pattern, main_cpp) or re.search(legacy_pattern, main_cpp)
     assert match, f"configure call not found for {var_name}"
     return int(match.group(1))
+
+
+def get_define_value(name: str) -> str | None:
+    """Rendered value of a CORE define, or None when absent.
+
+    Values are codegen expressions (IntLiteral); they are compared rendered.
+    A value-less define (e.g. USE_BK72XX_BLE) is present but renders as the
+    string "None", while an absent define returns the None object — easy to
+    conflate in assertions, so use this helper for valued defines only.
+    """
+    from esphome.core import CORE
+
+    for define in CORE.defines:
+        if define.name == name:
+            return str(define.value)
+    return None

tests/unit_tests/test_cpp_helpers.py

@@ -4,6 +4,7 @@ from unittest.mock import Mock
 import pytest
 
 from esphome import const, cpp_helpers as ch
+from esphome.core import CoroPriority, coroutine_with_priority
 from esphome.cpp_helpers import ComponentSourcePool, register_component_source
 
 
@@ -167,3 +168,53 @@ def test_register_component_source_overflow_suppressed_in_testing_mode(
         idx = register_component_source("overflow_component")
     assert idx == 0
     assert "Too many unique component source names" not in caplog.text
+
+
+def _define_value(name: str) -> str | None:
+    for define in ch.CORE.defines:
+        if define.name == name:
+            # Values are codegen expressions (IntLiteral); compare rendered.
+            return str(define.value)
+    return None
+
+
+def test_slot_counter_emits_requested_count() -> None:
+    """Each request bumps the count; the self-scheduled FINAL job emits it."""
+    request = ch.slot_counter("TEST_SLOT_COUNT")
+    request()
+    request()
+    ch.CORE.flush_tasks()
+    assert _define_value("TEST_SLOT_COUNT") == "2"
+
+
+def test_slot_counter_without_requests_emits_nothing() -> None:
+    """No requests, no job, no define — the guarded storage compiles out."""
+    ch.slot_counter("TEST_SLOT_COUNT_UNUSED")
+    ch.CORE.flush_tasks()
+    assert _define_value("TEST_SLOT_COUNT_UNUSED") is None
+
+
+def test_slot_counter_request_from_final_job_still_emits() -> None:
+    """The FIRST request for a define may come from a FINAL job: its emit job
+    is scheduled mid-drain and flush_tasks() loops until the heap is empty.
+    Later requests do not get this guarantee — see the companion test."""
+    request = ch.slot_counter("TEST_SLOT_COUNT_LATE")
+
+    @coroutine_with_priority(CoroPriority.FINAL)
+    async def late_requester() -> None:
+        request()
+
+    ch.CORE.add_job(late_requester)
+    ch.CORE.flush_tasks()
+    assert _define_value("TEST_SLOT_COUNT_LATE") == "1"
+
+
+def test_slot_counter_request_after_emit_raises() -> None:
+    """The boundary of FINAL-time requests: once the define was emitted, a
+    further request would silently undersize the storage, so it fails loudly."""
+    request = ch.slot_counter("TEST_SLOT_COUNT_TOO_LATE")
+    request()
+    ch.CORE.flush_tasks()
+    assert _define_value("TEST_SLOT_COUNT_TOO_LATE") == "1"
+    with pytest.raises(ValueError, match="TEST_SLOT_COUNT_TOO_LATE"):
+        request()