From ea1f7bbb5df26fbc18e8d46f72a6de43a84616dd Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sat, 11 Jul 2026 12:41:47 +0200 Subject: [PATCH] server: refactor server_stream (#25541) * server: refactoring, remove spipe from server_http_res * wip * remove non-thread-safe rd.stop() call * move server_res_spipe * nits * improve server_stream_create_spipe * server-stream: update dev docs for the improved API --------- Co-authored-by: Pascal --- tools/server/README-dev.md | 8 +- tools/server/server-context.cpp | 27 ++---- tools/server/server-http.cpp | 19 +---- tools/server/server-http.h | 11 +-- tools/server/server-stream.cpp | 146 ++++++++++++++------------------ tools/server/server-stream.h | 47 +++++----- 6 files changed, 99 insertions(+), 159 deletions(-) diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index 882adca09b..a5619485f0 100644 --- a/tools/server/README-dev.md +++ b/tools/server/README-dev.md @@ -126,15 +126,15 @@ It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`. The feature lives entirely in `server-stream.{h,cpp}` and rests on three types: -- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session. +- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` sets the flag the producer polls. One conv maps to at most one live session. - `stream_session_manager`: a file-static singleton (`g_stream_sessions`) inside `server-stream.cpp`, owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. Exposed to main only through `server_stream_session_manager_start/stop`. - `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation. -The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, `server_stream_session_attach_pipe`, `server_stream_aware_should_stop`, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager and consumer types stay in the `.cpp`. +The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, the `server_res_spipe` response base, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager, consumer and the `server_stream_create_spipe` factory stay in the `.cpp`. -Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `server_stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker. +Producer side: `server_res_generator` extends `server_res_spipe`, which keeps all spipe logic out of the generic `server_http_res`. `set_req` attaches a producer when the header is present, and the wrapped `next` tees each chunk into the ring before the socket, so a chunk lost to a dead wire is already buffered. While attached, `should_stop` ignores peer disconnect: only a `DELETE` stops generation. On an early peer drop, `on_complete` drains the tail into the ring on the http worker. -Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free. +Lifetime safety: the session holds no back reference to the response, so `spipe` is a plain `unique_ptr` touched only by the http worker. `cancel` raises an atomic the producer polls; the producer finalizes the session from its destructor, which also runs `~server_response_reader::stop()` to cancel the generation at the queue level. A `DELETE` stops work by raising the flag and letting the worker unwind. Consumer side: `GET /v1/stream/?from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400. diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 98d0cca1cc..4826ca5d03 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3979,11 +3979,9 @@ server_context_meta server_context::get_meta() const { }; } - - // generator-like API for HTTP response generation // may have bypass_sleep = true if the task does not use ctx_server -struct server_res_generator : server_http_res { +struct server_res_generator : server_res_spipe { server_response_reader rd; server_res_generator(server_queue & queue_tasks, server_response & queue_results, int sleep_idle_seconds, bool bypass_sleep = false) : rd(queue_tasks, queue_results, HTTP_POLLING_SECONDS) { @@ -3993,15 +3991,6 @@ struct server_res_generator : server_http_res { queue_tasks.wait_until_no_sleep(); } } - ~server_res_generator() override { - // cleanup() must run while rd is still alive (rd is destroyed after this body returns) - if (spipe) { - spipe->cleanup(); - } - } - void stop() override { - rd.stop(); - } void ok(const json & response_data) { status = 200; data = safe_json_to_str(response_data); @@ -4039,6 +4028,8 @@ std::unique_ptr server_routes::handle_completions_impl( auto & rd = res->rd; auto & params = this->params; + res->set_req(&req); // will also set spipe if needed + int32_t sse_ping_interval = params.sse_ping_interval; try { @@ -4181,7 +4172,7 @@ std::unique_ptr server_routes::handle_completions_impl( } res->status = 200; res->content_type = "text/event-stream"; - res->next = [res_this = res.get(), res_type, sse_ping_interval, &req](std::string & output) -> bool { + res->set_next([res_this = res.get(), res_type, sse_ping_interval](std::string & output) -> bool { static auto format_error = [](task_response_type res_type, const json & res_json) { if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { return format_anthropic_sse({ @@ -4193,7 +4184,9 @@ std::unique_ptr server_routes::handle_completions_impl( } }; - auto effective_should_stop = server_stream_aware_should_stop(res_this, req.should_stop); + auto effective_should_stop = [&res_this]() { + return res_this->should_stop(); + }; try { if (effective_should_stop()) { @@ -4284,13 +4277,9 @@ std::unique_ptr server_routes::handle_completions_impl( // terminate on exception return false; } - }; + }); } - // attach a producer pipe to the response when X-Conversation-Id is present. - // the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook. - server_stream_session_attach_pipe(*res, req.headers); - return res; } diff --git a/tools/server/server-http.cpp b/tools/server/server-http.cpp index 87eee5fd4d..48f903dfcb 100644 --- a/tools/server/server-http.cpp +++ b/tools/server/server-http.cpp @@ -1,7 +1,6 @@ #include "common.h" #include "http.h" #include "server-http.h" -#include "server-stream.h" #include "server-common.h" #include "ui.h" @@ -530,33 +529,20 @@ static void process_handler_response(server_http_req_ptr && request, server_http std::string chunk; const bool has_next = response->next(chunk); if (!chunk.empty()) { - // mirror into the ring buffer first, the session must reflect every SSE chunk - // whether or not the wire write below succeeds - if (response->spipe) { - response->spipe->write(chunk.data(), chunk.size()); - } if (!sink.write(chunk.data(), chunk.size())) { - // peer is gone, stop the wire path here return false; } SRV_DBG("http: streamed chunk: %s\n", chunk.c_str()); } if (!has_next) { - // producer reached its natural end on the wire, a later close() skips the drain - if (response->spipe) { - response->spipe->done(); - } sink.done(); SRV_DBG("%s", "http: stream ended\n"); } return has_next; }; const auto on_complete = [request = q_ptr, response = r_ptr](bool) mutable { - // on a dropped peer, close() drains the rest of the generation into the ring buffer - if (response->spipe) { - response->spipe->close(); - } - response.reset(); // spipe destructor finalizes the session if attached + response->on_complete(); + response.reset(); request.reset(); }; res.set_chunked_content_provider(content_type, chunked_content_provider, on_complete); @@ -564,6 +550,7 @@ static void process_handler_response(server_http_req_ptr && request, server_http res.status = response->status; set_headers(res, response->headers); res.set_content(response->data, response->content_type); + response->on_complete(); } } diff --git a/tools/server/server-http.h b/tools/server/server-http.h index 3508131836..032b08d0d2 100644 --- a/tools/server/server-http.h +++ b/tools/server/server-http.h @@ -11,7 +11,6 @@ #include struct common_params; -struct stream_pipe_producer; // defined in server-stream.h // generator-like API for HTTP response generation // this object response with one of the 2 modes: @@ -25,19 +24,13 @@ struct server_http_res { std::string data; std::map headers; - // if set, the stream survives a client disconnect: the producer pipe keeps draining into the - // ring buffer and finalizes the session on destruction, so no explicit on_stream_end is needed. - // shared_ptr (not unique_ptr) so the forward-declared type is safe to delete here. - std::shared_ptr spipe; - std::function next = nullptr; bool is_stream() const { return next != nullptr; } - // called when the session is cancelled (e.g. DELETE /v1/stream/). - // server_res_generator overrides this to stop its reader; the default is a no-op. - virtual void stop() {} + // fired before req and res are destroyed + virtual void on_complete() {} virtual ~server_http_res() = default; }; diff --git a/tools/server/server-stream.cpp b/tools/server/server-stream.cpp index 553ac26b1e..19db04988d 100644 --- a/tools/server/server-stream.cpp +++ b/tools/server/server-stream.cpp @@ -96,8 +96,6 @@ struct stream_session { size_t dropped_prefix() const; // bytes evicted from the front due to cap int64_t completed_at() const; // 0 while alive, unix seconds after finalize - void set_stop_producer(std::function fn); - void cancel(); private: @@ -109,7 +107,6 @@ private: bool done; std::atomic cancelled; // polled lock-free by the should_stop closure, no mu int64_t completed_ts; - std::function stop_producer; }; stream_session::stream_session(std::string conversation_id_, size_t max_bytes_) : conversation_id(std::move(conversation_id_)) @@ -217,26 +214,10 @@ int64_t stream_session::completed_at() const { return completed_ts; } -void stream_session::set_stop_producer(std::function fn) { - std::lock_guard lock(mu); - stop_producer = std::move(fn); -} - void stream_session::cancel() { - // flip cancelled first so the producer-side server_stream_aware_should_stop can break out of the - // recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task - // posted by rd.stop() will eventually notify, but we do not want to depend on that timing) + // the should_stop closure on both the producer and any HTTP reader polls is_cancelled() + // so flipping this is the only signal needed to unwind both sides cancelled.store(true, std::memory_order_release); - // copy the hook under the lock then invoke outside, the producer side may grab queue locks - // and we do not want to hold our mu across that path - std::function fn; - { - std::lock_guard lock(mu); - fn = stop_producer; - } - if (fn) { - fn(); - } } bool stream_session::is_cancelled() const { @@ -325,8 +306,10 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i s = it->second; sessions.erase(it); } - // signal the producer side first so the inference is cancelled at the queue level, - // then finalize, which wakes any pending HTTP reader and lets the drain exit naturally + // cancel first so the producer's on_complete() drain loop and any pending HTTP reader + // observe is_cancelled() and stop pulling further output, then finalize to wake readers + // blocked in read_from(). note: this does not interrupt the underlying generation itself, + // which keeps running to its own natural stop condition (EOS/max_tokens) s->cancel(); s->finalize(); } @@ -431,65 +414,15 @@ stream_pipe_producer::stream_pipe_producer(stream_session_ptr session) } stream_pipe_producer::~stream_pipe_producer() { - cleanup(); session_->finalize(); } -void stream_pipe_producer::cleanup() { - if (!alive_) { - return; - } - alive_->store(false, std::memory_order_release); - session_->set_stop_producer(nullptr); - alive_.reset(); -} - bool stream_pipe_producer::write(const char * data, size_t len) { return session_->append(data, len); } -void stream_pipe_producer::done() { - done_ = true; -} - -void stream_pipe_producer::close() { - // httplib bails its content provider the moment is_peer_alive() goes false, so pump the rest - // of the generation into the ring buffer here. a DELETE flips is_cancelled and cuts it short - if (done_ || session_->is_cancelled()) { - SRV_TRC("stream_pipe close: skip drain (done=%d cancelled=%d) conv=%s\n", - done_ ? 1 : 0, session_->is_cancelled() ? 1 : 0, session_->conversation_id.c_str()); - return; - } - SRV_TRC("stream_pipe close: draining conv=%s\n", session_->conversation_id.c_str()); - size_t drained = 0; - std::string chunk; - while (true) { - chunk.clear(); - bool has_next = res_->next(chunk); - if (!chunk.empty()) { - write(chunk.data(), chunk.size()); - drained += chunk.size(); - } - if (!has_next) { - break; - } - } - SRV_TRC("stream_pipe close: drain ended conv=%s bytes=%zu\n", session_->conversation_id.c_str(), drained); -} - -std::shared_ptr stream_pipe_producer::create(stream_session_ptr session, - server_http_res & res) { - auto alive = std::make_shared>(true); - auto * res_ptr = &res; - session->set_stop_producer([alive, res_ptr]() { - if (alive->load(std::memory_order_acquire)) { - res_ptr->stop(); - } - }); - auto pipe = std::shared_ptr(new stream_pipe_producer(std::move(session))); - pipe->alive_ = std::move(alive); - pipe->res_ = res_ptr; - return pipe; +stream_pipe_producer * stream_pipe_producer::create(stream_session_ptr session) { + return new stream_pipe_producer(std::move(session)); } // stream_pipe_consumer @@ -661,21 +594,68 @@ std::string server_stream_conv_id_from_headers(const std::map & headers) { +static stream_pipe_producer * server_stream_create_spipe(const std::map & headers) { std::string conversation_id = server_stream_conv_id_from_headers(headers); SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0); if (conversation_id.empty()) { - return; + return nullptr; } auto session = g_stream_sessions.create_or_replace(conversation_id); - res.spipe = stream_pipe_producer::create(session, res); + return stream_pipe_producer::create(session); } -std::function server_stream_aware_should_stop(server_http_res * res, std::function fallback) { - return [res, fallback = std::move(fallback)]() -> bool { - if (res->spipe) { - return res->spipe->is_cancelled(); +// +// server_res_spipe +// + +void server_res_spipe::set_req(const server_http_req * req) { + this->req = req; + // optionally attach spipe to the response when X-Conversation-Id is present + spipe.reset(server_stream_create_spipe(req->headers)); +} + +bool server_res_spipe::conn_alive() { + GGML_ASSERT(req != nullptr); + return !req->should_stop(); +} + +bool server_res_spipe::should_stop() { + if (spipe) { + // note: if DELETE /v1/stream/ is called, is_cancelled() will be true + return spipe->is_cancelled(); + } else { + return !conn_alive(); + } +} + +void server_res_spipe::on_complete() { + if (!spipe || next_finished) { + return; + } + std::string chunk; + while (!spipe->is_cancelled()) { + chunk.clear(); + bool has_next = next_orig(chunk); + if (!chunk.empty()) { + spipe->write(chunk.data(), chunk.size()); } - return fallback(); + if (!has_next) { + break; + } + } +} + +void server_res_spipe::set_next(std::function next_fn) { + next_orig = std::move(next_fn); + next = [this](std::string & out) { + bool has_next = next_orig(out); + if (spipe) { + // if spipe is set, tee-style pipe input to both HTTP and spipe + spipe->write(out.data(), out.size()); + } + if (!has_next) { + next_finished = true; + } + return has_next; }; } diff --git a/tools/server/server-stream.h b/tools/server/server-stream.h index c0c3e924fa..9753140dd6 100644 --- a/tools/server/server-stream.h +++ b/tools/server/server-stream.h @@ -30,36 +30,15 @@ protected: // producer end: writes chunks into the ring buffer and owns the session lifetime, finalizing it // on destruction. -// -// lifetime safety: holds a shared_ptr> alive also captured by the session's -// stop_producer hook. cleanup() sets alive=false and clears the hook; it must run while the -// response the hook calls stop() on is still alive. ~server_res_generator() does this explicitly. struct stream_pipe_producer : stream_pipe { ~stream_pipe_producer() override; bool write(const char * data, size_t len); - // mark the natural end on the wire so a later close() is a no-op - void done(); - - // on a peer drop, pump the response next() into the ring buffer until done. runs on the http - // worker from on_complete, no-op after done() or cancel - void close(); - - // disarm the stop hook and drop the alive guard, must run while the response the hook - // references is still alive. idempotent, the destructor calls it too - void cleanup(); - - // res.stop() is invoked when the session is cancelled, the alive guard ensures stop() is not - // called after cleanup() has run - static std::shared_ptr create(stream_session_ptr session, server_http_res & res); + static stream_pipe_producer * create(stream_session_ptr session); private: explicit stream_pipe_producer(stream_session_ptr session); - - bool done_ = false; - std::shared_ptr> alive_; - server_http_res * res_ = nullptr; }; void server_stream_session_manager_start(); @@ -73,10 +52,22 @@ server_http_context::handler_t server_stream_make_delete_handler(); // extract the X-Conversation-Id header value (case-insensitive), empty when absent std::string server_stream_conv_id_from_headers(const std::map & headers); -// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to res -void server_stream_session_attach_pipe(server_http_res & res, const std::map & headers); +// implement tee-style pipe (spipe) for "stream replay" functionality +struct server_res_spipe : server_http_res { +private: + // if set, the stream survives a client disconnect: + // connection kept alive, output is forwarded to spipe and reuse later + std::unique_ptr spipe; + // if spipe is set, use this next_orig to implement tee-style pipe + std::function next_orig; + const server_http_req * req = nullptr; + // set once next_orig reports no more data, so on_complete() doesn't re-drain a finished stream + bool next_finished = false; -// should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit -// DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it -// delegates to fallback, the legacy non-resumable flow -std::function server_stream_aware_should_stop(server_http_res * res, std::function fallback); +public: + void set_req(const server_http_req * req); + bool conn_alive(); + bool should_stop(); + void on_complete() override; + void set_next(std::function next_fn); +};