Skip to content

Commit bbebeec

Browse files
server-stream: follow-up on SSE Replay Buffer (#23226) (#25047)
* server-stream : pimpl * server-stream: prefix free functions with server_stream_ address review from ggerganov: scope the public stream functions under the server_stream_ prefix, matching server_stream_session_manager_start/stop. * server-stream: guard session and manager state with the mutex address review from ggerganov: make done, completed_ts and the GC running flag plain members under their mutex and set the condvar predicates under the lock. keep cancelled atomic for the lock-free should_stop poll. * server-stream: trim comments to the non-obvious address review from ggerganov: drop comments that restate the code, keep the concurrency, lifetime and ordering rationale. de-stale a few comments left by the pimpl: g_stream_sessions is now internal and the /v1/streams listing is gone. * server-stream: update dev docs for the pimpl and prefix reflect server_stream_session_manager_start/stop and the server_stream_ prefix, note the manager is now a file-static singleton hidden in the .cpp * server-stream: move stream traces to debug level keep the bring-up traces for diagnostics but off the default log: skip drain, draining, drain ended, DELETE evict, attach_pipe, and the router stream resume proxy. * server-stream: align router stream resume proxy trace with upstream the child-side bring-up traces are already SRV_TRC on master, move the router stream resume proxy trace to the same level. * server-stream: move stream_read_status enum to the cpp it is only used by the hidden session and consumer types, so it belongs with them behind the pimpl boundary, not on the public header surface. --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
1 parent 230ea9d commit bbebeec

6 files changed

Lines changed: 180 additions & 193 deletions

File tree

tools/server/README-dev.md

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ The core architecture consists of the following components:
5757
- `server_tokens`: Unified representation of token sequences (supports both text and multimodal tokens); used by `server_task` and `server_slot`.
5858
- `server_prompt_checkpoint`: For recurrent (e.g., RWKV) and SWA models, stores snapshots of KV cache state. Enables reuse when subsequent requests share the same prompt prefix, saving redundant computation.
5959
- `server_models`: Standalone component for managing multiple backend instances (used in router mode). It is completely independent of `server_context`.
60-
- `stream_session_manager`: Process wide owner of resumable SSE stream sessions (`g_stream_sessions`), keyed by conversation id. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below.
60+
- `stream_session_manager`: process wide owner of resumable SSE stream sessions, keyed by conversation id. A file-static singleton inside `server-stream.cpp`, driven through `server_stream_session_manager_start/stop`. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below.
6161

6262
```mermaid
6363
graph TD
@@ -127,10 +127,12 @@ It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`.
127127
The feature lives entirely in `server-stream.{h,cpp}` and rests on three types:
128128

129129
- `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.
130-
- `stream_session_manager` (`g_stream_sessions`): 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.
130+
- `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`.
131131
- `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.
132132

133-
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, `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.
133+
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`.
134+
135+
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.
134136

135137
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.
136138

@@ -144,7 +146,7 @@ Routes:
144146

145147
Router mode binds the same paths to proxy handlers. A `conv_id -> child` map (`conv_models`), populated when a POST is routed, resolves the owning child in one lookup with no polling. The lookup groups ids per child; GET and DELETE proxy straight to the owner. This loopback REST hop is expected to move to a websocket IPC later, swapping only the transport.
146148

147-
Lifecycle: `g_stream_sessions.start_gc()` runs in main after common init, `stop_gc()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin.
149+
Lifecycle: `server_stream_session_manager_start()` runs in main after common init, `server_stream_session_manager_stop()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin.
148150

149151
| Constant | Value | Role |
150152
| --- | --- | --- |

tools/server/server-context.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4188,7 +4188,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
41884188
}
41894189
};
41904190

4191-
auto effective_should_stop = stream_aware_should_stop(res_this, req.should_stop);
4191+
auto effective_should_stop = server_stream_aware_should_stop(res_this, req.should_stop);
41924192

41934193
try {
41944194
if (effective_should_stop()) {
@@ -4284,7 +4284,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
42844284

42854285
// attach a producer pipe to the response when X-Conversation-Id is present.
42864286
// the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook.
4287-
stream_session_attach_pipe(*res, req.headers);
4287+
server_stream_session_attach_pipe(*res, req.headers);
42884288

42894289
return res;
42904290
}

tools/server/server-models.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1681,7 +1681,7 @@ void server_models_routes::init_routes() {
16811681
}
16821682
// remember which child serves this conversation so the stream routes can route straight
16831683
// to it without polling, keyed on the exact conv id from the header
1684-
std::string conv_id = stream_conv_id_from_headers(req.headers);
1684+
std::string conv_id = server_stream_conv_id_from_headers(req.headers);
16851685
if (!conv_id.empty()) {
16861686
models.conv_models.remember(conv_id, name);
16871687
}
@@ -1896,7 +1896,7 @@ void server_models_routes::init_routes() {
18961896
if (!from.empty()) {
18971897
child_path += "?from=" + from;
18981898
}
1899-
SRV_INF("proxying stream resume to model %s on port %d, path=%s\n",
1899+
SRV_TRC("proxying stream resume to model %s on port %d, path=%s\n",
19001900
owner->name.c_str(), owner->port, child_path.c_str());
19011901
auto proxy = std::make_unique<server_http_proxy>(
19021902
"GET",

tools/server/server-stream.cpp

Lines changed: 147 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,111 @@
66
#include <chrono>
77
#include <memory>
88
#include <utility>
9+
#include <shared_mutex>
10+
11+
enum class stream_read_status {
12+
OK,
13+
OFFSET_LOST,
14+
};
915

1016
namespace {
1117
constexpr int64_t STREAM_SESSION_TTL_SECONDS = 300;
1218
constexpr size_t STREAM_SESSION_MAX_BYTES = 4 * 1024 * 1024;
1319
constexpr int64_t STREAM_SESSION_GC_INTERVAL_SECONDS = 60;
1420
constexpr int64_t STREAM_READ_WAKE_INTERVAL_MS = 200;
1521

16-
// returns unix time in seconds
1722
int64_t now_seconds() {
1823
return std::chrono::duration_cast<std::chrono::seconds>(
1924
std::chrono::system_clock::now().time_since_epoch()
2025
).count();
2126
}
2227
}
2328

29+
// owns all live sessions keyed by conversation_id, one conv = at most one live session.
30+
// a periodic GC evicts expired ones
31+
class stream_session_manager {
32+
public:
33+
stream_session_manager();
34+
~stream_session_manager();
35+
36+
stream_session_manager(const stream_session_manager &) = delete;
37+
stream_session_manager & operator=(const stream_session_manager &) = delete;
38+
39+
// install a new session, evicting and cancelling any previous one. conversation_id must be non empty
40+
stream_session_ptr create_or_replace(const std::string & conversation_id);
41+
42+
stream_session_ptr get(const std::string & conversation_id);
43+
44+
std::vector<stream_session_ptr> list_all() const;
45+
46+
void evict(const std::string & conversation_id);
47+
48+
void evict_and_cancel(const std::string & conversation_id);
49+
50+
void start_gc();
51+
void stop_gc();
52+
53+
private:
54+
void gc_loop();
55+
56+
mutable std::shared_mutex map_mu;
57+
std::unordered_map<std::string, stream_session_ptr> sessions; // key: conversation_id
58+
std::thread gc_thread;
59+
bool running;
60+
std::mutex gc_wake_mu;
61+
std::condition_variable gc_wake_cv;
62+
};
63+
64+
// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc
65+
static stream_session_manager g_stream_sessions;
66+
67+
void server_stream_session_manager_start() {
68+
g_stream_sessions.start_gc();
69+
}
70+
71+
void server_stream_session_manager_stop() {
72+
g_stream_sessions.stop_gc();
73+
}
74+
75+
struct stream_session {
76+
std::string conversation_id;
77+
int64_t started_ts; // unix seconds at construction
78+
79+
stream_session(std::string conversation_id_, size_t max_bytes_);
80+
stream_session(const stream_session &) = delete;
81+
stream_session & operator=(const stream_session &) = delete;
82+
83+
bool append(const char * data, size_t len);
84+
85+
void finalize();
86+
87+
// drain from offset into sink, blocking for more bytes or finalize. OFFSET_LOST if offset
88+
// fell below the dropped prefix
89+
stream_read_status read_from(size_t offset,
90+
const std::function<bool(const char *, size_t)> & sink,
91+
const std::function<bool()> & should_stop);
92+
93+
bool is_done() const;
94+
bool is_cancelled() const;
95+
size_t total_size() const; // bytes that ever entered the session
96+
size_t dropped_prefix() const; // bytes evicted from the front due to cap
97+
int64_t completed_at() const; // 0 while alive, unix seconds after finalize
98+
99+
void set_stop_producer(std::function<void()> fn);
100+
101+
void cancel();
102+
103+
private:
104+
mutable std::mutex mu;
105+
std::condition_variable cv;
106+
std::vector<char> buffer;
107+
size_t prefix_dropped;
108+
size_t cap_bytes;
109+
bool done;
110+
std::atomic<bool> cancelled; // polled lock-free by the should_stop closure, no mu
111+
int64_t completed_ts;
112+
std::function<void()> stop_producer;
113+
};
24114
stream_session::stream_session(std::string conversation_id_, size_t max_bytes_)
25115
: conversation_id(std::move(conversation_id_))
26116
, started_ts(now_seconds())
@@ -38,7 +128,7 @@ bool stream_session::append(const char * data, size_t len) {
38128
}
39129
{
40130
std::lock_guard<std::mutex> lock(mu);
41-
if (done.load(std::memory_order_relaxed)) {
131+
if (done) {
42132
return false;
43133
}
44134
if (len >= cap_bytes) {
@@ -62,11 +152,14 @@ bool stream_session::append(const char * data, size_t len) {
62152
}
63153

64154
void stream_session::finalize() {
65-
bool was_done = done.exchange(true, std::memory_order_acq_rel);
66-
if (was_done) {
67-
return;
155+
{
156+
std::lock_guard<std::mutex> lock(mu);
157+
if (done) {
158+
return;
159+
}
160+
done = true;
161+
completed_ts = now_seconds();
68162
}
69-
completed_ts.store(now_seconds(), std::memory_order_release);
70163
cv.notify_all();
71164
}
72165

@@ -96,7 +189,7 @@ stream_read_status stream_session::read_from(size_t offset,
96189
lock.lock();
97190
continue;
98191
}
99-
if (done.load(std::memory_order_acquire)) {
192+
if (done) {
100193
return stream_read_status::OK;
101194
}
102195
// wait for new bytes, finalize, or a periodic wake to re check should_stop
@@ -105,7 +198,8 @@ stream_read_status stream_session::read_from(size_t offset,
105198
}
106199

107200
bool stream_session::is_done() const {
108-
return done.load(std::memory_order_acquire);
201+
std::lock_guard<std::mutex> lock(mu);
202+
return done;
109203
}
110204

111205
size_t stream_session::total_size() const {
@@ -119,7 +213,8 @@ size_t stream_session::dropped_prefix() const {
119213
}
120214

121215
int64_t stream_session::completed_at() const {
122-
return completed_ts.load(std::memory_order_acquire);
216+
std::lock_guard<std::mutex> lock(mu);
217+
return completed_ts;
123218
}
124219

125220
void stream_session::set_stop_producer(std::function<void()> fn) {
@@ -128,7 +223,7 @@ void stream_session::set_stop_producer(std::function<void()> fn) {
128223
}
129224

130225
void stream_session::cancel() {
131-
// flip cancelled first so the producer-side stream_aware_should_stop can break out of the
226+
// flip cancelled first so the producer-side server_stream_aware_should_stop can break out of the
132227
// recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task
133228
// posted by rd.stop() will eventually notify, but we do not want to depend on that timing)
134229
cancelled.store(true, std::memory_order_release);
@@ -237,18 +332,24 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i
237332
}
238333

239334
void stream_session_manager::start_gc() {
240-
if (running.exchange(true)) {
241-
return;
335+
{
336+
std::lock_guard<std::mutex> lock(gc_wake_mu);
337+
if (running) {
338+
return;
339+
}
340+
running = true;
242341
}
243342
gc_thread = std::thread([this] { gc_loop(); });
244343
}
245344

246345
void stream_session_manager::stop_gc() {
247-
bool was_running = running.exchange(false);
346+
bool was_running;
347+
{
348+
std::lock_guard<std::mutex> lock(gc_wake_mu);
349+
was_running = running;
350+
running = false;
351+
}
248352
if (was_running) {
249-
{
250-
std::lock_guard<std::mutex> lock(gc_wake_mu);
251-
}
252353
gc_wake_cv.notify_all();
253354
if (gc_thread.joinable()) {
254355
gc_thread.join();
@@ -270,15 +371,15 @@ void stream_session_manager::stop_gc() {
270371
}
271372

272373
void stream_session_manager::gc_loop() {
273-
while (running.load(std::memory_order_acquire)) {
374+
while (true) {
274375
{
275376
std::unique_lock<std::mutex> lock(gc_wake_mu);
276377
gc_wake_cv.wait_for(lock,
277378
std::chrono::seconds(STREAM_SESSION_GC_INTERVAL_SECONDS),
278-
[this] { return !running.load(std::memory_order_acquire); });
279-
}
280-
if (!running.load(std::memory_order_acquire)) {
281-
return;
379+
[this] { return !running; });
380+
if (!running) {
381+
return;
382+
}
282383
}
283384
int64_t cutoff = now_seconds() - STREAM_SESSION_TTL_SECONDS;
284385
std::vector<stream_session_ptr> to_drop;
@@ -301,10 +402,19 @@ void stream_session_manager::gc_loop() {
301402
}
302403
}
303404

304-
// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc
305-
stream_session_manager g_stream_sessions;
405+
// stream_pipe
406+
407+
// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session
408+
struct stream_pipe_consumer : stream_pipe {
409+
stream_read_status read(size_t & offset,
410+
const std::function<bool(const char *, size_t)> & sink,
411+
const std::function<bool()> & should_stop);
412+
413+
static std::shared_ptr<stream_pipe_consumer> create(stream_session_ptr session);
306414

307-
// stream_pipe ---------------------------------------------------------------------------------
415+
private:
416+
explicit stream_pipe_consumer(stream_session_ptr session);
417+
};
308418

309419
stream_pipe::stream_pipe(stream_session_ptr session)
310420
: session_(std::move(session)) {
@@ -408,12 +518,10 @@ static server_http_res_ptr make_error_response(int status, const std::string & m
408518
return res;
409519
}
410520

411-
server_http_context::handler_t make_stream_get_handler() {
521+
server_http_context::handler_t server_stream_make_get_handler() {
412522
return [](const server_http_req & req) -> server_http_res_ptr {
413-
// GET /v1/stream/<conv_id>?from=N replays the SSE bytes already buffered for the
414-
// session, blocks for more bytes when the session is still running, returns when
415-
// the session is finalized. the body is streamed back as text/event-stream so the
416-
// browser EventSource can attach to it like a fresh request
523+
// GET /v1/stream/<conv_id>?from=N replays buffered SSE bytes then blocks for live
524+
// bytes until the session finalizes, streamed as text/event-stream for EventSource
417525
std::string conv_id = req.get_param("conv_id");
418526
if (conv_id.empty()) {
419527
return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
@@ -459,11 +567,10 @@ server_http_context::handler_t make_stream_get_handler() {
459567
};
460568
}
461569

462-
server_http_context::handler_t make_streams_lookup_handler() {
570+
server_http_context::handler_t server_stream_make_lookup_handler() {
463571
return [](const server_http_req & req) -> server_http_res_ptr {
464-
// POST /v1/streams/lookup with body {"conversation_ids": ["X", "Y", ...]} returns the
465-
// matching sessions, only for ids the caller already knows. each id matches the exact key
466-
// and any "<id>::<model>" variant, so one lookup covers every per model session for a conv
572+
// POST /v1/streams/lookup returns the matching sessions, only for ids the caller already
573+
// knows. each id matches the exact key and any "<id>::<model>" per model variant
467574
std::vector<std::string> requested;
468575
try {
469576
json body = json::parse(req.body);
@@ -518,11 +625,10 @@ server_http_context::handler_t make_streams_lookup_handler() {
518625
};
519626
}
520627

521-
server_http_context::handler_t make_stream_delete_handler() {
628+
server_http_context::handler_t server_stream_make_delete_handler() {
522629
return [](const server_http_req & req) -> server_http_res_ptr {
523-
// DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer hook
524-
// wired by handle_completions_impl and evicts the buffer. idempotent, a session that
525-
// already finalized or was never created returns 204 either way
630+
// DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer and evicts
631+
// the buffer. idempotent, returns 204 even if the session was already gone
526632
std::string conv_id = req.get_param("conv_id");
527633
if (conv_id.empty()) {
528634
return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
@@ -536,7 +642,7 @@ server_http_context::handler_t make_stream_delete_handler() {
536642
};
537643
}
538644

539-
std::string stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {
645+
std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {
540646
// case-insensitive scan for x-conversation-id
541647
static constexpr char target[] = "x-conversation-id";
542648
static constexpr size_t target_len = sizeof(target) - 1;
@@ -555,8 +661,8 @@ std::string stream_conv_id_from_headers(const std::map<std::string, std::string>
555661
return std::string();
556662
}
557663

558-
void stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
559-
std::string conversation_id = stream_conv_id_from_headers(headers);
664+
void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
665+
std::string conversation_id = server_stream_conv_id_from_headers(headers);
560666
SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0);
561667
if (conversation_id.empty()) {
562668
return;
@@ -565,7 +671,7 @@ void stream_session_attach_pipe(server_http_res & res, const std::map<std::strin
565671
res.spipe = stream_pipe_producer::create(session, res);
566672
}
567673

568-
std::function<bool()> stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
674+
std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
569675
return [res, fallback = std::move(fallback)]() -> bool {
570676
if (res->spipe) {
571677
return res->spipe->is_cancelled();

0 commit comments

Comments
 (0)