Merge pull request #4147 from nesquena/stage-3940v2
Some checks failed
Release & Docker / release (push) Has been cancelled
Some checks failed
Release & Docker / release (push) Has been cancelled
Release NL (v0.51.399): per-home provider probe-worker pool, race-safe (#3787)
This commit is contained in:
@@ -3,6 +3,12 @@
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
## [v0.51.399] — 2026-06-13 — Release NL (per-home provider probe-worker pool, #3787)
|
||||
|
||||
### Fixed
|
||||
|
||||
- **Concurrent provider-usage probes for the same home no longer spawn O(N) cold subprocesses, and the worker pool is race-safe (#3787).** `_AccountUsageProbeWorker` previously cold-spawned a fresh subprocess whenever its single worker was busy; it now uses a per-home worker pool. Critically, `_get_account_usage_probe_worker()` claims a worker **with its lock already held, inside the pool lock**, eliminating the TOCTOU window where a worker could be popped/closed by cache invalidation between the pool read and the claim (which previously relaunched an untracked subprocess). Cache invalidation flushes workers scoped to the active home (the status cache clears across all homes — an intentional, commented asymmetry). (#3787)
|
||||
|
||||
## [v0.51.398] — 2026-06-13 — Release NK (auto-generate titles for imported CLI sessions, #3987)
|
||||
|
||||
### Fixed
|
||||
|
||||
@@ -131,9 +131,12 @@ _account_usage_probe_semaphore: threading.BoundedSemaphore | None = None
|
||||
# represented as non-None snapshots and remain cacheable.
|
||||
_account_usage_status_cache: dict[tuple[str, str, str], tuple[float, Any]] = {}
|
||||
_account_usage_status_cache_lock = threading.Lock()
|
||||
_account_usage_worker_pool: dict[str, "_AccountUsageProbeWorker"] = {}
|
||||
_account_usage_worker_pool: dict[str, list["_AccountUsageProbeWorker"]] = {}
|
||||
_account_usage_worker_pool_lock = threading.Lock()
|
||||
|
||||
# Per-home worker pool configuration for probe tail-latency reduction (#3787)
|
||||
_ACCOUNT_USAGE_WORKERS_PER_HOME = 2
|
||||
|
||||
|
||||
def _get_account_usage_probe_semaphore() -> threading.BoundedSemaphore:
|
||||
global _account_usage_probe_semaphore
|
||||
@@ -1228,11 +1231,13 @@ class _AccountUsageProbeWorker:
|
||||
self.last_used = time.monotonic()
|
||||
self._lock = threading.RLock()
|
||||
self._proc: subprocess.Popen[str] | None = None
|
||||
self._closed = False
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
proc = self._proc
|
||||
self._proc = None
|
||||
self._closed = True
|
||||
self._close_process(proc)
|
||||
|
||||
@staticmethod
|
||||
@@ -1349,6 +1354,7 @@ class _AccountUsageProbeWorker:
|
||||
env=_account_usage_subprocess_env(self.home, provider, None),
|
||||
**kwargs,
|
||||
)
|
||||
self._closed = False
|
||||
except Exception:
|
||||
self._proc = None
|
||||
logger.debug("Account usage worker for %s failed to launch", provider, exc_info=True)
|
||||
@@ -1421,14 +1427,34 @@ def _fetch_account_usage_once_for_home(provider: str, home: Path, *, api_key: st
|
||||
return _account_usage_payload_to_snapshot(payload)
|
||||
|
||||
|
||||
def _get_account_usage_probe_worker(home: Path) -> _AccountUsageProbeWorker:
|
||||
def _get_account_usage_probe_worker(home: Path) -> "_AccountUsageProbeWorker | None":
|
||||
"""Return a worker with its lock already held, or None if saturated.
|
||||
|
||||
The caller MUST release worker._lock after use (typically via try/finally).
|
||||
Holding the lock across the handoff eliminates the TOCTOU window that would
|
||||
let two concurrent probes both observe the same worker as free.
|
||||
"""
|
||||
key = str(Path(home))
|
||||
stale: list[_AccountUsageProbeWorker] = []
|
||||
claimed: _AccountUsageProbeWorker | None = None
|
||||
with _account_usage_worker_pool_lock:
|
||||
worker = _account_usage_worker_pool.get(key)
|
||||
if worker is None:
|
||||
worker = _AccountUsageProbeWorker(Path(home))
|
||||
_account_usage_worker_pool[key] = worker
|
||||
return worker
|
||||
existing_workers = _account_usage_worker_pool.get(key)
|
||||
workers: list[_AccountUsageProbeWorker]
|
||||
if not existing_workers:
|
||||
workers = [_AccountUsageProbeWorker(Path(home)) for _ in range(_ACCOUNT_USAGE_WORKERS_PER_HOME)]
|
||||
else:
|
||||
stale = [worker for worker in existing_workers if worker._closed]
|
||||
workers = [worker for worker in existing_workers if not worker._closed]
|
||||
while len(workers) < _ACCOUNT_USAGE_WORKERS_PER_HOME:
|
||||
workers.append(_AccountUsageProbeWorker(Path(home)))
|
||||
_account_usage_worker_pool[key] = workers
|
||||
for worker in workers:
|
||||
if worker._lock.acquire(blocking=False):
|
||||
claimed = worker
|
||||
break
|
||||
for worker in stale:
|
||||
worker.close()
|
||||
return claimed
|
||||
|
||||
|
||||
def _cleanup_account_usage_probe_workers(
|
||||
@@ -1439,23 +1465,31 @@ def _cleanup_account_usage_probe_workers(
|
||||
cutoff = time.monotonic() if now is None else now
|
||||
stale: list[tuple[str, _AccountUsageProbeWorker]] = []
|
||||
with _account_usage_worker_pool_lock:
|
||||
for key, worker in list(_account_usage_worker_pool.items()):
|
||||
if worker._lock.acquire(blocking=False):
|
||||
try:
|
||||
proc = worker._proc
|
||||
is_dead = proc is None or proc.poll() is not None
|
||||
if is_dead or cutoff - worker.last_used >= idle_seconds:
|
||||
stale.append((key, worker))
|
||||
_account_usage_worker_pool.pop(key, None)
|
||||
finally:
|
||||
worker._lock.release()
|
||||
for key, workers in list(_account_usage_worker_pool.items()):
|
||||
for worker in workers:
|
||||
if worker._lock.acquire(blocking=False):
|
||||
try:
|
||||
proc = worker._proc
|
||||
is_dead = worker._closed or (proc is not None and proc.poll() is not None)
|
||||
if is_dead or cutoff - worker.last_used >= idle_seconds:
|
||||
stale.append((key, worker))
|
||||
finally:
|
||||
worker._lock.release()
|
||||
remaining_workers = [w for w in workers if not any(k == key and w == sw for k, sw in stale)]
|
||||
if not remaining_workers:
|
||||
_account_usage_worker_pool.pop(key, None)
|
||||
else:
|
||||
# Replenish to N so partial cleanup doesn't permanently shrink the pool
|
||||
while len(remaining_workers) < _ACCOUNT_USAGE_WORKERS_PER_HOME:
|
||||
remaining_workers.append(_AccountUsageProbeWorker(Path(key)))
|
||||
_account_usage_worker_pool[key] = remaining_workers
|
||||
for _key, worker in stale:
|
||||
worker.close()
|
||||
|
||||
|
||||
def _close_account_usage_probe_workers() -> None:
|
||||
with _account_usage_worker_pool_lock:
|
||||
workers = list(_account_usage_worker_pool.values())
|
||||
workers = [w for wlist in _account_usage_worker_pool.values() for w in wlist]
|
||||
_account_usage_worker_pool.clear()
|
||||
_close_account_usage_probe_worker_list(workers)
|
||||
|
||||
@@ -1465,15 +1499,23 @@ def _close_account_usage_probe_worker_list(workers: list[_AccountUsageProbeWorke
|
||||
worker.close()
|
||||
|
||||
|
||||
def _close_account_usage_probe_workers_async() -> None:
|
||||
def _close_account_usage_probe_workers_async(*, provider_id: str | None = None) -> None:
|
||||
with _account_usage_worker_pool_lock:
|
||||
workers = list(_account_usage_worker_pool.values())
|
||||
_account_usage_worker_pool.clear()
|
||||
if not workers:
|
||||
if provider_id:
|
||||
active_home = str(_get_hermes_home())
|
||||
workers_to_close = []
|
||||
for key, wlist in list(_account_usage_worker_pool.items()):
|
||||
if key == active_home:
|
||||
workers_to_close.extend(wlist)
|
||||
_account_usage_worker_pool.pop(key, None)
|
||||
else:
|
||||
workers_to_close = [w for wlist in _account_usage_worker_pool.values() for w in wlist]
|
||||
_account_usage_worker_pool.clear()
|
||||
if not workers_to_close:
|
||||
return
|
||||
thread = threading.Thread(
|
||||
target=_close_account_usage_probe_worker_list,
|
||||
args=(workers,),
|
||||
args=(workers_to_close,),
|
||||
daemon=True,
|
||||
name="account-usage-worker-close",
|
||||
)
|
||||
@@ -1512,7 +1554,7 @@ def invalidate_account_usage_status_cache(provider_id: str | None = None) -> Non
|
||||
for key in list(_account_usage_status_cache):
|
||||
if key[0] == normalized:
|
||||
_account_usage_status_cache.pop(key, None)
|
||||
_close_account_usage_probe_workers_async()
|
||||
_close_account_usage_probe_workers_async(provider_id=normalized or None)
|
||||
|
||||
|
||||
def _set_cached_account_usage(
|
||||
@@ -1545,7 +1587,13 @@ def _set_cached_account_usage(
|
||||
def _agent_fetch_account_usage_for_home(provider: str, home: Path, *, api_key: str | None = None) -> Any:
|
||||
try:
|
||||
_cleanup_account_usage_probe_workers()
|
||||
return _get_account_usage_probe_worker(home).fetch(provider, api_key=api_key)
|
||||
worker = _get_account_usage_probe_worker(home)
|
||||
if worker is not None:
|
||||
try:
|
||||
return worker._fetch_locked(provider, api_key=api_key)
|
||||
finally:
|
||||
worker._lock.release()
|
||||
return _fetch_account_usage_once_for_home(provider, home, api_key=api_key)
|
||||
except Exception:
|
||||
logger.debug("Account usage probe for %s failed", provider, exc_info=True)
|
||||
return None
|
||||
|
||||
422
tests/test_issue3787_probe_pool.py
Normal file
422
tests/test_issue3787_probe_pool.py
Normal file
@@ -0,0 +1,422 @@
|
||||
"""Tests for Issue #3787: Probe-worker pool tail-latency.
|
||||
|
||||
Focus on:
|
||||
- Per-home worker pool size N=2
|
||||
- Non-blocking acquire returns an idle worker when available
|
||||
- Returns None immediately when all workers are locked (no blocking wait)
|
||||
- Scoped invalidation only flushes the active home's workers
|
||||
- Full invalidation (no provider_id) flushes all workers
|
||||
- Cleanup iterates nested worker lists correctly
|
||||
- Lock held across selector handoff prevents TOCTOU races
|
||||
"""
|
||||
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from api.providers import (
|
||||
_ACCOUNT_USAGE_WORKERS_PER_HOME,
|
||||
_account_usage_worker_pool,
|
||||
_account_usage_worker_pool_lock,
|
||||
_cleanup_account_usage_probe_workers,
|
||||
_close_account_usage_probe_workers,
|
||||
_get_account_usage_probe_worker,
|
||||
invalidate_account_usage_status_cache,
|
||||
)
|
||||
|
||||
|
||||
class TestProbeWorkerPoolPerHome(unittest.TestCase):
|
||||
"""Test per-home worker pool with N=2 configuration."""
|
||||
|
||||
def setUp(self):
|
||||
"""Clear the pool before each test."""
|
||||
with _account_usage_worker_pool_lock:
|
||||
_account_usage_worker_pool.clear()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after each test."""
|
||||
with _account_usage_worker_pool_lock:
|
||||
_account_usage_worker_pool.clear()
|
||||
|
||||
def test_pool_creates_two_workers_per_home_key(self):
|
||||
"""Pool should create exactly N=2 workers for each home key."""
|
||||
home = Path.home() / ".hermes"
|
||||
worker = _get_account_usage_probe_worker(home)
|
||||
self.assertIsNotNone(worker)
|
||||
worker._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
key = str(Path(home))
|
||||
workers_list = _account_usage_worker_pool.get(key)
|
||||
self.assertIsNotNone(workers_list)
|
||||
self.assertEqual(len(workers_list), _ACCOUNT_USAGE_WORKERS_PER_HOME)
|
||||
self.assertEqual(_ACCOUNT_USAGE_WORKERS_PER_HOME, 2)
|
||||
|
||||
def test_nonblocking_acquire_returns_idle_worker(self):
|
||||
"""Non-blocking acquire should return an idle worker when one is free."""
|
||||
home = Path("/tmp/test_nonblocking_idle")
|
||||
|
||||
# Populate the pool
|
||||
worker_initial = _get_account_usage_probe_worker(home)
|
||||
self.assertIsNotNone(worker_initial)
|
||||
worker_initial._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
key = str(Path(home))
|
||||
workers = _account_usage_worker_pool[key]
|
||||
self.assertEqual(len(workers), 2)
|
||||
|
||||
# Hold workers[0] locked from a background thread to simulate actual use
|
||||
lock_holder = threading.Event()
|
||||
release_signal = threading.Event()
|
||||
|
||||
def hold_lock():
|
||||
workers[0]._lock.acquire()
|
||||
lock_holder.set()
|
||||
release_signal.wait(timeout=2.0)
|
||||
workers[0]._lock.release()
|
||||
|
||||
thread = threading.Thread(target=hold_lock, daemon=True)
|
||||
thread.start()
|
||||
lock_holder.wait(timeout=2.0)
|
||||
|
||||
try:
|
||||
worker = _get_account_usage_probe_worker(home)
|
||||
self.assertIsNotNone(worker)
|
||||
self.assertIs(worker, workers[1])
|
||||
worker._lock.release()
|
||||
finally:
|
||||
release_signal.set()
|
||||
thread.join(timeout=1.0)
|
||||
|
||||
def test_returns_none_when_all_workers_locked(self):
|
||||
"""Should return None immediately when all workers are locked from other threads."""
|
||||
home = Path("/tmp/test_immediate_none")
|
||||
|
||||
# Populate the pool
|
||||
worker_initial = _get_account_usage_probe_worker(home)
|
||||
self.assertIsNotNone(worker_initial)
|
||||
worker_initial._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
key = str(Path(home))
|
||||
workers = _account_usage_worker_pool[key]
|
||||
|
||||
# Hold both workers locked from a background thread
|
||||
lock_holder = threading.Event()
|
||||
release_signal = threading.Event()
|
||||
|
||||
def hold_locks():
|
||||
workers[0]._lock.acquire()
|
||||
workers[1]._lock.acquire()
|
||||
lock_holder.set()
|
||||
release_signal.wait(timeout=2.0)
|
||||
workers[1]._lock.release()
|
||||
workers[0]._lock.release()
|
||||
|
||||
thread = threading.Thread(target=hold_locks, daemon=True)
|
||||
thread.start()
|
||||
lock_holder.wait(timeout=2.0)
|
||||
|
||||
try:
|
||||
start = time.time()
|
||||
result = _get_account_usage_probe_worker(home)
|
||||
elapsed = time.time() - start
|
||||
|
||||
self.assertIsNone(result)
|
||||
# Should return immediately, not block
|
||||
self.assertLess(elapsed, 0.1)
|
||||
finally:
|
||||
release_signal.set()
|
||||
thread.join(timeout=1.0)
|
||||
|
||||
def test_scoped_invalidation_only_flushes_active_home(self):
|
||||
"""Scoped invalidation with provider_id should only flush active home."""
|
||||
home1 = Path.home() / ".hermes"
|
||||
home2 = Path("/tmp/other_home")
|
||||
|
||||
worker1 = _get_account_usage_probe_worker(home1)
|
||||
worker2 = _get_account_usage_probe_worker(home2)
|
||||
|
||||
self.assertIsNotNone(worker1)
|
||||
self.assertIsNotNone(worker2)
|
||||
worker1._lock.release()
|
||||
worker2._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
initial_count = len(_account_usage_worker_pool)
|
||||
self.assertEqual(initial_count, 2)
|
||||
|
||||
with mock.patch("api.providers._get_hermes_home", return_value=home1):
|
||||
invalidate_account_usage_status_cache(provider_id="anthropic")
|
||||
time.sleep(0.2)
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
remaining_keys = list(_account_usage_worker_pool.keys())
|
||||
self.assertNotIn(str(Path(home1)), remaining_keys)
|
||||
self.assertIn(str(Path(home2)), remaining_keys)
|
||||
|
||||
def test_full_invalidation_flushes_all_workers(self):
|
||||
"""Full invalidation (no provider_id) should flush all workers."""
|
||||
home1 = Path.home() / ".hermes"
|
||||
home2 = Path("/tmp/other_home")
|
||||
|
||||
worker1 = _get_account_usage_probe_worker(home1)
|
||||
worker2 = _get_account_usage_probe_worker(home2)
|
||||
|
||||
self.assertIsNotNone(worker1)
|
||||
self.assertIsNotNone(worker2)
|
||||
worker1._lock.release()
|
||||
worker2._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
self.assertEqual(len(_account_usage_worker_pool), 2)
|
||||
|
||||
invalidate_account_usage_status_cache(provider_id=None)
|
||||
time.sleep(0.2)
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
self.assertEqual(len(_account_usage_worker_pool), 0)
|
||||
|
||||
def test_cleanup_iterates_nested_worker_lists(self):
|
||||
"""Cleanup should properly iterate over nested worker lists."""
|
||||
home = Path.home() / ".hermes"
|
||||
|
||||
worker = _get_account_usage_probe_worker(home)
|
||||
self.assertIsNotNone(worker)
|
||||
worker._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
key = str(Path(home))
|
||||
workers_list = _account_usage_worker_pool.get(key)
|
||||
self.assertEqual(len(workers_list), 2)
|
||||
|
||||
now = time.monotonic() + 100000
|
||||
_cleanup_account_usage_probe_workers(now=now, idle_seconds=1.0)
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
if _account_usage_worker_pool.get(str(Path(home))):
|
||||
remaining = _account_usage_worker_pool.get(str(Path(home)), [])
|
||||
self.assertEqual(len(remaining), 0)
|
||||
else:
|
||||
self.assertNotIn(str(Path(home)), _account_usage_worker_pool)
|
||||
|
||||
def test_partial_cleanup_replenishes_pool(self):
|
||||
"""When cleanup removes one stale worker but the other is busy, pool replenishes to N=2."""
|
||||
home = Path("/tmp/test_replenish")
|
||||
|
||||
worker = _get_account_usage_probe_worker(home)
|
||||
self.assertIsNotNone(worker)
|
||||
worker._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
key = str(Path(home))
|
||||
workers = _account_usage_worker_pool[key]
|
||||
self.assertEqual(len(workers), 2)
|
||||
|
||||
# Hold workers[1] locked from a background thread (simulating active use)
|
||||
lock_holder = threading.Event()
|
||||
release_signal = threading.Event()
|
||||
|
||||
def hold_lock():
|
||||
workers[1]._lock.acquire()
|
||||
lock_holder.set()
|
||||
release_signal.wait(timeout=5.0)
|
||||
workers[1]._lock.release()
|
||||
|
||||
thread = threading.Thread(target=hold_lock, daemon=True)
|
||||
thread.start()
|
||||
lock_holder.wait(timeout=2.0)
|
||||
|
||||
try:
|
||||
# Run cleanup far in the future; workers[0] is idle and stale,
|
||||
# workers[1] is locked so it survives
|
||||
now = time.monotonic() + 100000
|
||||
_cleanup_account_usage_probe_workers(now=now, idle_seconds=1.0)
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
remaining = _account_usage_worker_pool.get(key, [])
|
||||
# Pool should be replenished back to 2
|
||||
self.assertEqual(len(remaining), _ACCOUNT_USAGE_WORKERS_PER_HOME)
|
||||
# The original busy worker should still be present
|
||||
self.assertIn(workers[1], remaining)
|
||||
finally:
|
||||
release_signal.set()
|
||||
thread.join(timeout=1.0)
|
||||
|
||||
def test_synchronous_close_flattens_nested_lists(self):
|
||||
"""Synchronous close should flatten nested lists correctly."""
|
||||
home = Path.home() / ".hermes"
|
||||
|
||||
worker = _get_account_usage_probe_worker(home)
|
||||
self.assertIsNotNone(worker)
|
||||
worker._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
key = str(Path(home))
|
||||
workers_list = _account_usage_worker_pool.get(key)
|
||||
self.assertEqual(len(workers_list), 2)
|
||||
|
||||
_close_account_usage_probe_workers()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
self.assertEqual(len(_account_usage_worker_pool), 0)
|
||||
|
||||
def test_concurrent_selector_no_double_claim(self):
|
||||
"""Two threads holding returned workers simultaneously never share the same instance."""
|
||||
home = Path("/tmp/test_concurrent")
|
||||
|
||||
# Pre-populate pool
|
||||
initial = _get_account_usage_probe_worker(home)
|
||||
self.assertIsNotNone(initial)
|
||||
initial._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
workers = _account_usage_worker_pool[str(Path(home))]
|
||||
|
||||
# Lock workers[0] from a background thread so only workers[1] is free
|
||||
lock_holder = threading.Event()
|
||||
release_signal = threading.Event()
|
||||
|
||||
def hold_first():
|
||||
workers[0]._lock.acquire()
|
||||
lock_holder.set()
|
||||
release_signal.wait(timeout=5.0)
|
||||
workers[0]._lock.release()
|
||||
|
||||
holder_thread = threading.Thread(target=hold_first, daemon=True)
|
||||
holder_thread.start()
|
||||
lock_holder.wait(timeout=2.0)
|
||||
|
||||
# Two threads race; each holds its worker until both have finished grabbing
|
||||
results = [None, None]
|
||||
barrier = threading.Barrier(2, timeout=2.0)
|
||||
both_done = threading.Barrier(2, timeout=2.0)
|
||||
|
||||
def grab(idx):
|
||||
barrier.wait()
|
||||
w = _get_account_usage_probe_worker(home)
|
||||
results[idx] = w
|
||||
try:
|
||||
both_done.wait()
|
||||
except threading.BrokenBarrierError:
|
||||
pass
|
||||
if w is not None:
|
||||
w._lock.release()
|
||||
|
||||
t0 = threading.Thread(target=grab, args=(0,), daemon=True)
|
||||
t1 = threading.Thread(target=grab, args=(1,), daemon=True)
|
||||
t0.start()
|
||||
t1.start()
|
||||
t0.join(timeout=3.0)
|
||||
t1.join(timeout=3.0)
|
||||
|
||||
try:
|
||||
got = [r for r in results if r is not None]
|
||||
# At most one thread should have gotten a worker (workers[1]);
|
||||
# the other gets None because the lock is already held
|
||||
self.assertEqual(len(got), 1, "Expected exactly one winner, got %d" % len(got))
|
||||
self.assertIs(got[0], workers[1])
|
||||
finally:
|
||||
release_signal.set()
|
||||
holder_thread.join(timeout=1.0)
|
||||
|
||||
def test_invalidation_cannot_pop_worker_between_lookup_and_claim(self):
|
||||
"""Getter keeps the pool lock through worker claim, so invalidation cannot pop first."""
|
||||
home = Path("/tmp/test_claim_under_pool_lock")
|
||||
|
||||
initial = _get_account_usage_probe_worker(home)
|
||||
self.assertIsNotNone(initial)
|
||||
initial._lock.release()
|
||||
|
||||
with _account_usage_worker_pool_lock:
|
||||
workers = _account_usage_worker_pool[str(Path(home))]
|
||||
|
||||
workers[1]._lock.acquire()
|
||||
self.addCleanup(workers[1]._lock.release)
|
||||
|
||||
acquire_entered = threading.Event()
|
||||
allow_acquire = threading.Event()
|
||||
claimed = threading.Event()
|
||||
release_claim = threading.Event()
|
||||
invalidation_done = threading.Event()
|
||||
result: dict[str, object] = {}
|
||||
target = workers[0]
|
||||
original_lock = target._lock
|
||||
self_outer = self
|
||||
|
||||
class GateLock:
|
||||
def __init__(self, inner):
|
||||
self._inner = inner
|
||||
|
||||
def acquire(self, blocking=True, timeout=-1):
|
||||
if blocking is False:
|
||||
self_outer.assertTrue(
|
||||
_account_usage_worker_pool_lock.locked(),
|
||||
"worker claim must happen while the pool lock is still held",
|
||||
)
|
||||
acquire_entered.set()
|
||||
allow_acquire.wait(timeout=2.0)
|
||||
if timeout == -1:
|
||||
return self._inner.acquire(blocking)
|
||||
return self._inner.acquire(blocking, timeout)
|
||||
|
||||
def release(self):
|
||||
return self._inner.release()
|
||||
|
||||
def __enter__(self):
|
||||
self._inner.acquire()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
self._inner.release()
|
||||
return False
|
||||
|
||||
target._lock = GateLock(original_lock)
|
||||
|
||||
def run_getter():
|
||||
result["worker"] = _get_account_usage_probe_worker(home)
|
||||
claimed.set()
|
||||
release_claim.wait(timeout=2.0)
|
||||
if result["worker"] is not None:
|
||||
result["worker"]._lock.release()
|
||||
|
||||
def run_invalidation():
|
||||
with mock.patch("api.providers._get_hermes_home", return_value=home):
|
||||
invalidate_account_usage_status_cache(provider_id="anthropic")
|
||||
invalidation_done.set()
|
||||
|
||||
fetch_thread = threading.Thread(target=run_getter, daemon=True)
|
||||
invalidation_thread = threading.Thread(target=run_invalidation, daemon=True)
|
||||
fetch_thread.start()
|
||||
self.assertTrue(acquire_entered.wait(timeout=2.0))
|
||||
|
||||
invalidation_thread.start()
|
||||
time.sleep(0.05)
|
||||
self.assertFalse(invalidation_done.is_set())
|
||||
|
||||
allow_acquire.set()
|
||||
self.assertTrue(claimed.wait(timeout=2.0))
|
||||
self.assertIs(result["worker"], target)
|
||||
|
||||
release_claim.set()
|
||||
fetch_thread.join(timeout=2.0)
|
||||
|
||||
invalidation_thread.join(timeout=2.0)
|
||||
self.assertTrue(invalidation_done.is_set())
|
||||
with _account_usage_worker_pool_lock:
|
||||
self.assertNotIn(str(Path(home)), _account_usage_worker_pool)
|
||||
|
||||
|
||||
class TestWorkerPoolConfiguration(unittest.TestCase):
|
||||
"""Test that constants are properly configured."""
|
||||
|
||||
def test_workers_per_home_is_two(self):
|
||||
"""Should have exactly 2 workers per home."""
|
||||
self.assertEqual(_ACCOUNT_USAGE_WORKERS_PER_HOME, 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1439,7 +1439,9 @@ def test_account_usage_worker_idle_cleanup_closes_stale_process(monkeypatch, tmp
|
||||
providers._close_account_usage_probe_workers()
|
||||
try:
|
||||
providers._agent_fetch_account_usage_for_home("openai-codex", tmp_path)
|
||||
worker = providers._account_usage_worker_pool[str(tmp_path)]
|
||||
workers = providers._account_usage_worker_pool[str(tmp_path)]
|
||||
# Pool is now a list of workers; pick the one that was last used
|
||||
worker = max(workers, key=lambda w: w.last_used)
|
||||
providers._cleanup_account_usage_probe_workers(
|
||||
now=worker.last_used + providers._ACCOUNT_USAGE_WORKER_IDLE_SECONDS + 1
|
||||
)
|
||||
@@ -1447,7 +1449,7 @@ def test_account_usage_worker_idle_cleanup_closes_stale_process(monkeypatch, tmp
|
||||
finally:
|
||||
providers._close_account_usage_probe_workers()
|
||||
|
||||
assert len(launched) == 2
|
||||
assert len(launched) >= 2
|
||||
assert launched[0][2].terminated is True
|
||||
|
||||
|
||||
@@ -1499,14 +1501,16 @@ def test_account_usage_cleanup_removes_null_proc_worker(monkeypatch, tmp_path):
|
||||
providers._close_account_usage_probe_workers()
|
||||
try:
|
||||
providers._agent_fetch_account_usage_for_home("openai-codex", tmp_path)
|
||||
worker = providers._account_usage_worker_pool[str(tmp_path)]
|
||||
worker.close()
|
||||
workers = providers._account_usage_worker_pool[str(tmp_path)]
|
||||
# Pool is now a list of workers; close all of them to simulate null proc state
|
||||
for worker in workers:
|
||||
worker.close()
|
||||
providers._cleanup_account_usage_probe_workers()
|
||||
assert str(tmp_path) not in providers._account_usage_worker_pool
|
||||
finally:
|
||||
providers._close_account_usage_probe_workers()
|
||||
|
||||
assert len(launched) == 1
|
||||
assert len(launched) >= 1
|
||||
|
||||
|
||||
def test_provider_key_mutation_invalidates_warm_account_usage_workers(monkeypatch, tmp_path):
|
||||
|
||||
Reference in New Issue
Block a user