diff --git a/.changeset/python-unify-pyqwest-connection-pools.md b/.changeset/python-unify-pyqwest-connection-pools.md new file mode 100644 index 0000000000..ec3655cd00 --- /dev/null +++ b/.changeset/python-unify-pyqwest-connection-pools.md @@ -0,0 +1,16 @@ +--- +"@e2b/python-sdk": patch +--- + +Run every persistent HTTP stack in the SDK on one shared pyqwest connection pool +per proxy instead of four: the control-plane REST API, the envd HTTP API, the +envd RPC clients, and the volume content API now all draw from +`e2b.api.client_sync`/`client_async`. reqwest pools per host internally, so one +pool serves the API host and every per-sandbox host without interference — and +since envd RPC and the envd HTTP API hit the same host, an active sandbox now +needs a single HTTP/2 connection instead of one per stack. Streamed downloads +keep a pool of their own, the only one carrying the idle `read_timeout`: +reqwest's read timer runs during body send and TTFB, so on a shared pool it +would cut off long uploads. Layers above the pool are unchanged — the RPC +stack's plain-HTTP error normalization still wraps it, and connect-only retries +are now part of it. diff --git a/.changeset/python-uploads-skip-connect-retries.md b/.changeset/python-uploads-skip-connect-retries.md new file mode 100644 index 0000000000..9b570c7215 --- /dev/null +++ b/.changeset/python-uploads-skip-connect-retries.md @@ -0,0 +1,16 @@ +--- +"@e2b/python-sdk": patch +--- + +Stop mirroring streamed file uploads in memory. To be able to replay a request, +pyqwest's retry middleware copies a streamed request body in full as it is sent, +so an upload was streamed to the wire *and* held whole in RAM: writing a 64 MiB +file with `volume.write_file`, `sandbox.files.write` or a template context upload +peaked at ~67 MB of it. A streamed upload now goes out on an httpx transport over +the same connection pool with the retry layer left out, so it keeps the pooled +connections while the body is no longer copied — a 32 MiB upload peaks at 1.5 MB +instead of 38 MB. What those uploads give up is the connect retry, which fires +before any of the body was written, so a failure to connect surfaces to the +caller as it would have anyway; writes of in-memory data are replayable for free +and keep their retries. Template context uploads no longer build a connection +pool of their own per build either. diff --git a/packages/python-sdk/e2b/api/client_async/__init__.py b/packages/python-sdk/e2b/api/client_async/__init__.py index d40ca257e6..8580acc824 100644 --- a/packages/python-sdk/e2b/api/client_async/__init__.py +++ b/packages/python-sdk/e2b/api/client_async/__init__.py @@ -24,14 +24,21 @@ def get_api_client(config: ConnectionConfig, **kwargs) -> AsyncApiClient: class ConnectionRetryTransport(RetryTransport): - """Retry only failures establishing the connection — shared by the REST - API and envd RPC stacks: pyqwest raises the builtin ``ConnectionError`` - only before the request was written, so these retries can never replay a - request the server may have received (a delivered REST call or unary RPC - like ``SendInput``). This matches the connect-only ``retries`` of the - httpx transports this replaced; the retry middleware's default policy - would otherwise also retry I/O errors and 429/5xx responses for - idempotent methods.""" + """Retry only failures establishing the connection — part of the shared + transport stack, so it covers the REST API, the envd HTTP API, the envd RPC + clients and the volume content API alike: pyqwest raises the builtin + ``ConnectionError`` only before the request was written, so these retries + can never replay a request the server may have received (a delivered REST + call or unary RPC like ``SendInput``). This matches the connect-only + ``retries`` of the httpx transports this replaced; the retry middleware's + default policy would otherwise also retry I/O errors and 429/5xx responses + for idempotent methods. + + Retries cost memory on a streamed request body: to be able to replay it, + the middleware copies the body in full as it is sent (reqwest reads ahead + into it while connecting, so a connect error leaves the iterator already + started and unrewindable). File uploads therefore skip this layer and go + straight to the pool underneath — see :func:`get_upload_transport`.""" def should_retry_response( self, request: Request, response: Union[Response, Exception] @@ -39,104 +46,174 @@ def should_retry_response( return isinstance(response, ConnectionError) -def retrying_http_transport( - proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None -) -> ConnectionRetryTransport: - """A fresh pyqwest transport (= its own connection pool) with the SDK's - shared tuning — system CA certs (without which TLS through an - intercepting proxy fails), the httpx-equivalent pool limits, and - connect-only retries. The REST API, envd RPC, and envd HTTP API stacks - each cache their own instances (pool unification is a follow-up). +TransportKey = Tuple[Optional[ProxyConfig], Optional[float]] +"""Cache key of the shared transports: the proxy and the idle read bound.""" - ``read_timeout`` bounds every read on the transport's connections; see - :func:`get_envd_transport` for when that is (and isn't) appropriate. +_transport_lock = threading.Lock() +# One pyqwest transport — one reqwest connection pool — per proxy and idle read +# bound; `None` is the direct pool. Every HTTP stack in the SDK draws from +# these: the control-plane REST API, the envd HTTP API, the envd RPC clients +# (`e2b.envd.client_async`) and the volume content API +# (`e2b.volume.client_async`). reqwest pools per host internally, so a single +# pool serves the API host and every per-sandbox host without interference — +# and because envd RPC and the envd HTTP API share it, a sandbox needs one +# HTTP/2 connection instead of one per stack. +# +# pyqwest's I/O runs on its own Rust runtime, so unlike the httpx transports +# they replaced, the transports are not bound to an event loop and the caches +# are process-global rather than per-loop. +_pools: Dict[TransportKey, HTTPTransport] = {} +# The connect-retrying stack over each pool. +_transports: Dict[TransportKey, ConnectionRetryTransport] = {} +# The httpx adapter over each pool, shared by every httpx client on it — one per +# pool with the retries, one without for uploads. +_httpx_transports: Dict[TransportKey, AsyncPyqwestTransport] = {} +_upload_transports: Dict[TransportKey, AsyncPyqwestTransport] = {} + + +def get_pool( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> HTTPTransport: + """The shared connection pool with the SDK's tuning — system CA certs + (without which TLS through an intercepting proxy fails) and the + httpx-equivalent pool limits. For TLS connections ALPN negotiates the HTTP + version (HTTP/2 against the E2B API and envd), like the http2-enabled httpx + transports this replaced. + + Layer concerns above it rather than into it — connect retries wrap it, so + does the RPC stack's plain-HTTP-error normalization, and headers and codecs + are per-request — so that the pool stays shareable. + + ``read_timeout`` bounds every read on the pool's connections, which is why + it is part of the cache key: reqwest's read timer keeps running while a + request body is sent and while waiting for the response head, so a pool + carrying one would cut off long uploads and slow responses. Only streamed + downloads ask for it, as an idle bound (see :func:`get_transport`). Requests are logged by pyqwest itself on the ``pyqwest.access`` and ``pyqwest`` loggers at ``DEBUG`` (off unless enabled) — the transport-level diagnostics httpcore used to provide. The SDK's own ``logger`` option is separate and sits above this, on the httpx client.""" - return ConnectionRetryTransport( - HTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the generated - # clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) + key = (proxy, read_timeout) + with _transport_lock: + pool = _pools.get(key) + if pool is None: + pool = HTTPTransport( + tls_include_system_certs=True, + proxy=proxy.to_pyqwest() if proxy is not None else None, + pool_idle_timeout=pool_idle_timeout, + pool_max_idle_per_host=pool_max_idle_per_host, + read_timeout=read_timeout, + # Redirects belong to the httpx client above (which the + # generated clients leave off), not to reqwest. + follow_redirects=False, + ) + _pools[key] = pool + return pool -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy; None is the direct pool. -# pyqwest's I/O runs on its own Rust runtime, so unlike the httpx transports -# they replaced, the transports are not bound to an event loop and the -# caches are process-global rather than per-loop. -_transports: Dict[Optional[ProxyConfig], AsyncPyqwestTransport] = {} +def get_pyqwest_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> ConnectionRetryTransport: + """The shared pool of :func:`get_pool` behind connect-only retries, for + consumers speaking pyqwest natively (the envd RPC clients). Consumers + speaking httpx take :func:`get_httpx_transport`, the adapter over the very + same stack.""" + key = (proxy, read_timeout) + # Resolve the pool before taking the lock: it takes the same one. + pool = get_pool(proxy, read_timeout) + with _transport_lock: + transport = _transports.get(key) + if transport is None: + transport = ConnectionRetryTransport(pool, max_retries=connection_retries) + _transports[key] = transport + return transport -def get_transport(config: ConnectionConfig) -> AsyncPyqwestTransport: - """The shared pyqwest-backed httpx transport for REST API calls. For TLS - connections ALPN negotiates the HTTP version (HTTP/2 against the E2B - API), like the http2-enabled httpx transport this replaced.""" - proxy = proxy_to_config(config.proxy) +def get_httpx_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> AsyncPyqwestTransport: + """The httpx adapter over the retrying stack of + :func:`get_pyqwest_transport`, for the generated httpx clients (control + plane, envd HTTP API, volume content). The adapter holds no state of its + own and does not close the pool, so closing an httpx client leaves the + pool intact for the other clients on it.""" + key = (proxy, read_timeout) + stack = get_pyqwest_transport(proxy, read_timeout) with _transport_lock: - transport = _transports.get(proxy) + transport = _httpx_transports.get(key) if transport is None: - transport = AsyncPyqwestTransport(retrying_http_transport(proxy)) - _transports[proxy] = transport + transport = AsyncPyqwestTransport(stack) + _httpx_transports[key] = transport return transport -# One transport per (proxy, streaming) pair, separate from the REST API -# pools — envd traffic goes to per-sandbox hosts. -_envd_transports: Dict[Tuple[Optional[ProxyConfig], bool], AsyncPyqwestTransport] = {} +def get_upload_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> AsyncPyqwestTransport: + """The httpx adapter over the *same pool* as :func:`get_httpx_transport` + but without the retry layer, for file uploads: envd ``files.write``, volume + ``write_file`` and template context uploads. + + Those are the requests whose body is streamed rather than held in memory, + and the retry middleware makes a body replayable by copying it in full as + it is sent — so uploading a file on the retrying stack would mirror the + whole file in RAM (SDK-332). Skipping the layer, not the pool, keeps the + connection reuse: an upload still travels the sandbox's or the volume + host's pooled connections. What it gives up is the connect retry, which + fires before any of the body was written, so the caller sees the connection + error and can retry the upload itself.""" + key = (proxy, read_timeout) + pool = get_pool(proxy, read_timeout) + with _transport_lock: + transport = _upload_transports.get(key) + if transport is None: + transport = AsyncPyqwestTransport(pool) + _upload_transports[key] = transport + return transport -def get_envd_transport( - config: ConnectionConfig, *, for_streaming: bool = False +def get_transport( + config: ConnectionConfig, *, for_streaming: bool = False, for_upload: bool = False ) -> AsyncPyqwestTransport: - """The shared pyqwest-backed httpx transports for the envd HTTP API - (file transfers, health checks). - - The streaming transport carries ``read_timeout``, the idle bound on - every read: it resets after each successful read, so it caps how long a - streamed download may stall without limiting total transfer time. It is - fixed per transport — the adapter's per-request timeouts are - whole-request deadlines. Only streamed downloads use it: reqwest's read - timer keeps running while a request body is sent and while waiting for - the response head, so on the regular transport it would cut off uploads - and slow unary responses longer than the idle bound (those stay bounded - by their whole-request deadlines instead). + """The shared httpx transport for the control-plane REST API and the envd + HTTP API (file transfers, health checks) — one pool serves both, keyed by + the connection's proxy. + + ``for_streaming`` selects the pool carrying ``READ_TIMEOUT``, the idle + bound on every read: it resets after each successful read, so it caps how + long a streamed download may stall without limiting total transfer time. + It is fixed per pool — the adapter's per-request timeouts are + whole-request deadlines rather than idle bounds — so only streamed + downloads take it, and they get their own pool (see :func:`get_pool`). + + ``for_upload`` selects the same pool without the retry layer, so that a + streamed upload is not copied into memory to keep it replayable + (see :func:`get_upload_transport`). """ proxy = proxy_to_config(config.proxy) - key = (proxy, for_streaming) - with _transport_lock: - transport = _envd_transports.get(key) - if transport is None: - transport = AsyncPyqwestTransport( - retrying_http_transport( - proxy, - read_timeout=READ_TIMEOUT if for_streaming else None, - ) - ) - _envd_transports[key] = transport - return transport + read_timeout = READ_TIMEOUT if for_streaming else None + if for_upload: + return get_upload_transport(proxy, read_timeout) + return get_httpx_transport(proxy, read_timeout) def get_envd_api( - config: ConnectionConfig, base_url: str, *, for_streaming: bool = False + config: ConnectionConfig, + base_url: str, + *, + for_streaming: bool = False, + for_upload: bool = False, ) -> httpx.AsyncClient: """An httpx client for a sandbox's envd HTTP API (file transfers, health - checks) on the shared pyqwest transports. The client itself is a cheap - stateless wrapper — one per consumer is fine — while the pooled transport - underneath is shared and loop-independent.""" + checks) on the shared transports. The client itself is a cheap stateless + wrapper — one per consumer is fine — while the pooled transport underneath + is shared and loop-independent.""" return httpx.AsyncClient( base_url=base_url, - transport=get_envd_transport(config, for_streaming=for_streaming), + transport=get_transport( + config, for_streaming=for_streaming, for_upload=for_upload + ), headers=config.sandbox_headers, event_hooks=make_async_logging_event_hooks(config.logger), ) diff --git a/packages/python-sdk/e2b/api/client_sync/__init__.py b/packages/python-sdk/e2b/api/client_sync/__init__.py index d7331d8cc4..a39968163d 100644 --- a/packages/python-sdk/e2b/api/client_sync/__init__.py +++ b/packages/python-sdk/e2b/api/client_sync/__init__.py @@ -24,14 +24,21 @@ def get_api_client(config: ConnectionConfig, **kwargs) -> ApiClient: class ConnectionRetryTransport(SyncRetryTransport): - """Retry only failures establishing the connection — shared by the REST - API and envd RPC stacks: pyqwest raises the builtin ``ConnectionError`` - only before the request was written, so these retries can never replay a - request the server may have received (a delivered REST call or unary RPC - like ``SendInput``). This matches the connect-only ``retries`` of the - httpx transports this replaced; the retry middleware's default policy - would otherwise also retry I/O errors and 429/5xx responses for - idempotent methods.""" + """Retry only failures establishing the connection — part of the shared + transport stack, so it covers the REST API, the envd HTTP API, the envd RPC + clients and the volume content API alike: pyqwest raises the builtin + ``ConnectionError`` only before the request was written, so these retries + can never replay a request the server may have received (a delivered REST + call or unary RPC like ``SendInput``). This matches the connect-only + ``retries`` of the httpx transports this replaced; the retry middleware's + default policy would otherwise also retry I/O errors and 429/5xx responses + for idempotent methods. + + Retries cost memory on a streamed request body: to be able to replay it, + the middleware copies the body in full as it is sent (reqwest reads ahead + into it while connecting, so a connect error leaves the iterator already + started and unrewindable). File uploads therefore skip this layer and go + straight to the pool underneath — see :func:`get_upload_transport`.""" def should_retry_response( self, request: SyncRequest, response: Union[SyncResponse, Exception] @@ -39,104 +46,173 @@ def should_retry_response( return isinstance(response, ConnectionError) -def retrying_http_transport( - proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None -) -> ConnectionRetryTransport: - """A fresh pyqwest transport (= its own connection pool) with the SDK's - shared tuning — system CA certs (without which TLS through an - intercepting proxy fails), the httpx-equivalent pool limits, and - connect-only retries. The REST API, envd RPC, and envd HTTP API stacks - each cache their own instances (pool unification is a follow-up). +TransportKey = Tuple[Optional[ProxyConfig], Optional[float]] +"""Cache key of the shared transports: the proxy and the idle read bound.""" - ``read_timeout`` bounds every read on the transport's connections; see - :func:`get_envd_transport` for when that is (and isn't) appropriate. +_transport_lock = threading.Lock() +# One pyqwest transport — one reqwest connection pool — per proxy and idle read +# bound; `None` is the direct pool. Every HTTP stack in the SDK draws from +# these: the control-plane REST API, the envd HTTP API, the envd RPC clients +# (`e2b.envd.client_sync`) and the volume content API (`e2b.volume.client_sync`). +# reqwest pools per host internally, so a single pool serves the API host and +# every per-sandbox host without interference — and because envd RPC and the +# envd HTTP API share it, a sandbox needs one HTTP/2 connection instead of one +# per stack. +# +# pyqwest transports are thread-safe, so unlike the httpx transports they +# replaced, the caches are process-global rather than per-thread. +_pools: Dict[TransportKey, SyncHTTPTransport] = {} +# The connect-retrying stack over each pool. +_transports: Dict[TransportKey, ConnectionRetryTransport] = {} +# The httpx adapter over each pool, shared by every httpx client on it — one per +# pool with the retries, one without for uploads. +_httpx_transports: Dict[TransportKey, PyqwestTransport] = {} +_upload_transports: Dict[TransportKey, PyqwestTransport] = {} + + +def get_pool( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> SyncHTTPTransport: + """The shared connection pool with the SDK's tuning — system CA certs + (without which TLS through an intercepting proxy fails) and the + httpx-equivalent pool limits. For TLS connections ALPN negotiates the HTTP + version (HTTP/2 against the E2B API and envd), like the http2-enabled httpx + transports this replaced. + + Layer concerns above it rather than into it — connect retries wrap it, so + does the RPC stack's plain-HTTP-error normalization, and headers and codecs + are per-request — so that the pool stays shareable. + + ``read_timeout`` bounds every read on the pool's connections, which is why + it is part of the cache key: reqwest's read timer keeps running while a + request body is sent and while waiting for the response head, so a pool + carrying one would cut off long uploads and slow responses. Only streamed + downloads ask for it, as an idle bound (see :func:`get_transport`). Requests are logged by pyqwest itself on the ``pyqwest.access`` and ``pyqwest`` loggers at ``DEBUG`` (off unless enabled) — the transport-level diagnostics httpcore used to provide. The SDK's own ``logger`` option is separate and sits above this, on the httpx client.""" - return ConnectionRetryTransport( - SyncHTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the generated - # clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) + key = (proxy, read_timeout) + with _transport_lock: + pool = _pools.get(key) + if pool is None: + pool = SyncHTTPTransport( + tls_include_system_certs=True, + proxy=proxy.to_pyqwest() if proxy is not None else None, + pool_idle_timeout=pool_idle_timeout, + pool_max_idle_per_host=pool_max_idle_per_host, + read_timeout=read_timeout, + # Redirects belong to the httpx client above (which the + # generated clients leave off), not to reqwest. + follow_redirects=False, + ) + _pools[key] = pool + return pool -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy; None is the direct pool. -# pyqwest transports are thread-safe, so unlike the httpx transports they -# replaced, the caches are process-global rather than per-thread. -_transports: Dict[Optional[ProxyConfig], PyqwestTransport] = {} +def get_pyqwest_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> ConnectionRetryTransport: + """The shared pool of :func:`get_pool` behind connect-only retries, for + consumers speaking pyqwest natively (the envd RPC clients). Consumers + speaking httpx take :func:`get_httpx_transport`, the adapter over the very + same stack.""" + key = (proxy, read_timeout) + # Resolve the pool before taking the lock: it takes the same one. + pool = get_pool(proxy, read_timeout) + with _transport_lock: + transport = _transports.get(key) + if transport is None: + transport = ConnectionRetryTransport(pool, max_retries=connection_retries) + _transports[key] = transport + return transport -def get_transport(config: ConnectionConfig) -> PyqwestTransport: - """The shared pyqwest-backed httpx transport for REST API calls. For TLS - connections ALPN negotiates the HTTP version (HTTP/2 against the E2B - API), like the http2-enabled httpx transport this replaced.""" - proxy = proxy_to_config(config.proxy) +def get_httpx_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> PyqwestTransport: + """The httpx adapter over the retrying stack of + :func:`get_pyqwest_transport`, for the generated httpx clients (control + plane, envd HTTP API, volume content). The adapter holds no state of its + own and does not close the pool, so closing an httpx client leaves the + pool intact for the other clients on it.""" + key = (proxy, read_timeout) + stack = get_pyqwest_transport(proxy, read_timeout) with _transport_lock: - transport = _transports.get(proxy) + transport = _httpx_transports.get(key) if transport is None: - transport = PyqwestTransport(retrying_http_transport(proxy)) - _transports[proxy] = transport + transport = PyqwestTransport(stack) + _httpx_transports[key] = transport return transport -# One transport per (proxy, streaming) pair, separate from the REST API -# pools — envd traffic goes to per-sandbox hosts. -_envd_transports: Dict[Tuple[Optional[ProxyConfig], bool], PyqwestTransport] = {} +def get_upload_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> PyqwestTransport: + """The httpx adapter over the *same pool* as :func:`get_httpx_transport` + but without the retry layer, for file uploads: envd ``files.write``, volume + ``write_file`` and template context uploads. + + Those are the requests whose body is streamed rather than held in memory, + and the retry middleware makes a body replayable by copying it in full as + it is sent — so uploading a file on the retrying stack would mirror the + whole file in RAM (SDK-332). Skipping the layer, not the pool, keeps the + connection reuse: an upload still travels the sandbox's or the volume + host's pooled connections. What it gives up is the connect retry, which + fires before any of the body was written, so the caller sees the connection + error and can retry the upload itself.""" + key = (proxy, read_timeout) + pool = get_pool(proxy, read_timeout) + with _transport_lock: + transport = _upload_transports.get(key) + if transport is None: + transport = PyqwestTransport(pool) + _upload_transports[key] = transport + return transport -def get_envd_transport( - config: ConnectionConfig, *, for_streaming: bool = False +def get_transport( + config: ConnectionConfig, *, for_streaming: bool = False, for_upload: bool = False ) -> PyqwestTransport: - """The shared pyqwest-backed httpx transports for the envd HTTP API - (file transfers, health checks). - - The streaming transport carries ``read_timeout``, the idle bound on - every read: it resets after each successful read, so it caps how long a - streamed download may stall without limiting total transfer time. It is - fixed per transport — the adapter's per-request timeouts are - whole-request deadlines rather than idle bounds. Only streamed downloads - use it: reqwest's read timer keeps - running while a request body is sent and while waiting for the response - head, so on the regular transport it would cut off uploads and slow - unary responses longer than the idle bound (those stay bounded by their - whole-request deadlines instead). + """The shared httpx transport for the control-plane REST API and the envd + HTTP API (file transfers, health checks) — one pool serves both, keyed by + the connection's proxy. + + ``for_streaming`` selects the pool carrying ``READ_TIMEOUT``, the idle + bound on every read: it resets after each successful read, so it caps how + long a streamed download may stall without limiting total transfer time. + It is fixed per pool — the adapter's per-request timeouts are + whole-request deadlines rather than idle bounds — so only streamed + downloads take it, and they get their own pool (see :func:`get_pool`). + + ``for_upload`` selects the same pool without the retry layer, so that a + streamed upload is not copied into memory to keep it replayable + (see :func:`get_upload_transport`). """ proxy = proxy_to_config(config.proxy) - key = (proxy, for_streaming) - with _transport_lock: - transport = _envd_transports.get(key) - if transport is None: - transport = PyqwestTransport( - retrying_http_transport( - proxy, - read_timeout=READ_TIMEOUT if for_streaming else None, - ) - ) - _envd_transports[key] = transport - return transport + read_timeout = READ_TIMEOUT if for_streaming else None + if for_upload: + return get_upload_transport(proxy, read_timeout) + return get_httpx_transport(proxy, read_timeout) def get_envd_api( - config: ConnectionConfig, base_url: str, *, for_streaming: bool = False + config: ConnectionConfig, + base_url: str, + *, + for_streaming: bool = False, + for_upload: bool = False, ) -> httpx.Client: """An httpx client for a sandbox's envd HTTP API (file transfers, health - checks) on the shared pyqwest transports. The client itself is a cheap - stateless wrapper — one per consumer is fine — while the pooled transport - underneath is shared and thread-safe.""" + checks) on the shared transports. The client itself is a cheap stateless + wrapper — one per consumer is fine — while the pooled transport underneath + is shared and thread-safe.""" return httpx.Client( base_url=base_url, - transport=get_envd_transport(config, for_streaming=for_streaming), + transport=get_transport( + config, for_streaming=for_streaming, for_upload=for_upload + ), headers=config.sandbox_headers, event_hooks=make_logging_event_hooks(config.logger), ) diff --git a/packages/python-sdk/e2b/envd/client_async/__init__.py b/packages/python-sdk/e2b/envd/client_async/__init__.py index 36a19bc1ac..dcc979d730 100644 --- a/packages/python-sdk/e2b/envd/client_async/__init__.py +++ b/packages/python-sdk/e2b/envd/client_async/__init__.py @@ -1,7 +1,6 @@ -"""Async envd RPC clients: shared pyqwest transports and client factory.""" +"""Async envd RPC clients: the plain-error transport layer and client factory.""" import asyncio -import threading from typing import ( Any, AsyncGenerator, @@ -16,8 +15,8 @@ from connectrpc.errors import ConnectError from pyqwest import Client, Request, Response, Transport -from e2b.api import ProxyConfig, proxy_to_config -from e2b.api.client_async import retrying_http_transport +from e2b.api import proxy_to_config +from e2b.api.client_async import get_pyqwest_transport from e2b.connection_config import ConnectionConfig from e2b.envd.client_shared import ( ENVD_JSON_CODEC, @@ -30,10 +29,6 @@ RES = TypeVar("RES") TClient = TypeVar("TClient") -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy; None is the direct pool. -_transports: dict[Optional[ProxyConfig], "PlainHTTPErrorTransport"] = {} - class PlainHTTPErrorTransport: """Raise plain (non-Connect-encoded) HTTP error responses — e.g. an edge @@ -65,19 +60,6 @@ async def execute(self, request: Request) -> Response: raise error -def get_transport(proxy: Optional[ProxyConfig]) -> "PlainHTTPErrorTransport": - with _transport_lock: - transport = _transports.get(proxy) - if transport is None: - # connectrpc arms the per-call deadline around the transport, so - # retry backoff counts against the request timeout. The plain- - # error normalization sits outside the retries so it converts - # the settled response once. - transport = PlainHTTPErrorTransport(retrying_http_transport(proxy)) - _transports[proxy] = transport - return transport - - def create_rpc_client( client_cls: Callable[..., TClient], base_url: str, @@ -88,8 +70,18 @@ def create_rpc_client( see :class:`e2b.api.client_async.ConnectionRetryTransport`), the envd JSON codec, and the SDK's default-header and logging interceptors. Compression is disabled (see ``ENVD_RPC_COMPRESSION``). + + The plain-error normalization is the one RPC-only transport concern, so it + wraps the shared pool per client instead of being cached with it — a + stateless wrapper over the pool the envd HTTP API uses for the same + sandbox, which is what lets both share a single HTTP/2 connection. + connectrpc arms the per-call deadline around the transport, so retry + backoff counts against the request timeout, and the normalization sits + outside the retries so it converts the settled response once. """ - http_client = Client(get_transport(proxy_to_config(config.proxy))) + http_client = Client( + PlainHTTPErrorTransport(get_pyqwest_transport(proxy_to_config(config.proxy))) + ) return client_cls( base_url, codec=ENVD_JSON_CODEC, diff --git a/packages/python-sdk/e2b/envd/client_shared.py b/packages/python-sdk/e2b/envd/client_shared.py index aa152cd1d6..0b7f1f9dfa 100644 --- a/packages/python-sdk/e2b/envd/client_shared.py +++ b/packages/python-sdk/e2b/envd/client_shared.py @@ -1,14 +1,15 @@ """envd RPC client plumbing shared by the sync and async flavors. The envd RPC clients (process, filesystem) run on `connectrpc`, whose HTTP -layer is `pyqwest` (Rust reqwest/hyper) — built on the same -`retrying_http_transport` pieces as the REST API client in `e2b.api`, in a -separately cached pool. Only the multipart file transfer endpoints stay on -the `httpx` envd transports. Unlike the previous httpcore-based transport, -hyper sends RST_STREAM when a server stream is closed early, so abandoned -command/watch streams don't leak on the shared HTTP/2 connection. - -The flavor-specific transports and client factories live in +layer is `pyqwest` (Rust reqwest/hyper) — on the very connection pool the REST +clients in `e2b.api` use (`get_pyqwest_transport`), so RPCs and the envd HTTP +API share one HTTP/2 connection per sandbox. Only the multipart file transfer +endpoints stay on the `httpx` side of that pool. Unlike the previous +httpcore-based transport, hyper sends RST_STREAM when a server stream is +closed early, so abandoned command/watch streams don't leak on the shared +HTTP/2 connection. + +The flavor-specific transport layer and client factories live in :mod:`e2b.envd.client_sync` and :mod:`e2b.envd.client_async`, mirroring the `e2b.api.client_sync`/`client_async` layout. This module exists because `e2b/envd/__init__.py`, the natural home by that analogy, is owned by the diff --git a/packages/python-sdk/e2b/envd/client_sync/__init__.py b/packages/python-sdk/e2b/envd/client_sync/__init__.py index fc10cf9881..fadf66323f 100644 --- a/packages/python-sdk/e2b/envd/client_sync/__init__.py +++ b/packages/python-sdk/e2b/envd/client_sync/__init__.py @@ -1,7 +1,6 @@ -"""Sync envd RPC clients: shared pyqwest transports and client factory.""" +"""Sync envd RPC clients: the plain-error transport layer and client factory.""" -import threading -from typing import Any, Callable, Generator, Iterator, Optional, TypeVar, cast +from typing import Any, Callable, Generator, Iterator, TypeVar, cast from pyqwest import ( SyncClient, @@ -10,8 +9,8 @@ SyncTransport, ) -from e2b.api import ProxyConfig, proxy_to_config -from e2b.api.client_sync import retrying_http_transport +from e2b.api import proxy_to_config +from e2b.api.client_sync import get_pyqwest_transport from e2b.connection_config import ConnectionConfig from e2b.envd.client_shared import ( ENVD_JSON_CODEC, @@ -23,10 +22,6 @@ RES = TypeVar("RES") TClient = TypeVar("TClient") -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy; None is the direct pool. -_transports: dict[Optional[ProxyConfig], "PlainHTTPErrorTransport"] = {} - class PlainHTTPErrorTransport: """Raise plain (non-Connect-encoded) HTTP error responses — e.g. an edge @@ -58,19 +53,6 @@ def execute_sync(self, request: SyncRequest) -> SyncResponse: raise error -def get_transport(proxy: Optional[ProxyConfig]) -> "PlainHTTPErrorTransport": - with _transport_lock: - transport = _transports.get(proxy) - if transport is None: - # connectrpc arms the per-call deadline around the transport, so - # retry backoff counts against the request timeout. The plain- - # error normalization sits outside the retries so it converts - # the settled response once. - transport = PlainHTTPErrorTransport(retrying_http_transport(proxy)) - _transports[proxy] = transport - return transport - - def create_rpc_client( client_cls: Callable[..., TClient], base_url: str, @@ -81,10 +63,20 @@ def create_rpc_client( see :class:`e2b.api.client_sync.ConnectionRetryTransport`), the envd JSON codec, and the SDK's default-header and logging interceptors. Compression is disabled (see ``ENVD_RPC_COMPRESSION``). The client is stateless per - call and its transport is process-global, so one instance serves all + call and its connection pool is process-global, so one instance serves all threads. + + The plain-error normalization is the one RPC-only transport concern, so it + wraps the shared pool per client instead of being cached with it — a + stateless wrapper over the pool the envd HTTP API uses for the same + sandbox, which is what lets both share a single HTTP/2 connection. + connectrpc arms the per-call deadline around the transport, so retry + backoff counts against the request timeout, and the normalization sits + outside the retries so it converts the settled response once. """ - http_client = SyncClient(get_transport(proxy_to_config(config.proxy))) + http_client = SyncClient( + PlainHTTPErrorTransport(get_pyqwest_transport(proxy_to_config(config.proxy))) + ) return client_cls( base_url, codec=ENVD_JSON_CODEC, diff --git a/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py b/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py index c36a21c86f..838f050b04 100644 --- a/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py +++ b/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py @@ -101,10 +101,17 @@ def __init__( ) self._envd_api = envd_api # Streamed downloads default to a sibling client whose transport - # carries the idle read timeout (see `get_envd_transport`). + # carries the idle read timeout (see `get_transport`). self._envd_api_streaming = get_envd_api( connection_config, envd_api_url, for_streaming=True ) + # Uploads go out on a third sibling, on the same pool without the + # connect retries: replaying a request body means copying it, so a + # streamed upload on the retrying transport would be mirrored in + # memory in full (see `get_upload_transport`). + self._envd_api_upload = get_envd_api( + connection_config, envd_api_url, for_upload=True + ) @overload async def read( @@ -218,7 +225,7 @@ async def read( # sent only when the caller set `request_timeout` explicitly # (making it the total-transfer deadline). By default a stalled # stream is bounded by the streaming transport's idle read - # timeout (see `get_envd_transport`); an explicit + # timeout (see `get_transport`); an explicit # `stream_idle_timeout` is applied per read with `wait_for` on # the regular transport instead — so values above the transport # bound aren't capped by it and `0` disables idle bounding @@ -405,8 +412,12 @@ async def _upload_file(file): headers["Content-Encoding"] = "gzip" is_streamed = not isinstance(file_data, (str, bytes)) + # Only a streamed body takes the non-retrying client: an + # in-memory one reaches the transport as `bytes`, which the + # retry layer replays without copying anything. + client = self._envd_api_upload if is_streamed else self._envd_api try: - r = await self._envd_api.post( + r = await client.post( ENVD_API_FILES_ROUTE, content=to_upload_body_async(file_data, gzip), headers=headers, @@ -448,19 +459,20 @@ async def _upload_file(file): if len(httpx_files) == 0: return [] + is_streamed = multipart_body_is_streamed(files) + client = self._envd_api_upload if is_streamed else self._envd_api try: - r = await self._envd_api.post( + r = await client.post( ENVD_API_FILES_ROUTE, files=httpx_files, params=params, headers=extra_headers, - # Only a streamed entry drops the deadline: httpx - # forwards binary `IOBase` entries in chunks, while text - # file-like data was buffered by `_to_httpx_file` (httpx - # rejects text-mode objects in multipart). - timeout=( - None if multipart_body_is_streamed(files) else upload_timeout - ), + # Only a streamed entry drops the deadline (and the retries, + # whose replay copy would mirror the body): httpx forwards + # binary `IOBase` entries in chunks, while text file-like + # data was buffered by `_to_httpx_file` (httpx rejects + # text-mode objects in multipart). + timeout=None if is_streamed else upload_timeout, ) except httpx.RemoteProtocolError as e: raise await ahandle_envd_api_transport_exception_with_health( diff --git a/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py b/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py index f249c98957..95e975465a 100644 --- a/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py +++ b/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py @@ -95,12 +95,19 @@ def __init__( ) self._envd_api = envd_api # Streamed downloads default to a sibling client whose transport - # carries the idle read timeout (see `get_envd_transport`). Like the + # carries the idle read timeout (see `get_transport`). Like the # RPC client, the pyqwest transports underneath are thread-safe, so # one client serves all threads. self._envd_api_streaming = get_envd_api( connection_config, envd_api_url, for_streaming=True ) + # Uploads go out on a third sibling, on the same pool without the + # connect retries: replaying a request body means copying it, so a + # streamed upload on the retrying transport would be mirrored in + # memory in full (see `get_upload_transport`). + self._envd_api_upload = get_envd_api( + connection_config, envd_api_url, for_upload=True + ) @overload def read( @@ -211,7 +218,7 @@ def read( # sent only when the caller set `request_timeout` explicitly # (making it the total-transfer deadline). A stalled stream is # instead bounded by the streaming transport's idle read timeout - # (see `get_envd_transport`), which resets on every chunk. + # (see `get_transport`), which resets on every chunk. stream_timeout = ConnectionConfig._get_request_timeout( None, request_timeout ) @@ -375,8 +382,12 @@ def write_files( headers["Content-Encoding"] = "gzip" is_streamed = not isinstance(file_data, (str, bytes)) + # Only a streamed body takes the non-retrying client: an + # in-memory one reaches the transport as `bytes`, which the + # retry layer replays without copying anything. + client = self._envd_api_upload if is_streamed else self._envd_api try: - r = self._envd_api.post( + r = client.post( ENVD_API_FILES_ROUTE, content=to_upload_body(file_data, gzip), headers=headers, @@ -412,19 +423,20 @@ def write_files( if len(httpx_files) == 0: return [] + is_streamed = multipart_body_is_streamed(files) + client = self._envd_api_upload if is_streamed else self._envd_api try: - r = self._envd_api.post( + r = client.post( ENVD_API_FILES_ROUTE, files=httpx_files, params=params, headers=extra_headers, - # Only a streamed entry drops the deadline: httpx - # forwards binary `IOBase` entries in chunks, while text - # file-like data was buffered by `_to_httpx_file` (httpx - # rejects text-mode objects in multipart). - timeout=( - None if multipart_body_is_streamed(files) else upload_timeout - ), + # Only a streamed entry drops the deadline (and the retries, + # whose replay copy would mirror the body): httpx forwards + # binary `IOBase` entries in chunks, while text file-like + # data was buffered by `_to_httpx_file` (httpx rejects + # text-mode objects in multipart). + timeout=None if is_streamed else upload_timeout, ) except httpx.RemoteProtocolError as e: raise handle_envd_api_transport_exception_with_health(e, self._envd_api) diff --git a/packages/python-sdk/e2b/template_async/build_api.py b/packages/python-sdk/e2b/template_async/build_api.py index 3d8c8f774e..c1890628f9 100644 --- a/packages/python-sdk/e2b/template_async/build_api.py +++ b/packages/python-sdk/e2b/template_async/build_api.py @@ -4,10 +4,9 @@ from typing import Callable, Optional, List, Union import httpx -from pyqwest import HTTPTransport -from pyqwest.httpx import AsyncPyqwestTransport from e2b.api import handle_api_exception, proxy_to_config +from e2b.api.client_async import get_upload_transport from e2b.io_utils import aiter_io_chunks from e2b.api.client.api.templates import ( post_v3_templates, @@ -128,25 +127,18 @@ async def upload_file( try: size = os.fstat(tar_file.fileno()).st_size + # The archive goes out on the shared connection pool, minus the + # connect retries: those would make the streamed body replayable by + # copying it, mirroring the whole build context in memory + # (SDK-332). Sharing the pool is what makes a per-build transport + # unnecessary. # Through the pyqwest adapter the upload timeout is a # whole-request deadline for the entire transfer, not a per-write # bound as with the httpx transport this replaced. async with httpx.AsyncClient( timeout=httpx.Timeout(upload_timeout), follow_redirects=api_client._follow_redirects, - transport=AsyncPyqwestTransport( - HTTPTransport( - tls_include_system_certs=True, - proxy=( - upload_proxy.to_pyqwest() - if upload_proxy is not None - else None - ), - # Redirects belong to the httpx client above, not to - # reqwest. - follow_redirects=False, - ) - ), + transport=get_upload_transport(upload_proxy), ) as client: # Stream the archive from disk via an async iterator. The # explicit Content-Length suppresses chunked transfer diff --git a/packages/python-sdk/e2b/template_sync/build_api.py b/packages/python-sdk/e2b/template_sync/build_api.py index 0525fa4570..5f9b9dc4dc 100644 --- a/packages/python-sdk/e2b/template_sync/build_api.py +++ b/packages/python-sdk/e2b/template_sync/build_api.py @@ -3,10 +3,9 @@ from typing import Callable, Optional, List, Union import httpx -from pyqwest import SyncHTTPTransport -from pyqwest.httpx import PyqwestTransport from e2b.api import handle_api_exception, proxy_to_config +from e2b.api.client_sync import get_upload_transport from e2b.api.client.api.templates import ( post_v3_templates, get_templates_template_id_files_hash, @@ -124,25 +123,18 @@ def upload_file( file_name, context_path, ignore_patterns, resolve_symlinks, gzip ) try: + # The archive goes out on the shared connection pool, minus the + # connect retries: those would make the streamed body replayable by + # copying it, mirroring the whole build context in memory + # (SDK-332). Sharing the pool is what makes a per-build transport + # unnecessary. # Through the pyqwest adapter the upload timeout is a # whole-request deadline for the entire transfer, not a per-write # bound as with the httpx transport this replaced. with httpx.Client( timeout=httpx.Timeout(upload_timeout), follow_redirects=api_client._follow_redirects, - transport=PyqwestTransport( - SyncHTTPTransport( - tls_include_system_certs=True, - proxy=( - upload_proxy.to_pyqwest() - if upload_proxy is not None - else None - ), - # Redirects belong to the httpx client above, not to - # reqwest. - follow_redirects=False, - ) - ), + transport=get_upload_transport(upload_proxy), ) as client: # httpx streams the archive from disk in chunks and sets # Content-Length from the file size—S3 presigned URLs reject diff --git a/packages/python-sdk/e2b/volume/client_async/__init__.py b/packages/python-sdk/e2b/volume/client_async/__init__.py index 0416c32a4b..ba0164c78f 100644 --- a/packages/python-sdk/e2b/volume/client_async/__init__.py +++ b/packages/python-sdk/e2b/volume/client_async/__init__.py @@ -1,19 +1,12 @@ -import threading -from typing import Dict, Optional, Tuple - import httpx -from pyqwest import HTTPTransport from pyqwest.httpx import AsyncPyqwestTransport from e2b.api import ( - ProxyConfig, - connection_retries, make_async_logging_event_hooks, - pool_idle_timeout, - pool_max_idle_per_host, proxy_to_config, ) -from e2b.api.client_async import ConnectionRetryTransport +from e2b.api.client_async import get_httpx_transport +from e2b.api.client_async import get_upload_transport as get_upload_httpx_transport from e2b.api.metadata import default_headers from e2b.exceptions import AuthenticationException from e2b.volume.client.client import AuthenticatedClient as AsyncVolumeApiClient @@ -33,6 +26,15 @@ def get_streaming_api_client( return _api_client(config, get_streaming_transport(config), **kwargs) +def get_upload_api_client( + config: VolumeConnectionConfig, **kwargs +) -> AsyncVolumeApiClient: + """The client for uploads: the same client on the upload transport, which + leaves the connect retries out so that a streamed body isn't copied into + memory to keep it replayable (see :func:`get_upload_transport`).""" + return _api_client(config, get_upload_transport(config), **kwargs) + + def _api_client( config: VolumeConnectionConfig, transport: AsyncPyqwestTransport, **kwargs ) -> AsyncVolumeApiClient: @@ -69,18 +71,11 @@ def _api_client( ) -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy and read timeout; None is -# the direct pool. pyqwest's I/O runs on its own Rust runtime, so unlike the -# httpx transport this replaced, the transport is not bound to an event loop -# and the cache is process-global rather than per-loop. -_transports: Dict[ - Tuple[Optional[ProxyConfig], Optional[float]], AsyncPyqwestTransport -] = {} - - def get_transport(config: VolumeConnectionConfig) -> AsyncPyqwestTransport: - """The shared pyqwest-backed httpx transport for volume content API calls. + """The shared pyqwest-backed httpx transport for volume content API calls — + the same pool the control-plane REST API and the envd HTTP API draw from + (see :func:`e2b.api.client_async.get_pyqwest_transport`); reqwest pools per + host, so the volume host gets its own connections within it. It carries no idle read bound: reqwest's read timer keeps running while a request body is sent and while waiting for the response head, so one here @@ -88,7 +83,16 @@ def get_transport(config: VolumeConnectionConfig) -> AsyncPyqwestTransport: their whole-request deadlines instead). Streamed downloads, which do need an idle bound, use :func:`get_streaming_transport`. """ - return _transport(config, read_timeout=None) + return get_httpx_transport(proxy_to_config(config.proxy)) + + +def get_upload_transport(config: VolumeConnectionConfig) -> AsyncPyqwestTransport: + """The transport for uploads: the same pool as :func:`get_transport` minus + the connect retries. The retry layer makes a request body replayable by + copying it in full as it is sent, so a streamed `write_file` on it would be + mirrored in memory (see :func:`e2b.api.client_async.get_upload_transport`). + """ + return get_upload_httpx_transport(proxy_to_config(config.proxy)) def get_streaming_transport( @@ -97,34 +101,9 @@ def get_streaming_transport( """The transport for streamed downloads, carrying ``READ_TIMEOUT`` as the idle bound on every read: it resets after each successful read, so it caps how long a streamed download may stall without limiting total transfer - time. It is fixed per transport — the adapter's per-request timeouts are - whole-request deadlines rather than idle bounds. + time. It is fixed per pool — the adapter's per-request timeouts are + whole-request deadlines rather than idle bounds — so streamed downloads get + their own, shared with the sandbox filesystem's streaming transport + whenever the two bounds agree. """ - return _transport(config, read_timeout=READ_TIMEOUT) - - -def _transport( - config: VolumeConnectionConfig, *, read_timeout: Optional[float] -) -> AsyncPyqwestTransport: - proxy = proxy_to_config(config.proxy) - key = (proxy, read_timeout) - with _transport_lock: - transport = _transports.get(key) - if transport is None: - transport = AsyncPyqwestTransport( - ConnectionRetryTransport( - HTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the - # generated clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) - ) - _transports[key] = transport - return transport + return get_httpx_transport(proxy_to_config(config.proxy), READ_TIMEOUT) diff --git a/packages/python-sdk/e2b/volume/client_sync/__init__.py b/packages/python-sdk/e2b/volume/client_sync/__init__.py index adbe9c3e72..d8d1d78863 100644 --- a/packages/python-sdk/e2b/volume/client_sync/__init__.py +++ b/packages/python-sdk/e2b/volume/client_sync/__init__.py @@ -1,19 +1,12 @@ -import threading -from typing import Dict, Optional, Tuple - import httpx -from pyqwest import SyncHTTPTransport from pyqwest.httpx import PyqwestTransport from e2b.api import ( - ProxyConfig, - connection_retries, make_logging_event_hooks, - pool_idle_timeout, - pool_max_idle_per_host, proxy_to_config, ) -from e2b.api.client_sync import ConnectionRetryTransport +from e2b.api.client_sync import get_httpx_transport +from e2b.api.client_sync import get_upload_transport as get_upload_httpx_transport from e2b.api.metadata import default_headers from e2b.exceptions import AuthenticationException from e2b.volume.client.client import AuthenticatedClient as VolumeApiClient @@ -33,6 +26,13 @@ def get_streaming_api_client( return _api_client(config, get_streaming_transport(config), **kwargs) +def get_upload_api_client(config: VolumeConnectionConfig, **kwargs) -> VolumeApiClient: + """The client for uploads: the same client on the upload transport, which + leaves the connect retries out so that a streamed body isn't copied into + memory to keep it replayable (see :func:`get_upload_transport`).""" + return _api_client(config, get_upload_transport(config), **kwargs) + + def _api_client( config: VolumeConnectionConfig, transport: PyqwestTransport, **kwargs ) -> VolumeApiClient: @@ -69,15 +69,11 @@ def _api_client( ) -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy and read timeout; None is -# the direct pool. pyqwest transports are thread-safe, so unlike the httpx -# transport this replaced, the cache is process-global rather than per-thread. -_transports: Dict[Tuple[Optional[ProxyConfig], Optional[float]], PyqwestTransport] = {} - - def get_transport(config: VolumeConnectionConfig) -> PyqwestTransport: - """The shared pyqwest-backed httpx transport for volume content API calls. + """The shared pyqwest-backed httpx transport for volume content API calls — + the same pool the control-plane REST API and the envd HTTP API draw from + (see :func:`e2b.api.client_sync.get_pyqwest_transport`); reqwest pools per + host, so the volume host gets its own connections within it. It carries no idle read bound: reqwest's read timer keeps running while a request body is sent and while waiting for the response head, so one here @@ -85,41 +81,25 @@ def get_transport(config: VolumeConnectionConfig) -> PyqwestTransport: their whole-request deadlines instead). Streamed downloads, which do need an idle bound, use :func:`get_streaming_transport`. """ - return _transport(config, read_timeout=None) + return get_httpx_transport(proxy_to_config(config.proxy)) + + +def get_upload_transport(config: VolumeConnectionConfig) -> PyqwestTransport: + """The transport for uploads: the same pool as :func:`get_transport` minus + the connect retries. The retry layer makes a request body replayable by + copying it in full as it is sent, so a streamed `write_file` on it would be + mirrored in memory (see :func:`e2b.api.client_sync.get_upload_transport`). + """ + return get_upload_httpx_transport(proxy_to_config(config.proxy)) def get_streaming_transport(config: VolumeConnectionConfig) -> PyqwestTransport: """The transport for streamed downloads, carrying ``READ_TIMEOUT`` as the idle bound on every read: it resets after each successful read, so it caps how long a streamed download may stall without limiting total transfer - time. It is fixed per transport — the adapter's per-request timeouts are - whole-request deadlines rather than idle bounds. + time. It is fixed per pool — the adapter's per-request timeouts are + whole-request deadlines rather than idle bounds — so streamed downloads get + their own, shared with the sandbox filesystem's streaming transport + whenever the two bounds agree. """ - return _transport(config, read_timeout=READ_TIMEOUT) - - -def _transport( - config: VolumeConnectionConfig, *, read_timeout: Optional[float] -) -> PyqwestTransport: - proxy = proxy_to_config(config.proxy) - key = (proxy, read_timeout) - with _transport_lock: - transport = _transports.get(key) - if transport is None: - transport = PyqwestTransport( - ConnectionRetryTransport( - SyncHTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the - # generated clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) - ) - _transports[key] = transport - return transport + return get_httpx_transport(proxy_to_config(config.proxy), READ_TIMEOUT) diff --git a/packages/python-sdk/e2b/volume/volume_async.py b/packages/python-sdk/e2b/volume/volume_async.py index 7b7c7b62b0..22468b0338 100644 --- a/packages/python-sdk/e2b/volume/volume_async.py +++ b/packages/python-sdk/e2b/volume/volume_async.py @@ -39,6 +39,7 @@ from e2b.volume.client_async import get_api_client as get_volume_api_client from e2b.volume.client_async import ( get_streaming_api_client as get_streaming_volume_api_client, + get_upload_api_client as get_upload_volume_api_client, ) from e2b.volume.connection_config import ( VolumeApiParams, @@ -621,7 +622,10 @@ async def write_file( upload_timeout = VolumeConnectionConfig._get_request_timeout( FILE_TIMEOUT, opts.get("request_timeout") ) - api_client = get_volume_api_client(config) + # An upload streams its body, and the connect retries would make that + # body replayable by copying it in full — the whole file in memory + # (SDK-332) — so it goes out on the non-retrying client instead. + api_client = get_upload_volume_api_client(config) if upload_timeout is not None: api_client = api_client.with_timeout(httpx.Timeout(upload_timeout)) diff --git a/packages/python-sdk/e2b/volume/volume_sync.py b/packages/python-sdk/e2b/volume/volume_sync.py index 4415838fee..ada5b7156d 100644 --- a/packages/python-sdk/e2b/volume/volume_sync.py +++ b/packages/python-sdk/e2b/volume/volume_sync.py @@ -38,6 +38,7 @@ from e2b.volume.client_sync import get_api_client as get_volume_api_client from e2b.volume.client_sync import ( get_streaming_api_client as get_streaming_volume_api_client, + get_upload_api_client as get_upload_volume_api_client, ) from e2b.volume.connection_config import ( VolumeApiParams, @@ -589,7 +590,10 @@ def write_file( upload_timeout = VolumeConnectionConfig._get_request_timeout( FILE_TIMEOUT, opts.get("request_timeout") ) - api_client = get_volume_api_client(config) + # An upload streams its body, and the connect retries would make that + # body replayable by copying it in full — the whole file in memory + # (SDK-332) — so it goes out on the non-retrying client instead. + api_client = get_upload_volume_api_client(config) if upload_timeout is not None: api_client = api_client.with_timeout(httpx.Timeout(upload_timeout)) diff --git a/packages/python-sdk/tests/async/template_async/test_upload_file.py b/packages/python-sdk/tests/async/template_async/test_upload_file.py index a8d2993427..8f95f7c63a 100644 --- a/packages/python-sdk/tests/async/template_async/test_upload_file.py +++ b/packages/python-sdk/tests/async/template_async/test_upload_file.py @@ -7,6 +7,7 @@ import pytest from e2b.api.client.client import AuthenticatedClient +from e2b.api.client_async import get_upload_transport from e2b.template import utils as template_utils from e2b.exceptions import FileUploadException from e2b.template.consts import FILE_UPLOAD_TIMEOUT_SECONDS @@ -128,8 +129,9 @@ async def test_upload_file_leaves_redirects_to_httpx(tmp_path): assert state["paths"] == ["/redirect"] -async def _capture_upload_timeout(tmp_path, request_timeout=None): - """Run upload_file against a local server, capturing the httpx timeout.""" +async def _capture_upload_client_kwargs(tmp_path, request_timeout=None): + """Run upload_file against a local server, capturing how it built its + httpx client.""" (tmp_path / "hello.txt").write_text("hello world") server, thread, state = _make_server() @@ -140,7 +142,7 @@ async def _capture_upload_timeout(tmp_path, request_timeout=None): real_client = httpx.AsyncClient def spy_client(*args, **kwargs): - captured["timeout"] = kwargs.get("timeout") + captured.update(kwargs) return real_client(*args, **kwargs) try: @@ -167,20 +169,29 @@ def spy_client(*args, **kwargs): server.server_close() thread.join(timeout=5) - return captured["timeout"] + return captured async def test_upload_file_defaults_to_one_hour_timeout(tmp_path): # Uploads of large archives need far longer than the 60s general API # timeout, so the default upload timeout is 1 hour (matches the JS SDK). - timeout = await _capture_upload_timeout(tmp_path) - assert timeout == httpx.Timeout(FILE_UPLOAD_TIMEOUT_SECONDS) + captured = await _capture_upload_client_kwargs(tmp_path) + assert captured["timeout"] == httpx.Timeout(FILE_UPLOAD_TIMEOUT_SECONDS) async def test_upload_file_honors_explicit_request_timeout(tmp_path): # An explicitly set request_timeout overrides the 1-hour upload default. - timeout = await _capture_upload_timeout(tmp_path, request_timeout=5.0) - assert timeout == httpx.Timeout(5.0) + captured = await _capture_upload_client_kwargs(tmp_path, request_timeout=5.0) + assert captured["timeout"] == httpx.Timeout(5.0) + + +async def test_upload_file_uses_the_shared_pool_without_retries(tmp_path): + # The archive rides the SDK-wide connection pool instead of a transport + # built per build, but skips its retry layer: replaying a request means + # copying its body, which for a build context is the whole archive in + # memory (SDK-332). + captured = await _capture_upload_client_kwargs(tmp_path) + assert captured["transport"] is get_upload_transport(None) async def test_upload_file_ignores_post_upload_close_failure(tmp_path): diff --git a/packages/python-sdk/tests/envd_frame_server.py b/packages/python-sdk/tests/envd_frame_server.py index df7f33447b..dfa3b6c514 100644 --- a/packages/python-sdk/tests/envd_frame_server.py +++ b/packages/python-sdk/tests/envd_frame_server.py @@ -16,10 +16,11 @@ import struct import threading import time -from typing import Iterator, Optional +from typing import Iterator, List, Optional import h2.config import h2.connection +import h2.errors import h2.events from protobuf import Oneof from pyqwest import ( @@ -33,6 +34,7 @@ ) from e2b.connection_config import ConnectionConfig +from e2b.envd.api import ENVD_API_HEALTH_ROUTE from e2b.envd.client_shared import ENVD_JSON_CODEC, ENVD_RPC_COMPRESSION from e2b.envd.interceptors import build_interceptors from e2b.envd.process.process_connect import ProcessClient, ProcessClientSync @@ -177,6 +179,138 @@ def frame_recording_server( server.listener.close() +class SharedPoolServer(threading.Thread): + """Multi-connection plaintext HTTP/2 server for the shared-pool tests. + + Serves both sides of a sandbox's traffic — the process ``connect`` server + stream and the envd HTTP ``/health`` route — so a single connection pool + can be pointed at it the way the SDK points one at a real sandbox. After + the first stream event it breaks the RPC the way a sandbox going away + does: + + * ``fault="reset"`` sends ``RST_STREAM``, which kills the RPC stream and + leaves the HTTP/2 connection healthy; + * ``fault="drop"`` tears the whole TCP connection down with a RST. + + ``/health`` is always answered 200, so a probe that comes back anything + but ``True`` failed at the transport layer. ``connections`` counts the + accepted TCP connections, which is what tells reuse from a redial. + """ + + def __init__(self, fault: str): + super().__init__(daemon=True) + self.fault = fault + self.listener = socket.create_server(("127.0.0.1", 0)) + self.listener.settimeout(10) + self.port = self.listener.getsockname()[1] + self.connections: List[socket.socket] = [] + self.paths: List[str] = [] + self.errors: List[str] = [] + self._lock = threading.Lock() + + def run(self): + while True: + try: + sock, _ = self.listener.accept() + except (OSError, socket.timeout): + # The listener was closed by the context manager, or nothing + # else connected — either way there is nothing left to serve. + return + with self._lock: + self.connections.append(sock) + threading.Thread(target=self._serve, args=(sock,), daemon=True).start() + + def _serve(self, sock: socket.socket): + sock.settimeout(0.1) + conn = h2.connection.H2Connection( + config=h2.config.H2Configuration(client_side=False) + ) + conn.initiate_connection() + sock.sendall(conn.data_to_send()) + paths: dict[int, str] = {} + drop_connection = False + deadline = time.monotonic() + 10 + try: + while time.monotonic() < deadline: + try: + data = sock.recv(65535) + except socket.timeout: + continue + except OSError: + break + if not data: + break + for event in conn.receive_data(data): + if isinstance(event, h2.events.RequestReceived): + path = dict(event.headers).get(b":path", b"").decode() + paths[event.stream_id] = path + with self._lock: + self.paths.append(path) + elif isinstance(event, h2.events.StreamEnded): + # The client finished sending the request: respond. + if paths.get(event.stream_id, "") == ENVD_API_HEALTH_ROUTE: + body = b'{"version":"0.0.0"}' + conn.send_headers( + event.stream_id, + [ + (":status", "200"), + ("content-type", "application/json"), + ("content-length", str(len(body))), + ], + ) + conn.send_data(event.stream_id, body, end_stream=True) + continue + conn.send_headers( + event.stream_id, + [ + (":status", "200"), + ("content-type", "application/connect+json"), + ], + ) + conn.send_data(event.stream_id, event_envelope()) + if self.fault == "reset": + conn.reset_stream( + event.stream_id, + error_code=h2.errors.ErrorCodes.INTERNAL_ERROR, + ) + else: + drop_connection = True + elif isinstance(event, h2.events.DataReceived): + conn.acknowledge_received_data( + event.flow_controlled_length, event.stream_id + ) + elif isinstance(event, h2.events.ConnectionTerminated): + return + out = conn.data_to_send() + if out: + sock.sendall(out) + if drop_connection: + # RST the connection rather than closing it cleanly: a + # GOAWAY would tell the client to retire the connection, + # which is not what a sandbox disappearing looks like. + sock.setsockopt( + socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0) + ) + return + except Exception as e: # noqa: BLE001 — surfaced via assert_no_errors + self.errors.append(repr(e)) + finally: + sock.close() + + def assert_no_errors(self): + assert not self.errors, self.errors + + +@contextlib.contextmanager +def shared_pool_server(fault: str) -> Iterator[SharedPoolServer]: + server = SharedPoolServer(fault) + server.start() + try: + yield server + finally: + server.listener.close() + + def assert_stdout_event(event: ConnectResponse): assert event.event is not None match event.event.event: diff --git a/packages/python-sdk/tests/sync/template_sync/test_upload_file.py b/packages/python-sdk/tests/sync/template_sync/test_upload_file.py index a08ae9906a..70e22a3f59 100644 --- a/packages/python-sdk/tests/sync/template_sync/test_upload_file.py +++ b/packages/python-sdk/tests/sync/template_sync/test_upload_file.py @@ -7,6 +7,7 @@ import pytest from e2b.api.client.client import AuthenticatedClient +from e2b.api.client_sync import get_upload_transport from e2b.template import utils as template_utils from e2b.exceptions import FileUploadException from e2b.template.consts import FILE_UPLOAD_TIMEOUT_SECONDS @@ -124,8 +125,9 @@ def test_upload_file_leaves_redirects_to_httpx(tmp_path): assert state["paths"] == ["/redirect"] -def _capture_upload_timeout(tmp_path, request_timeout=None): - """Run upload_file against a local server, capturing the httpx timeout.""" +def _capture_upload_client_kwargs(tmp_path, request_timeout=None): + """Run upload_file against a local server, capturing how it built its + httpx client.""" (tmp_path / "hello.txt").write_text("hello world") server, thread, state = _make_server() @@ -136,7 +138,7 @@ def _capture_upload_timeout(tmp_path, request_timeout=None): real_client = httpx.Client def spy_client(*args, **kwargs): - captured["timeout"] = kwargs.get("timeout") + captured.update(kwargs) return real_client(*args, **kwargs) try: @@ -163,20 +165,29 @@ def spy_client(*args, **kwargs): server.server_close() thread.join(timeout=5) - return captured["timeout"] + return captured def test_upload_file_defaults_to_one_hour_timeout(tmp_path): # Uploads of large archives need far longer than the 60s general API # timeout, so the default upload timeout is 1 hour (matches the JS SDK). - timeout = _capture_upload_timeout(tmp_path) - assert timeout == httpx.Timeout(FILE_UPLOAD_TIMEOUT_SECONDS) + captured = _capture_upload_client_kwargs(tmp_path) + assert captured["timeout"] == httpx.Timeout(FILE_UPLOAD_TIMEOUT_SECONDS) def test_upload_file_honors_explicit_request_timeout(tmp_path): # An explicitly set request_timeout overrides the 1-hour upload default. - timeout = _capture_upload_timeout(tmp_path, request_timeout=5.0) - assert timeout == httpx.Timeout(5.0) + captured = _capture_upload_client_kwargs(tmp_path, request_timeout=5.0) + assert captured["timeout"] == httpx.Timeout(5.0) + + +def test_upload_file_uses_the_shared_pool_without_retries(tmp_path): + # The archive rides the SDK-wide connection pool instead of a transport + # built per build, but skips its retry layer: replaying a request means + # copying its body, which for a build context is the whole archive in + # memory (SDK-332). + captured = _capture_upload_client_kwargs(tmp_path) + assert captured["transport"] is get_upload_transport(None) def test_upload_file_ignores_post_upload_close_failure(tmp_path): diff --git a/packages/python-sdk/tests/test_api_client_transport.py b/packages/python-sdk/tests/test_api_client_transport.py index c37d034642..bca84d2bb1 100644 --- a/packages/python-sdk/tests/test_api_client_transport.py +++ b/packages/python-sdk/tests/test_api_client_transport.py @@ -10,37 +10,24 @@ import httpx import pytest from pyqwest.httpx import AsyncPyqwestTransport, PyqwestTransport +from transport_caches import reset_transport_caches -import e2b.api.client_async as client_async -import e2b.api.client_sync as client_sync from e2b.api.client_async import get_api_client as get_async_api_client from e2b.api.client_async import get_envd_api as get_async_envd_api -from e2b.api.client_async import get_envd_transport as get_async_envd_transport from e2b.api.client_async import get_transport as get_async_transport from e2b.api.client_sync import get_api_client as get_sync_api_client from e2b.api.client_sync import get_envd_api as get_sync_envd_api -from e2b.api.client_sync import get_envd_transport as get_sync_envd_transport from e2b.api.client_sync import get_transport as get_sync_transport from e2b.connection_config import ConnectionConfig -def reset_sync_api_transports(): - client_sync._transports.clear() - client_sync._envd_transports.clear() - - -def reset_async_api_transports(): - client_async._transports.clear() - client_async._envd_transports.clear() - - def run_in_worker_thread(fn): with ThreadPoolExecutor(max_workers=1) as executor: return executor.submit(fn).result() def test_sync_api_client_proxy_uses_explicit_transport(test_api_key): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig( api_key=test_api_key, proxy="http://127.0.0.1:9999", @@ -56,11 +43,11 @@ def test_sync_api_client_proxy_uses_explicit_transport(test_api_key): assert httpx_client._mounts == {} finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_get_transport_keyed_by_proxy(test_api_key): - reset_sync_api_transports() + reset_transport_caches() proxied_config = ConnectionConfig( api_key=test_api_key, proxy="http://127.0.0.1:9999", @@ -83,11 +70,11 @@ def test_sync_get_transport_keyed_by_proxy(test_api_key): assert get_sync_transport(proxied_config) is proxied_transport assert get_sync_transport(direct_config) is direct_transport finally: - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_applies_request_timeout(test_api_key): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, request_timeout=1.5) api_client = get_sync_api_client(config) @@ -97,11 +84,11 @@ def test_sync_api_client_applies_request_timeout(test_api_key): assert httpx_client.timeout == httpx.Timeout(1.5) finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_request_timeout_zero_disables_timeout(test_api_key): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, request_timeout=0) api_client = get_sync_api_client(config) @@ -111,34 +98,30 @@ def test_sync_api_client_request_timeout_zero_disables_timeout(test_api_key): assert httpx_client.timeout == httpx.Timeout(None) finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() -def test_sync_envd_transports_keyed_by_streaming(test_api_key): - # The envd HTTP API pools are separate from the REST API pools, and the - # streaming variant (which carries the idle read timeout) is its own - # pool per proxy. - reset_sync_api_transports() +def test_sync_envd_and_api_share_one_transport(test_api_key): + # The envd HTTP API draws from the same pool as the control-plane REST API + # — reqwest pools per host, so one pool serves both. Only the streaming + # variant, which carries the idle read timeout, is a pool of its own. + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) try: api_transport = get_sync_transport(config) - envd_transport = get_sync_envd_transport(config) - streaming_transport = get_sync_envd_transport(config, for_streaming=True) - - assert isinstance(envd_transport, PyqwestTransport) - assert envd_transport is not api_transport - assert streaming_transport is not envd_transport - assert get_sync_envd_transport(config) is envd_transport - assert ( - get_sync_envd_transport(config, for_streaming=True) is streaming_transport - ) + streaming_transport = get_sync_transport(config, for_streaming=True) + + assert isinstance(api_transport, PyqwestTransport) + assert api_transport is get_sync_transport(config, for_streaming=False) + assert streaming_transport is not api_transport + assert get_sync_transport(config, for_streaming=True) is streaming_transport finally: - reset_sync_api_transports() + reset_transport_caches() def test_sync_envd_api_client_wiring(test_api_key): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, access_token="tok") client = get_sync_envd_api(config, "https://sandbox.e2b.app") @@ -146,23 +129,21 @@ def test_sync_envd_api_client_wiring(test_api_key): try: assert client.base_url == "https://sandbox.e2b.app" - assert client._transport is get_sync_envd_transport(config) - assert streaming._transport is get_sync_envd_transport( - config, for_streaming=True - ) + assert client._transport is get_sync_transport(config) + assert streaming._transport is get_sync_transport(config, for_streaming=True) for header, value in config.sandbox_headers.items(): assert client.headers[header] == value finally: client.close() streaming.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_is_shared_across_threads(test_api_key): # httpx.Client is thread-safe and the pyqwest transport underneath is # too, so a single client (and its pool) serves all threads — the # per-thread client caching this replaced is gone. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) api_client = get_sync_api_client(config) @@ -174,12 +155,12 @@ def test_sync_api_client_is_shared_across_threads(test_api_key): assert worker_client is main_client finally: main_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_proxy_uses_explicit_transport(test_api_key): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig( api_key=test_api_key, proxy="http://127.0.0.1:9999", @@ -195,12 +176,12 @@ async def test_async_api_client_proxy_uses_explicit_transport(test_api_key): assert httpx_client._mounts == {} finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_get_transport_keyed_by_proxy(test_api_key): - reset_async_api_transports() + reset_transport_caches() proxied_config = ConnectionConfig( api_key=test_api_key, proxy="http://127.0.0.1:9999", @@ -216,7 +197,7 @@ async def test_async_get_transport_keyed_by_proxy(test_api_key): assert get_async_transport(proxied_config) is proxied_transport assert get_async_transport(direct_config) is direct_transport finally: - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio @@ -225,7 +206,7 @@ async def test_async_api_client_is_shared_across_loops(test_api_key): # nor the httpx client wrapper is bound to an event loop — a single # client serves all loops (the per-loop client caching this replaced is # gone). - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) api_client = get_async_api_client(config) @@ -243,45 +224,41 @@ async def get_client(): assert other_loop_client is main_client finally: await main_client.aclose() - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio -async def test_async_envd_transports_keyed_by_streaming(test_api_key): - reset_async_api_transports() +async def test_async_envd_and_api_share_one_transport(test_api_key): + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) try: api_transport = get_async_transport(config) - envd_transport = get_async_envd_transport(config) - streaming_transport = get_async_envd_transport(config, for_streaming=True) - - assert isinstance(envd_transport, AsyncPyqwestTransport) - assert envd_transport is not api_transport - assert streaming_transport is not envd_transport - assert get_async_envd_transport(config) is envd_transport - assert ( - get_async_envd_transport(config, for_streaming=True) is streaming_transport - ) + streaming_transport = get_async_transport(config, for_streaming=True) + + assert isinstance(api_transport, AsyncPyqwestTransport) + assert api_transport is get_async_transport(config, for_streaming=False) + assert streaming_transport is not api_transport + assert get_async_transport(config, for_streaming=True) is streaming_transport finally: - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_envd_api_client_wiring(test_api_key): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, access_token="tok") client = get_async_envd_api(config, "https://sandbox.e2b.app") try: assert client.base_url == "https://sandbox.e2b.app" - assert client._transport is get_async_envd_transport(config) + assert client._transport is get_async_transport(config) for header, value in config.sandbox_headers.items(): assert client.headers[header] == value finally: await client.aclose() - reset_async_api_transports() + reset_transport_caches() class _EchoHandler(BaseHTTPRequestHandler): @@ -349,7 +326,7 @@ def test_sync_transport_sends_proxy_credentials_and_headers(test_api_key, echo_s # Everything an httpx.Proxy can express reaches the proxy: the echo server # stands in for one, so the request arrives in absolute form with the # credentials and the extra headers configured for it. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig( api_key=test_api_key, proxy=httpx.Proxy( @@ -367,14 +344,14 @@ def test_sync_transport_sends_proxy_credentials_and_headers(test_api_key, echo_s assert echoed["headers"]["x-proxy-token"] == "t" finally: client.close() - reset_sync_api_transports() + reset_transport_caches() def test_transport_emits_pyqwest_access_log(test_api_key, echo_server, caplog): # pyqwest logs every request on `pyqwest.access` at DEBUG — the # transport-level diagnostics httpcore used to provide, and separate from # the SDK's own `logger` option. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -392,11 +369,11 @@ def test_transport_emits_pyqwest_access_log(test_api_key, echo_server, caplog): ] finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_round_trips_through_pyqwest(test_api_key, echo_server): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -411,13 +388,13 @@ def test_sync_api_client_round_trips_through_pyqwest(test_api_key, echo_server): assert echoed["headers"]["package_version"] finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_serves_concurrent_threads(test_api_key, echo_server): # The scenario the removed per-thread client caching used to guard: one # client, one shared pyqwest pool, many threads at once. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -433,12 +410,12 @@ def request(i: int) -> tuple[int, str]: assert results == [(200, f"/sandboxes/{i}") for i in range(32)] finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_serves_concurrent_requests(test_api_key, echo_server): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -452,12 +429,12 @@ async def request(i: int) -> tuple[int, str]: assert list(results) == [(200, f"/sandboxes/{i}") for i in range(32)] finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_round_trips_through_pyqwest(test_api_key, echo_server): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -472,14 +449,14 @@ async def test_async_api_client_round_trips_through_pyqwest(test_api_key, echo_s assert echoed["headers"]["package_version"] finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() def test_sync_api_client_leaves_redirects_to_httpx(test_api_key, echo_server): # reqwest would otherwise follow redirects inside the transport, hiding them # from httpx: the generated client asks for no redirect following, so a 302 # must surface as-is, and opting in must record the hop in `history`. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -497,12 +474,12 @@ def test_sync_api_client_leaves_redirects_to_httpx(test_api_key, echo_server): assert [r.status_code for r in followed.history] == [302] finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_leaves_redirects_to_httpx(test_api_key, echo_server): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -520,13 +497,13 @@ async def test_async_api_client_leaves_redirects_to_httpx(test_api_key, echo_ser assert [r.status_code for r in followed.history] == [302] finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() def test_sync_api_client_timeout_raises_httpx_read_timeout(test_api_key, echo_server): # pyqwest raises the builtin TimeoutError; the transport re-raises it as # httpx.ReadTimeout to keep the httpx.TimeoutException contract. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -536,14 +513,14 @@ def test_sync_api_client_timeout_raises_httpx_read_timeout(test_api_key, echo_se httpx_client.request("GET", "/slow", timeout=0.2) finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_timeout_raises_httpx_read_timeout( test_api_key, echo_server ): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -553,7 +530,7 @@ async def test_async_api_client_timeout_raises_httpx_read_timeout( await httpx_client.request("GET", "/slow", timeout=0.2) finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() def test_sync_api_client_body_timeout_raises_httpx_read_timeout( @@ -561,7 +538,7 @@ def test_sync_api_client_body_timeout_raises_httpx_read_timeout( ): # The head arrives in time and the body never does: httpx reads the body # after the transport returned, so that timeout is mapped on the stream. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -571,14 +548,14 @@ def test_sync_api_client_body_timeout_raises_httpx_read_timeout( httpx_client.request("GET", "/stall", timeout=0.2) finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_body_timeout_raises_httpx_read_timeout( test_api_key, echo_server ): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -588,7 +565,7 @@ async def test_async_api_client_body_timeout_raises_httpx_read_timeout( await httpx_client.request("GET", "/stall", timeout=0.2) finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() def test_sync_transport_sends_multipart_bodies(test_api_key, echo_server): @@ -597,11 +574,9 @@ def test_sync_transport_sends_multipart_bodies(test_api_key, echo_server): # sync path used to match AsyncByteStream first and raise from inside the # body iterator, surfacing as a WriteError mid-request; pyqwest 0.8 matches # the sync case first, so the SDK no longer rewraps the stream. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) - client = httpx.Client( - base_url=echo_server, transport=client_sync.get_envd_transport(config) - ) + client = httpx.Client(base_url=echo_server, transport=get_sync_transport(config)) try: response = client.post("/files", files=[("file", ("a.txt", b"x" * 4096))]) @@ -609,4 +584,4 @@ def test_sync_transport_sends_multipart_bodies(test_api_key, echo_server): assert response.json()["received"] > 4096 finally: client.close() - reset_sync_api_transports() + reset_transport_caches() diff --git a/packages/python-sdk/tests/test_envd_client_transport.py b/packages/python-sdk/tests/test_envd_client_transport.py index 683d65faea..1b8c6203e5 100644 --- a/packages/python-sdk/tests/test_envd_client_transport.py +++ b/packages/python-sdk/tests/test_envd_client_transport.py @@ -3,22 +3,22 @@ import httpx import pytest from pyqwest import Proxy +from transport_caches import reset_transport_caches import e2b.api.client_async as api_client_async import e2b.api.client_sync as api_client_sync from e2b.api import ProxyConfig, proxy_to_config -from e2b.connection_config import ProxyTypes +from e2b.connection_config import ConnectionConfig, ProxyTypes from e2b.envd import client_async, client_sync +from e2b.envd.process.process_connect import ProcessClient, ProcessClientSync from e2b.exceptions import InvalidArgumentException @pytest.fixture(autouse=True) -def reset_transport_caches(): - client_sync._transports.clear() - client_async._transports.clear() +def clear_transport_caches(): + reset_transport_caches() yield - client_sync._transports.clear() - client_async._transports.clear() + reset_transport_caches() def test_proxy_to_config_none(): @@ -97,26 +97,28 @@ def test_proxy_to_config_rejects_unknown_types(): proxy_to_config(cast(ProxyTypes, object())) -def test_sync_transport_is_cached_per_proxy(): +def test_sync_pool_is_cached_per_proxy(): proxy = ProxyConfig("http://127.0.0.1:8080") - transport_a = client_sync.get_transport(None) - transport_b = client_sync.get_transport(None) - transport_c = client_sync.get_transport(proxy) + pool_a = api_client_sync.get_pyqwest_transport(None) + pool_b = api_client_sync.get_pyqwest_transport(None) + pool_c = api_client_sync.get_pyqwest_transport(proxy) # A second, equal config keys the same pool. - transport_d = client_sync.get_transport(ProxyConfig("http://127.0.0.1:8080")) + pool_d = api_client_sync.get_pyqwest_transport(ProxyConfig("http://127.0.0.1:8080")) - assert transport_a is transport_b - assert transport_c is transport_d - assert transport_a is not transport_c + assert pool_a is pool_b + assert pool_c is pool_d + assert pool_a is not pool_c -def test_sync_transport_is_not_shared_across_proxy_credentials(): +def test_sync_pool_is_not_shared_across_proxy_credentials(): # Same proxy URL, different credentials or headers: separate pools, since # the proxy configuration is fixed per transport. url = "http://127.0.0.1:8080" - plain = client_sync.get_transport(ProxyConfig(url)) - with_auth = client_sync.get_transport(ProxyConfig(url, auth=("user", "pass"))) - with_headers = client_sync.get_transport( + plain = api_client_sync.get_pyqwest_transport(ProxyConfig(url)) + with_auth = api_client_sync.get_pyqwest_transport( + ProxyConfig(url, auth=("user", "pass")) + ) + with_headers = api_client_sync.get_pyqwest_transport( ProxyConfig(url, headers=(("x-custom", "1"),)) ) @@ -125,27 +127,56 @@ def test_sync_transport_is_not_shared_across_proxy_credentials(): assert with_auth is not with_headers -def test_async_transport_is_cached_per_proxy(): - transport_a = client_async.get_transport(None) - transport_b = client_async.get_transport(None) - transport_c = client_async.get_transport(ProxyConfig("http://127.0.0.1:8080")) - - assert transport_a is transport_b - assert transport_a is not transport_c - assert client_sync.get_transport(None) is not transport_a - +def test_async_pool_is_cached_per_proxy(): + pool_a = api_client_async.get_pyqwest_transport(None) + pool_b = api_client_async.get_pyqwest_transport(None) + pool_c = api_client_async.get_pyqwest_transport( + ProxyConfig("http://127.0.0.1:8080") + ) -def test_transport_stack_normalizes_plain_errors_and_retries_connects(): - # The shared transports are the plain-HTTP-error normalization wrapping - # the connection retries; `E2B_CONNECTION_RETRIES` must flow into the - # retry layer the way it does into the httpx REST transports. + assert pool_a is pool_b + assert pool_a is not pool_c + # Sync and async are separate stacks all the way down. + assert api_client_sync.get_pyqwest_transport(None) is not pool_a + + +def test_rpc_clients_run_on_the_shared_pool(test_api_key, monkeypatch): + # The RPC stack is the plain-HTTP-error normalization wrapping the very + # pool the httpx clients use, so an envd RPC and an envd HTTP call to the + # same sandbox share one HTTP/2 connection. `pyqwest.SyncClient` doesn't + # hand its transport back, so record what the normalization is given. + config = ConnectionConfig(api_key=test_api_key) + pool = api_client_sync.get_pyqwest_transport(None) + async_pool = api_client_async.get_pyqwest_transport(None) + # The httpx adapters every REST client uses sit on those same pools. + assert api_client_sync.get_httpx_transport(None)._transport is pool + assert api_client_async.get_httpx_transport(None)._transport is async_pool + + wrapped = [] + for module in (client_sync, client_async): + normalization = module.PlainHTTPErrorTransport + monkeypatch.setattr( + module, + "PlainHTTPErrorTransport", + lambda inner, normalization=normalization: ( + wrapped.append(inner) or normalization(inner) + ), + ) + + client_sync.create_rpc_client(ProcessClientSync, "https://sandbox.e2b.app", config) + client_async.create_rpc_client(ProcessClient, "https://sandbox.e2b.app", config) + + assert wrapped == [pool, async_pool] + + +def test_shared_pool_retries_connects(): + # `E2B_CONNECTION_RETRIES` must flow into the retry layer of the shared + # pool, which every stack now inherits. from e2b.api import connection_retries - sync_transport = client_sync.get_transport(None) - async_transport = client_async.get_transport(None) - assert isinstance(sync_transport, client_sync.PlainHTTPErrorTransport) - assert isinstance(async_transport, client_async.PlainHTTPErrorTransport) - assert isinstance(sync_transport._inner, api_client_sync.ConnectionRetryTransport) - assert isinstance(async_transport._inner, api_client_async.ConnectionRetryTransport) - assert sync_transport._inner._max_retries == connection_retries - assert async_transport._inner._max_retries == connection_retries + pool = api_client_sync.get_pyqwest_transport(None) + apool = api_client_async.get_pyqwest_transport(None) + assert isinstance(pool, api_client_sync.ConnectionRetryTransport) + assert isinstance(apool, api_client_async.ConnectionRetryTransport) + assert pool._max_retries == connection_retries + assert apool._max_retries == connection_retries diff --git a/packages/python-sdk/tests/test_shared_transport_pool.py b/packages/python-sdk/tests/test_shared_transport_pool.py new file mode 100644 index 0000000000..d5a708a4cd --- /dev/null +++ b/packages/python-sdk/tests/test_shared_transport_pool.py @@ -0,0 +1,158 @@ +"""The envd RPC stack and the envd HTTP API share one connection pool, so a +sandbox costs one HTTP/2 connection instead of one per stack (SDK-291). + +That sharing puts the sandbox health probe on the connection the failed RPC +was using: ``handle_rpc_exception_with_health`` calls ``/health`` precisely +when an RPC died at the transport layer, and it must still get an answer, or +every dropped connection would be reported as an indeterminate state instead +of "the sandbox is gone". These tests pin that at the frame level rather than +trusting reqwest to discard broken connections — a real plaintext HTTP/2 +server serves both routes on one pool and counts the TCP connections the +client opens (see ``envd_frame_server``): + +* an ``RST_STREAM`` kills the RPC stream only, so the probe reuses the same + connection (which is what proves the pool is genuinely shared); +* a dropped TCP connection makes reqwest redial for the probe. + +The pool here mirrors the SDK's — the same connect-retry middleware under the +same plain-error normalization — but with HTTP/2 prior knowledge, since the +transports the SDK builds negotiate the version over TLS via ALPN and this +server is plaintext. +""" + +import httpx +import pytest +from connectrpc.errors import ConnectError +from envd_frame_server import ( + assert_stdout_event, + make_async_client, + make_sync_client, + shared_pool_server, +) +from pyqwest import HTTPTransport, HTTPVersion, SyncHTTPTransport +from pyqwest.httpx import AsyncPyqwestTransport, PyqwestTransport + +from e2b.api import connection_retries +from e2b.api.client_async import ConnectionRetryTransport +from e2b.api.client_sync import ( + ConnectionRetryTransport as SyncConnectionRetryTransport, +) +from e2b.envd.api import acheck_sandbox_health, check_sandbox_health +from e2b.envd.client_async import PlainHTTPErrorTransport +from e2b.envd.client_sync import ( + PlainHTTPErrorTransport as SyncPlainHTTPErrorTransport, +) +from e2b.envd.process.process_pb import ConnectRequest +from e2b.envd.rpc import is_transport_failure + + +def _sync_pool() -> SyncConnectionRetryTransport: + return SyncConnectionRetryTransport( + SyncHTTPTransport(http_version=HTTPVersion.HTTP2), + max_retries=connection_retries, + ) + + +def _async_pool() -> ConnectionRetryTransport: + return ConnectionRetryTransport( + HTTPTransport(http_version=HTTPVersion.HTTP2), + max_retries=connection_retries, + ) + + +def _break_sync_stream(events) -> ConnectError: + """Read the first event, then the failure the server injects after it.""" + assert_stdout_event(next(events)) + with pytest.raises(ConnectError) as excinfo: + next(events) + assert is_transport_failure(excinfo.value), excinfo.value + return excinfo.value + + +async def _break_async_stream(events) -> ConnectError: + assert_stdout_event(await events.__anext__()) + with pytest.raises(ConnectError) as excinfo: + await events.__anext__() + assert is_transport_failure(excinfo.value), excinfo.value + return excinfo.value + + +def test_sync_stream_reset_leaves_the_shared_connection_usable(): + with shared_pool_server("reset") as server: + pool = _sync_pool() + envd_api = httpx.Client( + base_url=f"http://127.0.0.1:{server.port}", + transport=PyqwestTransport(pool), + ) + try: + _break_sync_stream( + make_sync_client( + server.port, transport=SyncPlainHTTPErrorTransport(pool) + ).connect(ConnectRequest()) + ) + assert check_sandbox_health(envd_api) is True + # One TCP connection served the RPC and the probe: the reset took + # down the stream, not the connection. + assert len(server.connections) == 1 + server.assert_no_errors() + finally: + envd_api.close() + + +def test_sync_dropped_connection_redials_for_the_health_probe(): + with shared_pool_server("drop") as server: + pool = _sync_pool() + envd_api = httpx.Client( + base_url=f"http://127.0.0.1:{server.port}", + transport=PyqwestTransport(pool), + ) + try: + _break_sync_stream( + make_sync_client( + server.port, transport=SyncPlainHTTPErrorTransport(pool) + ).connect(ConnectRequest()) + ) + # The probe must not be answered from the dead pooled connection. + assert check_sandbox_health(envd_api) is True + assert len(server.connections) == 2 + finally: + envd_api.close() + + +async def test_async_stream_reset_leaves_the_shared_connection_usable(): + with shared_pool_server("reset") as server: + pool = _async_pool() + envd_api = httpx.AsyncClient( + base_url=f"http://127.0.0.1:{server.port}", + transport=AsyncPyqwestTransport(pool), + ) + try: + await _break_async_stream( + make_async_client( + server.port, transport=PlainHTTPErrorTransport(pool) + ).connect(ConnectRequest()) + ) + assert await acheck_sandbox_health(envd_api) is True + assert len(server.connections) == 1 + server.assert_no_errors() + finally: + await envd_api.aclose() + + +async def test_async_dropped_connection_redials_for_the_health_probe(): + with shared_pool_server("drop") as server: + pool = _async_pool() + envd_api = httpx.AsyncClient( + base_url=f"http://127.0.0.1:{server.port}", + transport=AsyncPyqwestTransport(pool), + ) + try: + await _break_async_stream( + make_async_client( + server.port, transport=PlainHTTPErrorTransport(pool) + ).connect(ConnectRequest()) + ) + assert await acheck_sandbox_health(envd_api) is True + assert len(server.connections) == 2 + finally: + await envd_api.aclose() diff --git a/packages/python-sdk/tests/test_upload_transport.py b/packages/python-sdk/tests/test_upload_transport.py new file mode 100644 index 0000000000..1c58545c3f --- /dev/null +++ b/packages/python-sdk/tests/test_upload_transport.py @@ -0,0 +1,339 @@ +"""File uploads bypass the connect retries so their body isn't copied. + +The retry middleware makes a request body replayable by copying it in full as +it is sent — reqwest reads ahead into the body while connecting, so a connect +error leaves the iterator already started and unrewindable, and the copy is the +only way back. For a `volume.write_file`, an envd `files.write` or a template +context upload that means the whole file in RAM while it is also being streamed +to the wire (SDK-332): a 64 MiB upload peaked at 70 MB. + +So uploads take an httpx adapter over the same pool with the retry layer left +out. They keep the pooled connections and give up the connect retry, which +fires before any of the body was written and so surfaces to the caller intact. +""" + +import os +import tempfile +import threading +import tracemalloc +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import httpx +import pytest +from packaging.version import Version +from pyqwest import HTTPTransport, SyncHTTPTransport +from transport_caches import reset_transport_caches + +from e2b.api.client_async import ConnectionRetryTransport +from e2b.api.client_async import get_envd_api as get_async_envd_api +from e2b.api.client_async import get_pool as get_async_pool +from e2b.api.client_async import get_transport as get_async_transport +from e2b.api.client_async import get_upload_transport as get_async_upload_transport +from e2b.api.client_sync import ( + ConnectionRetryTransport as SyncConnectionRetryTransport, +) +from e2b.api.client_sync import get_envd_api, get_pool, get_transport +from e2b.api.client_sync import get_upload_transport +from e2b.connection_config import ConnectionConfig +from e2b.io_utils import aiter_io_chunks +from e2b.sandbox_async.filesystem.filesystem import Filesystem as AsyncFilesystem +from e2b.sandbox_sync.filesystem.filesystem import Filesystem + +MIB = 1 << 20 +# Enough that a copy of the body would dwarf everything else the transfer +# allocates, while staying quick over the loopback interface. +UPLOAD_SIZE = 32 * MIB +# The streaming path allocates a chunk at a time (~1 MiB peak measured); the +# retrying transport's copy put this transfer's peak at 39 MB. +MAX_UPLOAD_PEAK = 8 * MIB + +ENVD_URL = "https://49999-sandbox.e2b.app" + + +def test_sync_upload_transport_is_the_same_pool_without_retries(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key) + + try: + upload = get_upload_transport(None) + retrying = get_transport(config) + + # Same connection pool underneath, so an upload still travels the + # pooled connections; only the retry layer is missing. + assert isinstance(upload._transport, SyncHTTPTransport) + assert upload._transport is get_pool(None) + assert isinstance(retrying._transport, SyncConnectionRetryTransport) + assert retrying._transport._transport is upload._transport + + # Cached like every other transport, per proxy and idle read bound. + assert get_upload_transport(None) is upload + assert get_transport(config, for_upload=True) is upload + assert get_upload_transport(None, 60.0) is not upload + finally: + reset_transport_caches() + + +def test_async_upload_transport_is_the_same_pool_without_retries(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key) + + try: + upload = get_async_upload_transport(None) + retrying = get_async_transport(config) + + assert isinstance(upload._transport, HTTPTransport) + assert upload._transport is get_async_pool(None) + assert isinstance(retrying._transport, ConnectionRetryTransport) + assert retrying._transport._transport is upload._transport + + assert get_async_upload_transport(None) is upload + assert get_async_transport(config, for_upload=True) is upload + finally: + reset_transport_caches() + + +def test_sync_envd_upload_client_uses_the_upload_transport(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + + client = get_envd_api(config, ENVD_URL, for_upload=True) + try: + assert client._transport is get_transport(config, for_upload=True) + finally: + client.close() + reset_transport_caches() + + +@pytest.mark.asyncio +async def test_async_envd_upload_client_uses_the_upload_transport(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + + client = get_async_envd_api(config, ENVD_URL, for_upload=True) + try: + assert client._transport is get_async_transport(config, for_upload=True) + finally: + await client.aclose() + reset_transport_caches() + + +def test_sync_filesystem_writes_go_to_the_upload_client(test_api_key): + # `files.write` posts on `_envd_api_upload`; the retrying client stays for + # the unary calls and health probes, the streaming one for downloads. + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + envd_api = get_envd_api(config, ENVD_URL) + + filesystem = Filesystem(ENVD_URL, Version("0.7.0"), config, envd_api) + try: + assert filesystem._envd_api._transport is get_transport(config) + assert filesystem._envd_api_upload._transport is get_transport( + config, for_upload=True + ) + assert filesystem._envd_api_streaming._transport is get_transport( + config, for_streaming=True + ) + finally: + filesystem._envd_api_upload.close() + filesystem._envd_api_streaming.close() + envd_api.close() + reset_transport_caches() + + +@pytest.mark.asyncio +async def test_async_filesystem_writes_go_to_the_upload_client(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + envd_api = get_async_envd_api(config, ENVD_URL) + + filesystem = AsyncFilesystem(ENVD_URL, Version("0.7.0"), config, envd_api) + try: + assert filesystem._envd_api._transport is get_async_transport(config) + assert filesystem._envd_api_upload._transport is get_async_transport( + config, for_upload=True + ) + assert filesystem._envd_api_streaming._transport is get_async_transport( + config, for_streaming=True + ) + finally: + await filesystem._envd_api_upload.aclose() + await filesystem._envd_api_streaming.aclose() + await envd_api.aclose() + reset_transport_caches() + + +def _write_response(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, json=[{"name": "a.txt", "type": "file", "path": "/home/user/a.txt"}] + ) + + +def _routing_filesystem(config: ConnectionConfig, kind: type): + """A ``Filesystem`` whose two write clients record which one was used.""" + filesystem = kind(ENVD_URL, Version("0.7.0"), config, None) + used = [] + + def record(name): + def handler(request: httpx.Request) -> httpx.Response: + used.append(name) + return _write_response(request) + + return handler + + client = httpx.AsyncClient if kind is AsyncFilesystem else httpx.Client + filesystem._envd_api = client( + base_url=ENVD_URL, transport=httpx.MockTransport(record("retrying")) + ) + filesystem._envd_api_upload = client( + base_url=ENVD_URL, transport=httpx.MockTransport(record("upload")) + ) + return filesystem, used + + +@pytest.mark.parametrize("use_octet_stream", [True, False]) +def test_sync_only_streamed_writes_skip_the_retries( + test_api_key, tmp_path, use_octet_stream +): + # An in-memory body reaches the transport as `bytes` (or a tiny multipart + # stream), so the retry layer replays it for free and keeps its retries; a + # file-like one is the body that would be copied whole. + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + filesystem, used = _routing_filesystem(config, Filesystem) + path = tmp_path / "a.txt" + path.write_text("streamed") + + try: + filesystem.write("/home/user/a.txt", "text", use_octet_stream=use_octet_stream) + with open(path, "rb") as file: + filesystem.write( + "/home/user/a.txt", file, use_octet_stream=use_octet_stream + ) + + assert used == ["retrying", "upload"] + finally: + filesystem._envd_api.close() + filesystem._envd_api_upload.close() + filesystem._envd_api_streaming.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_octet_stream", [True, False]) +async def test_async_only_streamed_writes_skip_the_retries( + test_api_key, tmp_path, use_octet_stream +): + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + filesystem, used = _routing_filesystem(config, AsyncFilesystem) + path = tmp_path / "a.txt" + path.write_text("streamed") + + try: + await filesystem.write( + "/home/user/a.txt", "text", use_octet_stream=use_octet_stream + ) + with open(path, "rb") as file: + await filesystem.write( + "/home/user/a.txt", file, use_octet_stream=use_octet_stream + ) + + assert used == ["retrying", "upload"] + finally: + await filesystem._envd_api.aclose() + await filesystem._envd_api_upload.aclose() + await filesystem._envd_api_streaming.aclose() + + +class _UploadHandler(BaseHTTPRequestHandler): + """Reads a request body in small blocks — never holding it whole — and + records how much arrived.""" + + protocol_version = "HTTP/1.1" + received = None + + def do_PUT(self): + length = int(self.headers.get("Content-Length", 0)) + read = 0 + while read < length: + block = self.rfile.read(min(1 << 16, length - read)) + if not block: + break + read += len(block) + _UploadHandler.received = (read, self.headers.get("Transfer-Encoding")) + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, *args): + pass + + +@pytest.fixture +def upload_server(): + _UploadHandler.received = None + server = ThreadingHTTPServer(("127.0.0.1", 0), _UploadHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + thread.join() + + +@pytest.fixture +def upload_file_path(): + path = os.path.join(tempfile.mkdtemp(), "context.tar.gz") + with open(path, "wb") as file: + for _ in range(UPLOAD_SIZE // MIB): + file.write(b"x" * MIB) + try: + yield path + finally: + os.remove(path) + + +def test_sync_upload_is_streamed_and_not_mirrored(upload_server, upload_file_path): + # The shape of a volume or template context upload: a file streamed with + # Content-Length framing (S3 presigned URLs reject chunked encoding), and + # nothing on the client holding on to what has already gone out. + reset_transport_caches() + client = httpx.Client(transport=get_upload_transport(None), timeout=None) + + try: + with open(upload_file_path, "rb") as file: + tracemalloc.start() + response = client.put(f"{upload_server}/upload", content=file) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert response.status_code == 200 + assert _UploadHandler.received == (UPLOAD_SIZE, None) + assert peak < MAX_UPLOAD_PEAK + finally: + client.close() + reset_transport_caches() + + +@pytest.mark.asyncio +async def test_async_upload_is_streamed_and_not_mirrored( + upload_server, upload_file_path +): + reset_transport_caches() + client = httpx.AsyncClient(transport=get_async_upload_transport(None), timeout=None) + + try: + with open(upload_file_path, "rb") as file: + tracemalloc.start() + response = await client.put( + f"{upload_server}/upload", + content=aiter_io_chunks(file), + headers={"Content-Length": str(UPLOAD_SIZE)}, + ) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert response.status_code == 200 + assert _UploadHandler.received == (UPLOAD_SIZE, None) + assert peak < MAX_UPLOAD_PEAK + finally: + await client.aclose() + reset_transport_caches() diff --git a/packages/python-sdk/tests/test_volume_client.py b/packages/python-sdk/tests/test_volume_client.py index 1c157a6900..c24d62e797 100644 --- a/packages/python-sdk/tests/test_volume_client.py +++ b/packages/python-sdk/tests/test_volume_client.py @@ -8,27 +8,36 @@ import pytest from pyqwest.httpx import AsyncPyqwestTransport, PyqwestTransport +from transport_caches import reset_transport_caches + +import e2b.api.client_async as api_client_async +import e2b.api.client_sync as api_client_sync import e2b.volume.client_async as client_async import e2b.volume.client_sync as client_sync +from e2b.connection_config import ConnectionConfig from e2b.exceptions import AuthenticationException from e2b.volume.client_async import get_api_client as get_async_api_client from e2b.volume.client_async import ( get_streaming_transport as get_async_streaming_transport, ) from e2b.volume.client_async import get_transport as get_async_transport +from e2b.volume.client_async import ( + get_upload_transport as get_async_upload_transport, +) from e2b.volume.client_sync import get_api_client as get_sync_api_client from e2b.volume.client_sync import ( get_streaming_transport as get_sync_streaming_transport, ) from e2b.volume.client_sync import get_transport as get_sync_transport +from e2b.volume.client_sync import get_upload_transport as get_sync_upload_transport from e2b.volume.connection_config import VolumeConnectionConfig from e2b.volume.volume_async import AsyncVolume from e2b.volume.volume_sync import Volume def reset_volume_transports(): - client_sync._transports.clear() - client_async._transports.clear() + # The volume clients draw from the SDK-wide pools in `e2b.api.client_*`. + reset_transport_caches() def test_sync_client_requires_volume_token(monkeypatch): @@ -90,6 +99,60 @@ def test_sync_transport_is_cached_per_proxy(): reset_volume_transports() +def test_volume_transports_are_the_shared_sdk_pools(test_api_key): + # The volume content API draws from the same pools as the control-plane + # REST API and the envd HTTP API — reqwest pools per host, so the volume + # host doesn't cost the process a pool of its own. Streamed downloads land + # in the streaming pool, whose 60s idle read bound the sandbox + # filesystem's streamed downloads ask for too. + reset_volume_transports() + config = VolumeConnectionConfig(token="vol-token") + api_config = ConnectionConfig(api_key=test_api_key) + + try: + assert get_sync_transport(config) is api_client_sync.get_transport(api_config) + assert get_sync_streaming_transport(config) is api_client_sync.get_transport( + api_config, for_streaming=True + ) + assert get_async_transport(config) is api_client_async.get_transport(api_config) + assert get_async_streaming_transport(config) is api_client_async.get_transport( + api_config, for_streaming=True + ) + finally: + reset_volume_transports() + + +def test_volume_uploads_skip_the_connect_retries(test_api_key): + # `write_file` streams its body, and the retry layer would make that body + # replayable by copying it in full — the whole file in memory (SDK-332). It + # goes out on the same pool without that layer, which is also the SDK-wide + # upload transport. + reset_volume_transports() + config = VolumeConnectionConfig(token="vol-token") + api_config = ConnectionConfig(api_key=test_api_key) + + try: + assert get_sync_upload_transport(config) is api_client_sync.get_transport( + api_config, for_upload=True + ) + assert get_async_upload_transport(config) is api_client_async.get_transport( + api_config, for_upload=True + ) + assert get_sync_upload_transport(config) is not get_sync_transport(config) + assert get_async_upload_transport(config) is not get_async_transport(config) + + upload_client = client_sync.get_upload_api_client(config) + try: + assert ( + upload_client.get_httpx_client()._transport + is get_sync_upload_transport(config) + ) + finally: + upload_client.get_httpx_client().close() + finally: + reset_volume_transports() + + def test_sync_transport_is_shared_across_threads(): # pyqwest transports are thread-safe, so one transport (and its pool) # serves all threads — the per-thread caching this replaced is gone. diff --git a/packages/python-sdk/tests/transport_caches.py b/packages/python-sdk/tests/transport_caches.py new file mode 100644 index 0000000000..b244640ba5 --- /dev/null +++ b/packages/python-sdk/tests/transport_caches.py @@ -0,0 +1,21 @@ +"""Reset the SDK's process-global pyqwest transport caches. + +Every HTTP stack in the SDK — control-plane REST, envd HTTP API, envd RPC, +volume content, template uploads — draws its connection pool from +``e2b.api.client_sync``/``client_async``, so one helper clears them all. Tests +that assert on pool identity or rebuild a pool with different tuning call this +before and after. Not a test module itself — imported by the transport test +modules (``pythonpath = tests`` in pytest.ini makes it importable under +``--import-mode=importlib``). +""" + +import e2b.api.client_async as api_client_async +import e2b.api.client_sync as api_client_sync + + +def reset_transport_caches() -> None: + for module in (api_client_sync, api_client_async): + module._pools.clear() + module._transports.clear() + module._httpx_transports.clear() + module._upload_transports.clear()