From cc2732a210690a0311d96b2918f5ebd69f68220e Mon Sep 17 00:00:00 2001 From: yasinlex Date: Wed, 22 Jul 2026 20:57:14 +0700 Subject: [PATCH] fix: peer gate race condition + async client connection leak MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- provider_adapters/openai_compatible.py | 237 +++++++++++++++---------- 1 file changed, 143 insertions(+), 94 deletions(-) diff --git a/provider_adapters/openai_compatible.py b/provider_adapters/openai_compatible.py index 5050824..22f45a9 100644 --- a/provider_adapters/openai_compatible.py +++ b/provider_adapters/openai_compatible.py @@ -211,14 +211,56 @@ def call(request: dict) -> dict: _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 def make_async_call_provider( @@ -252,7 +294,7 @@ async def call(request: dict) -> dict: offer = request.get("offer") or {} peer_id = offer.get("peer_id") cap = offer.get("max_concurrency") - gate = _peer_gate(peer_id, cap) if ( + gate = await _peer_gate(peer_id, cap) if ( peer_id and isinstance(cap, int) and cap > 0) else None if gate is not None: try: @@ -341,7 +383,8 @@ async def stream_openai_compatible( body["stream"] = True rules = (provider_rules or {}).get(request.get("provider_id")) or {} - if client is None: + _owns_client = client is None + if _owns_client: import httpx client = httpx.AsyncClient() @@ -365,92 +408,98 @@ def _timeout_err() -> dict: return first_token_timeout_err(first_timeout_s, _latency()) try: - async with AsyncExitStack() as stack: - try: - resp = await before_first_output(stack.enter_async_context( - client.stream("POST", url, json=body, headers=headers, - timeout=timeout)), first_timeout_s, t0, _saw_output) - except (asyncio.TimeoutError, TimeoutError): - return _timeout_err() - if not (200 <= resp.status_code < 300): - raw = (await resp.aread()).decode("utf-8", "replace")[:500] - kind = _classify_from_map(raw, rules.get("error_map")) \ - or _classify_status(resp.status_code, raw) - return _err(kind, resp.status_code, _latency(), raw) - - lines = resp.aiter_lines().__aiter__() - while True: + try: + async with AsyncExitStack() as stack: try: - line = await before_first_output( - lines.__anext__(), first_timeout_s, t0, _saw_output) - except StopAsyncIteration: - break + resp = await before_first_output(stack.enter_async_context( + client.stream("POST", url, json=body, headers=headers, + timeout=timeout)), first_timeout_s, t0, _saw_output) except (asyncio.TimeoutError, TimeoutError): - if not saw_output: - return _timeout_err() - raise - if not line or not line.startswith("data:"): - continue - data = line[len("data:"):].strip() - if data == "[DONE]": - break - try: - chunk = json.loads(data) - except ValueError: - continue - if raw_model is None: - raw_model = chunk.get("model") - if chunk.get("usage"): - usage = chunk["usage"] - for choice in chunk.get("choices") or []: - delta = choice.get("delta") or {} - if choice.get("finish_reason"): - finish_reason = choice["finish_reason"] - content = delta.get("content") - if content: - saw_output = True - text_parts.append(content) - await emit(content) - emitted = True - for tc in delta.get("tool_calls") or []: - saw_output = True - idx = tc.get("index", 0) - acc = tool_calls_acc.setdefault(idx, { - "id": None, "type": "function", - "function": {"name": "", "arguments": ""}}) - if tc.get("id"): - acc["id"] = tc["id"] - fn = tc.get("function") or {} - if fn.get("name"): - acc["function"]["name"] = fn["name"] - if fn.get("arguments"): - acc["function"]["arguments"] += fn["arguments"] - except Exception as exc: # noqa: BLE001 — classified below - partial = "".join(text_parts) - if emitted: - return _err("stream_interrupted", 0, _latency(), - f"{type(exc).__name__}: {exc} (partial: {partial[:200]!r})") - return _err("network_error", 0, _latency(), f"{type(exc).__name__}: {exc}") - - tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)] or None - text = "".join(text_parts) - if not text.strip() and not tool_calls: - return _err("bad_response", 200, _latency(), "empty assistant content") - return { - "ok": True, - "latency_ms": _latency(), - "response": { - "text": text, - "tool_calls": tool_calls, - "finish_reason": finish_reason, - "tokens_in": usage.get("prompt_tokens"), - "tokens_out": usage.get("completion_tokens"), - "tokens_total": usage.get("total_tokens"), - "tokens_cached": _cached_tokens(usage), - "cost_reported": usage.get("cost"), - "raw_model": raw_model, - }, - } + return _timeout_err() + if not (200 <= resp.status_code < 300): + raw = (await resp.aread()).decode("utf-8", "replace")[:500] + kind = _classify_from_map(raw, rules.get("error_map")) \ + or _classify_status(resp.status_code, raw) + return _err(kind, resp.status_code, _latency(), raw) + + lines = resp.aiter_lines().__aiter__() + while True: + try: + line = await before_first_output( + lines.__anext__(), first_timeout_s, t0, _saw_output) + except StopAsyncIteration: + break + except (asyncio.TimeoutError, TimeoutError): + if not saw_output: + return _timeout_err() + raise + if not line or not line.startswith("data:"): + continue + data = line[len("data:"):].strip() + if data == "[DONE]": + break + try: + chunk = json.loads(data) + except ValueError: + continue + if raw_model is None: + raw_model = chunk.get("model") + if chunk.get("usage"): + usage = chunk["usage"] + for choice in chunk.get("choices") or []: + delta = choice.get("delta") or {} + if choice.get("finish_reason"): + finish_reason = choice["finish_reason"] + content = delta.get("content") + if content: + saw_output = True + text_parts.append(content) + await emit(content) + emitted = True + for tc in delta.get("tool_calls") or []: + saw_output = True + idx = tc.get("index", 0) + acc = tool_calls_acc.setdefault(idx, { + "id": None, "type": "function", + "function": {"name": "", "arguments": ""}}) + if tc.get("id"): + acc["id"] = tc["id"] + fn = tc.get("function") or {} + if fn.get("name"): + acc["function"]["name"] = fn["name"] + if fn.get("arguments"): + acc["function"]["arguments"] += fn["arguments"] + except Exception as exc: # noqa: BLE001 — classified below + partial = "".join(text_parts) + if emitted: + return _err("stream_interrupted", 0, _latency(), + f"{type(exc).__name__}: {exc} (partial: {partial[:200]!r})") + return _err("network_error", 0, _latency(), f"{type(exc).__name__}: {exc}") + + tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)] or None + text = "".join(text_parts) + if not text.strip() and not tool_calls: + return _err("bad_response", 200, _latency(), "empty assistant content") + return { + "ok": True, + "latency_ms": _latency(), + "response": { + "text": text, + "tool_calls": tool_calls, + "finish_reason": finish_reason, + "tokens_in": usage.get("prompt_tokens"), + "tokens_out": usage.get("completion_tokens"), + "tokens_total": usage.get("total_tokens"), + "tokens_cached": _cached_tokens(usage), + "cost_reported": usage.get("cost"), + "raw_model": raw_model, + }, + } + finally: + # Close the client if we created it, to prevent connection leaks. + # Caller-owned clients are their responsibility. + if _owns_client: + await client.aclose() def _parse_openai_response(