Merge pull request #4147 from nesquena/stage-3940v2
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:
nesquena-hermes
2026-06-13 15:15:34 -07:00
committed by GitHub
4 changed files with 510 additions and 30 deletions

View File

@@ -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

View File

@@ -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

View 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()

View File

@@ -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):