test(e2e): race real processes on one rate limiter bucket so a missing lock fails

Eight spawned processes start together and drain a bucket that never refills. With the flock they get exactly its burst between them, without it every attempt is granted
This commit is contained in:
ryan-crabbe-berri 2026-10-30 20:47:15 -07:00
parent 3185f73572
commit 68a24fa3dc

View file

@ -1,11 +1,22 @@
from __future__ import annotations
import multiprocessing
from collections.abc import Mapping
from dataclasses import dataclass
from multiprocessing.synchronize import Barrier
from pathlib import Path
from typing import Final
from types import MappingProxyType
from typing import Final, NoReturn
from cli_harness.rate_limiter import ProviderConfig, RateBucket, RateLimiter
_RACERS: Final = 8
_ATTEMPTS_PER_RACER: Final = 100
_BURST: Final = 400
_NO_REFILL: Final[Mapping[RateBucket, ProviderConfig]] = MappingProxyType(
{RateBucket.OPENAI: ProviderConfig(rate_per_sec=1e-9, burst=_BURST)}
)
@dataclass
class _FakeTime:
@ -27,3 +38,41 @@ def test_two_processes_on_one_host_share_a_single_budget_per_bucket(tmp_path: Pa
assert first_process.acquire(RateBucket.OPENAI) == 0.0
assert 0.5 <= second_process.acquire(RateBucket.OPENAI) < 0.51
assert second_process.acquire(RateBucket.BEDROCK_MANTLE) == 0.0
def test_processes_racing_for_one_bucket_never_get_more_than_its_burst(tmp_path: Path) -> None:
grants: Final = tmp_path / "grants"
_race(_RACERS, state_dir=tmp_path / "state", grants=grants)
assert len(grants.read_bytes()) == _BURST
def _race(racers: int, *, state_dir: Path, grants: Path) -> None:
context: Final = multiprocessing.get_context("spawn")
start_line: Final = context.Barrier(racers)
processes: Final = tuple(
context.Process(target=_take_tokens_until_told_to_wait, args=(state_dir, grants, start_line))
for _ in range(racers)
)
for process in processes:
process.start()
for process in processes:
process.join()
def _take_tokens_until_told_to_wait(state_dir: Path, grants: Path, start_line: Barrier) -> None:
limiter: Final = RateLimiter(_NO_REFILL, state_dir, _frozen_clock, _exit_instead_of_sleeping)
with grants.open("ab", buffering=0) as log:
_ = start_line.wait()
for _ in range(_ATTEMPTS_PER_RACER):
_ = limiter.acquire(RateBucket.OPENAI)
_ = log.write(b".")
def _frozen_clock() -> float:
return 0.0
def _exit_instead_of_sleeping(_seconds: float) -> NoReturn:
raise SystemExit(0)