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
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

View File

@ -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

View File

@ -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,
)

View File

@ -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:

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)
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()

View File

@ -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()

View File

@ -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,
)

View File

@ -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")