Merge pull request #396 from mcdonc/perf/stagger-periodic-polls

perf: stagger periodic poll sweeps to avoid thundering herd
This commit is contained in:
Yiorgis Gozadinos 2026-06-01 17:23:22 +03:00 committed by GitHub
commit e3ea207358
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 44 additions and 3 deletions

View file

@ -1,5 +1,6 @@
import asyncio import asyncio
import logging import logging
import random
from datetime import UTC, datetime from datetime import UTC, datetime
from haiku.rag.config import SourceConfig from haiku.rag.config import SourceConfig
@ -15,6 +16,8 @@ from haiku.rag.telemetry import get_context, logfire
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_STAGGER_FRACTION = 0.25
def _enqueue_extra() -> dict | None: def _enqueue_extra() -> dict | None:
"""Per-job context the worker can't reconstruct from config alone. """Per-job context the worker can't reconstruct from config alone.
@ -84,6 +87,17 @@ class BasePoller:
await asyncio.gather(self._task, return_exceptions=True) await asyncio.gather(self._task, return_exceptions=True)
self._task = None self._task = None
async def _stagger_start(self) -> bool:
"""Sleep a random fraction of the interval so pollers sharing an
interval don't sweep in lockstep. Returns True if stop was
signalled during the wait."""
jitter = random.uniform(0, self.config.poll_interval_s * _STAGGER_FRACTION)
try:
await asyncio.wait_for(self._stop.wait(), timeout=jitter)
return True
except TimeoutError:
return False
async def _sweep_once(self) -> bool: async def _sweep_once(self) -> bool:
"""One discover() sweep. Returns True on success, False if the """One discover() sweep. Returns True on success, False if the
breaker is open, the source has pending work already queued, or the breaker is open, the source has pending work already queued, or the

View file

@ -64,6 +64,8 @@ class FSPoller(BasePoller):
"""Periodic full sweep. Catches files modified while the watcher """Periodic full sweep. Catches files modified while the watcher
wasn't running (gaps between starts, races, FS events the OS dropped). wasn't running (gaps between starts, races, FS events the OS dropped).
Sweep behaviour is unit-tested via `_sweep_once()` directly.""" Sweep behaviour is unit-tested via `_sweep_once()` directly."""
if await self._stagger_start():
return
while not self._stop.is_set(): while not self._stop.is_set():
try: try:
await asyncio.wait_for( await asyncio.wait_for(

View file

@ -1,5 +1,4 @@
import asyncio import asyncio
import logging
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from haiku.rag.ingester.pollers.base import BasePoller from haiku.rag.ingester.pollers.base import BasePoller
@ -7,8 +6,6 @@ from haiku.rag.ingester.pollers.base import BasePoller
if TYPE_CHECKING: if TYPE_CHECKING:
from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker
logger = logging.getLogger(__name__)
class PeriodicPoller(BasePoller): class PeriodicPoller(BasePoller):
"""Runs `source.discover()` on a fixed interval. Used for HTTP, S3, WebDAV """Runs `source.discover()` on a fixed interval. Used for HTTP, S3, WebDAV
@ -38,6 +35,8 @@ class PeriodicPoller(BasePoller):
# immediately instead of waiting one full interval. The sweep # immediately instead of waiting one full interval. The sweep
# behaviour itself is exercised via `_sweep_once()` unit tests. # behaviour itself is exercised via `_sweep_once()` unit tests.
await self._sweep_once() await self._sweep_once()
if await self._stagger_start():
return
while not self._stop.is_set(): while not self._stop.is_set():
try: try:
await asyncio.wait_for( await asyncio.wait_for(

View file

@ -114,6 +114,32 @@ def _periodic(source, config, jobs, sync, **kwargs):
) )
# --- _stagger_start ---
@pytest.mark.asyncio
async def test_stagger_start_sleeps_fraction_of_interval(jobs, sync, fs_config, monkeypatch):
"""_stagger_start should sleep for a random fraction of poll_interval_s
and return False (not stopped)."""
monkeypatch.setattr("random.uniform", lambda a, b: b) # max jitter
source = _StubSource("src", [])
poller = _periodic(source, fs_config, jobs, sync)
# poll_interval_s=0.05, so max jitter = 0.05 * 0.25 = 0.0125s
stopped = await poller._stagger_start()
assert stopped is False
@pytest.mark.asyncio
async def test_stagger_start_returns_true_when_stopped(jobs, sync, fs_config, monkeypatch):
"""If _stop is set before the jitter elapses, _stagger_start returns True."""
monkeypatch.setattr("random.uniform", lambda a, b: 10.0) # long jitter
source = _StubSource("src", [])
poller = _periodic(source, fs_config, jobs, sync)
poller._stop.set()
stopped = await poller._stagger_start()
assert stopped is True
# --- _sweep_once / event handling on the base class via PeriodicPoller --- # --- _sweep_once / event handling on the base class via PeriodicPoller ---