mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
b602209ca6
* test(integration): hand every harness server a port from one per-worker block
About 45 places in the integration harness picked a port by binding port 0,
reading the number back, closing the socket and starting a server on that
number later. Another process could take the port in between, which surfaced
as "Address already in use" flakes. Earlier fixes retried around the race in
single call sites without removing it.
tests/integration/_support/ports.py now owns every port the harness picks
ahead of a server. Each pytest process claims a 200-port block between 10000
and 32767 by taking an exclusive flock on a per-block lock file in the system
temp directory, so parallel runs on one machine never share a block and the
lock goes away with the process. The range is above the fixed service ports
and below the Linux and macOS ephemeral ranges. Ports go out in order, are
never reused in a run, skip anything an outside process already holds, and
a process that exhausts its block claims the next free one.
The three _free_port copies, worker_endpoints' port-0 binds,
free_loopback_port, the owned proxy's lost-port-race retry and the owned
Redis retry loop are gone. Tests that restart a server on the same port now
allocate that port up front instead of reading it back from the first
server. TQ012 in scripts/check_test_quality.py rejects a port-0 bind read
back with getsockname anywhere under tests/integration.
* ci(integration): record listening sockets and the ephemeral port range
* Revert "ci(integration): record listening sockets and the ephemeral port range"
This reverts commit d3fafc3da2.
330 lines
13 KiB
Python
330 lines
13 KiB
Python
import asyncio
|
|
import threading
|
|
import time
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
from typing import Final
|
|
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
|
|
|
|
|
from tests.integration._support.async_server import LoopbackServer, close_writer
|
|
from tests.integration._support.ports import allocate_port
|
|
|
|
OUTAGE_SECONDS: Final = 10.0
|
|
|
|
|
|
class DatabaseRelay:
|
|
def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None:
|
|
self.port: Final = allocate_port()
|
|
self._upstream_host: Final = upstream_host
|
|
self._upstream_port: Final = upstream_port
|
|
self._trigger: Final = trigger
|
|
self._server: Final = LoopbackServer(self.port, self._serve)
|
|
self._loop: Final = self._server.loop
|
|
self._armed: Final = threading.Event()
|
|
self.tripped: Final = threading.Event()
|
|
self.refused = 0
|
|
self.reconnected: Final = threading.Event()
|
|
self._tripped_at = 0.0
|
|
self._writers: tuple[asyncio.StreamWriter, ...] = ()
|
|
|
|
def arm(self) -> None:
|
|
self._armed.set()
|
|
|
|
def start(self) -> None:
|
|
self._server.start()
|
|
|
|
def stop(self) -> None:
|
|
self._server.stop()
|
|
|
|
def _drop_all(self) -> None:
|
|
for writer in self._writers:
|
|
writer.close()
|
|
self._writers = ()
|
|
|
|
async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
|
|
if self.tripped.is_set() and time.monotonic() - self._tripped_at < OUTAGE_SECONDS:
|
|
self.refused += 1
|
|
client_writer.close()
|
|
return
|
|
if self.tripped.is_set():
|
|
self.reconnected.set()
|
|
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
|
|
self._writers = (*self._writers, client_writer, server_writer)
|
|
|
|
async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None:
|
|
try:
|
|
while chunk := await reader.read(65536):
|
|
if inspect and self._armed.is_set() and not self.tripped.is_set() and self._trigger in chunk:
|
|
self._tripped_at = time.monotonic()
|
|
self.tripped.set()
|
|
self._drop_all()
|
|
return
|
|
writer.write(chunk)
|
|
await writer.drain()
|
|
except (ConnectionError, asyncio.IncompleteReadError):
|
|
return
|
|
finally:
|
|
await close_writer(writer)
|
|
|
|
await asyncio.gather(
|
|
forward(client_reader, server_writer, True),
|
|
forward(server_reader, client_writer, False),
|
|
)
|
|
|
|
|
|
class HeldStatementRelay:
|
|
def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None:
|
|
self.port: Final = allocate_port()
|
|
self._upstream_host: Final = upstream_host
|
|
self._upstream_port: Final = upstream_port
|
|
self._trigger: Final = trigger
|
|
self._server: Final = LoopbackServer(self.port, self._serve)
|
|
self._loop: Final = self._server.loop
|
|
self._released: Final = asyncio.Event()
|
|
self.held: Final = threading.Event()
|
|
|
|
def release(self) -> None:
|
|
if not self._loop.is_closed():
|
|
self._loop.call_soon_threadsafe(self._released.set)
|
|
|
|
def start(self) -> None:
|
|
self._server.start()
|
|
|
|
def stop(self) -> None:
|
|
self.release()
|
|
self._server.stop()
|
|
|
|
def _holds(self, window: bytes) -> bool:
|
|
return not self.held.is_set() and self._trigger in window
|
|
|
|
async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
|
|
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
|
|
|
|
async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None:
|
|
tail = b"" # rebind-ok: carries the previous read's end so a trigger split across reads still matches
|
|
try:
|
|
while chunk := await reader.read(65536):
|
|
window: Final = tail + chunk
|
|
if inspect and self._holds(window):
|
|
self.held.set()
|
|
await self._released.wait()
|
|
tail = window[-(len(self._trigger) - 1) :]
|
|
writer.write(chunk)
|
|
await writer.drain()
|
|
except (ConnectionError, asyncio.IncompleteReadError):
|
|
return
|
|
finally:
|
|
await close_writer(writer)
|
|
|
|
await asyncio.gather(
|
|
forward(client_reader, server_writer, True),
|
|
forward(server_reader, client_writer, False),
|
|
)
|
|
|
|
|
|
class TriggerScanner:
|
|
def __init__(self, trigger: bytes) -> None:
|
|
self._trigger: Final = trigger
|
|
self._tail: bytes = b""
|
|
|
|
def feed(self, chunk: bytes) -> bool:
|
|
window: Final = self._tail + chunk
|
|
self._tail = window[-(len(self._trigger) - 1) :]
|
|
return self._trigger in window
|
|
|
|
|
|
class DroppedConnectionRelay:
|
|
def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None:
|
|
self.port: Final = allocate_port()
|
|
self._upstream_host: Final = upstream_host
|
|
self._upstream_port: Final = upstream_port
|
|
self._trigger: Final = trigger
|
|
self._server: Final = LoopbackServer(self.port, self._serve)
|
|
self._loop: Final = self._server.loop
|
|
self._armed: Final = threading.Event()
|
|
self.dropped: Final = threading.Event()
|
|
|
|
def arm(self) -> None:
|
|
self._armed.set()
|
|
|
|
def disarm(self) -> None:
|
|
self._armed.clear()
|
|
|
|
def start(self) -> None:
|
|
self._server.start()
|
|
|
|
def stop(self) -> None:
|
|
self._server.stop()
|
|
|
|
async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
|
|
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
|
|
|
|
async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None:
|
|
scanner: Final = TriggerScanner(self._trigger)
|
|
try:
|
|
while chunk := await reader.read(65536):
|
|
matched: Final = scanner.feed(chunk)
|
|
if inspect and self._armed.is_set() and matched:
|
|
self.dropped.set()
|
|
client_writer.close()
|
|
return
|
|
writer.write(chunk)
|
|
await writer.drain()
|
|
except (ConnectionError, asyncio.IncompleteReadError):
|
|
return
|
|
finally:
|
|
await close_writer(writer)
|
|
|
|
await asyncio.gather(
|
|
forward(client_reader, server_writer, True),
|
|
forward(server_reader, client_writer, False),
|
|
)
|
|
|
|
|
|
SSL_REQUEST_CODE: Final = 80877103
|
|
GSSENC_REQUEST_CODE: Final = 80877104
|
|
|
|
|
|
class FrontendStatementScanner:
|
|
"""Counts the statements one Postgres client sends: a Sync closes every extended-protocol statement and a
|
|
Query is a simple-protocol statement, so the count moves with the number of statements, never their size."""
|
|
|
|
def __init__(self) -> None:
|
|
self._buffer = b""
|
|
self._started = False
|
|
|
|
def feed(self, chunk: bytes) -> int:
|
|
self._buffer += chunk # rebind-ok: a message can span two reads
|
|
statements = 0 # rebind-ok: summed over the messages this chunk completes
|
|
while (frame := self._next_frame()) is not None:
|
|
length, is_statement = frame
|
|
statements += int(is_statement)
|
|
self._buffer = self._buffer[length:] # rebind-ok: drops the parsed message
|
|
return statements
|
|
|
|
def _next_frame(self) -> tuple[int, bool] | None:
|
|
if not self._started:
|
|
if len(self._buffer) < 8:
|
|
return None
|
|
startup_length: Final = int.from_bytes(self._buffer[:4], "big")
|
|
if len(self._buffer) < startup_length:
|
|
return None
|
|
request_code: Final = int.from_bytes(self._buffer[4:8], "big")
|
|
self._started = request_code not in (SSL_REQUEST_CODE, GSSENC_REQUEST_CODE)
|
|
return startup_length, False
|
|
if len(self._buffer) < 5:
|
|
return None
|
|
frame_length: Final = 1 + int.from_bytes(self._buffer[1:5], "big")
|
|
if len(self._buffer) < frame_length:
|
|
return None
|
|
return frame_length, self._buffer[:1] in (b"S", b"Q")
|
|
|
|
|
|
class StatementCountingRelay:
|
|
def __init__(self, upstream_host: str, upstream_port: int) -> None:
|
|
self.port: Final = allocate_port()
|
|
self._upstream_host: Final = upstream_host
|
|
self._upstream_port: Final = upstream_port
|
|
self._loop: Final = asyncio.new_event_loop()
|
|
self.statements = 0
|
|
self._ready: Final = threading.Event()
|
|
self._thread: Final = threading.Thread(target=self._run, daemon=True)
|
|
|
|
def start(self) -> None:
|
|
self._thread.start()
|
|
assert self._ready.wait(10), "Database relay did not start"
|
|
|
|
def stop(self) -> None:
|
|
self._loop.call_soon_threadsafe(self._loop.stop)
|
|
self._thread.join(10)
|
|
|
|
def _run(self) -> None:
|
|
asyncio.set_event_loop(self._loop)
|
|
self._loop.run_until_complete(asyncio.start_server(self._serve, "127.0.0.1", self.port))
|
|
self._ready.set()
|
|
self._loop.run_forever()
|
|
|
|
async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
|
|
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
|
|
|
|
async def forward(
|
|
reader: asyncio.StreamReader, writer: asyncio.StreamWriter, scanner: FrontendStatementScanner | None
|
|
) -> None:
|
|
try:
|
|
while chunk := await reader.read(65536):
|
|
if scanner is not None:
|
|
self.statements += scanner.feed(chunk)
|
|
writer.write(chunk)
|
|
await writer.drain()
|
|
except (ConnectionError, asyncio.IncompleteReadError):
|
|
return
|
|
finally:
|
|
writer.close()
|
|
|
|
await asyncio.gather(
|
|
forward(client_reader, server_writer, FrontendStatementScanner()),
|
|
forward(server_reader, client_writer, None),
|
|
)
|
|
|
|
|
|
def plaintext_database_url(database_url: str) -> str:
|
|
parts: Final = urlsplit(database_url)
|
|
parameters: Final = tuple(
|
|
(name, value) for name, value in parse_qsl(parts.query, keep_blank_values=True) if name != "sslmode"
|
|
) + (("sslmode", "disable"),)
|
|
return urlunsplit(parts._replace(query=urlencode(parameters)))
|
|
|
|
|
|
def _relayed_url(database_url: str, port: int) -> str:
|
|
parts: Final = urlsplit(database_url)
|
|
credentials: Final = f"{parts.username}:{parts.password}@" if parts.username else ""
|
|
return urlunsplit(parts._replace(netloc=f"{credentials}127.0.0.1:{port}"))
|
|
|
|
|
|
@contextmanager
|
|
def database_relay(database_url: str, trigger: bytes) -> Generator[tuple[DatabaseRelay, str]]:
|
|
parts: Final = urlsplit(database_url)
|
|
assert parts.hostname is not None and parts.port is not None, database_url
|
|
relay: Final = DatabaseRelay(parts.hostname, parts.port, trigger)
|
|
try:
|
|
relay.start()
|
|
yield relay, _relayed_url(database_url, relay.port)
|
|
finally:
|
|
relay.stop()
|
|
|
|
|
|
@contextmanager
|
|
def held_statement_relay(database_url: str, trigger: bytes) -> Generator[tuple[HeldStatementRelay, str]]:
|
|
parts: Final = urlsplit(database_url)
|
|
assert parts.hostname is not None and parts.port is not None, database_url
|
|
relay: Final = HeldStatementRelay(parts.hostname, parts.port, trigger)
|
|
try:
|
|
relay.start()
|
|
yield relay, _relayed_url(database_url, relay.port)
|
|
finally:
|
|
relay.stop()
|
|
|
|
|
|
@contextmanager
|
|
def dropped_connection_relay(database_url: str, trigger: bytes) -> Generator[tuple[DroppedConnectionRelay, str]]:
|
|
parts: Final = urlsplit(database_url)
|
|
assert parts.hostname is not None and parts.port is not None, database_url
|
|
relay: Final = DroppedConnectionRelay(parts.hostname, parts.port, trigger)
|
|
try:
|
|
relay.start()
|
|
yield relay, _relayed_url(database_url, relay.port)
|
|
finally:
|
|
relay.stop()
|
|
|
|
|
|
@contextmanager
|
|
def statement_counting_relay(database_url: str) -> Generator[tuple[StatementCountingRelay, str]]:
|
|
parts: Final = urlsplit(database_url)
|
|
assert parts.hostname is not None and parts.port is not None, database_url
|
|
relay: Final = StatementCountingRelay(parts.hostname, parts.port)
|
|
relay.start()
|
|
try:
|
|
yield relay, plaintext_database_url(_relayed_url(database_url, relay.port))
|
|
finally:
|
|
relay.stop()
|