jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-07-17 02:07:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactor rate limit helper to track time in seconds (#113898)

Browse Source 
* 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
This commit is contained in:
J. Nick Koston 2024-03-20 13:49:37 -10:00 committed by GitHub
parent b311fe2a0f
commit b574220247
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 53 additions and 68 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
homeassistant/helpers/event.py
View File

@ -141,7 +141,7 @@ class TrackTemplate:
template: Template template: Template
variables: TemplateVarsType variables: TemplateVarsType
rate_limit: timedelta | None = None rate_limit: float | None = None
@dataclass(slots=True) @dataclass(slots=True)
@ -1077,7 +1077,7 @@ class TrackTemplateResultInfo:
def _render_template_if_ready( def _render_template_if_ready(
self, self,
track_template_: TrackTemplate, track_template_: TrackTemplate,
now: datetime, now: float,
event: Event[EventStateChangedData] | None, event: Event[EventStateChangedData] | None,
) -> bool | TrackTemplateResult: ) -> bool | TrackTemplateResult:
"""Re-render the template if conditions match. """Re-render the template if conditions match.
@ -1185,7 +1185,7 @@ class TrackTemplateResultInfo:
""" """
updates: list[TrackTemplateResult] = [] updates: list[TrackTemplateResult] = []
info_changed = False 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 block_updates = False
super_template = self._track_templates[0] if self._has_super_template else None 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], event: Event[EventStateChangedData],
info: RenderInfo, info: RenderInfo,
track_template_: TrackTemplate, track_template_: TrackTemplate,
) -> timedelta | None: ) -> float | None:
"""Determine the rate limit for an event.""" """Determine the rate limit for an event."""
# Specifically referenced entities are excluded # Specifically referenced entities are excluded
# from the rate limit # from the rate limit
@ -1937,7 +1937,7 @@ def _rate_limit_for_event(
if track_template_.rate_limit is not None: if track_template_.rate_limit is not None:
return track_template_.rate_limit return track_template_.rate_limit
rate_limit: timedelta | None = info.rate_limit rate_limit: float | None = info.rate_limit
return rate_limit return rate_limit

17
homeassistant/helpers/ratelimit.py
View File

@ -4,12 +4,11 @@ from __future__ import annotations
import asyncio import asyncio
from collections.abc import Callable, Hashable from collections.abc import Callable, Hashable
from datetime import datetime, timedelta
import logging import logging
import time
from typing import TypeVarTuple from typing import TypeVarTuple
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
import homeassistant.util.dt as dt_util
_Ts = TypeVarTuple("_Ts") _Ts = TypeVarTuple("_Ts")
@ -25,7 +24,7 @@ class KeyedRateLimit:
) -> None: ) -> None:
"""Initialize ratelimit tracker.""" """Initialize ratelimit tracker."""
self.hass = hass self.hass = hass
self._last_triggered: dict[Hashable, datetime] = {} self._last_triggered: dict[Hashable, float] = {}
self._rate_limit_timers: dict[Hashable, asyncio.TimerHandle] = {} self._rate_limit_timers: dict[Hashable, asyncio.TimerHandle] = {}
@callback @callback
@ -34,10 +33,10 @@ class KeyedRateLimit:
return bool(self._rate_limit_timers and key in self._rate_limit_timers) return bool(self._rate_limit_timers and key in self._rate_limit_timers)
@callback @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.""" """Call when the action we are tracking was triggered."""
self.async_cancel_timer(key) self.async_cancel_timer(key)
self._last_triggered[key] = now or dt_util.utcnow() self._last_triggered[key] = now or time.time()
@callback @callback
def async_cancel_timer(self, key: Hashable) -> None: def async_cancel_timer(self, key: Hashable) -> None:
@ -58,11 +57,11 @@ class KeyedRateLimit:
def async_schedule_action( def async_schedule_action(
self, self,
key: Hashable, key: Hashable,
rate_limit: timedelta | None, rate_limit: float | None,
now: datetime, now: float,
action: Callable[[*_Ts], None], action: Callable[[*_Ts], None],
*args: *_Ts, *args: *_Ts,
) -> datetime | None: ) -> float | None:
"""Check rate limits and schedule an action if we hit the limit. """Check rate limits and schedule an action if we hit the limit.
If the rate limit is hit: If the rate limit is hit:
@ -97,7 +96,7 @@ class KeyedRateLimit:
if key not in self._rate_limit_timers: if key not in self._rate_limit_timers:
self._rate_limit_timers[key] = self.hass.loop.call_later( self._rate_limit_timers[key] = self.hass.loop.call_later(
(next_call_time - now).total_seconds(), next_call_time - now,
action, action,
*args, *args,
) )

6
homeassistant/helpers/template.py
View File

@ -131,8 +131,8 @@ _T = TypeVar("_T")
_R = TypeVar("_R") _R = TypeVar("_R")
_P = ParamSpec("_P") _P = ParamSpec("_P")
ALL_STATES_RATE_LIMIT = timedelta(minutes=1) ALL_STATES_RATE_LIMIT = 60 # seconds
DOMAIN_STATES_RATE_LIMIT = timedelta(seconds=1) DOMAIN_STATES_RATE_LIMIT = 1 # seconds
_render_info: ContextVar[RenderInfo | None] = ContextVar("_render_info", default=None) _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: collections.abc.Set[str] = set()
self.domains_lifecycle: collections.abc.Set[str] = set() self.domains_lifecycle: collections.abc.Set[str] = set()
self.entities: 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 self.has_time = False
def __repr__(self) -> str: def __repr__(self) -> str:

2
tests/components/template/test_init.py
View File

@ -248,7 +248,7 @@ async def test_reload_sensors_that_reference_other_template_sensors(
next_time = dt_util.utcnow() + timedelta(seconds=1.2) next_time = dt_util.utcnow() + timedelta(seconds=1.2)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() await hass.async_block_till_done()

2
tests/components/template/test_sensor.py
View File

@ -974,7 +974,7 @@ async def test_self_referencing_entity_picture_loop(
assert len(hass.states.async_all()) == 1 assert len(hass.states.async_all()) == 1
next_time = dt_util.utcnow() + timedelta(seconds=1.2) next_time = dt_util.utcnow() + timedelta(seconds=1.2)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() await hass.async_block_till_done()

50
tests/helpers/test_event.py
View File

@ -1768,7 +1768,7 @@ async def test_track_template_result_complex(hass: HomeAssistant) -> None:
info = async_track_template_result( info = async_track_template_result(
hass, hass,
[TrackTemplate(template_complex, None, timedelta(seconds=0))], [TrackTemplate(template_complex, None, 0)],
specific_run_callback, specific_run_callback,
) )
await hass.async_block_till_done() await hass.async_block_till_done()
@ -2179,7 +2179,7 @@ async def test_track_template_result_iterator(hass: HomeAssistant) -> None:
hass, hass,
), ),
None, None,
timedelta(seconds=0), 0,
) )
], ],
iterator_callback, iterator_callback,
@ -2210,7 +2210,7 @@ async def test_track_template_result_iterator(hass: HomeAssistant) -> None:
hass, hass,
), ),
None, None,
timedelta(seconds=0), 0,
) )
], ],
filter_callback, filter_callback,
@ -2417,7 +2417,7 @@ async def test_track_template_rate_limit(hass: HomeAssistant) -> None:
info = async_track_template_result( info = async_track_template_result(
hass, hass,
[TrackTemplate(template_refresh, None, timedelta(seconds=0.1))], [TrackTemplate(template_refresh, None, 0.1)],
refresh_listener, refresh_listener,
) )
await hass.async_block_till_done() 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] assert refresh_runs == [0, 1]
next_time = dt_util.utcnow() + timedelta(seconds=0.125) next_time = dt_util.utcnow() + timedelta(seconds=0.125)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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] assert refresh_runs == [0, 1, 2]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2) next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -2485,7 +2485,7 @@ async def test_track_template_rate_limit_super(hass: HomeAssistant) -> None:
hass, hass,
[ [
TrackTemplate(template_availability, None), TrackTemplate(template_availability, None),
TrackTemplate(template_refresh, None, timedelta(seconds=0.1)), TrackTemplate(template_refresh, None, 0.1),
], ],
refresh_listener, refresh_listener,
has_super_template=True, has_super_template=True,
@ -2508,7 +2508,7 @@ async def test_track_template_rate_limit_super(hass: HomeAssistant) -> None:
assert refresh_runs == [0, 1] assert refresh_runs == [0, 1]
next_time = dt_util.utcnow() + timedelta(seconds=0.125) next_time = dt_util.utcnow() + timedelta(seconds=0.125)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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] assert refresh_runs == [0, 1, 4]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2) next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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( info = async_track_template_result(
hass, hass,
[ [
TrackTemplate(template_availability, None, timedelta(seconds=0.1)), TrackTemplate(template_availability, None, 0.1),
TrackTemplate(template_refresh, None, timedelta(seconds=0.1)), TrackTemplate(template_refresh, None, 0.1),
], ],
refresh_listener, refresh_listener,
has_super_template=True, has_super_template=True,
@ -2581,7 +2581,7 @@ async def test_track_template_rate_limit_super_2(hass: HomeAssistant) -> None:
assert refresh_runs == [1] assert refresh_runs == [1]
next_time = dt_util.utcnow() + timedelta(seconds=0.125) next_time = dt_util.utcnow() + timedelta(seconds=0.125)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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] assert refresh_runs == [1]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2) next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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( info = async_track_template_result(
hass, hass,
[ [
TrackTemplate(template_availability, None, timedelta(seconds=0.1)), TrackTemplate(template_availability, None, 0.1),
TrackTemplate(template_refresh, None), TrackTemplate(template_refresh, None),
], ],
refresh_listener, refresh_listener,
@ -2654,7 +2654,7 @@ async def test_track_template_rate_limit_super_3(hass: HomeAssistant) -> None:
assert refresh_runs == [1, 2] assert refresh_runs == [1, 2]
next_time = dt_util.utcnow() + timedelta(seconds=0.125) next_time = dt_util.utcnow() + timedelta(seconds=0.125)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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] assert refresh_runs == [1, 2]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2) next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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( info = async_track_template_result(
hass, hass,
[TrackTemplate(template_refresh, None, timedelta(seconds=0.1))], [TrackTemplate(template_refresh, None, 0.1)],
refresh_listener, refresh_listener,
) )
await hass.async_block_till_done() 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] assert refresh_runs == [0, 1]
next_time = dt_util.utcnow() + timedelta(seconds=0.125) next_time = dt_util.utcnow() + timedelta(seconds=0.125)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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) next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 2)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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( info = async_track_template_result(
hass, hass,
[TrackTemplate(template_refresh, None, timedelta(seconds=5))], [TrackTemplate(template_refresh, None, 5)],
refresh_listener, refresh_listener,
) )
await hass.async_block_till_done() 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( info = async_track_template_result(
hass, hass,
[TrackTemplate(template_refresh, None, timedelta(seconds=5))], [TrackTemplate(template_refresh, None, 5)],
refresh_listener, refresh_listener,
) )
await hass.async_block_till_done() 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( info = async_track_template_result(
hass, hass,
[ [
TrackTemplate(template_one, None, timedelta(seconds=0.1)), TrackTemplate(template_one, None, 0.1),
TrackTemplate(template_five, None, timedelta(seconds=5)), TrackTemplate(template_five, None, 5),
], ],
refresh_listener, refresh_listener,
) )
@ -3001,7 +3001,7 @@ async def test_track_two_templates_with_different_rate_limits(
assert refresh_runs[template_five] == [0, 1] assert refresh_runs[template_five] == [0, 1]
next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 1) next_time = dt_util.utcnow() + timedelta(seconds=0.125 * 1)
with patch( 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) async_fire_time_changed(hass, next_time)
await hass.async_block_till_done() 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_1, None),
TrackTemplate(template_2, None), TrackTemplate(template_2, None),
TrackTemplate(template_3, None), TrackTemplate(template_3, None),
TrackTemplate(template_4, None, timedelta(seconds=0)), TrackTemplate(template_4, None, 0),
], ],
refresh_listener, refresh_listener,
) )

34
tests/helpers/test_ratelimit.py
View File

@ -1,11 +1,10 @@
"""Tests for ratelimit.""" """Tests for ratelimit."""
import asyncio import asyncio
from datetime import timedelta import time
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import ratelimit from homeassistant.helpers import ratelimit
from homeassistant.util import dt as dt_util
async def test_hit(hass: HomeAssistant) -> None: async def test_hit(hass: HomeAssistant) -> None:
@ -19,12 +18,10 @@ async def test_hit(hass: HomeAssistant) -> None:
refresh_called = True refresh_called = True
rate_limiter = ratelimit.KeyedRateLimit(hass) rate_limiter = ratelimit.KeyedRateLimit(hass)
rate_limiter.async_triggered("key1", dt_util.utcnow()) rate_limiter.async_triggered("key1", time.time())
assert ( assert (
rate_limiter.async_schedule_action( rate_limiter.async_schedule_action("key1", 0.001, time.time(), _refresh)
"key1", timedelta(seconds=0.001), dt_util.utcnow(), _refresh
)
is not None is not None
) )
@ -36,10 +33,7 @@ async def test_hit(hass: HomeAssistant) -> None:
assert refresh_called assert refresh_called
assert ( assert (
rate_limiter.async_schedule_action( rate_limiter.async_schedule_action("key2", 0.001, time.time(), _refresh) is None
"key2", timedelta(seconds=0.001), dt_util.utcnow(), _refresh
)
is None
) )
rate_limiter.async_remove() rate_limiter.async_remove()
@ -56,19 +50,13 @@ async def test_miss(hass: HomeAssistant) -> None:
rate_limiter = ratelimit.KeyedRateLimit(hass) rate_limiter = ratelimit.KeyedRateLimit(hass)
assert ( assert (
rate_limiter.async_schedule_action( rate_limiter.async_schedule_action("key1", 0.1, time.time(), _refresh) is None
"key1", timedelta(seconds=0.1), dt_util.utcnow(), _refresh
)
is None
) )
assert not refresh_called assert not refresh_called
assert not rate_limiter.async_has_timer("key1") assert not rate_limiter.async_has_timer("key1")
assert ( assert (
rate_limiter.async_schedule_action( rate_limiter.async_schedule_action("key1", 0.1, time.time(), _refresh) is None
"key1", timedelta(seconds=0.1), dt_util.utcnow(), _refresh
)
is None
) )
assert not refresh_called assert not refresh_called
assert not rate_limiter.async_has_timer("key1") assert not rate_limiter.async_has_timer("key1")
@ -86,20 +74,18 @@ async def test_no_limit(hass: HomeAssistant) -> None:
refresh_called = True refresh_called = True
rate_limiter = ratelimit.KeyedRateLimit(hass) rate_limiter = ratelimit.KeyedRateLimit(hass)
rate_limiter.async_triggered("key1", dt_util.utcnow()) rate_limiter.async_triggered("key1", time.time())
assert ( assert (
rate_limiter.async_schedule_action("key1", None, dt_util.utcnow(), _refresh) rate_limiter.async_schedule_action("key1", None, time.time(), _refresh) is None
is None
) )
assert not refresh_called assert not refresh_called
assert not rate_limiter.async_has_timer("key1") assert not rate_limiter.async_has_timer("key1")
rate_limiter.async_triggered("key1", dt_util.utcnow()) rate_limiter.async_triggered("key1", time.time())
assert ( assert (
rate_limiter.async_schedule_action("key1", None, dt_util.utcnow(), _refresh) rate_limiter.async_schedule_action("key1", None, time.time(), _refresh) is None
is None
) )
assert not refresh_called assert not refresh_called
assert not rate_limiter.async_has_timer("key1") assert not rate_limiter.async_has_timer("key1")