tests/test_tcp_idle_timeout.py
raw ยท 9064 bytes
import asyncio
import os
from .support import (
LISTEN_HOST,
BACKEND_PORT,
PROXY_PORT,
SkipTest,
run_tinyproxy_with_conf,
start_tracked_stream_server,
)
async def delayed_response_upload_handler(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
try:
try:
header = await reader.readuntil(b"\r\n\r\n")
except asyncio.IncompleteReadError as e:
# The test harness/proxy readiness check may open a TCP
# connection and close it without sending any request bytes.
# Ignore that. A partial HTTP request is still a real failure.
if not e.partial:
return
raise AssertionError(f"incomplete request headers: {e.partial!r}") from e
content_length = None
for line in header.decode("latin1", errors="replace").split("\r\n"):
if line.lower().startswith("content-length:"):
content_length = int(line.split(":", 1)[1].strip())
break
assert content_length is not None, "missing Content-Length"
received = 0
while received < content_length:
chunk = await reader.read(min(65536, content_length - received))
if not chunk:
break
received += len(chunk)
assert received == content_length, (
f"upload body truncated: received={received} expected={content_length}"
)
body = b"ok\n"
writer.write(
b"HTTP/1.1 200 OK\r\n"
+ f"Content-Length: {len(body)}\r\n".encode("ascii")
+ b"Connection: close\r\n"
+ b"\r\n"
+ body
)
await writer.drain()
finally:
writer.close()
await writer.wait_closed()
async def stalled_upload_handler(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
try:
try:
await reader.readuntil(b"\r\n\r\n")
except asyncio.IncompleteReadError as e:
# Ignore readiness/preflight connections from the test harness.
if not e.partial:
return
raise AssertionError(f"incomplete request headers: {e.partial!r}") from e
# Read a little bit so the request is definitely established,
# then stop consuming the upload body. This should eventually
# backpressure tinyproxy and trigger the stalled I/O timeout.
await reader.readexactly(64 * 1024)
await asyncio.sleep(10.0)
finally:
writer.close()
await writer.wait_closed()
async def echo_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
try:
while True:
data = await reader.read(65536)
if not data:
break
writer.write(data)
await writer.drain()
finally:
writer.close()
await writer.wait_closed()
async def close_writer(writer: asyncio.StreamWriter) -> None:
try:
writer.close()
await writer.wait_closed()
except OSError:
pass
async def read_one_or_eof(reader: asyncio.StreamReader, timeout: float = 3.0) -> bytes:
try:
return await asyncio.wait_for(reader.read(1), timeout=timeout)
except ConnectionResetError:
return b""
async def test_tcp_idle_timeout_closes_idle_connection() -> None:
proxy_bin = os.environ.get("TINYPROXY_BIN")
if not proxy_bin:
raise SkipTest("TINYPROXY_BIN is not set")
conf_text = (
f"listen"
f" tcp {LISTEN_HOST}:{PROXY_PORT}"
f" tcp {LISTEN_HOST}:{BACKEND_PORT}"
f" idle_timeout=1\n"
)
backend_server = await start_tracked_stream_server(
echo_handler,
LISTEN_HOST,
BACKEND_PORT,
backlog=128,
)
try:
async with run_tinyproxy_with_conf(
proxy_bin=proxy_bin,
conf_text=conf_text,
listen_host=LISTEN_HOST,
listen_port=PROXY_PORT,
proto="tcp",
):
reader, writer = await asyncio.open_connection(LISTEN_HOST, PROXY_PORT)
try:
payload = b"before idle timeout\n"
writer.write(payload)
await writer.drain()
got = await asyncio.wait_for(
reader.readexactly(len(payload)),
timeout=3.0,
)
assert got == payload, (
f"initial roundtrip mismatch: got={got!r} expected={payload!r}"
)
await asyncio.sleep(2.0)
got = await read_one_or_eof(reader, timeout=3.0)
assert got == b"", (
f"expected idle connection to close, got={got!r}"
)
finally:
await close_writer(writer)
finally:
await backend_server.close()
async def test_tcp_idle_timeout_keeps_active_connection_open() -> None:
proxy_bin = os.environ.get("TINYPROXY_BIN")
if not proxy_bin:
raise SkipTest("TINYPROXY_BIN is not set")
conf_text = (
f"listen"
f" tcp {LISTEN_HOST}:{PROXY_PORT}"
f" tcp {LISTEN_HOST}:{BACKEND_PORT}"
f" idle_timeout=2\n"
)
backend_server = await start_tracked_stream_server(
echo_handler,
LISTEN_HOST,
BACKEND_PORT,
backlog=128,
)
try:
async with run_tinyproxy_with_conf(
proxy_bin=proxy_bin,
conf_text=conf_text,
listen_host=LISTEN_HOST,
listen_port=PROXY_PORT,
proto="tcp",
):
reader, writer = await asyncio.open_connection(LISTEN_HOST, PROXY_PORT)
try:
for i in range(3):
payload = f"still-active-{i}\n".encode()
writer.write(payload)
await writer.drain()
got = await asyncio.wait_for(
reader.readexactly(len(payload)),
timeout=3.0,
)
assert got == payload, (
f"active roundtrip {i} mismatch: "
f"got={got!r} expected={payload!r}"
)
await asyncio.sleep(1.0)
finally:
await close_writer(writer)
finally:
await backend_server.close()
async def test_tcp_idle_timeout_keeps_long_one_way_upload_open() -> None:
proxy_bin = os.environ.get("TINYPROXY_BIN")
if not proxy_bin:
raise SkipTest("TINYPROXY_BIN is not set")
conf_text = (
f"listen"
f" tcp {LISTEN_HOST}:{PROXY_PORT}"
f" tcp {LISTEN_HOST}:{BACKEND_PORT}"
f" idle_timeout=2\n"
)
backend_server = await start_tracked_stream_server(
delayed_response_upload_handler,
LISTEN_HOST,
BACKEND_PORT,
backlog=128,
)
try:
async with run_tinyproxy_with_conf(
proxy_bin=proxy_bin,
conf_text=conf_text,
listen_host=LISTEN_HOST,
listen_port=PROXY_PORT,
proto="tcp",
):
reader, writer = await asyncio.open_connection(LISTEN_HOST, PROXY_PORT)
try:
body_size = 6 * 1024 * 1024
chunk_size = 64 * 1024
delay = 0.05
req = (
b"PUT /v2/test/blobs/uploads/fake?digest=sha256:deadbeef HTTP/1.1\r\n"
b"Host: fake-registry\r\n"
b"User-Agent: test-buildkit/repro\r\n"
b"Content-Type: application/octet-stream\r\n"
+ f"Content-Length: {body_size}\r\n".encode("ascii")
+ b"Connection: close\r\n"
+ b"\r\n"
)
writer.write(req)
await writer.drain()
chunk = b"x" * chunk_size
sent = 0
while sent < body_size:
n = min(chunk_size, body_size - sent)
writer.write(chunk[:n])
await writer.drain()
sent += n
await asyncio.sleep(delay)
resp = await asyncio.wait_for(reader.read(), timeout=5.0)
assert b"HTTP/1.1 200 OK" in resp, resp
assert resp.endswith(b"\r\n\r\nok\n"), resp
finally:
await close_writer(writer)
finally:
await backend_server.close()
async def test_tcp_stalled_timeout_should_close() -> None:
proxy_bin = os.environ.get("TINYPROXY_BIN")
if not proxy_bin:
raise SkipTest("TINYPROXY_BIN is not set")
conf_text = (
f"listen"
f" tcp {LISTEN_HOST}:{PROXY_PORT}"
f" tcp {LISTEN_HOST}:{BACKEND_PORT}"
f" idle_timeout=2\n"
)
backend_server = await start_tracked_stream_server(
stalled_upload_handler,
LISTEN_HOST,
BACKEND_PORT,
backlog=128,
)
try:
async with run_tinyproxy_with_conf(
proxy_bin=proxy_bin,
conf_text=conf_text,
listen_host=LISTEN_HOST,
listen_port=PROXY_PORT,
proto="tcp",
):
reader, writer = await asyncio.open_connection(LISTEN_HOST, PROXY_PORT)
try:
body_size = 512 * 1024 * 1024
chunk_size = 1024 * 1024
req = (
b"PUT /v2/test/blobs/uploads/stalled?digest=sha256:deadbeef HTTP/1.1\r\n"
b"Host: fake-registry\r\n"
b"User-Agent: test-buildkit/stalled-repro\r\n"
b"Content-Type: application/octet-stream\r\n"
+ f"Content-Length: {body_size}\r\n".encode("ascii")
+ b"Connection: close\r\n"
+ b"\r\n"
)
writer.write(req)
await writer.drain()
chunk = b"x" * chunk_size
sent = 0
closed = False
deadline = asyncio.get_running_loop().time() + 8.0
while asyncio.get_running_loop().time() < deadline:
try:
writer.write(chunk)
await asyncio.wait_for(writer.drain(), timeout=3.0)
sent += len(chunk)
except (
ConnectionResetError,
BrokenPipeError,
asyncio.TimeoutError,
):
closed = True
break
if not closed:
got = await read_one_or_eof(reader, timeout=1.0)
closed = got == b""
assert closed, (
f"expected stalled upload to be closed by idle timeout; "
f"sent={sent}"
)
finally:
await close_writer(writer)
finally:
await backend_server.close()
TESTS = [
("test_tcp_idle_timeout_closes_idle_connection", test_tcp_idle_timeout_closes_idle_connection),
("test_tcp_idle_timeout_keeps_active_connection_open", test_tcp_idle_timeout_keeps_active_connection_open),
("test_tcp_idle_timeout_keeps_long_one_way_upload_open", test_tcp_idle_timeout_keeps_long_one_way_upload_open),
("test_tcp_stalled_timeout_should_close", test_tcp_stalled_timeout_should_close),
]