Felipe97/llama-cpp-compiled
01.1k
1#include "server-stream.h"2#include "server-common.h"3#include "server-http.h"4#include "server-queue.h"5 6#include <chrono>7#include <memory>8#include <utility>9#include <shared_mutex>10 11enum class stream_read_status {12 OK,13 OFFSET_LOST,14};15 16namespace {17constexpr int64_t STREAM_SESSION_TTL_SECONDS = 300;18constexpr size_t STREAM_SESSION_MAX_BYTES = 4 * 1024 * 1024;19constexpr int64_t STREAM_SESSION_GC_INTERVAL_SECONDS = 60;20constexpr int64_t STREAM_READ_WAKE_INTERVAL_MS = 200;21 22int64_t now_seconds() {23 return std::chrono::duration_cast<std::chrono::seconds>(24 std::chrono::system_clock::now().time_since_epoch()25 ).count();26}27}28 29// owns all live sessions keyed by conversation_id, one conv = at most one live session.30// a periodic GC evicts expired ones31class stream_session_manager {32public: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 empty40 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 53private:54 void gc_loop();55 56 mutable std::shared_mutex map_mu;57 std::unordered_map<std::string, stream_session_ptr> sessions; // key: conversation_id58 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_gc65static stream_session_manager g_stream_sessions;66 67void server_stream_session_manager_start() {68 g_stream_sessions.start_gc();69}70 71void server_stream_session_manager_stop() {72 g_stream_sessions.stop_gc();73}74 75struct stream_session {76 std::string conversation_id;77 int64_t started_ts; // unix seconds at construction78 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 offset88 // fell below the dropped prefix89 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 session96 size_t dropped_prefix() const; // bytes evicted from the front due to cap97 int64_t completed_at() const; // 0 while alive, unix seconds after finalize98 99 void cancel();100 101private:102 mutable std::mutex mu;103 std::condition_variable cv;104 std::vector<char> buffer;105 size_t prefix_dropped;106 size_t cap_bytes;107 bool done;108 std::atomic<bool> cancelled; // polled lock-free by the should_stop closure, no mu109 int64_t completed_ts;110};111stream_session::stream_session(std::string conversation_id_, size_t max_bytes_)112 : conversation_id(std::move(conversation_id_))113 , started_ts(now_seconds())114 , prefix_dropped(0)115 , cap_bytes(max_bytes_)116 , done(false)117 , cancelled(false)118 , completed_ts(0) {119 buffer.reserve(64 * 1024);120}121 122bool stream_session::append(const char * data, size_t len) {123 if (len == 0) {124 return true;125 }126 {127 std::lock_guard<std::mutex> lock(mu);128 if (done) {129 return false;130 }131 if (len >= cap_bytes) {132 // single chunk bigger than the cap, keep only the tail that fits133 size_t skip = len - cap_bytes;134 prefix_dropped += buffer.size() + skip;135 buffer.clear();136 buffer.insert(buffer.end(), data + skip, data + len);137 } else {138 size_t needed = buffer.size() + len;139 if (needed > cap_bytes) {140 size_t to_drop = needed - cap_bytes;141 buffer.erase(buffer.begin(), buffer.begin() + to_drop);142 prefix_dropped += to_drop;143 }144 buffer.insert(buffer.end(), data, data + len);145 }146 }147 cv.notify_all();148 return true;149}150 151void stream_session::finalize() {152 {153 std::lock_guard<std::mutex> lock(mu);154 if (done) {155 return;156 }157 done = true;158 completed_ts = now_seconds();159 }160 cv.notify_all();161}162 163stream_read_status stream_session::read_from(size_t offset,164 const std::function<bool(const char *, size_t)> & sink,165 const std::function<bool()> & should_stop) {166 std::unique_lock<std::mutex> lock(mu);167 while (true) {168 if (should_stop && should_stop()) {169 return stream_read_status::OK;170 }171 if (offset < prefix_dropped) {172 return stream_read_status::OFFSET_LOST;173 }174 size_t logical_end = prefix_dropped + buffer.size();175 if (offset < logical_end) {176 size_t local_off = offset - prefix_dropped;177 size_t n = buffer.size() - local_off;178 // copy the available chunk under the lock, release before calling the sink179 std::vector<char> chunk(buffer.begin() + local_off, buffer.begin() + local_off + n);180 offset += n;181 lock.unlock();182 bool keep_going = sink(chunk.data(), chunk.size());183 if (!keep_going) {184 return stream_read_status::OK;185 }186 lock.lock();187 continue;188 }189 if (done) {190 return stream_read_status::OK;191 }192 // wait for new bytes, finalize, or a periodic wake to re check should_stop193 cv.wait_for(lock, std::chrono::milliseconds(STREAM_READ_WAKE_INTERVAL_MS));194 }195}196 197bool stream_session::is_done() const {198 std::lock_guard<std::mutex> lock(mu);199 return done;200}201 202size_t stream_session::total_size() const {203 std::lock_guard<std::mutex> lock(mu);204 return prefix_dropped + buffer.size();205}206 207size_t stream_session::dropped_prefix() const {208 std::lock_guard<std::mutex> lock(mu);209 return prefix_dropped;210}211 212int64_t stream_session::completed_at() const {213 std::lock_guard<std::mutex> lock(mu);214 return completed_ts;215}216 217void stream_session::cancel() {218 // the should_stop closure on both the producer and any HTTP reader polls is_cancelled()219 // so flipping this is the only signal needed to unwind both sides220 cancelled.store(true, std::memory_order_release);221}222 223bool stream_session::is_cancelled() const {224 return cancelled.load(std::memory_order_acquire);225}226 227stream_session_manager::stream_session_manager()228 : running(false) {229}230 231stream_session_manager::~stream_session_manager() {232 stop_gc();233}234 235stream_session_ptr stream_session_manager::create_or_replace(const std::string & conversation_id) {236 // evict any previous session on the same conv, this guarantees the invariant237 // "one conv = at most one live session" and propagates cancel to its producer238 stream_session_ptr previous;239 auto fresh = std::make_shared<stream_session>(conversation_id, STREAM_SESSION_MAX_BYTES);240 {241 std::unique_lock<std::shared_mutex> lock(map_mu);242 auto it = sessions.find(conversation_id);243 if (it != sessions.end()) {244 previous = it->second;245 it->second = fresh;246 } else {247 sessions.emplace(conversation_id, fresh);248 }249 }250 if (previous) {251 previous->cancel();252 previous->finalize();253 }254 return fresh;255}256 257stream_session_ptr stream_session_manager::get(const std::string & conversation_id) {258 std::shared_lock<std::shared_mutex> lock(map_mu);259 auto it = sessions.find(conversation_id);260 if (it == sessions.end()) {261 return nullptr;262 }263 return it->second;264}265 266std::vector<stream_session_ptr> stream_session_manager::list_all() const {267 std::vector<stream_session_ptr> out;268 std::shared_lock<std::shared_mutex> lock(map_mu);269 out.reserve(sessions.size());270 for (auto & kv : sessions) {271 out.push_back(kv.second);272 }273 return out;274}275 276void stream_session_manager::evict(const std::string & conversation_id) {277 stream_session_ptr s;278 {279 std::unique_lock<std::shared_mutex> lock(map_mu);280 auto it = sessions.find(conversation_id);281 if (it == sessions.end()) {282 return;283 }284 s = it->second;285 sessions.erase(it);286 }287 // finalize outside the map lock so any pending readers wake up and exit288 s->finalize();289}290 291void stream_session_manager::evict_and_cancel(const std::string & conversation_id) {292 stream_session_ptr s;293 {294 std::unique_lock<std::shared_mutex> lock(map_mu);295 auto it = sessions.find(conversation_id);296 if (it == sessions.end()) {297 std::string live;298 for (const auto & kv : sessions) {299 if (!live.empty()) live += ", ";300 live += kv.first;301 }302 SRV_WRN("stop on unknown stream session, conv_id=%s matched nothing, %zu live: [%s]\n",303 conversation_id.c_str(), sessions.size(), live.c_str());304 return;305 }306 s = it->second;307 sessions.erase(it);308 }309 // cancel first so the producer's on_complete() drain loop and any pending HTTP reader310 // observe is_cancelled() and stop pulling further output, then finalize to wake readers311 // blocked in read_from(). note: this does not interrupt the underlying generation itself,312 // which keeps running to its own natural stop condition (EOS/max_tokens)313 s->cancel();314 s->finalize();315}316 317void stream_session_manager::start_gc() {318 {319 std::lock_guard<std::mutex> lock(gc_wake_mu);320 if (running) {321 return;322 }323 running = true;324 }325 gc_thread = std::thread([this] { gc_loop(); });326}327 328void stream_session_manager::stop_gc() {329 bool was_running;330 {331 std::lock_guard<std::mutex> lock(gc_wake_mu);332 was_running = running;333 running = false;334 }335 if (was_running) {336 gc_wake_cv.notify_all();337 if (gc_thread.joinable()) {338 gc_thread.join();339 }340 }341 // finalize all live sessions so no reader ever hangs342 std::vector<stream_session_ptr> snapshot;343 {344 std::unique_lock<std::shared_mutex> lock(map_mu);345 snapshot.reserve(sessions.size());346 for (auto & kv : sessions) {347 snapshot.push_back(kv.second);348 }349 sessions.clear();350 }351 for (auto & s : snapshot) {352 s->finalize();353 }354}355 356void stream_session_manager::gc_loop() {357 while (true) {358 {359 std::unique_lock<std::mutex> lock(gc_wake_mu);360 gc_wake_cv.wait_for(lock,361 std::chrono::seconds(STREAM_SESSION_GC_INTERVAL_SECONDS),362 [this] { return !running; });363 if (!running) {364 return;365 }366 }367 int64_t cutoff = now_seconds() - STREAM_SESSION_TTL_SECONDS;368 std::vector<stream_session_ptr> to_drop;369 {370 std::unique_lock<std::shared_mutex> lock(map_mu);371 for (auto it = sessions.begin(); it != sessions.end(); ) {372 int64_t completed = it->second->completed_at();373 if (completed != 0 && completed <= cutoff) {374 to_drop.push_back(it->second);375 it = sessions.erase(it);376 } else {377 ++it;378 }379 }380 }381 // finalize outside the map lock, idempotent if the session was already done382 for (auto & s : to_drop) {383 s->finalize();384 }385 }386}387 388// stream_pipe389 390// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session391struct stream_pipe_consumer : stream_pipe {392 stream_read_status read(size_t & offset,393 const std::function<bool(const char *, size_t)> & sink,394 const std::function<bool()> & should_stop);395 396 static std::shared_ptr<stream_pipe_consumer> create(stream_session_ptr session);397 398private:399 explicit stream_pipe_consumer(stream_session_ptr session);400};401 402stream_pipe::stream_pipe(stream_session_ptr session)403 : session_(std::move(session)) {404}405 406bool stream_pipe::is_cancelled() const {407 return session_->is_cancelled();408}409 410// stream_pipe_producer411 412stream_pipe_producer::stream_pipe_producer(stream_session_ptr session)413 : stream_pipe(std::move(session)) {414}415 416stream_pipe_producer::~stream_pipe_producer() {417 session_->finalize();418}419 420bool stream_pipe_producer::write(const char * data, size_t len) {421 return session_->append(data, len);422}423 424stream_pipe_producer * stream_pipe_producer::create(stream_session_ptr session) {425 return new stream_pipe_producer(std::move(session));426}427 428// stream_pipe_consumer429 430stream_pipe_consumer::stream_pipe_consumer(stream_session_ptr session)431 : stream_pipe(std::move(session)) {432}433 434stream_read_status stream_pipe_consumer::read(size_t & offset,435 const std::function<bool(const char *, size_t)> & sink,436 const std::function<bool()> & should_stop) {437 return session_->read_from(offset, sink, should_stop);438}439 440std::shared_ptr<stream_pipe_consumer> stream_pipe_consumer::create(stream_session_ptr session) {441 return std::shared_ptr<stream_pipe_consumer>(new stream_pipe_consumer(std::move(session)));442}443 444// helper, builds the standard error response and assigns it to a brand new http_res445static server_http_res_ptr make_error_response(int status, const std::string & message, error_type type) {446 auto res = std::make_unique<server_http_res>();447 json err = format_error_response(message, type);448 res->status = json_value(err, "code", status);449 res->content_type = "application/json; charset=utf-8";450 res->data = safe_json_to_str({{"error", err}});451 return res;452}453 454server_http_context::handler_t server_stream_make_get_handler() {455 return [](const server_http_req & req) -> server_http_res_ptr {456 // GET /v1/stream?conv_id=<id>&from=N replays buffered SSE bytes then blocks for live457 // bytes until the session finalizes, streamed as text/event-stream for EventSource458 std::string conv_id = req.get_param("conv_id");459 if (conv_id.empty()) {460 return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);461 }462 auto session = g_stream_sessions.get(conv_id);463 if (!session) {464 return make_error_response(404, "Stream not found or expired", ERROR_TYPE_NOT_FOUND);465 }466 size_t from = 0;467 std::string from_str = req.get_param("from");468 if (!from_str.empty()) {469 try {470 from = static_cast<size_t>(std::stoull(from_str));471 } catch (const std::exception &) {472 return make_error_response(400, "Invalid 'from' offset", ERROR_TYPE_INVALID_REQUEST);473 }474 }475 if (from < session->dropped_prefix()) {476 return make_error_response(400, "Stream offset lost, please restart", ERROR_TYPE_INVALID_REQUEST);477 }478 auto res = std::make_unique<server_http_res>();479 res->status = 200;480 res->content_type = "text/event-stream";481 // the next closure reads from the ring buffer at the requested offset, blocks until482 // bytes arrive or the session finalizes. exit each call after draining the available483 // chunk so set_chunked_content_provider gets a chance to flush to the socket484 auto offset_ptr = std::make_shared<size_t>(from);485 // consumer pipe: read-only, does not finalize the session on destruction486 auto pipe = stream_pipe_consumer::create(session);487 res->next = [pipe, offset_ptr, &req](std::string & output) -> bool {488 bool got_any = false;489 pipe->read(*offset_ptr,490 [&](const char * d, size_t n) {491 output.append(d, n);492 *offset_ptr += n;493 got_any = true;494 return false;495 },496 req.should_stop);497 return got_any;498 };499 return res;500 };501}502 503server_http_context::handler_t server_stream_make_lookup_handler() {504 return [](const server_http_req & req) -> server_http_res_ptr {505 // POST /v1/streams/lookup returns the matching sessions, only for ids the caller already506 // knows. each id matches the exact key and any "<id>::<model>" per model variant507 std::vector<std::string> requested;508 try {509 json body = json::parse(req.body);510 if (body.contains("conversation_ids") && body["conversation_ids"].is_array()) {511 for (const auto & v : body["conversation_ids"]) {512 if (v.is_string()) {513 std::string id = v.get<std::string>();514 if (!id.empty()) {515 requested.push_back(std::move(id));516 }517 }518 }519 }520 } catch (const std::exception & e) {521 auto res = std::make_unique<server_http_res>();522 res->status = 400;523 res->content_type = "application/json; charset=utf-8";524 res->data = safe_json_to_str({{"error", {{"message", std::string("invalid body: ") + e.what()},525 {"type", "invalid_request_error"}}}});526 return res;527 }528 529 std::vector<stream_session_ptr> sessions;530 if (!requested.empty()) {531 auto all = g_stream_sessions.list_all();532 for (const auto & rid : requested) {533 const std::string with_sep = rid + "::";534 for (auto & s : all) {535 if (s->conversation_id == rid ||536 s->conversation_id.compare(0, with_sep.size(), with_sep) == 0) {537 sessions.push_back(s);538 }539 }540 }541 }542 543 json arr = json::array();544 for (auto & s : sessions) {545 arr.push_back({546 {"conversation_id", s->conversation_id},547 {"is_done", s->is_done()},548 {"total_bytes", s->total_size()},549 {"started_at", s->started_ts},550 {"completed_at", s->completed_at()},551 });552 }553 auto res = std::make_unique<server_http_res>();554 res->status = 200;555 res->content_type = "application/json; charset=utf-8";556 res->data = safe_json_to_str(arr);557 return res;558 };559}560 561server_http_context::handler_t server_stream_make_delete_handler() {562 return [](const server_http_req & req) -> server_http_res_ptr {563 // DELETE /v1/stream?conv_id=<id> is the explicit user Stop, cancels the producer and evicts564 // the buffer. idempotent, returns 204 even if the session was already gone565 std::string conv_id = req.get_param("conv_id");566 if (conv_id.empty()) {567 return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);568 }569 SRV_TRC("DELETE /v1/stream conv_id=%s -> evict_and_cancel\n", conv_id.c_str());570 g_stream_sessions.evict_and_cancel(conv_id);571 auto res = std::make_unique<server_http_res>();572 res->status = 204;573 res->content_type = "application/json";574 return res;575 };576}577 578std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {579 // case-insensitive scan for x-conversation-id580 static constexpr char target[] = "x-conversation-id";581 static constexpr size_t target_len = sizeof(target) - 1;582 for (const auto & [hk, hv] : headers) {583 if (hk.size() != target_len) continue;584 bool match = true;585 for (size_t i = 0; i < target_len; ++i) {586 char c = hk[i];587 if (c >= 'A' && c <= 'Z') c = char(c + 32);588 if (c != target[i]) { match = false; break; }589 }590 if (match) {591 return hv;592 }593 }594 return std::string();595}596 597static stream_pipe_producer * server_stream_create_spipe(const std::map<std::string, std::string> & headers) {598 std::string conversation_id = server_stream_conv_id_from_headers(headers);599 SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0);600 if (conversation_id.empty()) {601 return nullptr;602 }603 auto session = g_stream_sessions.create_or_replace(conversation_id);604 return stream_pipe_producer::create(session);605}606 607//608// server_res_spipe609//610 611void server_res_spipe::set_req(const server_http_req * req) {612 this->req = req;613 // optionally attach spipe to the response when X-Conversation-Id is present614 spipe.reset(server_stream_create_spipe(req->headers));615}616 617bool server_res_spipe::conn_alive() {618 GGML_ASSERT(req != nullptr);619 return !req->should_stop();620}621 622bool server_res_spipe::should_stop() {623 if (spipe) {624 // note: if DELETE /v1/stream is called for this conv, is_cancelled() will be true625 return spipe->is_cancelled();626 } else {627 return !conn_alive();628 }629}630 631void server_res_spipe::on_complete() {632 if (!spipe || next_finished) {633 return;634 }635 // an empty next_orig means set_next() never ran: the request failed before streaming636 // started, typically a params validation throw. evict the session installed by set_req()637 // so the failed request leaves nothing behind for discovery or replay638 if (!next_orig) {639 g_stream_sessions.evict(server_stream_conv_id_from_headers(req->headers));640 return;641 }642 std::string chunk;643 while (!spipe->is_cancelled()) {644 chunk.clear();645 bool has_next = next_orig(chunk);646 if (!chunk.empty()) {647 spipe->write(chunk.data(), chunk.size());648 }649 if (!has_next) {650 break;651 }652 }653}654 655void server_res_spipe::set_next(std::function<bool(std::string &)> next_fn) {656 next_orig = std::move(next_fn);657 next = [this](std::string & out) {658 bool has_next = next_orig(out);659 if (spipe) {660 // if spipe is set, tee-style pipe input to both HTTP and spipe661 spipe->write(out.data(), out.size());662 }663 if (!has_next) {664 next_finished = true;665 }666 return has_next;667 };668}669