Skip to content

Isolate long-polls and file transfers on separate HTTP/2 pools #825

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gautam-rl merged 11 commits into from
Jul 28, 2026
Merged

Isolate long-polls and file transfers on separate HTTP/2 pools #825

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
7fc11ba
fix: open extra H2 conns under stream pressure; isolate file transfer…
Jul 27, 2026
bfc71df
fix: isolate waits and file transfers on sharded H2 pools
Jul 27, 2026
6783ef6
chore: fix ruff import sorting for pool bulkhead changes
Jul 27, 2026
bb09391
chore: apply ruff format to bulkhead pool changes
Jul 27, 2026
76a9ecb
test: prove H2 bulkheads on the wire; lock sync lazy init
Jul 27, 2026
8bdcd53
chore: fix pyright on H2 bulkhead integration test
Jul 27, 2026
b6451ee
fix: round-robin H2 bulkhead shards per SDK client
Jul 28, 2026
5900470
fix: randomize per-client H2 shard start offsets
Jul 28, 2026
e6ef85a
refactor: encapsulate sharded H2 pools; raise defaults
Jul 28, 2026
e7042d7
chore: mark Registry test helpers; fix pyright on pool cleanup
Jul 28, 2026
a0f69b3
chore: ruff format _base_client.py
Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
261 changes: 259 additions & 2 deletions src/runloop_api_client/_base_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import asyncio
import inspect
import logging
import secrets
import weakref
import platform
import warnings
Expand Down Expand Up @@ -80,6 +81,8 @@
RAW_RESPONSE_HEADER,
OVERRIDE_CAST_TO_HEADER,
DEFAULT_CONNECTION_LIMITS,
DEFAULT_TRANSFER_POOL_SHARDS,
DEFAULT_BACKGROUND_POOL_SHARDS,
)
from ._streaming import Stream, SSEDecoder, AsyncStream, SSEBytesDecoder
from ._exceptions import (
Expand Down Expand Up @@ -174,6 +177,63 @@ async def aclose(self) -> None:
weakref.WeakKeyDictionary()
)

# Sharded H2 bulkheads: long-polls and file transfers stay off the API control-plane
# connection. Each shard index maps to its own shared transport (≈ one H2 connection).
# Per-client round-robin (random start offset) spreads concurrent requests across
# shards. Removable once httpcore respects stream capacity when opening connections.
_shared_sync_background_transports: dict[int, _SharedTransport] = {}
_shared_sync_transfer_transports: dict[int, _SharedTransport] = {}
_shared_async_background_transports: weakref.WeakKeyDictionary[
asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport]
] = weakref.WeakKeyDictionary()
_shared_async_transfer_transports: weakref.WeakKeyDictionary[
asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport]
] = weakref.WeakKeyDictionary()

_BACKGROUND_PATH_SUFFIXES = ("/wait_for_status",)
_TRANSFER_PATH_SUFFIXES = ("/upload_file", "/download_file")


def _is_background_path(path: str) -> bool:
return path.endswith(_BACKGROUND_PATH_SUFFIXES)


def _is_transfer_path(path: str) -> bool:
return path.endswith(_TRANSFER_PATH_SUFFIXES)


def _acquire_shared_sync_transport(bucket: dict[int, _SharedTransport], shard: int) -> _SharedTransport:
with _pool_lock:
existing = bucket.get(shard)
if existing is not None and existing.acquire():
return existing
transport = _SharedTransport(
httpx.HTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True),
)
bucket[shard] = transport
return transport


def _acquire_shared_async_transport(
by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport]],
loop: asyncio.AbstractEventLoop,
shard: int,
) -> _SharedAsyncTransport:
with _pool_lock:
bucket = by_loop.get(loop)
if bucket is None:
bucket = {}
by_loop[loop] = bucket
existing = bucket.get(shard)
if existing is not None and existing.acquire():
return existing
transport = _SharedAsyncTransport(
httpx.AsyncHTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True),
)
bucket[shard] = transport
return transport


# TODO: make base page type vars covariant
SyncPageT = TypeVar("SyncPageT", bound="BaseSyncPage[Any]")
AsyncPageT = TypeVar("AsyncPageT", bound="BaseAsyncPage[Any]")
Expand Down Expand Up @@ -929,8 +989,15 @@ def __del__(self) -> None:

class SyncAPIClient(BaseClient[httpx.Client, Stream[Any]]):
_client: httpx.Client
_background_clients: dict[int, httpx.Client]
_transfer_clients: dict[int, httpx.Client]
_default_stream_cls: type[Stream[Any]] | None = None
_uses_shared_pool: bool
_isolate_workload_pools: bool
_background_pool_shards: int
_transfer_pool_shards: int
_background_next: int
_transfer_next: int
_closed: bool

def __init__(
Expand All @@ -945,6 +1012,8 @@ def __init__(
custom_query: Mapping[str, object] | None = None,
_strict_response_validation: bool,
shared_http_pool: bool = True,
background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS,
transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS,
) -> None:
if not is_given(timeout):
# if the user passed in a custom http client with a non-default
Expand All @@ -964,6 +1033,11 @@ def __init__(
f"Invalid `http_client` argument; Expected an instance of `httpx.Client` but got {type(http_client)}"
)

if background_pool_shards < 1:
raise ValueError("background_pool_shards must be >= 1")
if transfer_pool_shards < 1:
raise ValueError("transfer_pool_shards must be >= 1")

super().__init__(
version=version,
# cast to a valid type because mypy doesn't understand our type narrowing
Expand All @@ -976,6 +1050,16 @@ def __init__(
)

self._closed = False
self._background_clients = {}
self._transfer_clients = {}
self._bulkhead_lock = threading.Lock()
self._background_pool_shards = background_pool_shards
self._transfer_pool_shards = transfer_pool_shards
# Random start avoids every short-lived SDK client pinning global shard 0.
self._background_next = secrets.randbelow(background_pool_shards)
self._transfer_next = secrets.randbelow(transfer_pool_shards)
# Custom http_client owns the full transport stack; don't invent sibling pools.
self._isolate_workload_pools = http_client is None

if http_client is not None:
self._client = http_client
Expand All @@ -1000,6 +1084,78 @@ def __init__(
)
self._uses_shared_pool = False

def _make_bulkhead_client(self, *, transport: httpx.BaseTransport | None) -> httpx.Client:
timeout = cast(Timeout, self.timeout)
if transport is not None:
return SyncHttpxClientWrapper(
base_url=self._base_url,
timeout=timeout,
transport=transport,
)
return SyncHttpxClientWrapper(
base_url=self._base_url,
timeout=timeout,
)

def _ensure_background_client(self, shard: int) -> httpx.Client:
existing = self._background_clients.get(shard)
if existing is not None:
return existing
with self._bulkhead_lock:
existing = self._background_clients.get(shard)
if existing is not None:
return existing
if self._uses_shared_pool:
transport: httpx.BaseTransport | None = _acquire_shared_sync_transport(
_shared_sync_background_transports, shard
)
else:
transport = None
client = self._make_bulkhead_client(transport=transport)
self._background_clients[shard] = client
return client

def _ensure_transfer_client(self, shard: int) -> httpx.Client:
existing = self._transfer_clients.get(shard)
if existing is not None:
return existing
with self._bulkhead_lock:
existing = self._transfer_clients.get(shard)
if existing is not None:
return existing
if self._uses_shared_pool:
transport: httpx.BaseTransport | None = _acquire_shared_sync_transport(
_shared_sync_transfer_transports, shard
)
else:
transport = None
client = self._make_bulkhead_client(transport=transport)
self._transfer_clients[shard] = client
return client

def _next_background_client(self) -> httpx.Client:
# Select under the lock; ensure afterward so _ensure_* can take the same lock.
with self._bulkhead_lock:
shard = self._background_next % self._background_pool_shards
self._background_next += 1
return self._ensure_background_client(shard)

@jrvb-rl jrvb-rl Jul 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we've some duplication here that would be cleaner if we pushed it into the _SharedTransport class. I think that would simplify the locking and acquire logic as well, and reduce the chance of errors later on.

@gautam-rl gautam-rl Jul 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done — acquire/refcount/shard map are in the Registry; client wrappers no longer duplicate that locking.


def _next_transfer_client(self) -> httpx.Client:
with self._bulkhead_lock:
shard = self._transfer_next % self._transfer_pool_shards
self._transfer_next += 1
return self._ensure_transfer_client(shard)

def _send_client_for_request(self, request: httpx.Request) -> httpx.Client:
if not self._isolate_workload_pools:
return self._client
path = request.url.path
if _is_background_path(path):
return self._next_background_client()
if _is_transfer_path(path):
return self._next_transfer_client()

@jrvb-rl jrvb-rl Jul 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how about just:

client = get_client_for_path(path)

and then we push all of the pool management down in to the wrapper library?

@gautam-rl gautam-rl Jul 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Switched to _get_client_for_path(path) as the single entry point; pool selection lives under that.

return self._client

def is_closed(self) -> bool:
return self._closed or self._client.is_closed

Expand All @@ -1014,6 +1170,12 @@ def close(self) -> None:
return
self._closed = True
self._client.close()
for client in self._background_clients.values():
client.close()
self._background_clients.clear()
for client in self._transfer_clients.values():
client.close()
self._transfer_clients.clear()

def __enter__(self: _T) -> _T:
return self
Expand Down Expand Up @@ -1114,7 +1276,7 @@ def request(

response = None
try:
response = self._client.send(
response = self._send_client_for_request(request).send(
request,
stream=stream or self._should_stream_response_body(request=request),
**kwargs,
Expand Down Expand Up @@ -1561,8 +1723,15 @@ def __del__(self) -> None:

class AsyncAPIClient(BaseClient[httpx.AsyncClient, AsyncStream[Any]]):
_client: httpx.AsyncClient
_background_clients: dict[int, httpx.AsyncClient]
_transfer_clients: dict[int, httpx.AsyncClient]
_default_stream_cls: type[AsyncStream[Any]] | None = None
_uses_shared_pool: bool
_isolate_workload_pools: bool
_background_pool_shards: int
_transfer_pool_shards: int
_background_next: int
_transfer_next: int
_closed: bool

def __init__(
Expand All @@ -1577,6 +1746,8 @@ def __init__(
custom_headers: Mapping[str, str] | None = None,
custom_query: Mapping[str, object] | None = None,
shared_http_pool: bool = True,
background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS,
transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS,
) -> None:
if not is_given(timeout):
# if the user passed in a custom http client with a non-default
Expand All @@ -1596,6 +1767,11 @@ def __init__(
f"Invalid `http_client` argument; Expected an instance of `httpx.AsyncClient` but got {type(http_client)}"
)

if background_pool_shards < 1:
raise ValueError("background_pool_shards must be >= 1")
if transfer_pool_shards < 1:
raise ValueError("transfer_pool_shards must be >= 1")

super().__init__(
version=version,
base_url=base_url,
Expand All @@ -1608,6 +1784,15 @@ def __init__(
)

self._closed = False
self._background_clients = {}
self._transfer_clients = {}
self._background_pool_shards = background_pool_shards
self._transfer_pool_shards = transfer_pool_shards
# Random start avoids every short-lived SDK client pinning global shard 0.
self._background_next = secrets.randbelow(background_pool_shards)
self._transfer_next = secrets.randbelow(transfer_pool_shards)
# Custom http_client owns the full transport stack; don't invent sibling pools.
self._isolate_workload_pools = http_client is None

if http_client is not None:
self._client = http_client
Expand Down Expand Up @@ -1646,6 +1831,72 @@ def __init__(
)
self._uses_shared_pool = False

def _make_bulkhead_client(self, *, transport: httpx.AsyncBaseTransport | None) -> httpx.AsyncClient:
timeout = cast(Timeout, self.timeout)
if transport is not None:
return AsyncHttpxClientWrapper(
base_url=self._base_url,
timeout=timeout,
transport=transport,
)
return AsyncHttpxClientWrapper(
base_url=self._base_url,
timeout=timeout,
)

def _ensure_background_client(self, shard: int) -> httpx.AsyncClient:
existing = self._background_clients.get(shard)
if existing is not None:
return existing
transport: httpx.AsyncBaseTransport | None = None
if self._uses_shared_pool:
try:
loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
except RuntimeError:
loop = None
if loop is not None:
transport = _acquire_shared_async_transport(_shared_async_background_transports, loop, shard)
client = self._make_bulkhead_client(transport=transport)
self._background_clients[shard] = client
return client

def _ensure_transfer_client(self, shard: int) -> httpx.AsyncClient:
existing = self._transfer_clients.get(shard)
if existing is not None:
return existing
transport: httpx.AsyncBaseTransport | None = None
if self._uses_shared_pool:
try:
loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
except RuntimeError:
loop = None
if loop is not None:
transport = _acquire_shared_async_transport(_shared_async_transfer_transports, loop, shard)
client = self._make_bulkhead_client(transport=transport)
self._transfer_clients[shard] = client
return client

def _next_background_client(self) -> httpx.AsyncClient:
# Single-threaded event loop: counter bump needs no lock when there is no await.
shard = self._background_next % self._background_pool_shards
self._background_next += 1
return self._ensure_background_client(shard)

def _next_transfer_client(self) -> httpx.AsyncClient:
shard = self._transfer_next % self._transfer_pool_shards
self._transfer_next += 1
return self._ensure_transfer_client(shard)

def _send_client_for_request(self, request: httpx.Request) -> httpx.AsyncClient:
if not self._isolate_workload_pools:
return self._client
path = request.url.path
if _is_background_path(path):
return self._next_background_client()
if _is_transfer_path(path):
return self._next_transfer_client()
return self._client

def is_closed(self) -> bool:
return self._closed or self._client.is_closed

Expand All @@ -1660,6 +1911,12 @@ async def close(self) -> None:
return
self._closed = True
await self._client.aclose()
for client in self._background_clients.values():
await client.aclose()
self._background_clients.clear()
for client in self._transfer_clients.values():
await client.aclose()
self._transfer_clients.clear()

async def __aenter__(self: _T) -> _T:
return self
Expand Down Expand Up @@ -1765,7 +2022,7 @@ async def request(

response = None
try:
response = await self._client.send(
response = await self._send_client_for_request(request).send(
request,
stream=stream or self._should_stream_response_body(request=request),
**kwargs,
Expand Down
Loading