diff --git a/src/server/server.cpp b/src/server/server.cpp index 08c0c06c5..5a7adfa40 100644 --- a/src/server/server.cpp +++ b/src/server/server.cpp @@ -172,6 +172,18 @@ HttpSession::HttpSession(tcp::socket socket, WebServer& server) header_print("🔗 ", "TCP connection established - Remote: " + socket_.remote_endpoint().address().to_string() + ":" + std::to_string(socket_.remote_endpoint().port())); } +///@brief destructor: a session never outlives its slot in the connection limit +HttpSession::~HttpSession() { + release_connection(); +} + +///@brief release this connection's slot (idempotent; every close path may call it) +void HttpSession::release_connection() { + if (!connection_released_.exchange(true)) { + server_.active_connections_.fetch_sub(1); + } +} + ///@brief start void HttpSession::start(bool cors) { read_request(cors); @@ -188,7 +200,7 @@ void HttpSession::close_connection() { header_print("🔒 ", "TCP connection closing - Remote endpoint unavailable"); } socket_.shutdown(tcp::socket::shutdown_both, ec); - server_.active_connections_.fetch_sub(1); + release_connection(); } ///@brief read request @@ -229,7 +241,7 @@ void HttpSession::read_request(bool cors) { } catch (...) { header_print("🔒 ", "TCP connection closed - Remote endpoint unavailable"); } - self->server_.active_connections_.fetch_sub(1); + self->release_connection(); }); } @@ -318,7 +330,7 @@ void HttpSession::write_response() { [self](beast::error_code ec, std::size_t) { self->socket_.shutdown(tcp::socket::shutdown_both, ec); // Decrement connection counter for non-keep-alive connections - self->server_.active_connections_.fetch_sub(1); + self->release_connection(); try { header_print("🔒 ", "TCP connection closed - Remote: " + self->socket_.remote_endpoint().address().to_string() + ":" + std::to_string(self->socket_.remote_endpoint().port())); } @@ -453,10 +465,10 @@ void HttpSession::send_chunk_data(const json& data, bool is_final) { } //socket_.shutdown(tcp::socket::shutdown_both, ec); - //server_.active_connections_.fetch_sub(1); boost::system::error_code ignore_ec; socket_.close(ignore_ec); + release_connection(); return; } @@ -467,7 +479,7 @@ void HttpSession::send_chunk_data(const json& data, bool is_final) { header_print("🔒 ", "Closing TCP connection (streaming, non-keep-alive)"); socket_.shutdown(tcp::socket::shutdown_both, ec); // Decrement connection counter for non-keep-alive connections - server_.active_connections_.fetch_sub(1); + release_connection(); } } diff --git a/src/server/server.hpp b/src/server/server.hpp index 14c351b8f..83a82fe67 100644 --- a/src/server/server.hpp +++ b/src/server/server.hpp @@ -265,6 +265,7 @@ class WebServer { class HttpSession : public std::enable_shared_from_this { public: HttpSession(tcp::socket socket, WebServer& server); + ~HttpSession(); void start(bool cors); void write_streaming_response(const json& data, bool is_final); void close_connection(); @@ -276,6 +277,8 @@ class HttpSession : public std::enable_shared_from_this { void write_response(); void send_chunk_data(const json& data, bool is_final); void start_client_disconnect_monitor(std::shared_ptr token); + /// Give this connection's slot back to the server's connection limit, exactly once. + void release_connection(); ///@brief socket tcp::socket socket_; @@ -292,6 +295,8 @@ class HttpSession : public std::enable_shared_from_this { ///@brief stream buffer std::shared_ptr stream_buf_; std::shared_ptr cancellation_token_; + ///@brief whether this session's slot in active_connections_ was already released + std::atomic connection_released_{false}; }; // Forward declarations diff --git a/src/test/openai_api_test/README.md b/src/test/openai_api_test/README.md index d4066706b..8cee11b16 100644 --- a/src/test/openai_api_test/README.md +++ b/src/test/openai_api_test/README.md @@ -40,6 +40,7 @@ Each case prints `PASS`, `FAIL` or `SKIP`. The exit code is non-zero if anything | `test_sampling.py` | The same `seed` samples the same text (with an unseeded request in between, streamed or not), `repeat_last_n: 0` turns penalties off, `-1` covers the whole reply, malformed values get 400 | | `test_continue_final_message.py` | `continue_final_message` continues a trailing assistant message instead of answering it: content is only the new text, whitespace as written (a model that reasons first may read the prefix as reasoning, and close it), `logprobs` cover exactly the continued tokens (the first entry scores the next token), stream matches non-stream, a cache-resumed continuation matches a cold one, a normal turn can follow, malformed combinations get 400 | | `test_errors.py` | Every malformed `tool_choice` / `response_format` gets HTTP 400 with an OpenAI error body, streaming too, and the server keeps serving | +| `test_connections.py` | Connections give their slot back however they end: clients hanging up mid-stream, before the first chunk or on a non-streaming request, finished requests, idle and malformed connections. Needs the default `--socket 10` | Some cases skip by design: diff --git a/src/test/openai_api_test/run_all.py b/src/test/openai_api_test/run_all.py index 58df824f4..ed344db0d 100644 --- a/src/test/openai_api_test/run_all.py +++ b/src/test/openai_api_test/run_all.py @@ -13,6 +13,7 @@ import sys import flm_api +import test_connections import test_continue_final_message import test_errors import test_logprobs @@ -23,7 +24,7 @@ import test_tool_calls MODULES = [test_response_shape, test_reasoning, test_tool_calls, test_response_format, test_logprobs, test_sampling, - test_continue_final_message, test_errors] + test_continue_final_message, test_errors, test_connections] def main() -> None: diff --git a/src/test/openai_api_test/test_connections.py b/src/test/openai_api_test/test_connections.py new file mode 100644 index 000000000..be406629f --- /dev/null +++ b/src/test/openai_api_test/test_connections.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +"""Connection slots: every connection must give its slot back, however it ends. + +The server accepts at most ``--socket`` connections at once (default 10) and +closes the rest on accept. A client that hangs up on a streaming reply used to +keep its slot forever, so after that many dropped streams the server turned +every connection away until it was restarted. These cases drop more clients +than the default limit allows, then check the limit is still exactly what it +should be: ``limit - 1`` idle sockets leave room for one request, ``limit`` +idle sockets do not. A leaked slot fails the first check, a slot released twice +fails the second. + +Needs a server started with the default ``--socket 10``. + +Usage:: + + python3 test_connections.py [--endpoint URL] [-k NAME] [-v] +""" + +from __future__ import annotations + +import json +import socket +import time +import urllib.parse +import urllib.request +from typing import Any + +import flm_api +from flm_api import Api, case, expect +from fixtures import user + +LIMIT = 10 # the default --socket +DROPS = LIMIT + 2 # more than the limit, so a leak per drop cannot go unnoticed +BARRIER_S = 180.0 # most we wait for the server to work through the dropped requests + + +def address(api: Api) -> tuple[str, int]: + parts = urllib.parse.urlsplit(api.url) + return parts.hostname or "127.0.0.1", parts.port or 80 + + +def send_chat(api: Api, body: dict[str, Any]) -> socket.socket: + payload = json.dumps(body).encode("utf-8") + sock = socket.create_connection(address(api), timeout=api.timeout) + sock.sendall(b"POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\n" + b"Content-Type: application/json\r\nContent-Length: " + str(len(payload)).encode() + + b"\r\n\r\n" + payload) + return sock + + +def drop_clients(api: Api, *, stream: bool, wait_for_reply: bool) -> None: + """Send DROPS long requests and hang up on each one.""" + body = api.body(user("Write a long story about a dragon."), max_tokens=200, stream=stream) + for _ in range(DROPS): + sock = send_chat(api, body) + if wait_for_reply: + sock.recv(4096) # the first streamed bytes, or the whole non-streaming reply + else: + time.sleep(0.05) + sock.close() + time.sleep(0.2) + drain(api) + + +def drain(api: Api) -> None: + """Wait until the server has unwound every dropped request. + + A dropped request keeps its slot until it is dequeued and cancelled, and the + server runs one request at a time, so a short request that completes proves + all the earlier ones are gone. (It may be turned away while the queue is full.) + """ + deadline = time.monotonic() + BARRIER_S + while time.monotonic() < deadline: + try: + if api.post(api.body(user("Say hi."), max_tokens=4)).status == 200: + time.sleep(0.5) + return + except OSError: + pass # turned away (limit reached) or reset: try again + time.sleep(1.0) + expect(False, f"server did not get through the dropped requests within {BARRIER_S:.0f}s") + + +def models_ok(api: Api) -> bool: + """GET /v1/models; False when the server turns the connection away.""" + url = api.url.replace("/chat/completions", "/models") + try: + with urllib.request.urlopen(url, timeout=5) as resp: + resp.read() + return resp.status == 200 + except OSError: + return False + + +def idle_sockets(api: Api, count: int) -> list[socket.socket]: + socks = [socket.create_connection(address(api), timeout=5) for _ in range(count)] + time.sleep(0.3) # let the server accept them all + return socks + + +def expect_exact_limit(api: Api) -> None: + """The limit still admits LIMIT connections in total, and no more.""" + expect(models_ok(api), "server turns away a lone request") + for idle, admitted in ((LIMIT - 1, True), (LIMIT, False)): + socks = idle_sockets(api, idle) + try: + got = models_ok(api) + finally: + for sock in socks: + sock.close() + time.sleep(0.5) + expect(got == admitted, f"with {idle} idle connections a request was " + f"{'served' if got else 'turned away'}, expected " + f"{'served' if admitted else 'turned away'} (leaked or double-released slots)") + + +@case +def dropped_streams_release_their_slot(api: Api) -> None: + """More hang-ups mid-stream than the limit allows leave the limit intact.""" + drop_clients(api, stream=True, wait_for_reply=True) + expect_exact_limit(api) + + +@case +def dropped_streams_before_first_chunk(api: Api) -> None: + """Hanging up straight after sending a streaming request (during prefill).""" + drop_clients(api, stream=True, wait_for_reply=False) + expect_exact_limit(api) + + +@case +def dropped_non_streaming_requests(api: Api) -> None: + """Hanging up on a non-streaming request while it generates.""" + drop_clients(api, stream=False, wait_for_reply=False) + expect_exact_limit(api) + + +@case +def finished_requests_release_their_slot(api: Api) -> None: + """Completed streaming and non-streaming requests release exactly once.""" + for stream in (True, False): + for _ in range(DROPS): + body = api.body(user("Say hi."), max_tokens=8) + if stream: + api.ok_stream(body) + else: + api.ok(body) + expect_exact_limit(api) + + +@case +def idle_and_malformed_connections(api: Api) -> None: + """Connections that open and close, or send garbage, release their slot.""" + host, port = address(api) + for _ in range(DROPS): + socket.create_connection((host, port), timeout=5).close() + for junk in (b"GARBAGE\r\n\r\n", b"POST /v1/chat/completions HTTP/1.1\r\nContent-Length: 5\r\n\r\nabc"): + for _ in range(DROPS): + sock = socket.create_connection((host, port), timeout=5) + sock.sendall(junk) + time.sleep(0.05) + sock.close() + time.sleep(1.0) + expect_exact_limit(api) + + +@case +def server_still_serves(api: Api) -> None: + """After all the hang-ups a normal request still succeeds.""" + flm_api.check_message_contract(api.ok(api.body(user("Say hi."), max_tokens=16))) + + +CASES = flm_api.collect(globals()) + +if __name__ == "__main__": + flm_api.main(CASES, __doc__.splitlines()[0])