fix(credential-pool): lock the quarantine read-modify-write of _entries
#71775 moved deferred single-use-token refreshes outside the pool lock (correct — they hold a cross-process flock plus network I/O). But _refresh_entry_impl's three terminal-auth-failure quarantine paths do a bare read-modify-write of self._entries. Those used to run with the caller holding self._lock; on the deferred path they run unlocked, so a concurrent mutation between the read and the write is silently lost. Wrap all three in 'with self._lock' (an RLock, so locked callers re-enter safely) and correct the _refresh_pending_entries docstring, which claimed the mutations were already self-locking. Post-merge gate-sweep finding on the #71775 salvage (#77714). Sibling to the acquire_lease re-select fix.
This commit is contained in:
kshitijk4poor
kshitij
parent
db0bd42119
commit
4075c8fd5a
|
|
@ -1458,17 +1458,22 @@ class CredentialPool:
|
|||
logger.debug(
|
||||
"Failed to clear terminal xAI OAuth state: %s", clear_exc
|
||||
)
|
||||
removed_ids = [
|
||||
item.id for item in self._entries
|
||||
if item.source == "device_code"
|
||||
]
|
||||
self._entries = [
|
||||
item for item in self._entries
|
||||
if item.source != "device_code"
|
||||
]
|
||||
if self._current_id == entry.id:
|
||||
self._current_id = None
|
||||
self._persist(removed_ids=removed_ids)
|
||||
# Read-modify-write of self._entries: must be atomic.
|
||||
# This runs on the DEFERRED refresh path (outside the
|
||||
# pool lock), so take it here. self._lock is an RLock,
|
||||
# so the still-locked callers re-enter safely.
|
||||
with self._lock:
|
||||
removed_ids = [
|
||||
item.id for item in self._entries
|
||||
if item.source == "device_code"
|
||||
]
|
||||
self._entries = [
|
||||
item for item in self._entries
|
||||
if item.source != "device_code"
|
||||
]
|
||||
if self._current_id == entry.id:
|
||||
self._current_id = None
|
||||
self._persist(removed_ids=removed_ids)
|
||||
return None
|
||||
# For openai-codex: same race as xAI/nous — another Hermes process
|
||||
# may have consumed the refresh token between our proactive sync
|
||||
|
|
@ -1528,17 +1533,22 @@ class CredentialPool:
|
|||
logger.debug(
|
||||
"Failed to clear terminal Codex OAuth state: %s", clear_exc
|
||||
)
|
||||
removed_ids = [
|
||||
item.id for item in self._entries
|
||||
if item.source == "device_code"
|
||||
]
|
||||
self._entries = [
|
||||
item for item in self._entries
|
||||
if item.source != "device_code"
|
||||
]
|
||||
if self._current_id == entry.id:
|
||||
self._current_id = None
|
||||
self._persist(removed_ids=removed_ids)
|
||||
# Read-modify-write of self._entries: must be atomic.
|
||||
# This runs on the DEFERRED refresh path (outside the
|
||||
# pool lock), so take it here. self._lock is an RLock,
|
||||
# so the still-locked callers re-enter safely.
|
||||
with self._lock:
|
||||
removed_ids = [
|
||||
item.id for item in self._entries
|
||||
if item.source == "device_code"
|
||||
]
|
||||
self._entries = [
|
||||
item for item in self._entries
|
||||
if item.source != "device_code"
|
||||
]
|
||||
if self._current_id == entry.id:
|
||||
self._current_id = None
|
||||
self._persist(removed_ids=removed_ids)
|
||||
return None
|
||||
# For nous: another process may have consumed the refresh token
|
||||
# between our proactive sync and the HTTP call. Re-sync from
|
||||
|
|
@ -1595,17 +1605,19 @@ class CredentialPool:
|
|||
auth_mod.NOUS_DEVICE_CODE_SOURCE,
|
||||
f"manual:{auth_mod.NOUS_DEVICE_CODE_SOURCE}",
|
||||
}
|
||||
removed_ids = [
|
||||
item.id for item in self._entries
|
||||
if item.source in singleton_sources
|
||||
]
|
||||
self._entries = [
|
||||
item for item in self._entries
|
||||
if item.source not in singleton_sources
|
||||
]
|
||||
if self._current_id == entry.id:
|
||||
self._current_id = None
|
||||
self._persist(removed_ids=removed_ids)
|
||||
# Atomic read-modify-write; see the note above.
|
||||
with self._lock:
|
||||
removed_ids = [
|
||||
item.id for item in self._entries
|
||||
if item.source in singleton_sources
|
||||
]
|
||||
self._entries = [
|
||||
item for item in self._entries
|
||||
if item.source not in singleton_sources
|
||||
]
|
||||
if self._current_id == entry.id:
|
||||
self._current_id = None
|
||||
self._persist(removed_ids=removed_ids)
|
||||
return None
|
||||
self._mark_exhausted(entry, None)
|
||||
return None
|
||||
|
|
@ -1716,9 +1728,12 @@ class CredentialPool:
|
|||
On failure the entry is silently skipped.
|
||||
"""
|
||||
for entry, sync_fn in pending:
|
||||
# _refresh_entry already merges the refreshed entry into the
|
||||
# pool internally (its mutation primitives are self-locking),
|
||||
# so no second _replace_entry is needed here.
|
||||
# _refresh_entry merges the refreshed entry into the pool
|
||||
# internally. Its mutation primitives (_replace_entry, _persist)
|
||||
# are self-locking, and the quarantine paths inside
|
||||
# _refresh_entry_impl take self._lock explicitly around their
|
||||
# read-modify-write of self._entries — required because this
|
||||
# call site runs OUTSIDE the pool lock.
|
||||
self._refresh_entry(entry, force=False)
|
||||
|
||||
def _available_entries(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,136 @@
|
|||
"""Codex/nous quarantine paths must mutate self._entries under the lock.
|
||||
|
||||
Post-merge gate-sweep finding on the #71775 salvage (#77714). That PR moved
|
||||
single-use-token refreshes OUTSIDE the pool lock to avoid stalling every
|
||||
consumer during cross-process flock + OAuth network I/O — correct in intent,
|
||||
but ``_refresh_entry_impl``'s three "terminal auth failure" quarantine paths
|
||||
do a bare read-modify-write of ``self._entries``:
|
||||
|
||||
removed_ids = [item.id for item in self._entries if ...]
|
||||
self._entries = [item for item in self._entries if ...]
|
||||
|
||||
Before #71775 those ran with the caller (``_available_entries``) holding the
|
||||
lock. On the deferred path they now run unlocked, so a concurrent mutation
|
||||
interleaved between the read and the write is silently lost.
|
||||
"""
|
||||
|
||||
import threading
|
||||
|
||||
from agent.credential_pool import CredentialPool, PooledCredential
|
||||
|
||||
|
||||
def _entry(entry_id: str, source: str) -> PooledCredential:
|
||||
return PooledCredential(
|
||||
id=entry_id,
|
||||
provider="anthropic",
|
||||
auth_type="oauth",
|
||||
access_token="tok",
|
||||
label=entry_id,
|
||||
source=source,
|
||||
priority=0,
|
||||
)
|
||||
|
||||
|
||||
def _bare_pool(entries):
|
||||
pool = CredentialPool.__new__(CredentialPool)
|
||||
pool._lock = threading.RLock()
|
||||
pool._entries = list(entries)
|
||||
pool._active_leases = {}
|
||||
pool._current_id = None
|
||||
pool._max_concurrent = 2
|
||||
pool._unmatched_rotation_streak = 0
|
||||
pool.provider = "anthropic"
|
||||
return pool
|
||||
|
||||
|
||||
def test_quarantine_read_modify_write_is_atomic():
|
||||
"""A concurrent mutation must not be lost across the quarantine filter.
|
||||
|
||||
The quarantine reads the surviving entries, then writes back a filtered
|
||||
list. If a concurrent writer lands between the read and the write and the
|
||||
section is unlocked, that write is clobbered. Under the lock the writer is
|
||||
serialized — it either lands fully before or fully after.
|
||||
"""
|
||||
pool = _bare_pool([_entry("dc1", "device_code")])
|
||||
survivor = _entry("keep", "manual")
|
||||
started = threading.Event()
|
||||
|
||||
def concurrent_add():
|
||||
started.set()
|
||||
with pool._lock: # blocks while the quarantine holds the lock
|
||||
pool._entries = pool._entries + [survivor]
|
||||
|
||||
t = threading.Thread(target=concurrent_add)
|
||||
|
||||
with pool._lock:
|
||||
_removed = [i.id for i in pool._entries if i.source == "device_code"]
|
||||
t.start()
|
||||
started.wait(timeout=2)
|
||||
# Give the writer a chance to (incorrectly) interleave.
|
||||
t.join(timeout=0.2)
|
||||
pool._entries = [i for i in pool._entries if i.source != "device_code"]
|
||||
|
||||
# Outside the lock the writer can now proceed; wait for it to finish.
|
||||
t.join(timeout=2)
|
||||
assert not t.is_alive(), "concurrent writer did not complete"
|
||||
|
||||
ids = {e.id for e in pool._entries}
|
||||
assert "dc1" not in ids, "the device_code entry should be quarantined"
|
||||
assert "keep" in ids, (
|
||||
"the concurrent append was LOST — the quarantine read-modify-write "
|
||||
"of self._entries is not atomic"
|
||||
)
|
||||
|
||||
|
||||
def test_quarantine_paths_hold_the_pool_lock():
|
||||
"""Static guard: every bare ``self._entries = [`` inside
|
||||
_refresh_entry_impl must sit under a ``with self._lock`` block.
|
||||
|
||||
The deferred-refresh call site runs outside the pool lock, so an
|
||||
unguarded rebind there is a lost-update window.
|
||||
"""
|
||||
import inspect
|
||||
import textwrap
|
||||
|
||||
src = textwrap.dedent(inspect.getsource(CredentialPool._refresh_entry_impl))
|
||||
lines = src.splitlines()
|
||||
|
||||
unguarded = []
|
||||
for idx, line in enumerate(lines):
|
||||
if "self._entries = [" not in line:
|
||||
continue
|
||||
indent = len(line) - len(line.lstrip())
|
||||
# Walk backwards for an enclosing `with self._lock` at lower indent.
|
||||
guarded = False
|
||||
for prev in range(idx - 1, -1, -1):
|
||||
p = lines[prev]
|
||||
if not p.strip():
|
||||
continue
|
||||
p_indent = len(p) - len(p.lstrip())
|
||||
if p_indent < indent:
|
||||
if "with self._lock" in p:
|
||||
guarded = True
|
||||
break
|
||||
if p.lstrip().startswith("def "):
|
||||
break
|
||||
if not guarded:
|
||||
unguarded.append(line.strip())
|
||||
|
||||
assert not unguarded, (
|
||||
"unguarded self._entries rebind(s) in _refresh_entry_impl — the "
|
||||
f"deferred refresh path runs outside the pool lock: {unguarded}"
|
||||
)
|
||||
|
||||
|
||||
def test_rlock_allows_locked_callers_to_reenter():
|
||||
"""The already-locked callers must still work after adding the lock.
|
||||
|
||||
self._lock is an RLock, so a caller holding it can re-enter the new
|
||||
quarantine block without deadlocking.
|
||||
"""
|
||||
pool = _bare_pool([_entry("dc1", "device_code")])
|
||||
|
||||
with pool._lock:
|
||||
acquired = pool._lock.acquire(timeout=1)
|
||||
assert acquired, "RLock must allow same-thread re-entry"
|
||||
pool._lock.release()
|
||||
Loading…
Reference in New Issue