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:
@@ -1168,6 +1168,22 @@ class Choreographer:
|
||||
# above. These migrate into spec.extra_preconditions in a later
|
||||
# task; until then, keep them imperative so concurrency invariants
|
||||
# stay enforced.
|
||||
#
|
||||
# The agent-wide guards below read the agent's OTHER tasks via plain
|
||||
# (unlocked) SELECTs, and claim()'s FOR UPDATE locks only the TARGET
|
||||
# row — neither serializes two concurrent claims by the SAME agent on
|
||||
# TWO DIFFERENT pending tasks. A transaction-scoped advisory lock keyed
|
||||
# by agent_id, acquired here BEFORE the guard reads and held until the
|
||||
# request transaction commits, makes the second concurrent claim's read
|
||||
# see the first's committed in_progress task and get rejected. The
|
||||
# in-process asyncio.Lock is lost on orchestrator-restart split-brain,
|
||||
# so this is the only DB-level guarantee of the one-task-per-agent
|
||||
# invariant. Coordinators (cell_pm / main_pm) are exempt — matching the
|
||||
# already_active/paused guard exemption below — so a PM can still plan +
|
||||
# delegate many roots in parallel (the PM coordinator concurrency
|
||||
# feature). A hash collision only causes benign false serialization.
|
||||
if role_str not in self._COORDINATOR_ROLES:
|
||||
await self.task.acquire_claim_lock(ctx.agent_id)
|
||||
if guard := await self._run_claim_guards(
|
||||
agent_id=ctx.agent_id,
|
||||
task=t,
|
||||
|
||||
+26
-1
@@ -11,7 +11,7 @@ from datetime import UTC, datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, cast
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
from sqlalchemy import and_, func, or_, select, update
|
||||
from sqlalchemy import and_, func, or_, select, text, update
|
||||
from sqlalchemy import inspect as sa_inspect
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.orm import InstanceState
|
||||
@@ -2341,6 +2341,31 @@ class TaskService(BaseService):
|
||||
self._background_tasks.add(bg_task)
|
||||
bg_task.add_done_callback(self._background_tasks.discard)
|
||||
|
||||
async def acquire_claim_lock(self, agent_id: UUID) -> None:
|
||||
"""Take a per-agent transaction-scoped advisory lock.
|
||||
|
||||
``claim``'s ``SELECT ... FOR UPDATE`` serializes concurrent claims of
|
||||
the SAME task, but the one-task-per-agent invariant is an agent-WIDE
|
||||
guard: two concurrent claims by the SAME agent on TWO DIFFERENT pending
|
||||
tasks lock different rows and both pass the already_active guard (each
|
||||
reads an empty in_progress set before either commits). A
|
||||
``pg_advisory_xact_lock`` keyed by the agent serializes the WHOLE
|
||||
claim+guard+start critical section per agent — held until the outer
|
||||
request transaction commits, the second concurrent claim blocks until
|
||||
the first commits, then its guard read sees the first's committed
|
||||
in_progress task and is rejected.
|
||||
|
||||
Transaction-scoped (not session-scoped) so it auto-releases on commit
|
||||
or rollback and cannot outlive the request. ``hashtextextended`` maps a
|
||||
UUID to a ``bigint`` key; a hash collision only causes benign false
|
||||
serialization (two different agents momentarily serializing), never a
|
||||
false negative — so it is not a correctness concern.
|
||||
"""
|
||||
await self.session.execute(
|
||||
text("SELECT pg_advisory_xact_lock(hashtextextended(CAST(:aid AS text), 0))"),
|
||||
{"aid": str(agent_id)},
|
||||
)
|
||||
|
||||
async def claim(
|
||||
self, task_id: UUID, agent_id: UUID, allow_reassign: bool = False
|
||||
) -> TaskTable | None:
|
||||
|
||||
Reference in New Issue
Block a user