ocr: one handshake per pod, not one per page

The gateway built an httpx.AsyncClient for every request, so each page paid a
fresh TCP and TLS handshake to the upstream. With surya_instances: 24 that is
24 streams reconnecting per page, against a round trip that already dominates
the cost of reading one.

One pooled client with keep-alive instead, so a document's pages reuse the
connection. Correctness-neutral; it only ever cost time.

Keyed by the event loop, because an AsyncClient binds to the loop that created
it and coderai runs more than one (the OCR pool keeps its own). The key is the
loop OBJECT held weakly, not id(): keying by id would leak an entry per dead
loop and, since CPython reuses addresses, could hand a new loop a client bound
to the closed one that used to live at that address — a fault that would only
appear under the load this change exists to handle.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PxYHWbCAjjjD7A1rhFKtaW

codai/ocr/gateway.py

@@ -33,6 +33,7 @@ that was not handed both.
 import os
 import secrets
 import threading
+import weakref
 from typing import Optional
 
 from fastapi import APIRouter, Request
@@ -51,8 +52,25 @@ PREFIX = "/v1/ocr/attach"
 #: takes minutes; a timeout here would look to the engine like a failed page.
 UPSTREAM_TIMEOUT_S = 3600.0
 
+#: Upstream connection pooling. A client per REQUEST meant a fresh TCP and TLS
+#: handshake for every page: with `surya_instances: 24` that is 24 streams each
+#: reconnecting per page, against a pod whose round trip is already the dominant
+#: cost. One client holds keep-alive connections open instead, which is the
+#: difference between a handshake per page and a handshake per pod.
+#:
+#: Keyed by event loop: an httpx.AsyncClient binds to the loop that created it,
+#: and coderai runs more than one (the OCR pool has a loop of its own), so a
+#: single global client would be used from a loop it does not belong to. The key
+#: is the loop OBJECT, held weakly — keying by id() would both leak an entry per
+#: dead loop and, because CPython reuses addresses, risk handing a new loop a
+#: client bound to the closed one that used to live there.
+_MAX_CONNECTIONS = 256
+_MAX_KEEPALIVE = 64
+_KEEPALIVE_EXPIRY_S = 60.0
+
 _lock = threading.Lock()
 _leases: dict = {}              # lease token -> {"model_id", "served_name"}
+_clients: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()  # loop -> client
 
 
 def issue(model_id: str, served_name: str) -> str:
@@ -107,6 +125,41 @@ def engine_url(token: str) -> str:
 
 
 # --------------------------------------------------------------------- routes #
+def _upstream_client():
+    """The shared client for this event loop, created on first use."""
+    import asyncio
+
+    import httpx
+
+    try:
+        key = asyncio.get_running_loop()
+    except RuntimeError:
+        key = _upstream_client            # a stable stand-in when there is no loop
+    with _lock:
+        got = _clients.get(key)
+        if got is not None and not got.is_closed:
+            return got
+        got = httpx.AsyncClient(
+            timeout=UPSTREAM_TIMEOUT_S,
+            limits=httpx.Limits(max_connections=_MAX_CONNECTIONS,
+                                max_keepalive_connections=_MAX_KEEPALIVE,
+                                keepalive_expiry=_KEEPALIVE_EXPIRY_S))
+        _clients[key] = got
+        return got
+
+
+async def close_upstream_clients() -> None:
+    """Close the pooled clients (process shutdown, or a test tidying up)."""
+    with _lock:
+        clients = list(_clients.values())
+        _clients.clear()
+    for c in clients:
+        try:
+            await c.aclose()
+        except Exception:
+            pass
+
+
 def _unknown_lease() -> JSONResponse:
     return JSONResponse(status_code=404,
                         content={"error": {"message": "unknown OCR engine lease",
@@ -165,8 +218,6 @@ async def chat_completions(token: str, request: Request):
     if got is None:
         return _unknown_lease()
 
-    import httpx
-
     from codai.models import placement
     body = await request.body()
 
@@ -188,8 +239,7 @@ async def chat_completions(token: str, request: Request):
         if chosen.api_key:
             headers["Authorization"] = f"Bearer {chosen.api_key}"
         url = chosen.base.rstrip("/") + "/chat/completions"
-        async with httpx.AsyncClient(timeout=UPSTREAM_TIMEOUT_S) as client:
-            reply = await client.post(url, content=body, headers=headers)
+        reply = await _upstream_client().post(url, content=body, headers=headers)
         # The reply goes back as it came: a re-serialised response is how
         # logprobs got lost in the first place.
         media = reply.headers.get("content-type", "application/json")

tests/test_ocr_fanout.py

@@ -35,6 +35,9 @@ from codai.ocr.base import OcrError                                # noqa: E402
 def _reset():
     dispatch.forget()
     gateway.revoke()
+    # The upstream client is pooled and cached per event loop; a cached one would
+    # carry a previous test's stub into the next test.
+    gateway._clients.clear()
 
 
 def _cand(key, kind="node", base="http://n/v1", priority=1, node_url="",
@@ -926,3 +929,107 @@ def test_the_internal_gate_exempts_only_the_attach_prefix(monkeypatch):
         reached.clear()
         asyncio.run(mw({"type": "http", "path": path, "headers": []}, None, _send))
         assert reached == [], f"{path} must still require the internal token"
+
+
+# --------------------------------------------- one handshake per pod, not per page #
+# The gateway built a new httpx.AsyncClient for every request, so each page paid a
+# fresh TCP and TLS handshake to the pod. With surya_instances: 24 that is 24
+# streams reconnecting per page against a round trip that already dominates.
+
+
+def test_the_upstream_client_is_reused_within_one_event_loop(monkeypatch):
+    """The point of the fix: a second page on the same loop reuses the connection
+    pool instead of building another one."""
+    _reset()
+    import asyncio
+    import httpx
+    built = []
+
+    class _Counting(_CapturingClient):
+        def __init__(self, *a, **kw):
+            built.append(kw)
+            super().__init__(*a, **kw)
+
+        @property
+        def is_closed(self):
+            return False
+
+    monkeypatch.setattr(httpx, "AsyncClient", _Counting)
+    _no_node_load(monkeypatch)
+    _fixed_candidates(monkeypatch, [_cand("node:a", base="http://a/v1")])
+    token = gateway.issue("m", "m")
+
+    async def _two_pages():
+        await gateway.chat_completions(token, _Req(SURYA_BODY))
+        await gateway.chat_completions(token, _Req(SURYA_BODY))
+
+    asyncio.run(_two_pages())
+    assert len(built) == 1, "a second page must not build a second client"
+
+
+def test_the_client_is_created_with_keep_alive_limits(monkeypatch):
+    """A pool with no keep-alive would reconnect anyway."""
+    _reset()
+    import asyncio
+    import httpx
+    built = []
+
+    class _Counting(_CapturingClient):
+        def __init__(self, *a, **kw):
+            built.append(kw)
+            super().__init__(*a, **kw)
+
+        @property
+        def is_closed(self):
+            return False
+
+    monkeypatch.setattr(httpx, "AsyncClient", _Counting)
+    _no_node_load(monkeypatch)
+    _fixed_candidates(monkeypatch, [_cand("node:a", base="http://a/v1")])
+    token = gateway.issue("m", "m")
+    asyncio.run(gateway.chat_completions(token, _Req(SURYA_BODY)))
+    assert built and built[0].get("limits") is not None
+    assert built[0].get("timeout") == gateway.UPSTREAM_TIMEOUT_S
+
+
+def test_a_client_is_not_shared_between_different_loops(monkeypatch):
+    """An AsyncClient binds to the loop that made it; using one from another loop
+    is an error waiting for load to find it."""
+    _reset()
+    import asyncio
+    import httpx
+    built = []
+
+    class _Counting(_CapturingClient):
+        def __init__(self, *a, **kw):
+            built.append(kw)
+            super().__init__(*a, **kw)
+
+        @property
+        def is_closed(self):
+            return False
+
+    monkeypatch.setattr(httpx, "AsyncClient", _Counting)
+    _no_node_load(monkeypatch)
+    _fixed_candidates(monkeypatch, [_cand("node:a", base="http://a/v1")])
+    token = gateway.issue("m", "m")
+    asyncio.run(gateway.chat_completions(token, _Req(SURYA_BODY)))
+    asyncio.run(gateway.chat_completions(token, _Req(SURYA_BODY)))
+    assert len(built) == 2, "each loop needs its own client"
+
+
+def test_a_dead_loop_does_not_keep_its_client_alive(monkeypatch):
+    """Keyed weakly by the loop object, so a finished loop's entry disappears
+    rather than accumulating — and cannot be handed to a new loop that happens to
+    reuse the address."""
+    _reset()
+    import asyncio
+    import gc
+    import httpx
+    monkeypatch.setattr(httpx, "AsyncClient", _CapturingClient)
+    _no_node_load(monkeypatch)
+    _fixed_candidates(monkeypatch, [_cand("node:a", base="http://a/v1")])
+    token = gateway.issue("m", "m")
+    asyncio.run(gateway.chat_completions(token, _Req(SURYA_BODY)))
+    gc.collect()
+    assert len(gateway._clients) == 0