litellm/tests/integration/_support/idle_gateway.py

from __future__ import annotations

import itertools
import re
import socket
import struct
import threading
import time
from collections.abc import Callable, Generator, Mapping
from contextlib import contextmanager, suppress
from dataclasses import dataclass
from enum import Enum
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from io import BufferedIOBase
from queue import SimpleQueue
from socketserver import BaseServer
from typing import Final

from integration._support.wire import Reply, Request

MARKER: Final = re.compile(rb"idle-[0-9a-f]{16}")
_LOOPBACK: Final = "127.0.0.1"


def _exactly(stream: BufferedIOBase, size: int) -> bytes | None:
    data: Final = stream.read(size)
    return data if len(data) == size else None


def _body(headers: Mapping[str, str], stream: BufferedIOBase) -> bytes | None:
    if headers.get("transfer-encoding", "").lower() == "chunked":
        return None
    return _exactly(stream, int(headers.get("content-length", "0")))


class Drop(Enum):
    CLOSE = "close"
    RESET = "reset"


@dataclass(frozen=True, slots=True)
class Arrival:
    marker: str
    target: str
    connection: int
    dropped: bool


@dataclass(frozen=True, slots=True)
class IdleGateway:
    """A keep-alive HTTP/1.1 peer that never closes an idle connection itself and, like a load balancer whose
    idle timer already fired, swallows a request that reuses a connection idle for `drop_after` seconds."""

    url: str
    port: int
    arrivals: SimpleQueue[Arrival]

    def pending(self) -> int:
        return self.arrivals.qsize()

    def drain(self) -> tuple[Arrival, ...]:
        return tuple(self.arrivals.get_nowait() for _ in range(self.arrivals.qsize()))


class _Server(ThreadingHTTPServer):
    daemon_threads = True
    block_on_close = False
    request_queue_size = 128

    def __init__(
        self,
        port: int,
        respond: Callable[[Request], Reply],
        drop_after: float | None,
        drop: Drop,
        arrivals: SimpleQueue[Arrival],
    ) -> None:
        super().__init__((_LOOPBACK, port), _Handler)
        self.respond = respond
        self.drop_after = drop_after
        self.drop = drop
        self.arrivals = arrivals
        self.connections = itertools.count(1)
        self.sockets: SimpleQueue[socket.socket] = SimpleQueue()
        self.failures: SimpleQueue[Exception] = SimpleQueue()

    def close_connections(self) -> None:
        for client in tuple(self.sockets.get_nowait() for _ in range(self.sockets.qsize())):
            with suppress(OSError):
                client.shutdown(socket.SHUT_RDWR)

    def handle_error(self, request: object, client_address: object) -> None:
        return


class _Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def __init__(
        self, request: socket.socket | tuple[bytes, socket.socket], client_address: object, server: BaseServer
    ) -> None:
        assert isinstance(server, _Server)
        assert isinstance(request, socket.socket)
        self.gateway = server
        self.client = request
        super().__init__(request, client_address, server)

    def setup(self) -> None:
        super().setup()
        self.connection_id = next(self.gateway.connections)
        self.idle_since: float | None = None
        self.arrived_at = time.monotonic()
        self.gateway.sockets.put(self.client)

    def parse_request(self) -> bool:
        self.arrived_at = time.monotonic()
        return super().parse_request()

    def _stale(self) -> bool:
        return (
            self.gateway.drop_after is not None
            and self.idle_since is not None
            and self.arrived_at - self.idle_since >= self.gateway.drop_after
        )

    def answer(self) -> None:
        headers: Final = {name.lower(): value for name, value in self.headers.items()}
        body: Final = _body(headers, self.rfile)
        if body is None:
            self.gateway.failures.put(ValueError(f"unreadable body for {self.path}: {headers}"))
            self.close_connection = True
            return
        found: Final = MARKER.search(body) or MARKER.search(self.path.encode())
        stale: Final = self._stale()
        self.gateway.arrivals.put(
            Arrival(found.group().decode() if found else "", self.path, self.connection_id, stale)
        )
        if stale:
            self._drop()
            return
        try:
            reply: Final = self.gateway.respond(Request(self.command, self.path, headers, body))
        except (ValueError, KeyError, TypeError) as error:
            self.gateway.failures.put(error)
            self._write(Reply(status=500))
            return
        self._write(reply)
        self.idle_since = time.monotonic()

    do_POST = answer
    do_GET = answer

    def _drop(self) -> None:
        self.close_connection = True
        match self.gateway.drop:
            case Drop.CLOSE:
                with suppress(OSError):
                    self.client.shutdown(socket.SHUT_RDWR)
            case Drop.RESET:
                self.client.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
                self.rfile.close()
                self.client.close()

    def _write(self, reply: Reply) -> None:
        self.send_response(reply.status)
        self.send_header("content-type", reply.content_type)
        for name, value in reply.headers.items():
            self.send_header(name, value)
        if reply.chunks is None:
            self.send_header("content-length", str(len(reply.body)))
        else:
            self.send_header("transfer-encoding", "chunked")
        self.end_headers()
        if reply.chunks is None:
            self.wfile.write(reply.body)
            return
        for index, chunk in enumerate(reply.chunks):
            if index and reply.pause_between_chunks:
                time.sleep(reply.pause_between_chunks)
            self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk))
        self.wfile.write(b"0\r\n\r\n")

    def log_message(self, format: str, *args: object) -> None:
        return


@contextmanager
def idle_gateway(
    respond: Callable[[Request], Reply],
    *,
    drop_after: float | None,
    drop: Drop = Drop.CLOSE,
    port: int = 0,
) -> Generator[IdleGateway]:
    arrivals: Final[SimpleQueue[Arrival]] = SimpleQueue()
    with _Server(port, respond, drop_after, drop, arrivals) as server:
        thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
        thread.start()
        try:
            yield IdleGateway(f"http://{_LOOPBACK}:{server.server_port}", server.server_port, arrivals)
        finally:
            server.shutdown()
            thread.join(timeout=6)
            server.close_connections()
        assert not thread.is_alive(), "Idle gateway survived cleanup"
        failure: Final = None if server.failures.empty() else server.failures.get_nowait()
        assert failure is None, f"Idle gateway failed to answer: {failure!r}"