Skip to content

Commit d943237

Browse files
authored
fix: long partitions behind a NAT idle timeout never return (#359)
## What Partition requests that take longer than a few minutes now return their response when the client is behind a NAT or load balancer that drops idle connections. The SDK's own HTTP clients now send TCP keepalive probes while they wait. That covers the default sync and async clients and the split-PDF client. ``` hi_res partition, response headers after ~5 min, client behind a NAT gateway with a 4 min idle timeout Before: request sent → no packets for 4 min → NAT drops the flow → response never arrives → ReadTimeout at the full read timeout (900 s) After: request sent → keepalive probe every 30 s after 60 s idle → flow stays open → 200 ``` It is on by default (60 s idle, 30 s interval, 5 probes). It can be tuned or turned off with `UNSTRUCTURED_CLIENT_TCP_KEEPALIVE`, `UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_IDLE`, `UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_INTERVAL` and `UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_COUNT`. The split-PDF `plan_created` and `batch_async_start` logs show the setting in use. ## Why Long-running partition requests can exceed a NAT gateway’s idle timeout, causing responses to be lost while the client waits for its read timeout. This change enables TCP keepalive to keep those connections active. Neither httpx nor the SDK turns on keepalive, and the OS keepalive settings only apply to sockets that ask for it. That makes the SDK, which opens these sockets, the place to fix it. Unchanged: - Clients passed in by the caller (`client=` / `async_client=`) are not modified. - Proxy routing from `HTTP(S)_PROXY` / `NO_PROXY` works as before. Connections made through a proxy do not get keepalive, because httpcore drops socket options for proxied connections. - If a future httpx release moves the helper used to rebuild proxy routing, the SDK falls back to building its clients exactly as it did before this change. ## Validation - Unit tests (416 passed, on top of 0.46.4), pylint and mypy pass. - Tests check that real sockets carry `SO_KEEPALIVE` and every idle, interval and count setting. They also cover the bounds on the environment values, proxy routing and parity of TLS, limits and timeouts. - Docker reproduction with a NAT router that silently drops idle flows after 20 s, and a server that answers after 45 s: - with keepalive: `200` at 45 s; - without keepalive: `ReadTimeout` at 90 s. - Released as 0.46.5: CHANGELOG entry and `_version.py` bump (0.46.4 is already published).
1 parent 24fadc4 commit d943237

8 files changed

Lines changed: 396 additions & 29 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,8 @@
1+
## 0.46.5
2+
3+
### Fixes
4+
* Long partitions behind a NAT or load balancer idle timeout could never return, leaving the client waiting for its full read timeout. SDK-owned HTTP clients now enable TCP keepalive, including split-PDF requests. Connections made through an HTTP(S)/SOCKS proxy are not covered. Configure probes with `UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_IDLE` (60 seconds), `UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_INTERVAL` (30 seconds), and `UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_COUNT` (5); set `UNSTRUCTURED_CLIENT_TCP_KEEPALIVE=0` or `false` to disable. Caller-supplied clients are unchanged.
5+
16
## 0.46.4
27

38
### Fixes

‎_test_unstructured_client/unit/test_split_collection_urls.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from pypdf import PdfWriter
77

88
from unstructured_client import UnstructuredClient
9+
from unstructured_client._hooks.custom import split_pdf_hook
910
from unstructured_client._hooks.custom.split_pdf_hook import _get_collection_url
1011
from unstructured_client.models import operations, shared
1112
from unstructured_client.utils import BackoffStrategy, RetryConfig
@@ -60,9 +61,9 @@ def respond(request):
6061
)
6162
# Split chunks create their own async HTTP client; keep every request offline.
6263
monkeypatch.setattr(
63-
httpx,
64-
"AsyncClient",
65-
lambda **kwargs: original_async_client(transport=transport, **kwargs),
64+
split_pdf_hook,
65+
"create_async_http_client",
66+
lambda **kwargs: original_async_client(**kwargs, transport=transport),
6667
)
6768
return sdk, sent, sync_client, async_client
6869

‎_test_unstructured_client/unit/test_split_pdf_hook.py‎

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -498,15 +498,15 @@ async def test_unit_run_tasks_pool_limits_configurable_via_env(
498498
monkeypatch.setenv("UNSTRUCTURED_CLIENT_KEEPALIVE_EXPIRY", "30.0")
499499

500500
captured: dict[str, httpx.Limits] = {}
501-
real_async_client = httpx.AsyncClient
501+
real_async_transport = httpx.AsyncHTTPTransport
502502

503-
def _capturing_async_client(*args, **kwargs):
503+
def _capturing_async_transport(*args, **kwargs):
504504
captured["limits"] = kwargs.get("limits")
505-
return real_async_client(*args, **kwargs)
505+
return real_async_transport(*args, **kwargs)
506506

507507
with patch(
508-
"unstructured_client._hooks.custom.split_pdf_hook.httpx.AsyncClient",
509-
side_effect=_capturing_async_client,
508+
"unstructured_client._hooks.custom.http_transport.httpx.AsyncHTTPTransport",
509+
side_effect=_capturing_async_transport,
510510
):
511511
await run_tasks(
512512
[partial(_request_mock, fails=False, content="ok")],
@@ -622,24 +622,21 @@ def test_unit_resolve_tls_config_client_cert_and_key_split(
622622

623623

624624
@pytest.mark.asyncio
625-
async def test_unit_run_tasks_forwards_tls_config_to_httpx_async_client(
625+
async def test_unit_run_tasks_forwards_tls_config_to_httpx_async_transport(
626626
monkeypatch: pytest.MonkeyPatch, tmp_path
627627
):
628-
"""run_tasks() actually wires verify+cert into httpx.AsyncClient(). End-to-
629-
end check that _resolve_tls_config is called and its result reaches the
630-
client construction. AsyncClient is fully mocked so we don't have to feed
631-
httpx a real cert chain."""
628+
"""Split requests apply TLS settings to the transport."""
632629
ca_bundle = tmp_path / "ca.pem"
633630
ca_bundle.write_text("-----BEGIN CERTIFICATE-----\nfake\n-----END CERTIFICATE-----\n")
634631
monkeypatch.setenv("SSL_CERT_FILE", str(ca_bundle))
635632
monkeypatch.setenv("UNSTRUCTURED_CLIENT_TLS_CLIENT_CERT", "/etc/ssl/client.pem")
636633

634+
monkeypatch.setenv("NO_PROXY", "*")
637635
captured: dict = {}
638636

639637
class _MockAsyncClient:
640638
def __init__(self, *args, **kwargs):
641-
captured["verify"] = kwargs.get("verify")
642-
captured["cert"] = kwargs.get("cert")
639+
captured["transport"] = kwargs.get("transport")
643640

644641
async def __aenter__(self):
645642
return self
@@ -648,16 +645,21 @@ async def __aexit__(self, *exc):
648645
return False
649646

650647
with patch(
651-
"unstructured_client._hooks.custom.split_pdf_hook.httpx.AsyncClient",
648+
"unstructured_client._hooks.custom.http_transport.httpx.AsyncClient",
652649
new=_MockAsyncClient,
653-
):
650+
), patch(
651+
"unstructured_client._hooks.custom.http_transport.httpx.AsyncHTTPTransport"
652+
) as transport:
654653
await run_tasks(
655654
[partial(_request_mock, fails=False, content="ok")],
656655
allow_failed=True,
657656
)
658657

659-
assert captured["verify"] == str(ca_bundle)
660-
assert captured["cert"] == "/etc/ssl/client.pem"
658+
kwargs = transport.call_args.kwargs
659+
assert captured["transport"] is transport.return_value
660+
assert kwargs["verify"] == str(ca_bundle)
661+
assert kwargs["cert"] == "/etc/ssl/client.pem"
662+
assert kwargs["socket_options"]
661663

662664

663665
@patch("unstructured_client._hooks.custom.form_utils.Path")
Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
"""TCP keepalive configuration and transport integration."""
2+
3+
import asyncio
4+
import socket
5+
from contextlib import contextmanager
6+
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
7+
from threading import Thread
8+
from unittest.mock import patch
9+
10+
import httpcore
11+
import httpx
12+
import pytest
13+
14+
from unstructured_client import UnstructuredClient
15+
from unstructured_client._hooks.custom import http_transport as transport
16+
17+
18+
@pytest.fixture(autouse=True)
19+
def clean_env(monkeypatch):
20+
for suffix in ("", "_IDLE", "_INTERVAL", "_COUNT"):
21+
monkeypatch.delenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE" + suffix, raising=False)
22+
monkeypatch.setenv("NO_PROXY", "*")
23+
24+
25+
DEFAULTS = transport.TCPKeepAlive(enabled=True)
26+
BASE_OPTIONS = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
27+
28+
29+
def _close(client):
30+
if isinstance(client, httpx.AsyncClient):
31+
asyncio.run(client.aclose())
32+
else:
33+
client.close()
34+
35+
36+
@contextmanager
37+
def local_server(handler):
38+
server = ThreadingHTTPServer(("127.0.0.1", 0), handler)
39+
thread = Thread(target=server.serve_forever, daemon=True)
40+
thread.start()
41+
try:
42+
yield server.server_port
43+
finally:
44+
server.shutdown()
45+
server.server_close()
46+
thread.join()
47+
48+
49+
@pytest.mark.parametrize("attrs_to_set,expected_options", [
50+
({"TCP_KEEPIDLE": 100, "TCP_KEEPINTVL": 101, "TCP_KEEPCNT": 102},
51+
[(socket.IPPROTO_TCP, 100, DEFAULTS.idle),
52+
(socket.IPPROTO_TCP, 101, DEFAULTS.interval),
53+
(socket.IPPROTO_TCP, 102, DEFAULTS.count)]),
54+
({"TCP_KEEPALIVE": 100, "TCP_KEEPINTVL": 101, "TCP_KEEPCNT": 102},
55+
[(socket.IPPROTO_TCP, 100, DEFAULTS.idle),
56+
(socket.IPPROTO_TCP, 101, DEFAULTS.interval),
57+
(socket.IPPROTO_TCP, 102, DEFAULTS.count)]),
58+
({"TCP_KEEPALIVE": 100}, [(socket.IPPROTO_TCP, 100, DEFAULTS.idle)]),
59+
({}, []),
60+
], ids=["linux", "macos", "macos_old", "windows"])
61+
def test_unit_platform_options(monkeypatch, attrs_to_set, expected_options):
62+
for name in ("TCP_KEEPIDLE", "TCP_KEEPALIVE", "TCP_KEEPINTVL", "TCP_KEEPCNT"):
63+
monkeypatch.delattr(socket, name, raising=False)
64+
for name, value in attrs_to_set.items():
65+
monkeypatch.setattr(socket, name, value, raising=False)
66+
assert transport.resolve_tcp_keepalive().socket_options() == BASE_OPTIONS + expected_options
67+
68+
69+
@pytest.mark.parametrize("value", ["0", "false", "False", " FALSE "])
70+
def test_unit_disabled(monkeypatch, caplog, value):
71+
monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE", value)
72+
for suffix in ("IDLE", "INTERVAL", "COUNT"):
73+
monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_" + suffix, "invalid")
74+
settings = transport.resolve_tcp_keepalive()
75+
assert settings.socket_options() == []
76+
assert settings.describe() == "off"
77+
assert not caplog.records
78+
79+
80+
@pytest.mark.parametrize("suffix,value,expected", [
81+
(suffix, value, expected)
82+
for suffix, maximum in (
83+
("IDLE", transport.MAX_IDLE),
84+
("INTERVAL", transport.MAX_INTERVAL),
85+
("COUNT", transport.MAX_COUNT),
86+
)
87+
for value, expected in (
88+
*[(invalid, getattr(DEFAULTS, suffix.lower()))
89+
for invalid in ("0", "-1", "bad", "1.5", "", str(maximum + 1))],
90+
("1", 1),
91+
(str(maximum), maximum),
92+
)
93+
])
94+
def test_unit_value_bounds(monkeypatch, caplog, suffix, value, expected):
95+
monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_" + suffix, value)
96+
assert getattr(transport.resolve_tcp_keepalive(), suffix.lower()) == expected
97+
assert ("using default" in caplog.text) == (value != str(expected))
98+
99+
100+
@pytest.mark.parametrize("asynchronous", [False, True])
101+
def test_unit_missing_proxy_helper_client_options(monkeypatch, asynchronous):
102+
monkeypatch.setattr(transport, "get_environment_proxies", None)
103+
if asynchronous:
104+
options = {"verify": False, "cert": "client.pem", "limits": httpx.Limits(max_connections=3)}
105+
with patch.object(httpx, "AsyncClient") as constructor:
106+
transport.create_async_http_client(timeout=42, **options)
107+
constructor.assert_called_once_with(timeout=42, **options)
108+
else:
109+
with patch.object(httpx, "Client") as constructor:
110+
transport.create_http_client()
111+
constructor.assert_called_once_with()
112+
113+
114+
def test_unit_custom_values(monkeypatch):
115+
for suffix, value in (("IDLE", "5"), ("INTERVAL", "6"), ("COUNT", "7")):
116+
monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_" + suffix, value)
117+
settings = transport.resolve_tcp_keepalive()
118+
assert settings == transport.TCPKeepAlive(True, 5, 6, 7)
119+
assert settings.describe() == "idle=5s interval=6s count=7"
120+
121+
122+
def test_unit_sdk_owned_clients():
123+
sdk = UnstructuredClient()
124+
options = transport.resolve_tcp_keepalive().socket_options()
125+
assert sdk.sdk_configuration.client._transport.base_transport._pool._socket_options == options
126+
assert sdk.sdk_configuration.async_client._transport._pool._socket_options == options
127+
_close(sdk.sdk_configuration.client)
128+
_close(sdk.sdk_configuration.async_client)
129+
130+
131+
def test_unit_supplied_clients_untouched():
132+
sync = httpx.Client()
133+
asynchronous = httpx.AsyncClient()
134+
sync_transport = sync._transport
135+
async_transport = asynchronous._transport
136+
UnstructuredClient(client=sync, async_client=asynchronous)
137+
assert sync._transport.base_transport is sync_transport
138+
assert asynchronous._transport is async_transport
139+
_close(sync)
140+
_close(asynchronous)
141+
142+
143+
@pytest.mark.parametrize("missing_proxy_helper", [False, True])
144+
@pytest.mark.parametrize("asynchronous", [False, True])
145+
def test_unit_environment_proxy_routes(monkeypatch, asynchronous, missing_proxy_helper):
146+
if missing_proxy_helper:
147+
monkeypatch.setattr(transport, "get_environment_proxies", None)
148+
monkeypatch.setenv("HTTPS_PROXY", "http://127.0.0.1:8123")
149+
monkeypatch.setenv("NO_PROXY", "excluded.example")
150+
client = transport.create_async_http_client() if asynchronous else transport.create_http_client()
151+
proxied = client._transport_for_url(httpx.URL("https://api.example"))
152+
assert isinstance(proxied._pool, httpcore.AsyncHTTPProxy if asynchronous else httpcore.HTTPProxy)
153+
assert proxied._pool._proxy_url.host == b"127.0.0.1"
154+
assert proxied._pool._socket_options is None
155+
expected = None if missing_proxy_helper else transport.resolve_tcp_keepalive().socket_options()
156+
assert client._transport._pool._socket_options == expected
157+
assert client._transport_for_url(httpx.URL("https://excluded.example")) is client._transport
158+
_close(client)
159+
160+
161+
@pytest.mark.parametrize("asynchronous", [False, True])
162+
def test_unit_real_socket_keepalive(asynchronous):
163+
class Handler(BaseHTTPRequestHandler):
164+
protocol_version = "HTTP/1.1"
165+
166+
def do_GET(self):
167+
self.send_response(200)
168+
self.send_header("Content-Length", "2")
169+
self.end_headers()
170+
self.wfile.write(b"ok")
171+
172+
def check(response):
173+
stream = response.extensions["network_stream"]
174+
sock = stream.get_extra_info("socket")
175+
for level, option, value in transport.resolve_tcp_keepalive().socket_options():
176+
if (level, option) == (socket.SOL_SOCKET, socket.SO_KEEPALIVE):
177+
assert sock.getsockopt(level, option)
178+
else:
179+
assert sock.getsockopt(level, option) == value
180+
181+
async def request():
182+
async with transport.create_async_http_client() as client:
183+
async with client.stream("GET", url) as response:
184+
check(response)
185+
await response.aread()
186+
187+
with local_server(Handler) as port:
188+
url = f"http://127.0.0.1:{port}"
189+
if asynchronous:
190+
asyncio.run(request())
191+
else:
192+
with transport.create_http_client() as client:
193+
with client.stream("GET", url) as response:
194+
check(response)
195+
response.read()
196+
197+
198+
@pytest.mark.parametrize("asynchronous", [False, True])
199+
def test_unit_https_proxy_receives_connect(monkeypatch, asynchronous):
200+
destinations = []
201+
202+
class ProxyHandler(BaseHTTPRequestHandler):
203+
def do_CONNECT(self):
204+
destinations.append(self.path)
205+
self.send_error(502)
206+
207+
async def request():
208+
async with transport.create_async_http_client() as client:
209+
await client.get("https://proxy-target.invalid/")
210+
211+
with local_server(ProxyHandler) as port:
212+
monkeypatch.setenv("HTTPS_PROXY", f"http://127.0.0.1:{port}")
213+
monkeypatch.setenv("NO_PROXY", "")
214+
with pytest.raises(httpx.ProxyError, match="502"):
215+
if asynchronous:
216+
asyncio.run(request())
217+
else:
218+
with transport.create_http_client() as client:
219+
client.get("https://proxy-target.invalid/")
220+
assert destinations == ["proxy-target.invalid:443"]

0 commit comments

Comments
 (0)