mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
test(integration): stop the wire peer from dropping requests whose bytes arrive late (#49889)
* test(integration): stop the wire peer from dropping requests whose bytes arrive late The scripted wire peer set a fixed 5 s socket timeout, so a request whose bytes reached it late on a loaded rig was dropped with no reply and the proxy answered 500 Server disconnected. Remove the timeout and, at teardown, stop reading on every accepted connection so idle handlers exit at once. A request still arriving when teardown starts never reaches the test's respond callback. * fix(tests): bound wire peer teardown on unread replies and answer complete requests * fix(tests): poll wire peer reads so teardown keeps queued requests and only cuts stuck calls * fix(tests): write wire peer replies in pieces so teardown never cuts a client still reading * fix(tests): give a wire peer reply the whole teardown grace when its write blocked before teardown --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tessellark
GitHub
parent
a12a6cd154
commit
f357872c98
2 changed files with 409 additions and 6 deletions
|
|
@ -1,16 +1,25 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import selectors
|
||||
import socket
|
||||
import ssl
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Generator, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from contextlib import contextmanager, suppress
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
from typing import BinaryIO, Final
|
||||
from typing import TYPE_CHECKING, BinaryIO, Final
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Buffer
|
||||
|
||||
_POLL_SECONDS: Final = 0.05
|
||||
_WRITE_PIECE_BYTES: Final = 64 * 1024
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -97,12 +106,108 @@ def _await_release(
|
|||
return gate.is_set()
|
||||
|
||||
|
||||
class _RequestReader(io.BufferedReader):
|
||||
cut = False
|
||||
|
||||
def readline(self, size: int | None = -1, /) -> bytes:
|
||||
line: Final = super().readline(size)
|
||||
self.cut = self.cut or not line.endswith(b"\n")
|
||||
return line
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Blocked:
|
||||
since: float
|
||||
writing: bool
|
||||
|
||||
|
||||
class _Connection(io.BufferedIOBase):
|
||||
blocked: _Blocked | None = None
|
||||
|
||||
def __init__(self, connection: socket.socket, closing: threading.Event) -> None:
|
||||
super().__init__()
|
||||
self.connection = connection
|
||||
self.closing = closing
|
||||
self.selector = selectors.DefaultSelector()
|
||||
self.selector.register(connection, selectors.EVENT_READ)
|
||||
|
||||
def readable(self) -> bool:
|
||||
return True
|
||||
|
||||
def writable(self) -> bool:
|
||||
return True
|
||||
|
||||
def readinto(self, buffer: Buffer, /) -> int:
|
||||
if not self._wait_for_bytes():
|
||||
return 0
|
||||
with self._blocking(writing=False):
|
||||
return self.connection.recv_into(buffer)
|
||||
|
||||
def write(self, buffer: Buffer, /) -> int:
|
||||
view: Final = memoryview(buffer).cast("B")
|
||||
for start in range(0, view.nbytes, _WRITE_PIECE_BYTES):
|
||||
with self._blocking(writing=True):
|
||||
self.connection.sendall(view[start : start + _WRITE_PIECE_BYTES])
|
||||
return view.nbytes
|
||||
|
||||
def close(self) -> None:
|
||||
self.selector.close()
|
||||
super().close()
|
||||
|
||||
def stalled(self, grace: float, teardown_started: float) -> _Blocked | None:
|
||||
blocked: Final = self.blocked
|
||||
if blocked is None or time.monotonic() - max(blocked.since, teardown_started) < grace:
|
||||
return None
|
||||
return blocked
|
||||
|
||||
def shut(self) -> None:
|
||||
with suppress(OSError):
|
||||
self.connection.shutdown(socket.SHUT_RDWR)
|
||||
|
||||
def _wait_for_bytes(self) -> bool:
|
||||
while not self.closing.is_set():
|
||||
if self._has_bytes(_POLL_SECONDS):
|
||||
return True
|
||||
return self._has_bytes(0)
|
||||
|
||||
def _has_bytes(self, timeout: float) -> bool:
|
||||
decrypted: Final = isinstance(self.connection, ssl.SSLSocket) and self.connection.pending() > 0
|
||||
return decrypted or bool(self.selector.select(timeout))
|
||||
|
||||
@contextmanager
|
||||
def _blocking(self, writing: bool) -> Iterator[None]:
|
||||
self.blocked = _Blocked(time.monotonic(), writing)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.blocked = None
|
||||
|
||||
|
||||
def _close_unless_stalled(
|
||||
close: Callable[[], None], connections: SimpleQueue[_Connection], grace: float, errors: SimpleQueue[Exception]
|
||||
) -> None:
|
||||
teardown_started: Final = time.monotonic()
|
||||
closer: Final = threading.Thread(target=close)
|
||||
closer.start()
|
||||
while closer.is_alive():
|
||||
for connection in tuple(connections.get_nowait() for _ in range(connections.qsize())):
|
||||
stalled: Final = connection.stalled(grace, teardown_started)
|
||||
if stalled is None:
|
||||
connections.put(connection)
|
||||
continue
|
||||
if stalled.writing:
|
||||
errors.put(TimeoutError(f"a reply write made no progress for {grace} s"))
|
||||
connection.shut()
|
||||
closer.join(timeout=_POLL_SECONDS)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def wire_server(
|
||||
respond: Callable[[Request], Reply],
|
||||
tls: ssl.SSLContext | None = None,
|
||||
port: int = 0,
|
||||
keep_alive: bool = False,
|
||||
teardown_grace: float = 10,
|
||||
) -> Generator[Wire, None, None]:
|
||||
"""Owned TCP peer; requests traverse the real HTTP client and serialization. With `keep_alive` the
|
||||
peer honours HTTP/1.1 persistent connections so `connections()` counts the client's TCP sessions."""
|
||||
|
|
@ -111,14 +216,20 @@ def wire_server(
|
|||
disconnected: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
connected: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
closed: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
streams: Final[SimpleQueue[_Connection]] = SimpleQueue()
|
||||
closing: Final = threading.Event()
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
timeout = 5
|
||||
request: socket.socket
|
||||
reader: _RequestReader
|
||||
|
||||
def setup(self) -> None:
|
||||
super().setup()
|
||||
stream: Final = _Connection(self.request, closing)
|
||||
self.reader = _RequestReader(stream)
|
||||
self.rfile = self.reader
|
||||
self.wfile = stream
|
||||
streams.put(stream)
|
||||
connected.put(f"{self.client_address[0]}:{self.client_address[1]}")
|
||||
|
||||
def finish(self) -> None:
|
||||
|
|
@ -130,6 +241,9 @@ def wire_server(
|
|||
self.wfile.flush()
|
||||
|
||||
def respond(self) -> None:
|
||||
if self.reader.cut:
|
||||
self.close_connection = True
|
||||
return
|
||||
headers: Final = {name.lower(): value for name, value in self.headers.items()}
|
||||
try:
|
||||
body: Final = read_body(headers, self.rfile)
|
||||
|
|
@ -221,7 +335,7 @@ def wire_server(
|
|||
self.socket = tls.wrap_socket(self.socket, server_side=True)
|
||||
|
||||
with OwnedHTTPServer(("127.0.0.1", port), Handler) as server:
|
||||
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
|
||||
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": _POLL_SECONDS})
|
||||
thread.start()
|
||||
try:
|
||||
yield Wire(
|
||||
|
|
@ -236,6 +350,6 @@ def wire_server(
|
|||
server.shutdown()
|
||||
thread.join(timeout=6)
|
||||
assert not thread.is_alive(), "Owned HTTP server survived cleanup"
|
||||
server.server_close()
|
||||
_close_unless_stalled(server.server_close, streams, teardown_grace, errors)
|
||||
failure: Final = None if errors.empty() else errors.get_nowait()
|
||||
assert failure is None, f"Owned HTTP peer failed: {failure!r}"
|
||||
|
|
|
|||
289
tests/unit/integration_support/test_wire.py
Normal file
289
tests/unit/integration_support/test_wire.py
Normal file
|
|
@ -0,0 +1,289 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import io
|
||||
import socket
|
||||
import ssl
|
||||
import time
|
||||
from collections.abc import Callable, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import ExitStack
|
||||
from http.client import HTTPConnection, HTTPResponse, HTTPSConnection
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import pytest
|
||||
from cryptography import x509
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import ec
|
||||
from cryptography.x509.oid import NameOID
|
||||
|
||||
from tests.integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
PAUSE_SECONDS: Final = 6
|
||||
HANDLER_READ_SECONDS: Final = 0.3
|
||||
FIRST: Final = b'{"call": "first"}'
|
||||
LATE: Final = b'{"call": "late"}'
|
||||
UNREAD: Final = b"x" * (32 * 1024 * 1024)
|
||||
|
||||
|
||||
def _self_signed_server_context(directory: Path) -> ssl.SSLContext:
|
||||
key: Final = ec.generate_private_key(ec.SECP256R1())
|
||||
name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "wire peer")])
|
||||
now: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)
|
||||
certificate: Final = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(name)
|
||||
.issuer_name(name)
|
||||
.public_key(key.public_key())
|
||||
.serial_number(1)
|
||||
.not_valid_before(now)
|
||||
.not_valid_after(now + datetime.timedelta(days=36500))
|
||||
.sign(key, hashes.SHA256())
|
||||
)
|
||||
(directory / "peer.pem").write_bytes(
|
||||
certificate.public_bytes(serialization.Encoding.PEM)
|
||||
+ key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
|
||||
)
|
||||
context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||||
context.load_cert_chain(directory / "peer.pem")
|
||||
return context
|
||||
|
||||
|
||||
def _trusting() -> ssl.SSLContext:
|
||||
context: Final = ssl.create_default_context()
|
||||
context.check_hostname = False
|
||||
context.verify_mode = ssl.CERT_NONE
|
||||
return context
|
||||
|
||||
|
||||
def _client(url: str) -> HTTPConnection:
|
||||
address: Final = urlsplit(url)
|
||||
if address.scheme == "http":
|
||||
return HTTPConnection(address.hostname or "", address.port)
|
||||
return HTTPSConnection(address.hostname or "", address.port, context=_trusting())
|
||||
|
||||
|
||||
def _tls_socket(url: str) -> ssl.SSLSocket:
|
||||
address: Final = urlsplit(url)
|
||||
return _trusting().wrap_socket(socket.create_connection((address.hostname or "", address.port or 0)))
|
||||
|
||||
|
||||
def _echo(request: Request) -> Reply:
|
||||
return Reply(body=request.body)
|
||||
|
||||
|
||||
def _answer_slowly(request: Request) -> Reply:
|
||||
time.sleep(HANDLER_READ_SECONDS)
|
||||
return Reply(body=request.body)
|
||||
|
||||
|
||||
def _answer_in_paused_chunks(request: Request) -> Reply:
|
||||
return Reply(chunks=(request.body[:1], request.body[1:]), pause_between_chunks=HANDLER_READ_SECONDS)
|
||||
|
||||
|
||||
def _hold_first(request: Request) -> Reply:
|
||||
if request.target == "/first":
|
||||
time.sleep(HANDLER_READ_SECONDS)
|
||||
return Reply(body=request.body)
|
||||
|
||||
|
||||
def _paced_pieces(client: socket.socket) -> Iterator[bytes]:
|
||||
while piece := client.recv(256 * 1024):
|
||||
yield piece
|
||||
time.sleep(0.02)
|
||||
|
||||
|
||||
def _read_slowly(client: socket.socket) -> bytes:
|
||||
return b"".join(_paced_pieces(client))
|
||||
|
||||
|
||||
def _read_after(client: socket.socket, delay: float) -> bytes:
|
||||
time.sleep(delay)
|
||||
with client.makefile("rb") as stream:
|
||||
return stream.read()
|
||||
|
||||
|
||||
def _post(target: str, body: bytes) -> bytes:
|
||||
return b"POST %s HTTP/1.1\r\nhost: peer\r\ncontent-length: %d\r\n\r\n%s" % (target.encode(), len(body), body)
|
||||
|
||||
|
||||
def _pause_then_send(connection: HTTPConnection, body: bytes) -> HTTPResponse:
|
||||
time.sleep(PAUSE_SECONDS)
|
||||
connection.request("POST", "/late", body)
|
||||
return connection.getresponse()
|
||||
|
||||
|
||||
def _send_headers_then_pause_before_the_body(connection: HTTPConnection, body: bytes) -> HTTPResponse:
|
||||
connection.putrequest("POST", "/late")
|
||||
connection.putheader("content-length", str(len(body)))
|
||||
connection.endheaders()
|
||||
time.sleep(PAUSE_SECONDS)
|
||||
connection.send(body)
|
||||
return connection.getresponse()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"send_late",
|
||||
(_pause_then_send, _send_headers_then_pause_before_the_body),
|
||||
ids=("pause-before-the-request", "pause-before-the-body"),
|
||||
)
|
||||
def test_a_kept_alive_connection_answers_a_request_whose_bytes_arrive_after_a_long_pause(
|
||||
send_late: Callable[[HTTPConnection, bytes], HTTPResponse],
|
||||
) -> None:
|
||||
with wire_server(_echo, keep_alive=True) as wire:
|
||||
address: Final = urlsplit(wire.url)
|
||||
connection: Final = HTTPConnection(address.hostname or "", address.port)
|
||||
connection.request("POST", "/first", FIRST)
|
||||
first: Final = connection.getresponse().read()
|
||||
late: Final = send_late(connection, LATE)
|
||||
assert (first, late.status, late.read()) == (FIRST, 200, LATE)
|
||||
assert tuple(request.body for request in wire.drain()) == (FIRST, LATE)
|
||||
assert wire.connections() == 1
|
||||
assert connection.sock.recv(1) == b""
|
||||
connection.close()
|
||||
|
||||
|
||||
def test_a_request_cut_off_by_teardown_never_reaches_respond() -> None:
|
||||
with wire_server(_echo, keep_alive=True) as wire:
|
||||
address: Final = urlsplit(wire.url)
|
||||
connection: Final = HTTPConnection(address.hostname or "", address.port)
|
||||
connection.request("POST", "/first", FIRST)
|
||||
first: Final = connection.getresponse().read()
|
||||
connection.sock.sendall(b"POST /cut HTTP/1.1\r\ncontent-type: application/json\r\ncontent-len")
|
||||
time.sleep(HANDLER_READ_SECONDS)
|
||||
assert (first, tuple(request.target for request in wire.drain())) == (FIRST, ("/first",))
|
||||
connection.close()
|
||||
|
||||
|
||||
@pytest.mark.timeout(60)
|
||||
def test_a_reply_its_client_stopped_reading_fails_teardown_instead_of_hanging() -> None:
|
||||
peer: Final = ExitStack()
|
||||
wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=0.5))
|
||||
address: Final = urlsplit(wire.url)
|
||||
with socket.create_connection((address.hostname or "", address.port or 0)) as client:
|
||||
client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n")
|
||||
while wire.received.empty():
|
||||
time.sleep(0.01)
|
||||
with pytest.raises(AssertionError, match="no progress"):
|
||||
peer.close()
|
||||
|
||||
|
||||
@pytest.mark.timeout(60)
|
||||
def test_a_long_reply_its_client_is_still_reading_at_teardown_is_not_cut() -> None:
|
||||
peer: Final = ExitStack()
|
||||
wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=0.5))
|
||||
address: Final = urlsplit(wire.url)
|
||||
with socket.create_connection((address.hostname or "", address.port or 0)) as client, ThreadPoolExecutor(1) as pool:
|
||||
client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n")
|
||||
while wire.received.empty():
|
||||
time.sleep(0.01)
|
||||
reply: Final = pool.submit(_read_slowly, client)
|
||||
time.sleep(1)
|
||||
peer.close()
|
||||
received: Final = reply.result()
|
||||
assert (received[: len(b"HTTP/1.1 200")], received[-len(UNREAD) :] == UNREAD) == (b"HTTP/1.1 200", True)
|
||||
|
||||
|
||||
@pytest.mark.timeout(60)
|
||||
def test_a_reply_its_client_starts_reading_after_teardown_starts_still_arrives_whole() -> None:
|
||||
peer: Final = ExitStack()
|
||||
wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=2))
|
||||
address: Final = urlsplit(wire.url)
|
||||
with socket.create_connection((address.hostname or "", address.port or 0)) as client, ThreadPoolExecutor(1) as pool:
|
||||
client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n")
|
||||
while wire.received.empty():
|
||||
time.sleep(0.01)
|
||||
time.sleep(2.5)
|
||||
reply: Final = pool.submit(_read_after, client, 0.3)
|
||||
peer.close()
|
||||
received: Final = reply.result()
|
||||
assert (received[: len(b"HTTP/1.1 200")], received[-len(UNREAD) :] == UNREAD) == (b"HTTP/1.1 200", True)
|
||||
|
||||
|
||||
def test_a_complete_request_sent_as_teardown_starts_is_still_answered() -> None:
|
||||
with wire_server(_echo, keep_alive=True) as wire:
|
||||
connection: Final = _client(wire.url)
|
||||
connection.request("POST", "/first", FIRST)
|
||||
first: Final = connection.getresponse().read()
|
||||
connection.request("POST", "/complete", LATE)
|
||||
late: Final = connection.getresponse()
|
||||
assert (first, late.status, late.read(), tuple(request.target for request in wire.drain())) == (
|
||||
FIRST,
|
||||
200,
|
||||
LATE,
|
||||
("/first", "/complete"),
|
||||
)
|
||||
connection.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("secure", (False, True), ids=("plain", "tls"))
|
||||
def test_a_reply_still_being_written_at_teardown_reaches_its_client_whole(secure: bool, tmp_path: Path) -> None:
|
||||
slow: Final = Reply(chunks=(b"first", b"second"), pause_between_chunks=HANDLER_READ_SECONDS)
|
||||
tls: Final = _self_signed_server_context(tmp_path) if secure else None
|
||||
with wire_server(lambda _: slow, tls=tls, keep_alive=True) as wire:
|
||||
connection: Final = _client(wire.url)
|
||||
connection.request("GET", "/slow")
|
||||
response: Final = connection.getresponse()
|
||||
head: Final = response.read(len(b"first"))
|
||||
assert head + response.read() == b"firstsecond"
|
||||
connection.close()
|
||||
|
||||
|
||||
def test_a_complete_request_queued_behind_a_busy_handler_is_still_answered() -> None:
|
||||
with wire_server(_hold_first, keep_alive=True) as wire:
|
||||
address: Final = urlsplit(wire.url)
|
||||
client: Final = socket.create_connection((address.hostname or "", address.port or 0))
|
||||
client.sendall(_post("/first", FIRST))
|
||||
while wire.received.empty():
|
||||
time.sleep(0.01)
|
||||
client.sendall(_post("/second", LATE))
|
||||
with client, client.makefile("rb") as stream:
|
||||
replies: Final = stream.read()
|
||||
assert (
|
||||
tuple(request.target for request in wire.drain()),
|
||||
replies.count(b"HTTP/1.1 200"),
|
||||
replies[-len(LATE) :],
|
||||
) == (
|
||||
("/first", "/second"),
|
||||
2,
|
||||
LATE,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("respond", (_answer_slowly, _answer_in_paused_chunks), ids=("slow-respond", "paused-chunks"))
|
||||
def test_a_reply_that_takes_longer_than_the_teardown_grace_still_reaches_its_client(
|
||||
respond: Callable[[Request], Reply],
|
||||
) -> None:
|
||||
with wire_server(respond, keep_alive=True, teardown_grace=HANDLER_READ_SECONDS / 2) as wire:
|
||||
connection: Final = _client(wire.url)
|
||||
connection.request("POST", "/slow", FIRST)
|
||||
while wire.received.empty():
|
||||
time.sleep(0.01)
|
||||
response: Final = connection.getresponse()
|
||||
assert (response.status, response.read()) == (200, FIRST)
|
||||
connection.close()
|
||||
|
||||
|
||||
@pytest.mark.timeout(60)
|
||||
def test_a_tls_request_larger_than_the_read_buffer_is_answered(tmp_path: Path) -> None:
|
||||
large: Final = b"x" * (3 * io.DEFAULT_BUFFER_SIZE // 2)
|
||||
with wire_server(_echo, tls=_self_signed_server_context(tmp_path), keep_alive=True) as wire:
|
||||
client: Final = _tls_socket(wire.url)
|
||||
client.sendall(_post("/large", large))
|
||||
response: Final = HTTPResponse(client)
|
||||
response.begin()
|
||||
assert (response.status, response.read()) == (200, large)
|
||||
client.close()
|
||||
|
||||
|
||||
@pytest.mark.timeout(60)
|
||||
def test_a_tls_record_cut_off_by_teardown_closes_its_connection_instead_of_hanging(tmp_path: Path) -> None:
|
||||
tls: Final = _self_signed_server_context(tmp_path)
|
||||
with wire_server(_echo, tls=tls, keep_alive=True, teardown_grace=HANDLER_READ_SECONDS) as wire:
|
||||
client: Final = _tls_socket(wire.url)
|
||||
socket.socket.sendall(client, b"\x17\x03\x03\x40\x00" + b"x" * 16)
|
||||
time.sleep(HANDLER_READ_SECONDS)
|
||||
assert wire.drain() == ()
|
||||
client.close()
|
||||
Loading…
Add table
Reference in a new issue