From b57422024704c42a67034fc394e0cfe8a654a9df Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 20 Mar 2024 13:49:37 -1000 Subject: [PATCH] Refactor rate limit helper to track time in seconds (#113898) * Refactor rate limit helper to track time in seconds Currently we created datetime and timedelta objects to enforce the rate limit. When the rate limit was being hit hard, this got expensive. We now use floats everywhere instead as they are much cheaper which is important when we are running up against a rate limit, which is by definition a hot path The rate limit helper is currently only used for templates and we do not have any code in the code base that directly passes in a rate limit so the impact to custom components is expected to be negligible if any * misesd two --- homeassistant/helpers/event.py | 10 ++--- homeassistant/helpers/ratelimit.py | 17 ++++---- homeassistant/helpers/template.py | 6 +-- tests/components/template/test_init.py | 2 +- tests/components/template/test_sensor.py | 2 +- tests/helpers/test_event.py | 50 ++++++++++++------------ tests/helpers/test_ratelimit.py | 34 +++++----------- 7 files changed, 53 insertions(+), 68 deletions(-) diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py index dd9c038340e..749c6d3e6e4 100644 --- a/homeassistant/helpers/event.py +++ b/homeassistant/helpers/event.py @@ -141,7 +141,7 @@ class TrackTemplate: template: Template variables: TemplateVarsType - rate_limit: timedelta | None = None + rate_limit: float | None = None @dataclass(slots=True) @@ -1077,7 +1077,7 @@ class TrackTemplateResultInfo: def _render_template_if_ready( self, track_template_: TrackTemplate, - now: datetime, + now: float, event: Event[EventStateChangedData] | None, ) -> bool | TrackTemplateResult: """Re-render the template if conditions match. @@ -1185,7 +1185,7 @@ class TrackTemplateResultInfo: """ updates: list[TrackTemplateResult] = [] info_changed = False - now = event.time_fired if not replayed and event else dt_util.utcnow() + now = event.time_fired_timestamp if not replayed and event else time.time() block_updates = False super_template = self._track_templates[0] if self._has_super_template else None @@ -1927,7 +1927,7 @@ def _rate_limit_for_event( event: Event[EventStateChangedData], info: RenderInfo, track_template_: TrackTemplate, -) -> timedelta | None: +) -> float | None: """Determine the rate limit for an event.""" # Specifically referenced entities are excluded # from the rate limit @@ -1937,7 +1937,7 @@ def _rate_limit_for_event( if track_template_.rate_limit is not None: return track_template_.rate_limit - rate_limit: timedelta | None = info.rate_limit + rate_limit: float | None = info.rate_limit return rate_limit diff --git a/homeassistant/helpers/ratelimit.py b/homeassistant/helpers/ratelimit.py index 430ec906bdb..516d4134f76 100644 --- a/homeassistant/helpers/ratelimit.py +++ b/homeassistant/helpers/ratelimit.py @@ -4,12 +4,11 @@ from __future__ import annotations import asyncio from collections.abc import Callable, Hashable -from datetime import datetime, timedelta import logging +import time from typing import TypeVarTuple from homeassistant.core import HomeAssistant, callback -import homeassistant.util.dt as dt_util _Ts = TypeVarTuple("_Ts") @@ -25,7 +24,7 @@ class KeyedRateLimit: ) -> None: """Initialize ratelimit tracker.""" self.hass = hass - self._last_triggered: dict[Hashable, datetime] = {} + self._last_triggered: dict[Hashable, float] = {} self._rate_limit_timers: dict[Hashable, asyncio.TimerHandle] = {} @callback @@ -34,10 +33,10 @@ class KeyedRateLimit: return bool(self._rate_limit_timers and key in self._rate_limit_timers) @callback - def async_triggered(self, key: Hashable, now: datetime | None = None) -> None: + def async_triggered(self, key: Hashable, now: float | None = None) -> None: """Call when the action we are tracking was triggered.""" self.async_cancel_timer(key) - self._last_triggered[key] = now or dt_util.utcnow() + self._last_triggered[key] = now or time.time() @callback def async_cancel_timer(self, key: Hashable) -> None: @@ -58,11 +57,11 @@ class KeyedRateLimit: def async_schedule_action( self, key: Hashable, - rate_limit: timedelta | None, - now: datetime, + rate_limit: float | None, + now: float, action: Callable[[*_Ts], None], *args: *_Ts, - ) -> datetime | None: + ) -> float | None: """Check rate limits and schedule an action if we hit the limit. If the rate limit is hit: @@ -97,7 +96,7 @@ class KeyedRateLimit: if key not in self._rate_limit_timers: self._rate_limit_timers[key] = self.hass.loop.call_later( - (next_call_time - now).total_seconds(), + next_call_time - now, action, *args, ) diff --git a/homeassistant/helpers/template.py b/homeassistant/helpers/template.py index 36d481d9e29..bb0d868bf3c 100644 --- a/homeassistant/helpers/template.py +++ b/homeassistant/helpers/template.py @@ -131,8 +131,8 @@ _T = TypeVar("_T") _R = TypeVar("_R") _P = ParamSpec("_P") -ALL_STATES_RATE_LIMIT = timedelta(minutes=1) -DOMAIN_STATES_RATE_LIMIT = timedelta(seconds=1) +ALL_STATES_RATE_LIMIT = 60 # seconds +DOMAIN_STATES_RATE_LIMIT = 1 # seconds _render_info: ContextVar[RenderInfo | None] = ContextVar("_render_info", default=None) @@ -374,7 +374,7 @@ class RenderInfo: self.domains: collections.abc.Set[str] = set() self.domains_lifecycle: collections.abc.Set[str] = set() self.entities: collections.abc.Set[str] = set() - self.rate_limit: timedelta | None = None + self.rate_limit: float | None = None self.has_time = False def __repr__(self) -> str: diff --git a/tests/components/template/test_init.py b/tests/components/template/test_init.py index ed0426b7e0e..991228623b1 100644 --- a/tests/components/template/test_init.py +++ b/tests/components/template/test_init.py @@ -248,7 +248,7 @@ async def test_reload_sensors_that_reference_other_template_sensors( next_time = dt_util.utcnow() + timedelta(seconds=1.2) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() diff --git a/tests/components/template/test_sensor.py b/tests/components/template/test_sensor.py index 8ce5596a446..c046e961fac 100644 --- a/tests/components/template/test_sensor.py +++ b/tests/components/template/test_sensor.py @@ -974,7 +974,7 @@ async def test_self_referencing_entity_picture_loop( assert len(hass.states.async_all()) == 1 next_time = dt_util.utcnow() + timedelta(seconds=1.2) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() diff --git a/tests/helpers/test_event.py b/tests/helpers/test_event.py index 9535c714bb3..cf5051e657a 100644 --- a/tests/helpers/test_event.py +++ b/tests/helpers/test_event.py @@ -1768,7 +1768,7 @@ async def test_track_template_result_complex(hass: HomeAssistant) -> None: info = async_track_template_result( hass, - [TrackTemplate(template_complex, None, timedelta(seconds=0))], + [TrackTemplate(template_complex, None, 0)], specific_run_callback, ) await hass.async_block_till_done() @@ -2179,7 +2179,7 @@ async def test_track_template_result_iterator(hass: HomeAssistant) -> None: hass, ), None, - timedelta(seconds=0), + 0, ) ], iterator_callback, @@ -2210,7 +2210,7 @@ async def test_track_template_result_iterator(hass: HomeAssistant) -> None: hass, ), None, - timedelta(seconds=0), + 0, ) ], filter_callback, @@ -2417,7 +2417,7 @@ async def test_track_template_rate_limit(hass: HomeAssistant) -> None: info = async_track_template_result( hass, - [TrackTemplate(template_refresh, None, timedelta(seconds=0.1))], + [TrackTemplate(template_refresh, None, 0.1)], refresh_listener, ) await hass.async_block_till_done() @@ -2435,7 +2435,7 @@ async def test_track_template_rate_limit(hass: HomeAssistant) -> None: assert refresh_runs == [0, 1] next_time = dt_util.utcnow() + timedelta(seconds=0.125) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2448,7 +2448,7 @@ async def test_track_template_rate_limit(hass: HomeAssistant) -> None: assert refresh_runs == [0, 1, 2] next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2485,7 +2485,7 @@ async def test_track_template_rate_limit_super(hass: HomeAssistant) -> None: hass, [ TrackTemplate(template_availability, None), - TrackTemplate(template_refresh, None, timedelta(seconds=0.1)), + TrackTemplate(template_refresh, None, 0.1), ], refresh_listener, has_super_template=True, @@ -2508,7 +2508,7 @@ async def test_track_template_rate_limit_super(hass: HomeAssistant) -> None: assert refresh_runs == [0, 1] next_time = dt_util.utcnow() + timedelta(seconds=0.125) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2525,7 +2525,7 @@ async def test_track_template_rate_limit_super(hass: HomeAssistant) -> None: assert refresh_runs == [0, 1, 4] next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2560,8 +2560,8 @@ async def test_track_template_rate_limit_super_2(hass: HomeAssistant) -> None: info = async_track_template_result( hass, [ - TrackTemplate(template_availability, None, timedelta(seconds=0.1)), - TrackTemplate(template_refresh, None, timedelta(seconds=0.1)), + TrackTemplate(template_availability, None, 0.1), + TrackTemplate(template_refresh, None, 0.1), ], refresh_listener, has_super_template=True, @@ -2581,7 +2581,7 @@ async def test_track_template_rate_limit_super_2(hass: HomeAssistant) -> None: assert refresh_runs == [1] next_time = dt_util.utcnow() + timedelta(seconds=0.125) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2597,7 +2597,7 @@ async def test_track_template_rate_limit_super_2(hass: HomeAssistant) -> None: assert refresh_runs == [1] next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2632,7 +2632,7 @@ async def test_track_template_rate_limit_super_3(hass: HomeAssistant) -> None: info = async_track_template_result( hass, [ - TrackTemplate(template_availability, None, timedelta(seconds=0.1)), + TrackTemplate(template_availability, None, 0.1), TrackTemplate(template_refresh, None), ], refresh_listener, @@ -2654,7 +2654,7 @@ async def test_track_template_rate_limit_super_3(hass: HomeAssistant) -> None: assert refresh_runs == [1, 2] next_time = dt_util.utcnow() + timedelta(seconds=0.125) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2671,7 +2671,7 @@ async def test_track_template_rate_limit_super_3(hass: HomeAssistant) -> None: assert refresh_runs == [1, 2] next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2701,7 +2701,7 @@ async def test_track_template_rate_limit_suppress_listener(hass: HomeAssistant) info = async_track_template_result( hass, - [TrackTemplate(template_refresh, None, timedelta(seconds=0.1))], + [TrackTemplate(template_refresh, None, 0.1)], refresh_listener, ) await hass.async_block_till_done() @@ -2733,7 +2733,7 @@ async def test_track_template_rate_limit_suppress_listener(hass: HomeAssistant) assert refresh_runs == [0, 1] next_time = dt_util.utcnow() + timedelta(seconds=0.125) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2760,7 +2760,7 @@ async def test_track_template_rate_limit_suppress_listener(hass: HomeAssistant) } next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -2801,7 +2801,7 @@ async def test_track_template_rate_limit_five(hass: HomeAssistant) -> None: info = async_track_template_result( hass, - [TrackTemplate(template_refresh, None, timedelta(seconds=5))], + [TrackTemplate(template_refresh, None, 5)], refresh_listener, ) await hass.async_block_till_done() @@ -2928,7 +2928,7 @@ async def test_specifically_referenced_entity_is_not_rate_limited( info = async_track_template_result( hass, - [TrackTemplate(template_refresh, None, timedelta(seconds=5))], + [TrackTemplate(template_refresh, None, 5)], refresh_listener, ) await hass.async_block_till_done() @@ -2976,8 +2976,8 @@ async def test_track_two_templates_with_different_rate_limits( info = async_track_template_result( hass, [ - TrackTemplate(template_one, None, timedelta(seconds=0.1)), - TrackTemplate(template_five, None, timedelta(seconds=5)), + TrackTemplate(template_one, None, 0.1), + TrackTemplate(template_five, None, 5), ], refresh_listener, ) @@ -3001,7 +3001,7 @@ async def test_track_two_templates_with_different_rate_limits( assert refresh_runs[template_five] == [0, 1] next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 1) with patch( - "homeassistant.helpers.ratelimit.dt_util.utcnow", return_value=next_time + "homeassistant.helpers.ratelimit.time.time", return_value=next_time.timestamp() ): async_fire_time_changed(hass, next_time) await hass.async_block_till_done() @@ -3194,7 +3194,7 @@ async def test_async_track_template_result_multiple_templates_mixing_domain( TrackTemplate(template_1, None), TrackTemplate(template_2, None), TrackTemplate(template_3, None), - TrackTemplate(template_4, None, timedelta(seconds=0)), + TrackTemplate(template_4, None, 0), ], refresh_listener, ) diff --git a/tests/helpers/test_ratelimit.py b/tests/helpers/test_ratelimit.py index 9bd18a7e882..a87493b1731 100644 --- a/tests/helpers/test_ratelimit.py +++ b/tests/helpers/test_ratelimit.py @@ -1,11 +1,10 @@ """Tests for ratelimit.""" import asyncio -from datetime import timedelta +import time from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import ratelimit -from homeassistant.util import dt as dt_util async def test_hit(hass: HomeAssistant) -> None: @@ -19,12 +18,10 @@ async def test_hit(hass: HomeAssistant) -> None: refresh_called = True rate_limiter = ratelimit.KeyedRateLimit(hass) - rate_limiter.async_triggered("key1", dt_util.utcnow()) + rate_limiter.async_triggered("key1", time.time()) assert ( - rate_limiter.async_schedule_action( - "key1", timedelta(seconds=0.001), dt_util.utcnow(), _refresh - ) + rate_limiter.async_schedule_action("key1", 0.001, time.time(), _refresh) is not None ) @@ -36,10 +33,7 @@ async def test_hit(hass: HomeAssistant) -> None: assert refresh_called assert ( - rate_limiter.async_schedule_action( - "key2", timedelta(seconds=0.001), dt_util.utcnow(), _refresh - ) - is None + rate_limiter.async_schedule_action("key2", 0.001, time.time(), _refresh) is None ) rate_limiter.async_remove() @@ -56,19 +50,13 @@ async def test_miss(hass: HomeAssistant) -> None: rate_limiter = ratelimit.KeyedRateLimit(hass) assert ( - rate_limiter.async_schedule_action( - "key1", timedelta(seconds=0.1), dt_util.utcnow(), _refresh - ) - is None + rate_limiter.async_schedule_action("key1", 0.1, time.time(), _refresh) is None ) assert not refresh_called assert not rate_limiter.async_has_timer("key1") assert ( - rate_limiter.async_schedule_action( - "key1", timedelta(seconds=0.1), dt_util.utcnow(), _refresh - ) - is None + rate_limiter.async_schedule_action("key1", 0.1, time.time(), _refresh) is None ) assert not refresh_called assert not rate_limiter.async_has_timer("key1") @@ -86,20 +74,18 @@ async def test_no_limit(hass: HomeAssistant) -> None: refresh_called = True rate_limiter = ratelimit.KeyedRateLimit(hass) - rate_limiter.async_triggered("key1", dt_util.utcnow()) + rate_limiter.async_triggered("key1", time.time()) assert ( - rate_limiter.async_schedule_action("key1", None, dt_util.utcnow(), _refresh) - is None + rate_limiter.async_schedule_action("key1", None, time.time(), _refresh) is None ) assert not refresh_called assert not rate_limiter.async_has_timer("key1") - rate_limiter.async_triggered("key1", dt_util.utcnow()) + rate_limiter.async_triggered("key1", time.time()) assert ( - rate_limiter.async_schedule_action("key1", None, dt_util.utcnow(), _refresh) - is None + rate_limiter.async_schedule_action("key1", None, time.time(), _refresh) is None ) assert not refresh_called assert not rate_limiter.async_has_timer("key1")