fix: peer gate race condition + async client connection leak#83
Open
yasinlex wants to merge 1 commit into
Open
fix: peer gate race condition + async client connection leak#83yasinlex wants to merge 1 commit into
yasinlex wants to merge 1 commit into
Conversation
Two issues in provider_adapters/openai_compatible.py:
1. _PEER_GATES race condition and memory leak:
- The previous _peer_gate() was synchronous and had a TOCTOU race:
two concurrent coroutines could both see None and create separate
semaphores for the same peer_id. The second write would silently
overwrite the first, dropping any coroutines already waiting on it
(permanent hang until the semaphore is GC'd).
- _PEER_GATES grew without bound as AntSeed peers were dynamically
discovered and retired — old entries were never evicted.
Fix: _peer_gate is now async, guarded by an asyncio.Lock, and stale
entries are evicted after a TTL (600s) when the dict exceeds 256
entries. The lock is lazy-initialised to avoid creating asyncio
primitives at import time (safe for test/sync contexts).
2. stream_openai_compatible() leaked httpx.AsyncClient:
When no external client was passed, a new AsyncClient was created but
never closed — each call leaked a TCP connection pool. Under sustained
streaming traffic this exhausts the file descriptor limit.
Fix: track whether we own the client via _owns_client and close it in
a finally block after the response is fully consumed.
📝 WalkthroughWalkthroughThe provider adapter now manages peer concurrency gates with a locked, TTL-evicted cache and awaits asynchronous gate acquisition. Streaming requests track client ownership and close internally created HTTP clients while preserving response parsing and error handling. ChangesProvider adapter safety
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@provider_adapters/openai_compatible.py`:
- Around line 213-263: Update _peer_gate to track the capacity associated with
each cached peer semaphore and recreate the semaphore when the requested cap
differs from the cached cap. Keep the existing semaphore when cap is unchanged,
preserve lock protection and timestamp updates, and ensure replacement entries
update the stored capacity so subsequent calls use the new limit.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7ecb1055-8211-460f-8804-5d07e1b0703f
📒 Files selected for processing (1)
provider_adapters/openai_compatible.py
Comment on lines
213
to
+263
| _PEER_GATES: dict[str, asyncio.Semaphore] = {} | ||
|
|
||
|
|
||
| def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore: | ||
| sem = _PEER_GATES.get(peer_id) | ||
| if sem is None: | ||
| sem = asyncio.Semaphore(cap) | ||
| _PEER_GATES[peer_id] = sem | ||
| return sem | ||
| _PEER_GATE_LOCK: asyncio.Lock | None = None | ||
|
|
||
| # Maximum idle time (seconds) before a peer gate is evicted to prevent | ||
| # unbounded memory growth when peers are dynamically discovered and retired. | ||
| _PEER_GATE_TTL: float = 600.0 | ||
| _PEER_GATE_TIMESTAMPS: dict[str, float] = {} | ||
|
|
||
|
|
||
| def _peer_gate_lock() -> asyncio.Lock: | ||
| """Lazy-initialised lock to avoid creating asyncio primitives at import time | ||
| (which may happen outside an event loop in test/synchronous contexts).""" | ||
| global _PEER_GATE_LOCK | ||
| if _PEER_GATE_LOCK is None: | ||
| _PEER_GATE_LOCK = asyncio.Lock() | ||
| return _PEER_GATE_LOCK | ||
|
|
||
|
|
||
| async def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore: | ||
| """Return a per-peer concurrency semaphore, creating one if needed. | ||
|
|
||
| Uses a lock to prevent a race where two concurrent coroutines both see | ||
| ``None`` and create separate semaphores for the same peer_id — the | ||
| second would overwrite the first, silently dropping any coroutines | ||
| already waiting on it. | ||
|
|
||
| Old entries are evicted after ``_PEER_GATE_TTL`` seconds of inactivity | ||
| so that the dict does not grow without bound when AntSeed peers are | ||
| dynamically discovered and later retired. | ||
| """ | ||
| import time as _time | ||
|
|
||
| lock = _peer_gate_lock() | ||
| async with lock: | ||
| # Evict stale entries first (best-effort, not on every call to | ||
| # avoid quadratic sweep cost — only when the lock is already held | ||
| # and the dict is getting large). | ||
| if len(_PEER_GATES) > 256: | ||
| now = _time.monotonic() | ||
| stale = [k for k, ts in _PEER_GATE_TIMESTAMPS.items() | ||
| if now - ts > _PEER_GATE_TTL] | ||
| for k in stale: | ||
| _PEER_GATES.pop(k, None) | ||
| _PEER_GATE_TIMESTAMPS.pop(k, None) | ||
|
|
||
| sem = _PEER_GATES.get(peer_id) | ||
| if sem is None: | ||
| sem = asyncio.Semaphore(cap) | ||
| _PEER_GATES[peer_id] = sem | ||
| _PEER_GATE_TIMESTAMPS[peer_id] = _time.monotonic() | ||
| return sem |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Solid async gate cache design — but cached semaphore ignores a changed cap for the same peer.
The lock-protected creation and TTL eviction are correctly implemented (no await inside the critical section, so no race). However, once a semaphore exists for a peer_id, a later call with a different cap (e.g., seller updates max_concurrency) is silently ignored — sem = _PEER_GATES.get(peer_id) short-circuits before the new cap is ever used, and the stale limit persists until the entry ages out (_PEER_GATE_TTL = 600s).
♻️ Proposed fix: recreate semaphore when cap changes
+_PEER_GATE_CAPS: dict[str, int] = {}
+
async def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore:
...
sem = _PEER_GATES.get(peer_id)
- if sem is None:
+ if sem is None or _PEER_GATE_CAPS.get(peer_id) != cap:
sem = asyncio.Semaphore(cap)
_PEER_GATES[peer_id] = sem
+ _PEER_GATE_CAPS[peer_id] = cap
_PEER_GATE_TIMESTAMPS[peer_id] = _time.monotonic()
return semNote: in-flight callers still holding the old semaphore continue enforcing the old cap briefly, same as the existing TTL-eviction transition — acceptable tradeoff.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| _PEER_GATES: dict[str, asyncio.Semaphore] = {} | |
| def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore: | |
| sem = _PEER_GATES.get(peer_id) | |
| if sem is None: | |
| sem = asyncio.Semaphore(cap) | |
| _PEER_GATES[peer_id] = sem | |
| return sem | |
| _PEER_GATE_LOCK: asyncio.Lock | None = None | |
| # Maximum idle time (seconds) before a peer gate is evicted to prevent | |
| # unbounded memory growth when peers are dynamically discovered and retired. | |
| _PEER_GATE_TTL: float = 600.0 | |
| _PEER_GATE_TIMESTAMPS: dict[str, float] = {} | |
| def _peer_gate_lock() -> asyncio.Lock: | |
| """Lazy-initialised lock to avoid creating asyncio primitives at import time | |
| (which may happen outside an event loop in test/synchronous contexts).""" | |
| global _PEER_GATE_LOCK | |
| if _PEER_GATE_LOCK is None: | |
| _PEER_GATE_LOCK = asyncio.Lock() | |
| return _PEER_GATE_LOCK | |
| async def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore: | |
| """Return a per-peer concurrency semaphore, creating one if needed. | |
| Uses a lock to prevent a race where two concurrent coroutines both see | |
| ``None`` and create separate semaphores for the same peer_id — the | |
| second would overwrite the first, silently dropping any coroutines | |
| already waiting on it. | |
| Old entries are evicted after ``_PEER_GATE_TTL`` seconds of inactivity | |
| so that the dict does not grow without bound when AntSeed peers are | |
| dynamically discovered and later retired. | |
| """ | |
| import time as _time | |
| lock = _peer_gate_lock() | |
| async with lock: | |
| # Evict stale entries first (best-effort, not on every call to | |
| # avoid quadratic sweep cost — only when the lock is already held | |
| # and the dict is getting large). | |
| if len(_PEER_GATES) > 256: | |
| now = _time.monotonic() | |
| stale = [k for k, ts in _PEER_GATE_TIMESTAMPS.items() | |
| if now - ts > _PEER_GATE_TTL] | |
| for k in stale: | |
| _PEER_GATES.pop(k, None) | |
| _PEER_GATE_TIMESTAMPS.pop(k, None) | |
| sem = _PEER_GATES.get(peer_id) | |
| if sem is None: | |
| sem = asyncio.Semaphore(cap) | |
| _PEER_GATES[peer_id] = sem | |
| _PEER_GATE_TIMESTAMPS[peer_id] = _time.monotonic() | |
| return sem | |
| _PEER_GATES: dict[str, asyncio.Semaphore] = {} | |
| _PEER_GATE_LOCK: asyncio.Lock | None = None | |
| # Maximum idle time (seconds) before a peer gate is evicted to prevent | |
| # unbounded memory growth when peers are dynamically discovered and retired. | |
| _PEER_GATE_TTL: float = 600.0 | |
| _PEER_GATE_TIMESTAMPS: dict[str, float] = {} | |
| _PEER_GATE_CAPS: dict[str, int] = {} | |
| def _peer_gate_lock() -> asyncio.Lock: | |
| """Lazy-initialised lock to avoid creating asyncio primitives at import time | |
| (which may happen outside an event loop in test/synchronous contexts).""" | |
| global _PEER_GATE_LOCK | |
| if _PEER_GATE_LOCK is None: | |
| _PEER_GATE_LOCK = asyncio.Lock() | |
| return _PEER_GATE_LOCK | |
| async def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore: | |
| """Return a per-peer concurrency semaphore, creating one if needed. | |
| Uses a lock to prevent a race where two concurrent coroutines both see | |
| ``None`` and create separate semaphores for the same peer_id — the | |
| second would overwrite the first, silently dropping any coroutines | |
| already waiting on it. | |
| Old entries are evicted after ``_PEER_GATE_TTL`` seconds of inactivity | |
| so that the dict does not grow without bound when AntSeed peers are | |
| dynamically discovered and later retired. | |
| """ | |
| import time as _time | |
| lock = _peer_gate_lock() | |
| async with lock: | |
| # Evict stale entries first (best-effort, not on every call to | |
| # avoid quadratic sweep cost — only when the lock is already held | |
| # and the dict is getting large). | |
| if len(_PEER_GATES) > 256: | |
| now = _time.monotonic() | |
| stale = [k for k, ts in _PEER_GATE_TIMESTAMPS.items() | |
| if now - ts > _PEER_GATE_TTL] | |
| for k in stale: | |
| _PEER_GATES.pop(k, None) | |
| _PEER_GATE_TIMESTAMPS.pop(k, None) | |
| _PEER_GATE_CAPS.pop(k, None) | |
| sem = _PEER_GATES.get(peer_id) | |
| if sem is None or _PEER_GATE_CAPS.get(peer_id) != cap: | |
| sem = asyncio.Semaphore(cap) | |
| _PEER_GATES[peer_id] = sem | |
| _PEER_GATE_CAPS[peer_id] = cap | |
| _PEER_GATE_TIMESTAMPS[peer_id] = _time.monotonic() | |
| return sem |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@provider_adapters/openai_compatible.py` around lines 213 - 263, Update
_peer_gate to track the capacity associated with each cached peer semaphore and
recreate the semaphore when the requested cap differs from the cached cap. Keep
the existing semaphore when cap is unchanged, preserve lock protection and
timestamp updates, and ensure replacement entries update the stored capacity so
subsequent calls use the new limit.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Two bugs in provider_adapters/openai_compatible.py that affect every LLM request routed through unhardcoded:
_PEER_GATES race condition and unbounded memory growth - two concurrent coroutines can both see None and create separate semaphores for the same peer_id, the second silently overwrites the first, dropping any coroutines already waiting on it. The dict only grows as AntSeed peers are dynamically discovered and retired, but old entries are never evicted.
stream_openai_compatible() leaks httpx.AsyncClient - when no external client is passed, a new AsyncClient is created but never closed, leaking a TCP connection pool on every streaming call.
Fix
Peer gates: _peer_gate() is now async, guarded by asyncio.Lock, with TTL-based eviction (600s) for stale entries when dict exceeds 256 entries. Lock is lazy-initialised to avoid creating asyncio primitives at import time.
AsyncClient: track ownership via _owns_client flag, close in finally block.
Summary by CodeRabbit