[F100] atomic Redis probe-failure counter via server-side Lua

increment_probe_failures / reset_probe_failures did a non-atomic
get_state (GET) -> mutate -> set (SET) in Python. A concurrent
activate() re-park writes a FRESH episode blob (probe_failures: 0 +
fresh activated_at / retry_after / affected_agents / kind); if the
stale increment's SET landed after the fresh activate's SET, the stale
blob overwrote the fresh episode metadata AND un-reset the counter
(clobbering the new episode).

Redis single-threads a Lua EVAL, so a server-side read-modify-write
is indivisible: activate's SET is serialized entirely before or after
the script, never interleaved between the script's GET and SET. The
two scripts mutate ONLY probe_failures, so every other episode field
survives the bump. activate stays a single atomic SET (a fresh episode
resetting the counter to 0 is correct semantics).
This commit is contained in:
Renn F
2026-06-28 21:16:31 +02:00
parent f4c67bedff
commit 3e768bbb89
3 changed files with 264 additions and 10 deletions
+29 -1
View File
@@ -9,6 +9,7 @@ visible to a fresh instance.
from __future__ import annotations
import json
from typing import Any
from unittest.mock import AsyncMock
@@ -22,7 +23,8 @@ from roboco.services.gateway.rate_limit_tracker import RateLimitStateTracker
def _make_redis_mock(initial_store: dict[str, Any] | None = None) -> AsyncMock:
"""Build an async Redis mock backed by a plain dict.
The mock supports ``get``, ``set``, and ``delete`` with the same
The mock supports ``get``, ``set``, ``delete`` and ``eval`` (server-side
Lua for the tracker's atomic probe-failure scripts) with the same
semantics as the real redis.asyncio.Redis client.
"""
# Use the dict AS-IS (no copy) so that two mocks sharing the same
@@ -44,10 +46,36 @@ def _make_redis_mock(initial_store: dict[str, Any] | None = None) -> AsyncMock:
async def _delete(key: str) -> int:
return 1 if store.pop(key, None) is not None else 0
async def _eval(script: str, _numkeys: int, *keys_and_args: Any) -> Any:
# Mirror the tracker's two atomic Lua scripts (see rate_limit_tracker.py)
# so the counter update is observable in-process. Single-threaded tests
# get the same result production gets from Redis' single-threaded Lua.
key = keys_and_args[0]
raw = store.get(key)
text = raw.decode() if isinstance(raw, bytes) else (str(raw) if raw else None)
if "roboco:increment_probe_failures" in script:
if text is None:
state: dict[str, Any] = {"probe_failures": 1}
else:
state = json.loads(text)
state["probe_failures"] = state.get("probe_failures", 0) + 1
store[key] = json.dumps(state)
return state["probe_failures"]
if "roboco:reset_probe_failures" in script:
if text is None:
state = {"probe_failures": 0}
else:
state = json.loads(text)
state["probe_failures"] = 0
store[key] = json.dumps(state)
return None
raise AssertionError(f"unknown eval script: {script[:80]}")
mock = AsyncMock()
mock.get = AsyncMock(side_effect=_get)
mock.set = AsyncMock(side_effect=_set)
mock.delete = AsyncMock(side_effect=_delete)
mock.eval = AsyncMock(side_effect=_eval)
# Stash the backing store so tests can inspect raw state
mock._store = store
return mock
@@ -0,0 +1,184 @@
"""The probe-failure counter update must be a single atomic Redis op.
``increment_probe_failures`` / ``reset_probe_failures`` used to do a non-atomic
read-modify-write: ``get_state`` (GET) → mutate the dict → ``SET`` the whole
blob back. A concurrent ``activate()`` (a re-park — grok 429 re-park, a 529
overload re-park) writes a FRESH episode blob (``probe_failures: 0`` + fresh
``activated_at`` / ``retry_after`` / ``affected_agents`` / ``kind``). If the
stale ``increment``'s SET lands AFTER the fresh ``activate``'s SET, the stale
blob overwrites the fresh episode metadata AND un-resets the counter (writes
back the old ``probe_failures`` + old metadata) — clobbering the new episode.
Redis single-threads a Lua ``EVAL``, so a server-side read-modify-write is
indivisible: ``activate``'s ``SET`` is serialized entirely before or entirely
after the script — never interleaved between the script's GET and SET. These
tests pin the WIRING (the increment/reset go through ``eval``, a single atomic
server-side call, NOT a separate ``get``+``set`` pair) and the field-preservation
(the script decodes, mutates ONLY ``probe_failures``, re-encodes — every other
episode field survives the bump).
"""
from __future__ import annotations
import json
from typing import Any
from unittest.mock import AsyncMock
import pytest
from roboco.services.gateway.rate_limit_tracker import RateLimitStateTracker
def _make_redis_mock(initial_store: dict[str, Any] | None = None) -> AsyncMock:
"""Async Redis mock backed by a dict, with server-side ``eval`` for the
tracker's two Lua scripts.
The ``eval`` impl mirrors the Lua (GET → decode → mutate only
``probe_failures`` → SET) so a single-threaded test observes the same result
production gets from Redis' atomic Lua execution. Real concurrency cannot be
simulated with a mock; the atomicity guarantee in production is Redis'
single-threaded Lua, which these tests pin by asserting the tracker routes
through ``eval`` (one atomic op) rather than a separate ``get``+``set``.
"""
store: dict[str, Any] = initial_store if initial_store is not None else {}
async def _get(key: str) -> bytes | None:
val = store.get(key)
if val is None:
return None
if isinstance(val, bytes):
return val
return str(val).encode()
async def _set(key: str, value: Any) -> None:
store[key] = value
async def _delete(key: str) -> int:
return 1 if store.pop(key, None) is not None else 0
async def _eval(script: str, _numkeys: int, *keys_and_args: Any) -> Any:
key = keys_and_args[0]
raw = store.get(key)
text = raw.decode() if isinstance(raw, bytes) else (str(raw) if raw else None)
if "roboco:increment_probe_failures" in script:
if text is None:
state: dict[str, Any] = {"probe_failures": 1}
else:
state = json.loads(text)
state["probe_failures"] = state.get("probe_failures", 0) + 1
store[key] = json.dumps(state)
return state["probe_failures"]
if "roboco:reset_probe_failures" in script:
if text is None:
state = {"probe_failures": 0}
else:
state = json.loads(text)
state["probe_failures"] = 0
store[key] = json.dumps(state)
return None
raise AssertionError(f"unknown eval script: {script[:80]}")
mock = AsyncMock()
mock.get = AsyncMock(side_effect=_get)
mock.set = AsyncMock(side_effect=_set)
mock.delete = AsyncMock(side_effect=_delete)
mock.eval = AsyncMock(side_effect=_eval)
mock._store = store
return mock
def _make_tracker(redis_mock: AsyncMock) -> RateLimitStateTracker:
tracker = RateLimitStateTracker(provider="anthropic", redis_url="redis://unused")
tracker._redis = redis_mock
return tracker
@pytest.mark.asyncio
async def test_increment_uses_atomic_eval_not_separate_get_set() -> None:
"""The increment must route through ``eval`` (one atomic server-side op) and
must NOT issue a separate ``set`` for the read-modify-write — the separate
SET is exactly the non-atomic write a concurrent ``activate`` can clobber."""
mock = _make_redis_mock()
tracker = _make_tracker(mock)
await tracker.activate(retry_after=60.0, affected_agents=["be-dev-1"])
# activate issued the only legitimate SET; reset call counts so the increment
# path's commands are isolated.
mock.set.reset_mock()
mock.get.reset_mock()
mock.eval.reset_mock()
count = await tracker.increment_probe_failures()
assert count == 1
mock.eval.assert_awaited_once()
# The atomic op is server-side — the tracker must not issue its own SET
# (a separate SET is the non-atomic write a racing activate clobbers).
mock.set.assert_not_awaited()
@pytest.mark.asyncio
async def test_reset_uses_atomic_eval_not_separate_get_set() -> None:
mock = _make_redis_mock()
tracker = _make_tracker(mock)
await tracker.activate()
await tracker.increment_probe_failures()
mock.set.reset_mock()
mock.get.reset_mock()
mock.eval.reset_mock()
await tracker.reset_probe_failures()
mock.eval.assert_awaited_once()
mock.set.assert_not_awaited()
assert (await tracker.get_state())["probe_failures"] == 0
@pytest.mark.asyncio
async def test_increment_preserves_episode_metadata() -> None:
"""The atomic script decodes, mutates ONLY ``probe_failures``, and re-encodes
— every other episode field (rate_limited / kind / activated_at / retry_after
/ affected_agents) survives the bump. This is the property a non-atomic
GET+SET that read a STALE blob would violate under a concurrent activate."""
retry_after = 120.0
bumps = 2
mock = _make_redis_mock()
tracker = _make_tracker(mock)
await tracker.activate(
retry_after=retry_after,
affected_agents=["be-dev-1", "fe-dev-1"],
kind="overloaded",
)
before = await tracker.get_state()
for _ in range(bumps):
await tracker.increment_probe_failures()
after = await tracker.get_state()
assert after["probe_failures"] == bumps
# Episode metadata untouched by the counter bump.
assert after["rate_limited"] is before["rate_limited"] is True
assert after["kind"] == "overloaded"
assert after["retry_after"] == retry_after
assert after["affected_agents"] == ["be-dev-1", "fe-dev-1"]
assert after["activated_at"] == before["activated_at"]
@pytest.mark.asyncio
async def test_increment_accumulates_across_calls() -> None:
mock = _make_redis_mock()
tracker = _make_tracker(mock)
await tracker.activate()
total = 4
counts = [await tracker.increment_probe_failures() for _ in range(total)]
assert counts == [1, 2, 3, total]
assert (await tracker.get_state())["probe_failures"] == total
@pytest.mark.asyncio
async def test_reset_zeroes_after_increments() -> None:
mock = _make_redis_mock()
tracker = _make_tracker(mock)
await tracker.activate()
for _ in range(3):
await tracker.increment_probe_failures()
await tracker.reset_probe_failures()
assert (await tracker.get_state())["probe_failures"] == 0