From 501a5879164a3740c904c2f3e123deddb996ef0a Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Wed, 10 Jun 2026 12:25:40 -0400 Subject: [PATCH 1/8] perf(providers): per-home probe-worker pool + scoped credential invalidation (#3787) --- api/providers.py | 78 ++++--- tests/test_issue3787_probe_pool.py | 319 +++++++++++++++++++++++++++++ 2 files changed, 372 insertions(+), 25 deletions(-) create mode 100644 tests/test_issue3787_probe_pool.py diff --git a/api/providers.py b/api/providers.py index 680f6f70..28b062a7 100644 --- a/api/providers.py +++ b/api/providers.py @@ -131,9 +131,13 @@ _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 +_ACCOUNT_USAGE_WORKER_WAIT_SECONDS = 0.5 + def _get_account_usage_probe_semaphore() -> threading.BoundedSemaphore: global _account_usage_probe_semaphore @@ -1421,14 +1425,21 @@ 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": key = str(Path(home)) 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 + workers = _account_usage_worker_pool.get(key) + if not workers: + workers = [_AccountUsageProbeWorker(Path(home)) for _ in range(_ACCOUNT_USAGE_WORKERS_PER_HOME)] + _account_usage_worker_pool[key] = workers + for w in workers: + if w._lock.acquire(blocking=False): + w._lock.release() + return w + if workers[0]._lock.acquire(timeout=_ACCOUNT_USAGE_WORKER_WAIT_SECONDS): + workers[0]._lock.release() + return workers[0] + return None def _cleanup_account_usage_probe_workers( @@ -1439,23 +1450,29 @@ 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 = proc is None or proc.poll() is not None + if is_dead or cutoff - worker.last_used >= idle_seconds: + stale.append((key, worker)) + finally: + worker._lock.release() + # Remove the entire key if all workers are stale + 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: + _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 +1482,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 +1537,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 +1570,10 @@ 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: + return worker.fetch(provider, api_key=api_key) + 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 diff --git a/tests/test_issue3787_probe_pool.py b/tests/test_issue3787_probe_pool.py new file mode 100644 index 00000000..819b4f87 --- /dev/null +++ b/tests/test_issue3787_probe_pool.py @@ -0,0 +1,319 @@ +"""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 when all workers are locked and timeout expires +- Scoped invalidation only flushes the active home's workers +- Full invalidation (no provider_id) flushes all workers +- Cleanup iterates nested worker lists correctly +""" + +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_WAIT_SECONDS, + _account_usage_worker_pool, + _account_usage_worker_pool_lock, + _AccountUsageProbeWorker, + _cleanup_account_usage_probe_workers, + _close_account_usage_probe_workers, + _close_account_usage_probe_workers_async, + _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) + + 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 available.""" + home = Path("/tmp/test_nonblocking_idle") + + # First, populate the pool with two workers + worker_initial = _get_account_usage_probe_worker(home) + self.assertIsNotNone(worker_initial) + + # Get reference to the actual workers in pool + with _account_usage_worker_pool_lock: + key = str(Path(home)) + workers = _account_usage_worker_pool[key] + self.assertEqual(len(workers), 2) + + # Hold first worker 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) # Wait until background thread has the lock + + try: + # Get a worker - should pick the idle one (workers[1]) + worker = _get_account_usage_probe_worker(home) + self.assertIsNotNone(worker) + # Worker should be the second one since first is locked + self.assertIs(worker, workers[1]) + finally: + release_signal.set() + thread.join(timeout=1.0) + + def test_returns_none_when_all_workers_locked_and_timeout_expires(self): + """Should return None when all workers locked and timeout expires.""" + home = Path("/tmp/test_timeout_none") + + # First populate the pool + worker_initial = _get_account_usage_probe_worker(home) + self.assertIsNotNone(worker_initial) + + # Get reference to the workers + with _account_usage_worker_pool_lock: + key = str(Path(home)) + workers = _account_usage_worker_pool[key] + self.assertEqual(len(workers), 2) + + # Hold both workers locked from background threads + 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) # Wait until background thread has both locks + + try: + # Now attempt to get a worker; should timeout and return None + start = time.time() + result = _get_account_usage_probe_worker(home) + elapsed = time.time() - start + + self.assertIsNone(result) + # Verify timeout was respected (within a small margin) + self.assertGreaterEqual(elapsed, _ACCOUNT_USAGE_WORKER_WAIT_SECONDS - 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") + + # Create workers for both homes + worker1 = _get_account_usage_probe_worker(home1) + worker2 = _get_account_usage_probe_worker(home2) + + self.assertIsNotNone(worker1) + self.assertIsNotNone(worker2) + + with _account_usage_worker_pool_lock: + initial_count = len(_account_usage_worker_pool) + self.assertEqual(initial_count, 2) + + # Mock _get_hermes_home to return home1 + with mock.patch("api.providers._get_hermes_home", return_value=home1): + invalidate_account_usage_status_cache(provider_id="anthropic") + # Give async thread time to complete + time.sleep(0.2) + + with _account_usage_worker_pool_lock: + remaining_keys = list(_account_usage_worker_pool.keys()) + # home1 should be flushed, home2 should remain + 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") + + # Create workers for both homes + worker1 = _get_account_usage_probe_worker(home1) + worker2 = _get_account_usage_probe_worker(home2) + + self.assertIsNotNone(worker1) + self.assertIsNotNone(worker2) + + with _account_usage_worker_pool_lock: + self.assertEqual(len(_account_usage_worker_pool), 2) + + # Full invalidation with no provider_id + invalidate_account_usage_status_cache(provider_id=None) + # Give async thread time to complete + 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" + + # Create workers + worker = _get_account_usage_probe_worker(home) + self.assertIsNotNone(worker) + + with _account_usage_worker_pool_lock: + key = str(Path(home)) + workers_list = _account_usage_worker_pool.get(key) + self.assertEqual(len(workers_list), 2) + + # Run cleanup with future timestamp (should mark workers as stale) + now = time.monotonic() + 100000 # Far in the future + _cleanup_account_usage_probe_workers(now=now, idle_seconds=1.0) + + # Pool should be cleared of stale workers + with _account_usage_worker_pool_lock: + if _account_usage_worker_pool.get(str(Path(home))): + # If key still exists, it should have no workers + remaining = _account_usage_worker_pool.get(str(Path(home)), []) + self.assertEqual(len(remaining), 0) + else: + # Key should be removed entirely if no workers + self.assertNotIn(str(Path(home)), _account_usage_worker_pool) + + def test_synchronous_close_flattens_nested_lists(self): + """Synchronous close should flatten nested lists correctly.""" + home = Path.home() / ".hermes" + + # Create workers + worker = _get_account_usage_probe_worker(home) + self.assertIsNotNone(worker) + + 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 all workers synchronously + _close_account_usage_probe_workers() + + # Pool should be empty + with _account_usage_worker_pool_lock: + self.assertEqual(len(_account_usage_worker_pool), 0) + + def test_multiple_gets_return_different_idle_workers(self): + """Multiple rapid gets should return different workers when available.""" + home = Path("/tmp/test_multiple_gets") + + # First populate the pool + worker_initial = _get_account_usage_probe_worker(home) + self.assertIsNotNone(worker_initial) + + # Get reference to the workers + with _account_usage_worker_pool_lock: + key = str(Path(home)) + workers = _account_usage_worker_pool[key] + self.assertEqual(len(workers), 2) + + # Hold first worker locked from background thread + lock_holder1 = threading.Event() + release_signal1 = threading.Event() + + def hold_lock1(): + workers[0]._lock.acquire() + lock_holder1.set() + release_signal1.wait(timeout=2.0) + workers[0]._lock.release() + + thread1 = threading.Thread(target=hold_lock1, daemon=True) + thread1.start() + lock_holder1.wait(timeout=2.0) + + try: + # First call should get workers[1] (idle) + worker1 = _get_account_usage_probe_worker(home) + self.assertIsNotNone(worker1) + self.assertIs(worker1, workers[1]) + + # Lock workers[1] too from another background thread, then release workers[0] + lock_holder2 = threading.Event() + release_signal2 = threading.Event() + + def hold_lock2(): + workers[1]._lock.acquire() + lock_holder2.set() + release_signal2.wait(timeout=2.0) + workers[1]._lock.release() + + thread2 = threading.Thread(target=hold_lock2, daemon=True) + thread2.start() + lock_holder2.wait(timeout=2.0) + + # Now release workers[0] but keep workers[1] locked + release_signal1.set() + thread1.join(timeout=1.0) + + try: + # Second call should timeout waiting for workers[1], then get workers[0] which is now free + # But the timeout in _get_account_usage_probe_worker is 0.5s, and workers[1] holds for 2s, + # so it should return workers[0] after the timeout + worker2 = _get_account_usage_probe_worker(home) + self.assertIsNotNone(worker2) + # Should get workers[0] since workers[1] is still locked + self.assertIs(worker2, workers[0]) + finally: + release_signal2.set() + thread2.join(timeout=1.0) + except Exception: + release_signal1.set() + thread1.join(timeout=1.0) + raise + + +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) + + def test_worker_wait_seconds_is_half(self): + """Should wait 0.5 seconds for worker availability.""" + self.assertEqual(_ACCOUNT_USAGE_WORKER_WAIT_SECONDS, 0.5) + + +if __name__ == "__main__": + unittest.main() From dd72a36456cff4adcb00613a3a288d1caf31fdce Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Wed, 10 Jun 2026 12:42:12 -0400 Subject: [PATCH 2/8] fix(tests): adapt existing probe pool tests to N=2 worker list structure --- tests/test_issue3787_probe_pool.py | 2 -- tests/test_provider_quota_status.py | 14 +++++++++----- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/tests/test_issue3787_probe_pool.py b/tests/test_issue3787_probe_pool.py index 819b4f87..4bb815da 100644 --- a/tests/test_issue3787_probe_pool.py +++ b/tests/test_issue3787_probe_pool.py @@ -20,10 +20,8 @@ from api.providers import ( _ACCOUNT_USAGE_WORKER_WAIT_SECONDS, _account_usage_worker_pool, _account_usage_worker_pool_lock, - _AccountUsageProbeWorker, _cleanup_account_usage_probe_workers, _close_account_usage_probe_workers, - _close_account_usage_probe_workers_async, _get_account_usage_probe_worker, invalidate_account_usage_status_cache, ) diff --git a/tests/test_provider_quota_status.py b/tests/test_provider_quota_status.py index b1a6a0be..c8a4fb11 100644 --- a/tests/test_provider_quota_status.py +++ b/tests/test_provider_quota_status.py @@ -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): From 40abc0f701bbc1c3dc4a560734fc0f3f26c9862c Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Wed, 10 Jun 2026 19:03:45 -0400 Subject: [PATCH 3/8] fix(providers): hold probe-worker lock across selector handoff --- api/providers.py | 16 ++- tests/test_issue3787_probe_pool.py | 170 +++++++++++++---------------- 2 files changed, 83 insertions(+), 103 deletions(-) diff --git a/api/providers.py b/api/providers.py index 28b062a7..f75b04c6 100644 --- a/api/providers.py +++ b/api/providers.py @@ -136,7 +136,6 @@ _account_usage_worker_pool_lock = threading.Lock() # Per-home worker pool configuration for probe tail-latency reduction (#3787) _ACCOUNT_USAGE_WORKERS_PER_HOME = 2 -_ACCOUNT_USAGE_WORKER_WAIT_SECONDS = 0.5 def _get_account_usage_probe_semaphore() -> threading.BoundedSemaphore: @@ -1426,6 +1425,12 @@ def _fetch_account_usage_once_for_home(provider: str, home: Path, *, api_key: st 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)) with _account_usage_worker_pool_lock: workers = _account_usage_worker_pool.get(key) @@ -1434,11 +1439,7 @@ def _get_account_usage_probe_worker(home: Path) -> "_AccountUsageProbeWorker | N _account_usage_worker_pool[key] = workers for w in workers: if w._lock.acquire(blocking=False): - w._lock.release() return w - if workers[0]._lock.acquire(timeout=_ACCOUNT_USAGE_WORKER_WAIT_SECONDS): - workers[0]._lock.release() - return workers[0] return None @@ -1572,7 +1573,10 @@ def _agent_fetch_account_usage_for_home(provider: str, home: Path, *, api_key: s _cleanup_account_usage_probe_workers() worker = _get_account_usage_probe_worker(home) if worker is not None: - return worker.fetch(provider, api_key=api_key) + 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) diff --git a/tests/test_issue3787_probe_pool.py b/tests/test_issue3787_probe_pool.py index 4bb815da..e07b2cfc 100644 --- a/tests/test_issue3787_probe_pool.py +++ b/tests/test_issue3787_probe_pool.py @@ -3,10 +3,11 @@ Focus on: - Per-home worker pool size N=2 - Non-blocking acquire returns an idle worker when available -- Returns None when all workers are locked and timeout expires +- 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 @@ -17,7 +18,6 @@ from unittest import mock from api.providers import ( _ACCOUNT_USAGE_WORKERS_PER_HOME, - _ACCOUNT_USAGE_WORKER_WAIT_SECONDS, _account_usage_worker_pool, _account_usage_worker_pool_lock, _cleanup_account_usage_probe_workers, @@ -45,6 +45,7 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase): 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)) @@ -54,20 +55,20 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase): 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 available.""" + """Non-blocking acquire should return an idle worker when one is free.""" home = Path("/tmp/test_nonblocking_idle") - # First, populate the pool with two workers + # Populate the pool worker_initial = _get_account_usage_probe_worker(home) self.assertIsNotNone(worker_initial) + worker_initial._lock.release() - # Get reference to the actual workers in pool with _account_usage_worker_pool_lock: key = str(Path(home)) workers = _account_usage_worker_pool[key] self.assertEqual(len(workers), 2) - # Hold first worker locked from a background thread to simulate actual use + # Hold workers[0] locked from a background thread to simulate actual use lock_holder = threading.Event() release_signal = threading.Event() @@ -79,33 +80,31 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase): thread = threading.Thread(target=hold_lock, daemon=True) thread.start() - lock_holder.wait(timeout=2.0) # Wait until background thread has the lock + lock_holder.wait(timeout=2.0) try: - # Get a worker - should pick the idle one (workers[1]) worker = _get_account_usage_probe_worker(home) self.assertIsNotNone(worker) - # Worker should be the second one since first is locked 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_and_timeout_expires(self): - """Should return None when all workers locked and timeout expires.""" - home = Path("/tmp/test_timeout_none") + 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") - # First populate the pool + # Populate the pool worker_initial = _get_account_usage_probe_worker(home) self.assertIsNotNone(worker_initial) + worker_initial._lock.release() - # Get reference to the workers with _account_usage_worker_pool_lock: key = str(Path(home)) workers = _account_usage_worker_pool[key] - self.assertEqual(len(workers), 2) - # Hold both workers locked from background threads + # Hold both workers locked from a background thread lock_holder = threading.Event() release_signal = threading.Event() @@ -119,17 +118,16 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase): thread = threading.Thread(target=hold_locks, daemon=True) thread.start() - lock_holder.wait(timeout=2.0) # Wait until background thread has both locks + lock_holder.wait(timeout=2.0) try: - # Now attempt to get a worker; should timeout and return None start = time.time() result = _get_account_usage_probe_worker(home) elapsed = time.time() - start self.assertIsNone(result) - # Verify timeout was respected (within a small margin) - self.assertGreaterEqual(elapsed, _ACCOUNT_USAGE_WORKER_WAIT_SECONDS - 0.1) + # Should return immediately, not block + self.assertLess(elapsed, 0.1) finally: release_signal.set() thread.join(timeout=1.0) @@ -139,26 +137,24 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase): home1 = Path.home() / ".hermes" home2 = Path("/tmp/other_home") - # Create workers for both homes 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) - # Mock _get_hermes_home to return home1 with mock.patch("api.providers._get_hermes_home", return_value=home1): invalidate_account_usage_status_cache(provider_id="anthropic") - # Give async thread time to complete time.sleep(0.2) with _account_usage_worker_pool_lock: remaining_keys = list(_account_usage_worker_pool.keys()) - # home1 should be flushed, home2 should remain self.assertNotIn(str(Path(home1)), remaining_keys) self.assertIn(str(Path(home2)), remaining_keys) @@ -167,19 +163,18 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase): home1 = Path.home() / ".hermes" home2 = Path("/tmp/other_home") - # Create workers for both homes 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) - # Full invalidation with no provider_id invalidate_account_usage_status_cache(provider_id=None) - # Give async thread time to complete time.sleep(0.2) with _account_usage_worker_pool_lock: @@ -189,116 +184,101 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase): """Cleanup should properly iterate over nested worker lists.""" home = Path.home() / ".hermes" - # Create workers 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) - # Run cleanup with future timestamp (should mark workers as stale) - now = time.monotonic() + 100000 # Far in the future + now = time.monotonic() + 100000 _cleanup_account_usage_probe_workers(now=now, idle_seconds=1.0) - # Pool should be cleared of stale workers with _account_usage_worker_pool_lock: if _account_usage_worker_pool.get(str(Path(home))): - # If key still exists, it should have no workers remaining = _account_usage_worker_pool.get(str(Path(home)), []) self.assertEqual(len(remaining), 0) else: - # Key should be removed entirely if no workers self.assertNotIn(str(Path(home)), _account_usage_worker_pool) def test_synchronous_close_flattens_nested_lists(self): """Synchronous close should flatten nested lists correctly.""" home = Path.home() / ".hermes" - # Create workers 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 all workers synchronously _close_account_usage_probe_workers() - # Pool should be empty with _account_usage_worker_pool_lock: self.assertEqual(len(_account_usage_worker_pool), 0) - def test_multiple_gets_return_different_idle_workers(self): - """Multiple rapid gets should return different workers when available.""" - home = Path("/tmp/test_multiple_gets") + def test_concurrent_selector_no_double_claim(self): + """Two threads holding returned workers simultaneously never share the same instance.""" + home = Path("/tmp/test_concurrent") - # First populate the pool - worker_initial = _get_account_usage_probe_worker(home) - self.assertIsNotNone(worker_initial) + # Pre-populate pool + initial = _get_account_usage_probe_worker(home) + self.assertIsNotNone(initial) + initial._lock.release() - # Get reference to the workers with _account_usage_worker_pool_lock: - key = str(Path(home)) - workers = _account_usage_worker_pool[key] - self.assertEqual(len(workers), 2) + workers = _account_usage_worker_pool[str(Path(home))] - # Hold first worker locked from background thread - lock_holder1 = threading.Event() - release_signal1 = threading.Event() + # Lock workers[0] from a background thread so only workers[1] is free + lock_holder = threading.Event() + release_signal = threading.Event() - def hold_lock1(): + def hold_first(): workers[0]._lock.acquire() - lock_holder1.set() - release_signal1.wait(timeout=2.0) + lock_holder.set() + release_signal.wait(timeout=5.0) workers[0]._lock.release() - thread1 = threading.Thread(target=hold_lock1, daemon=True) - thread1.start() - lock_holder1.wait(timeout=2.0) + 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: - # First call should get workers[1] (idle) - worker1 = _get_account_usage_probe_worker(home) - self.assertIsNotNone(worker1) - self.assertIs(worker1, workers[1]) - - # Lock workers[1] too from another background thread, then release workers[0] - lock_holder2 = threading.Event() - release_signal2 = threading.Event() - - def hold_lock2(): - workers[1]._lock.acquire() - lock_holder2.set() - release_signal2.wait(timeout=2.0) - workers[1]._lock.release() - - thread2 = threading.Thread(target=hold_lock2, daemon=True) - thread2.start() - lock_holder2.wait(timeout=2.0) - - # Now release workers[0] but keep workers[1] locked - release_signal1.set() - thread1.join(timeout=1.0) - - try: - # Second call should timeout waiting for workers[1], then get workers[0] which is now free - # But the timeout in _get_account_usage_probe_worker is 0.5s, and workers[1] holds for 2s, - # so it should return workers[0] after the timeout - worker2 = _get_account_usage_probe_worker(home) - self.assertIsNotNone(worker2) - # Should get workers[0] since workers[1] is still locked - self.assertIs(worker2, workers[0]) - finally: - release_signal2.set() - thread2.join(timeout=1.0) - except Exception: - release_signal1.set() - thread1.join(timeout=1.0) - raise + 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) class TestWorkerPoolConfiguration(unittest.TestCase): @@ -308,10 +288,6 @@ class TestWorkerPoolConfiguration(unittest.TestCase): """Should have exactly 2 workers per home.""" self.assertEqual(_ACCOUNT_USAGE_WORKERS_PER_HOME, 2) - def test_worker_wait_seconds_is_half(self): - """Should wait 0.5 seconds for worker availability.""" - self.assertEqual(_ACCOUNT_USAGE_WORKER_WAIT_SECONDS, 0.5) - if __name__ == "__main__": unittest.main() From b5f1965b3cf50df6ea5e2280a9102d8ace745b19 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Wed, 10 Jun 2026 19:22:45 -0400 Subject: [PATCH 4/8] fix(providers): replenish pool after partial cleanup removes stale workers --- api/providers.py | 4 ++- tests/test_issue3787_probe_pool.py | 43 ++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/api/providers.py b/api/providers.py index f75b04c6..1c52871b 100644 --- a/api/providers.py +++ b/api/providers.py @@ -1461,11 +1461,13 @@ def _cleanup_account_usage_probe_workers( stale.append((key, worker)) finally: worker._lock.release() - # Remove the entire key if all workers are stale 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() diff --git a/tests/test_issue3787_probe_pool.py b/tests/test_issue3787_probe_pool.py index e07b2cfc..54a0c56b 100644 --- a/tests/test_issue3787_probe_pool.py +++ b/tests/test_issue3787_probe_pool.py @@ -203,6 +203,49 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase): 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" From 0652cf9b162f4b885254ac962478a07fc2bfd1d9 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Wed, 10 Jun 2026 20:35:48 -0400 Subject: [PATCH 5/8] fix(providers): treat proc=None as idle, not dead, in cleanup --- api/providers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/providers.py b/api/providers.py index 1c52871b..4d989218 100644 --- a/api/providers.py +++ b/api/providers.py @@ -1456,7 +1456,7 @@ def _cleanup_account_usage_probe_workers( if worker._lock.acquire(blocking=False): try: proc = worker._proc - is_dead = proc is None or proc.poll() is not None + is_dead = 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: From 60570f1fa6d1b176aa9816af987636f28a274669 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Thu, 11 Jun 2026 05:52:24 -0400 Subject: [PATCH 6/8] Distinguish closed workers from idle ones in cleanup --- api/providers.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/api/providers.py b/api/providers.py index 4d989218..ff7b934d 100644 --- a/api/providers.py +++ b/api/providers.py @@ -1231,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 @@ -1456,7 +1458,7 @@ def _cleanup_account_usage_probe_workers( if worker._lock.acquire(blocking=False): try: proc = worker._proc - is_dead = proc is not None and proc.poll() is not None + 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: From bb8d5e958744f2c8f7a6942a5f14052f1a32e73d Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Sat, 13 Jun 2026 16:42:43 -0400 Subject: [PATCH 7/8] fix(providers): claim probe workers under the pool lock (#3787) --- api/providers.py | 26 ++++++--- tests/test_issue3787_probe_pool.py | 86 ++++++++++++++++++++++++++++++ 2 files changed, 105 insertions(+), 7 deletions(-) diff --git a/api/providers.py b/api/providers.py index ff7b934d..b94046ed 100644 --- a/api/providers.py +++ b/api/providers.py @@ -1354,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) @@ -1434,15 +1435,26 @@ def _get_account_usage_probe_worker(home: Path) -> "_AccountUsageProbeWorker | N 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: - workers = _account_usage_worker_pool.get(key) - if not workers: + 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)] - _account_usage_worker_pool[key] = workers - for w in workers: - if w._lock.acquire(blocking=False): - return w - return None + 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( diff --git a/tests/test_issue3787_probe_pool.py b/tests/test_issue3787_probe_pool.py index 54a0c56b..7b1d32a3 100644 --- a/tests/test_issue3787_probe_pool.py +++ b/tests/test_issue3787_probe_pool.py @@ -323,6 +323,92 @@ class TestProbeWorkerPoolPerHome(unittest.TestCase): 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.""" From c0ea19407b563bb2acd67c9247226e2a0201bddf Mon Sep 17 00:00:00 2001 From: nesquena-hermes Date: Sat, 13 Jun 2026 22:05:58 +0000 Subject: [PATCH 8/8] docs(changelog): stamp #3940 v2 (claim-under-pool-lock) as v0.51.399 (Release NL) --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9b216555..d589f4e8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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