Convert tracing context managers to classes to cut per-step overhead

Script and automation tracing wraps every executed step and evaluated condition in @contextmanager / @asynccontextmanager decorated generators (trace_path, trace_condition, trace_action). Each use allocates a generator and a _GeneratorContextManager and drives it through next() / throw(), which is noticeably more expensive than a plain context manager object: a microbenchmark measured the generator form at ~2.4x the cost of an equivalent class (958 vs 393 ns per use).

Reimplement the three as hand written context manager classes with __slots__. The behavior is identical (the exception handling order in trace_action, the breakpoint await, and trace_condition's reuse_by_child handling are all preserved); only the generator machinery is removed. This runs on the hot path of every script and automation step, and a four step script benchmark drops from ~123 to ~104 us per run (~15%).

homeassistant/helpers/condition.py

@@ -2,8 +2,7 @@
 
 import abc
 from collections import deque
-from collections.abc import Callable, Container, Coroutine, Generator, Iterable, Mapping
-from contextlib import contextmanager
+from collections.abc import Callable, Container, Coroutine, Iterable, Mapping
 from dataclasses import dataclass
 from datetime import datetime, time as dt_time, timedelta
 import functools as ft
@@ -1000,24 +999,44 @@ def condition_trace_update_result(**kwargs: Any) -> None:
     node.update_result(**kwargs)
 
 
-@contextmanager
-def trace_condition(variables: TemplateVarsType) -> Generator[TraceElement]:
-    """Trace condition evaluation."""
-    should_pop = True
-    trace_element = trace_stack_top(trace_stack_cv)
-    if trace_element and trace_element.reuse_by_child:
-        should_pop = False
-        trace_element.reuse_by_child = False
-    else:
-        trace_element = condition_trace_append(variables, trace_path_get())
-        trace_stack_push(trace_stack_cv, trace_element)
-    try:
-        yield trace_element
-    except Exception as ex:
-        trace_element.set_error(ex)
-        raise
-    finally:
-        if should_pop:
+class trace_condition:
+    """Trace condition evaluation.
+
+    Implemented as a context manager class rather than a
+    @contextmanager-decorated generator to avoid the per-use generator
+    allocation and iteration overhead on the condition evaluation hot path.
+    """
+
+    __slots__ = ("_should_pop", "_trace_element", "_variables")
+
+    _should_pop: bool
+    _trace_element: TraceElement
+
+    def __init__(self, variables: TemplateVarsType) -> None:
+        """Store the variables for the trace element."""
+        self._variables = variables
+
+    def __enter__(self) -> TraceElement:
+        """Start tracing the condition evaluation."""
+        should_pop = True
+        trace_element = trace_stack_top(trace_stack_cv)
+        if trace_element and trace_element.reuse_by_child:
+            should_pop = False
+            trace_element.reuse_by_child = False
+        else:
+            trace_element = condition_trace_append(self._variables, trace_path_get())
+            trace_stack_push(trace_stack_cv, trace_element)
+        self._should_pop = should_pop
+        self._trace_element = trace_element
+        return trace_element
+
+    def __exit__(
+        self, exc_type: object, exc_val: BaseException | None, exc_tb: object
+    ) -> None:
+        """Finish tracing the condition evaluation."""
+        if exc_val is not None and isinstance(exc_val, Exception):
+            self._trace_element.set_error(exc_val)
+        if self._should_pop:
             trace_stack_pop(trace_stack_cv)
 
 

homeassistant/helpers/script.py

@@ -1,8 +1,7 @@
 """Helpers to execute scripts."""
 
 import asyncio
-from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
-from contextlib import asynccontextmanager
+from collections.abc import Callable, Mapping, Sequence
 from contextvars import ContextVar
 from copy import copy
 from dataclasses import dataclass
@@ -186,79 +185,108 @@ def action_trace_append(variables: TemplateVarsType, path: str) -> TraceElement:
     return trace_element
 
 
-@asynccontextmanager
-async def trace_action(
-    hass: HomeAssistant,
-    script_run: _ScriptRun,
-    stop: asyncio.Future[None],
-    variables: TemplateVarsType,
-) -> AsyncGenerator[TraceElement]:
-    """Trace action execution."""
-    path = trace_path_get()
-    trace_element = action_trace_append(variables, path)
-    trace_stack_push(trace_stack_cv, trace_element)
+class trace_action:
+    """Trace action execution.
 
-    trace_id = trace_id_get()
-    if trace_id:
-        key = trace_id[0]
-        run_id = trace_id[1]
-        breakpoints = hass.data[DATA_SCRIPT_BREAKPOINTS]
-        if key in breakpoints and (
-            (
-                run_id in breakpoints[key]
-                and (
-                    path in breakpoints[key][run_id]
-                    or NODE_ANY in breakpoints[key][run_id]
+    Implemented as a context manager class rather than an
+    @asynccontextmanager-decorated generator to avoid the per-use async
+    generator allocation and iteration overhead on the script execution
+    hot path.
+    """
+
+    __slots__ = ("_hass", "_stop", "_trace_element", "_variables")
+
+    _trace_element: TraceElement
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        script_run: _ScriptRun,
+        stop: asyncio.Future[None],
+        variables: TemplateVarsType,
+    ) -> None:
+        """Store the data needed to trace the action."""
+        self._hass = hass
+        self._stop = stop
+        self._variables = variables
+
+    async def __aenter__(self) -> TraceElement:
+        """Start tracing the action, handling any configured breakpoint."""
+        hass = self._hass
+        stop = self._stop
+        path = trace_path_get()
+        trace_element = action_trace_append(self._variables, path)
+        self._trace_element = trace_element
+        trace_stack_push(trace_stack_cv, trace_element)
+
+        trace_id = trace_id_get()
+        if trace_id:
+            key = trace_id[0]
+            run_id = trace_id[1]
+            breakpoints = hass.data[DATA_SCRIPT_BREAKPOINTS]
+            if key in breakpoints and (
+                (
+                    run_id in breakpoints[key]
+                    and (
+                        path in breakpoints[key][run_id]
+                        or NODE_ANY in breakpoints[key][run_id]
+                    )
                 )
-            )
-            or (
-                RUN_ID_ANY in breakpoints[key]
-                and (
-                    path in breakpoints[key][RUN_ID_ANY]
-                    or NODE_ANY in breakpoints[key][RUN_ID_ANY]
+                or (
+                    RUN_ID_ANY in breakpoints[key]
+                    and (
+                        path in breakpoints[key][RUN_ID_ANY]
+                        or NODE_ANY in breakpoints[key][RUN_ID_ANY]
+                    )
+                )
+            ):
+                async_dispatcher_send_internal(
+                    hass, SCRIPT_BREAKPOINT_HIT, key, run_id, path
                 )
-            )
-        ):
-            async_dispatcher_send_internal(
-                hass, SCRIPT_BREAKPOINT_HIT, key, run_id, path
-            )
 
-            done = hass.loop.create_future()
+                done = hass.loop.create_future()
 
-            @callback
-            def async_continue_stop(
-                command: Literal["continue", "stop"] | None = None,
-            ) -> None:
-                if command == "stop":
-                    _set_result_unless_done(stop)
-                _set_result_unless_done(done)
+                @callback
+                def async_continue_stop(
+                    command: Literal["continue", "stop"] | None = None,
+                ) -> None:
+                    if command == "stop":
+                        _set_result_unless_done(stop)
+                    _set_result_unless_done(done)
 
-            signal = SCRIPT_DEBUG_CONTINUE_STOP.format(key, run_id)
-            remove_signal1 = async_dispatcher_connect(hass, signal, async_continue_stop)
-            remove_signal2 = async_dispatcher_connect(
-                hass, SCRIPT_DEBUG_CONTINUE_ALL, async_continue_stop
-            )
+                signal = SCRIPT_DEBUG_CONTINUE_STOP.format(key, run_id)
+                remove_signal1 = async_dispatcher_connect(
+                    hass, signal, async_continue_stop
+                )
+                remove_signal2 = async_dispatcher_connect(
+                    hass, SCRIPT_DEBUG_CONTINUE_ALL, async_continue_stop
+                )
 
-            await asyncio.wait([stop, done], return_when=asyncio.FIRST_COMPLETED)
-            remove_signal1()
-            remove_signal2()
+                await asyncio.wait([stop, done], return_when=asyncio.FIRST_COMPLETED)
+                remove_signal1()
+                remove_signal2()
 
-    try:
-        yield trace_element
-    except _AbortScript as ex:
-        trace_element.set_error(ex.__cause__ or ex)
-        raise
-    except _ConditionFail:
-        # Clear errors which may have been set when evaluating the condition
-        trace_element.set_error(None)
-        raise
-    except _StopScript:
-        raise
-    except Exception as ex:
-        trace_element.set_error(ex)
-        raise
-    finally:
-        trace_stack_pop(trace_stack_cv)
+        return trace_element
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
+    ) -> None:
+        """Finish tracing the action, recording any error."""
+        try:
+            if isinstance(exc_val, _AbortScript):
+                self._trace_element.set_error(exc_val.__cause__ or exc_val)
+            elif isinstance(exc_val, _ConditionFail):
+                # Clear errors which may have been set when evaluating the condition
+                self._trace_element.set_error(None)
+            elif isinstance(exc_val, _StopScript):
+                pass
+            elif isinstance(exc_val, Exception):
+                self._trace_element.set_error(exc_val)
+        finally:
+            trace_stack_pop(trace_stack_cv)
 
 
 def make_script_schema(

homeassistant/helpers/trace.py

@@ -302,17 +302,32 @@ def script_execution_get() -> str | None:
     return data.script_execution
 
 
-@contextmanager
-def trace_path(suffix: str | list[str]) -> Generator[None]:
+class trace_path:
     """Go deeper in the config tree.
 
-    Can not be used as a decorator on couroutine functions.
+    Implemented as a context manager class rather than a
+    @contextmanager-decorated generator to avoid the per-use generator
+    allocation and iteration overhead; this runs on every traced step and
+    condition.
+
+    Can not be used as a decorator on coroutine functions.
     """
-    count = trace_path_push(suffix)
-    try:
-        yield
-    finally:
-        trace_path_pop(count)
+
+    __slots__ = ("_count", "_suffix")
+
+    _count: int
+
+    def __init__(self, suffix: str | list[str]) -> None:
+        """Store the path suffix to push on enter."""
+        self._suffix = suffix
+
+    def __enter__(self) -> None:
+        """Go deeper in the config tree."""
+        self._count = trace_path_push(self._suffix)
+
+    def __exit__(self, *exc: object) -> None:
+        """Go back up in the config tree."""
+        trace_path_pop(self._count)
 
 
 def async_trace_path[*_Ts](