mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
[F074] per-agent advisory lock closes claim TOCTOU
_run_claim_guards read the agent's other tasks via unlocked SELECTs before claim() took its row lock, and claim()'s FOR UPDATE locked only the TARGET row — so two concurrent i_will_work_on by the SAME agent on TWO DIFFERENT pending tasks each locked their own row, each read an empty in_progress set, each passed already_active, each claimed+started → the agent ended with two in_progress tasks (the in-process asyncio Lock is lost on orchestrator-restart split-brain, so it wasn't a DB-level guarantee). Fix: TaskService.acquire_claim_lock takes a transaction-scoped pg_advisory_xact_lock keyed by hashtextextended(agent_id). The gate acquires it BEFORE the guard reads (for non-coordinator roles only) so the second concurrent claim's read sees the first's committed in_progress task and is rejected. Tx-scoped → auto-releases on commit/rollback, can't outlive the request. Coordinator exemption (the key logical-regression guard): cell_pm / main_pm do NOT take the lock — the PM coordinator concurrency feature lets a PM plan+delegate many roots in parallel, and a per-agent lock would serialize those claims and regress it. Matches the existing _COORDINATOR_ROLES already_active/paused guard exemption. A hash collision only causes benign false serialization, never a false negative. Tests: unit (dev acquires lock before guard read; coordinator does not) + real-PG integration (same-agent serializes, different-agent does not, releases on rollback).
This commit is contained in:
@@ -0,0 +1,112 @@
|
||||
"""F074 — real-Postgres proof that ``TaskService.acquire_claim_lock`` serializes
|
||||
concurrent claims by the SAME agent (the one-task-per-agent invariant) while NOT
|
||||
serializing claims by DIFFERENT agents.
|
||||
|
||||
The choreographer-level ordering + coordinator-exemption is covered by the unit
|
||||
suite (``test_choreographer_claim_lock.py``); this test pins the DB-level
|
||||
contract the unit suite mocks out: that ``pg_advisory_xact_lock`` keyed by
|
||||
``hashtextextended(agent_id)`` actually blocks a second transaction trying to
|
||||
acquire the same agent's lock until the first commits/rolls back, and that a
|
||||
different agent's lock is uncontended. Skips when Postgres is unreachable.
|
||||
|
||||
Each ``acquire_claim_lock`` call uses its own fresh session/engine. The
|
||||
blocking call's session is single-use: ``asyncio.wait_for`` cancelling an
|
||||
in-flight asyncpg query leaves the SQLAlchemy session mid-connection-checkout
|
||||
("provisioning a new connection"), so a throwaway session per timed acquire
|
||||
keeps the rest of the test on clean connections.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from roboco.services.task import TaskService
|
||||
from sqlalchemy.ext.asyncio import (
|
||||
AsyncEngine,
|
||||
AsyncSession,
|
||||
async_sessionmaker,
|
||||
create_async_engine,
|
||||
)
|
||||
|
||||
|
||||
async def _fresh_session(url: str) -> tuple[AsyncSession, AsyncEngine]:
|
||||
"""A session on a brand-new engine (caller disposes via ``_dispose``)."""
|
||||
engine = create_async_engine(url, future=True)
|
||||
factory = async_sessionmaker(
|
||||
bind=engine, class_=AsyncSession, expire_on_commit=False
|
||||
)
|
||||
return factory(), engine
|
||||
|
||||
|
||||
async def _dispose(session: AsyncSession, engine: AsyncEngine) -> None:
|
||||
with contextlib.suppress(Exception):
|
||||
await session.rollback()
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_advisory_lock_serializes_same_agent_not_different(
|
||||
_test_database_url: str,
|
||||
) -> None:
|
||||
"""TX1 holds agent A's lock (uncommitted). A second acquire for agent A
|
||||
must block past a short deadline; an acquire for a DIFFERENT agent B must
|
||||
return immediately. After TX1 rolls back, agent A's lock is releasable.
|
||||
Proves the SQL is valid against real PG, serializes per-agent, is
|
||||
collision-free across agents, and is transaction-scoped (releases on
|
||||
rollback)."""
|
||||
url = _test_database_url
|
||||
agent_a = uuid4()
|
||||
agent_b = uuid4()
|
||||
|
||||
# TX1 acquires agent A's tx-scoped advisory lock and holds it (no commit).
|
||||
holder, holder_engine = await _fresh_session(url)
|
||||
try:
|
||||
await TaskService(holder).acquire_claim_lock(agent_a)
|
||||
|
||||
# A second transaction acquiring the SAME agent's lock must block.
|
||||
blocked, blocked_engine = await _fresh_session(url)
|
||||
try:
|
||||
with pytest.raises(TimeoutError):
|
||||
await asyncio.wait_for(
|
||||
TaskService(blocked).acquire_claim_lock(agent_a), timeout=0.5
|
||||
)
|
||||
finally:
|
||||
await _dispose(blocked, blocked_engine)
|
||||
|
||||
# A different agent's lock is uncontended — acquires immediately.
|
||||
other, other_engine = await _fresh_session(url)
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
TaskService(other).acquire_claim_lock(agent_b), timeout=2.0
|
||||
)
|
||||
finally:
|
||||
await _dispose(other, other_engine)
|
||||
|
||||
# Rolling back TX1 ends its transaction → releases the tx-scoped lock.
|
||||
await holder.rollback()
|
||||
finally:
|
||||
await _dispose(holder, holder_engine)
|
||||
|
||||
# Agent A's lock is now free — a fresh transaction acquires it at once.
|
||||
reacquire, reacquire_engine = await _fresh_session(url)
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
TaskService(reacquire).acquire_claim_lock(agent_a), timeout=2.0
|
||||
)
|
||||
finally:
|
||||
await _dispose(reacquire, reacquire_engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_runs_without_error(_test_database_url: str) -> None:
|
||||
"""Sanity: the advisory-lock SQL executes against the migrated schema
|
||||
(``hashtextextended`` exists, the statement is valid) — pins that the
|
||||
one-liner hasn't been broken by a schema/PG-version drift."""
|
||||
session, engine = await _fresh_session(_test_database_url)
|
||||
try:
|
||||
await TaskService(session).acquire_claim_lock(uuid4())
|
||||
finally:
|
||||
await _dispose(session, engine)
|
||||
Reference in New Issue
Block a user