diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index 280bd9e19dca..700f08fd2cd4 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -63,3 +63,23 @@ install(TARGETS ${TARGET} RUNTIME) target_link_libraries(${TARGET} PRIVATE llama-server-impl) target_compile_features(${TARGET} PRIVATE cxx_std_17) + +# server-queue unit test: needs no model, so it stays out of the default build + +if (LLAMA_BUILD_TESTS AND NOT CMAKE_CROSSCOMPILING) + set(TARGET test-server-queue) + + add_executable(${TARGET} tests/test-server-queue.cpp) + target_link_libraries(${TARGET} PRIVATE server-context ${CMAKE_THREAD_LIBS_INIT}) + target_compile_features(${TARGET} PRIVATE cxx_std_17) + + add_test(NAME ${TARGET} COMMAND ${TARGET}) + + set(TARGET test-server-response) + + add_executable(${TARGET} tests/test-server-response.cpp) + target_link_libraries(${TARGET} PRIVATE server-context ${CMAKE_THREAD_LIBS_INIT}) + target_compile_features(${TARGET} PRIVATE cxx_std_17) + + add_test(NAME ${TARGET} COMMAND ${TARGET}) +endif() diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..ef78971d8b1f 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -431,6 +431,12 @@ struct server_slot { return; } + // only send_final_response() reads this, and only with n_probs > 0; otherwise every token + // copied a string and a vector into a list grown for the whole generation, then dropped + if (task->params.sampling.n_probs <= 0) { + return; + } + generated_token_probs.push_back(token); } diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d86..e72d53fe8ab8 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -387,84 +387,247 @@ void server_queue::cleanup_pending_task(int id_target) { // void server_response::add_waiting_task_id(int id_task) { - RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting_task_ids.size()); - std::unique_lock lock(mutex_results); - waiting_task_ids.insert(id_task); + + RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting.size()); + + waiting.emplace(id_task, std::make_shared()); + + // a reader may already be parked on these ids waiting for exactly this + condition_gone.notify_all(); } void server_response::add_waiting_task_ids(const std::unordered_set & id_tasks) { std::unique_lock lock(mutex_results); + // one waiter for the whole set: these ids belong to one reader + auto w = std::make_shared(); + for (const auto & id_task : id_tasks) { - RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting_task_ids.size()); - waiting_task_ids.insert(id_task); + RES_DBG("add task %d to waiting list. current waiting = %d (before add)\n", id_task, (int) waiting.size()); + waiting.emplace(id_task, w); } + + // a reader may already be parked on these ids waiting for exactly this + condition_gone.notify_all(); } void server_response::remove_waiting_task_id(int id_task) { - RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting_task_ids.size()); - std::unique_lock lock(mutex_results); - waiting_task_ids.erase(id_task); - // make sure to clean up all pending results - queue_results.erase( - std::remove_if(queue_results.begin(), queue_results.end(), [id_task](const server_task_result_ptr & res) { - return res->id == id_task; + + RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting.size()); + + auto it = waiting.find(id_task); + if (it == waiting.end()) { + return; + } + + // the waiter is shared with the reader's other ids, so drop only this task's results + auto & results = it->second->results; + results.erase( + std::remove_if(results.begin(), results.end(), [id_task](const pending & p) { + return p.res->id == id_task; }), - queue_results.end()); + results.end()); + + // a reader may be parked on this waiter; it has to repeat the lookup rather than wait out its + // deadline on a condition that nothing will fire again + auto w = it->second; + waiting.erase(it); + w->cv.notify_all(); + condition_gone.notify_all(); } void server_response::remove_waiting_task_ids(const std::unordered_set & id_tasks) { std::unique_lock lock(mutex_results); + std::vector removed; + + for (const auto & id_task : id_tasks) { + RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting.size()); + + auto it = waiting.find(id_task); + if (it == waiting.end()) { + continue; + } + + removed.push_back(it->second); + waiting.erase(it); + } + + // same as the single id form: wake anyone parked on a waiter that no longer serves these ids + for (const auto & w : removed) { + w->cv.notify_all(); + } + condition_gone.notify_all(); +} + +server_response::waiter_ptr server_response::find_waiter(const std::unordered_set & id_tasks) const { for (const auto & id_task : id_tasks) { - RES_DBG("remove task %d from waiting list. current waiting = %d (before remove)\n", id_task, (int) waiting_task_ids.size()); - waiting_task_ids.erase(id_task); + auto it = waiting.find(id_task); + if (it != waiting.end()) { + return it->second; + } } + + return nullptr; +} + +// The waiter that covers every requested id, or nullptr when they sit in more than one waiter or +// any of them is not registered yet. An absent id matters: it can be registered onto a different +// waiter while the reader waits, so no single waiter's condition covers the call. +server_response::waiter_ptr server_response::sole_waiter(const std::unordered_set & id_tasks) const { + waiter_ptr found = nullptr; + + for (const auto & id_task : id_tasks) { + auto it = waiting.find(id_task); + if (it == waiting.end()) { + return nullptr; + } + if (found == nullptr) { + found = it->second; + continue; + } + if (it->second != found) { + return nullptr; + } + } + + return found; +} + +// A waiter is shared by every id its reader registered in one call, so its queue can hold a +// sibling's result. Return only an id the caller asked for, and the oldest such result across +// every waiter the ids map to, which is what scanning the shared vector did. Each waiter's queue +// is already in arrival order, so its first match is its oldest and only the winners are compared. +server_task_result_ptr server_response::take_result(const std::unordered_set & id_tasks) { + auto first_match = [&](waiter * w) { + return std::find_if(w->results.begin(), w->results.end(), [&](const pending & p) { + return id_tasks.find(p.res->id) != id_tasks.end(); + }); + }; + + auto claim = [](waiter * w, std::deque::iterator it) { + server_task_result_ptr res = std::move(it->res); + w->results.erase(it); + return res; + }; + + // the ordinary case: every id the caller named shares one waiter, so no comparison is needed + if (auto w = sole_waiter(id_tasks)) { + auto it = first_match(w.get()); + return it == w->results.end() ? nullptr : claim(w.get(), it); + } + + waiter * best_w = nullptr; + std::deque::iterator best_it; + uint64_t best_seq = 0; + std::vector examined; + + for (const auto & id_task : id_tasks) { + auto it = waiting.find(id_task); + if (it == waiting.end()) { + continue; + } + + waiter * w = it->second.get(); + if (std::find(examined.begin(), examined.end(), w) != examined.end()) { + continue; // ids commonly share a waiter, so do not scan the same queue twice + } + examined.push_back(w); + + auto rit = first_match(w); + if (rit != w->results.end() && (best_w == nullptr || rit->seq < best_seq)) { + best_w = w; + best_it = rit; + best_seq = rit->seq; + } + } + + return best_w == nullptr ? nullptr : claim(best_w, best_it); } server_task_result_ptr server_response::recv(const std::unordered_set & id_tasks) { + std::unique_lock lock(mutex_results); + while (true) { - std::unique_lock lock(mutex_results); - condition_results.wait(lock, [&]{ - if (!running) { - RES_DBG("%s : queue result stop\n", "recv"); - std::terminate(); // we cannot return here since the caller is HTTP code - } - return !queue_results.empty(); - }); + if (!running) { + RES_DBG("%s : queue result stop\n", "recv"); + std::terminate(); // we cannot return here since the caller is HTTP code + } - for (size_t i = 0; i < queue_results.size(); i++) { - if (id_tasks.find(queue_results[i]->id) != id_tasks.end()) { - server_task_result_ptr res = std::move(queue_results[i]); - queue_results.erase(queue_results.begin() + i); - return res; - } + server_task_result_ptr res = take_result(id_tasks); + if (res != nullptr) { + return res; + } + + // The waiter can be absent, so this cannot assert. A cancel or a cleanup drops the ids + // between the caller posting them and arriving here, and recv() runs on the HTTP + // thread: aborting there turns one stuck request into a dead server for every other + // client. Before the per-waiter queues this waited on a condition that no longer fires + // for these ids, which blocks this one connection and nothing else, so that is what it + // does here too. The lookup is inside the loop rather than above it because a waiter + // re-added while we wait should be picked up instead of waited out. + // Only a waiter that covers every requested id has a condition that covers the whole + // receive. Anything else parks on the shared one, and send() notifies that for readers + // whose ids are at least partly registered, so a result cannot be missed. + auto w = sole_waiter(id_tasks); + if (w == nullptr) { + const bool deliverable = find_waiter(id_tasks) != nullptr; + + // registration and terminate() both fire condition_gone; the timeout is only a backstop + if (deliverable) { n_split_readers++; } + condition_gone.wait_for(lock, std::chrono::seconds(1)); + if (deliverable) { n_split_readers--; } + continue; } + + // bounded: a terminate() landing after the id left the map is still noticed here + w->cv.wait_for(lock, std::chrono::seconds(1)); } // should never reach here } server_task_result_ptr server_response::recv_with_timeout(const std::unordered_set & id_tasks, int timeout) { - while (true) { - std::unique_lock lock(mutex_results); + std::unique_lock lock(mutex_results); - for (int i = 0; i < (int) queue_results.size(); i++) { - if (id_tasks.find(queue_results[i]->id) != id_tasks.end()) { - server_task_result_ptr res = std::move(queue_results[i]); - queue_results.erase(queue_results.begin() + i); - return res; - } - } + // one deadline for the whole call: waiting for a registration and then for a result must not + // add up to twice the timeout the caller asked for + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout); - std::cv_status cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout)); + while (true) { if (!running) { RES_DBG("%s : queue result stop\n", __func__); std::terminate(); // we cannot return here since the caller is HTTP code } - if (cr_res == std::cv_status::timeout) { + + server_task_result_ptr res = take_result(id_tasks); + if (res != nullptr) { + return res; + } + + // Park on the shared condition unless one waiter covers every requested id: the ids may + // be spread over several waiters, or some of them may not be registered yet, and either + // way no one waiter's condition covers the call. add_waiting_task_id(s) fires the shared + // one, so a result that arrives during this call is still seen, which is what the single + // shared condition used to give; terminate() fires it too; and send() fires it while a + // reader that could already be served is parked there. + auto w = sole_waiter(id_tasks); + + const bool deliverable = w == nullptr && find_waiter(id_tasks) != nullptr; + + std::condition_variable & cv = w == nullptr ? condition_gone : w->cv; + + if (deliverable) { n_split_readers++; } + const std::cv_status st = cv.wait_until(lock, deadline); + if (deliverable) { n_split_readers--; } + + if (st == std::cv_status::timeout) { + if (!running) { + RES_DBG("%s : queue result stop\n", __func__); + std::terminate(); // we cannot return here since the caller is HTTP code + } return nullptr; } } @@ -481,31 +644,54 @@ void server_response::send(server_task_result_ptr && result) { RES_DBG("sending result for task id = %d\n", result->id); std::unique_lock lock(mutex_results); - for (const auto & id_task : waiting_task_ids) { - if (result->id == id_task) { - RES_DBG("task id = %d pushed to result queue\n", result->id); - queue_results.emplace_back(std::move(result)); - condition_results.notify_all(); - return; - } + auto it = waiting.find(result->id); + if (it == waiting.end()) { + return; + } + + RES_DBG("task id = %d pushed to result queue\n", result->id); + + auto & w = *it->second; + + w.results.push_back(pending{next_seq++, std::move(result)}); + + // notify_all, not notify_one: results are filtered by id, so waking a single waiter can wake + // one taking a disjoint subset of this reader's ids, which finds nothing and sleeps again + // while the reader whose result this is stays asleep. This is one reader's own condition, + // not the single global one the shared vector used, so it is still O(1) in the common case + // of one thread per reader. + w.cv.notify_all(); + + // normally zero: only a reader whose ids span several waiters parks on the shared condition + if (n_split_readers > 0) { + condition_gone.notify_all(); } } void server_response::broadcast(server_task_result_ptr && result) { std::unique_lock lock(mutex_results); - for (const auto & id_task : waiting_task_ids) { + for (const auto & [id_task, w] : waiting) { RES_DBG("task id = %d pushed to result queue\n", id_task); server_task_result_ptr res_copy(result->clone()); res_copy->id = id_task; // override id with target task id - queue_results.emplace_back(std::move(res_copy)); + w->results.push_back(pending{next_seq++, std::move(res_copy)}); + w->cv.notify_all(); + } + + if (n_split_readers > 0) { + condition_gone.notify_all(); } - condition_results.notify_all(); } void server_response::terminate() { + std::unique_lock lock(mutex_results); running = false; - condition_results.notify_all(); + for (const auto & [id_task, w] : waiting) { + (void) id_task; + w->cv.notify_all(); + } + condition_gone.notify_all(); } // diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index e17733a743f6..b688d7dcd677 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -5,10 +5,13 @@ #include #include #include +#include +#include #include #include -#include +#include #include +#include // struct for managing server tasks // in most cases, use server_response_reader to post new tasks and retrieve results @@ -155,14 +158,52 @@ struct server_response { private: bool running = true; - // for keeping track of all tasks waiting for the result - std::unordered_set waiting_task_ids; + // One waiter per reader, shared by every id it registered in one call. A single shared vector + // plus one cv instead costs N wakeups and N scans per token, N^2 per decode step. A send + // notifies this reader's own cv, so the wakeup is O(1) even though it is a notify_all. + // arrival order is global, not per waiter: a reader can name ids from several waiters and + // must still be served oldest first, which is what scanning the shared vector gave + struct pending { + uint64_t seq; + server_task_result_ptr res; + }; + + struct waiter { + std::condition_variable cv; + + std::deque results; + }; + + using waiter_ptr = std::shared_ptr; - // the main result queue (using ptr for polymorphism) - std::vector queue_results; + std::unordered_map waiting; + + // stamped onto every queued result so arrival order survives being split across waiters + uint64_t next_seq = 0; std::mutex mutex_results; - std::condition_variable condition_results; + + // parks a reader whose ids left the waiting list, so it honours its timeout, and a reader + // whose ids span several waiters, for which no single waiter's condition is enough + std::condition_variable condition_gone; + + // how many readers are parked on condition_gone even though a result could already be + // delivered to them, i.e. their ids are spread over several waiters or only partly + // registered. Normally zero, so send() pays one integer compare rather than a second notify: + // server_response_reader registers all of its ids in one call, before it ever receives. + size_t n_split_readers = 0; + + // ids registered together share one waiter, so the first hit is the right one. mutex_results held. + waiter_ptr find_waiter(const std::unordered_set & id_tasks) const; + + // pop the oldest queued result whose id the caller asked for. mutex_results held. + server_task_result_ptr take_result(const std::unordered_set & id_tasks); + + // The waiter that covers EVERY requested id, or nullptr when they are spread over several + // waiters or any of them is not registered. Only then does one waiter's condition cover the + // whole receive; an id that is absent now can be registered onto a different waiter while + // the reader waits. mutex_results held. + waiter_ptr sole_waiter(const std::unordered_set & id_tasks) const; public: // add the id_task to the list of tasks waiting for response diff --git a/tools/server/tests/test-server-queue.cpp b/tools/server/tests/test-server-queue.cpp new file mode 100644 index 000000000000..74c7a1676cf2 --- /dev/null +++ b/tools/server/tests/test-server-queue.cpp @@ -0,0 +1,56 @@ +#include "server-queue.h" + +#include +#include +#include +#include +#include +#include + +// recv() can be called with ids that are not in the waiting list: a cancel or a cleanup drops +// them between the caller posting and the caller arriving. That has to park this one connection +// and nothing else, which is what the unbounded wait did before the per-waiter queues. +// +// It must not assert. GGML_ASSERT is GGML_ABORT, and recv() runs on the HTTP thread, so a +// single dropped request would take the whole server down for every other client. The parent +// commit dies inside recv() here, in well under the 2.5 s this waits. +int main() { + server_response res; + + std::atomic returned{false}; + + std::thread parked([&] { + server_task_result_ptr r = res.recv(std::unordered_set{4242}); + (void) r; + returned.store(true); + }); + + std::this_thread::sleep_for(std::chrono::milliseconds(2500)); + + if (returned.load()) { + fprintf(stderr, "FAIL: recv() returned for ids that are not in the waiting list\n"); + return 1; + } + + // the timeout form already tolerated this, and must keep doing so promptly + const auto t0 = std::chrono::steady_clock::now(); + server_task_result_ptr none = res.recv_with_timeout(std::unordered_set{4243}, 1); + const auto waited = std::chrono::steady_clock::now() - t0; + + if (none != nullptr) { + fprintf(stderr, "FAIL: recv_with_timeout() invented a result\n"); + return 1; + } + if (waited > std::chrono::seconds(5)) { + fprintf(stderr, "FAIL: recv_with_timeout() did not honour its timeout\n"); + return 1; + } + + printf("OK: a dropped request parks its own caller and leaves the server up\n"); + + // parked is still inside recv() by design: terminate() would make it std::terminate(), + // which is the documented behaviour for an HTTP caller, so leave without joining it. + parked.detach(); + fflush(stdout); + _Exit(0); +} diff --git a/tools/server/tests/test-server-response.cpp b/tools/server/tests/test-server-response.cpp new file mode 100644 index 000000000000..594622d7588b --- /dev/null +++ b/tools/server/tests/test-server-response.cpp @@ -0,0 +1,537 @@ +// Unit test for server_response, the result queue between the decode loop and the HTTP threads. +// +// It exercises the public API directly, so the awkward cases are injected rather than waited +// for: a reader whose ids were dropped between posting and receiving, a send that races a +// cancel, per id and bulk teardown, broadcast, and concurrent registration and removal. +// +// Run with "leak " to measure what the queue retains over n reader lifecycles, each leaving +// one result queued at teardown, which is what a client disconnect during generation does. + +#include "server-queue.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef _WIN32 +# include +# include +#endif + +static int g_fail = 0; + +static void check(bool ok, const char * name, const char * detail = "") { + printf("%-46s %s %s\n", name, ok ? "PASS" : "FAIL", detail); + if (!ok) { g_fail++; } +} + +// minimal concrete result so we can put things on the queue without a model +struct fake_result : server_task_result { + int payload = 0; + bool stop = false; + fake_result(int id_, int payload_, bool stop_) : payload(payload_), stop(stop_) { id = id_; } + bool is_stop() override { return stop; } + json to_json() override { return json{{"payload", payload}}; } + server_task_result * clone() const override { return new fake_result(*this); } +}; + +static server_task_result_ptr mk(int id, int payload, bool stop = true) { + return server_task_result_ptr(new fake_result(id, payload, stop)); +} + +static int payload_of(const server_task_result_ptr & p) { + return p ? static_cast(p.get())->payload : -1; +} + +using ms = std::chrono::milliseconds; + +// --------------------------------------------------------------------------- + +// a result for an id that has left the waiting list is dropped, silently, no crash +static void t_send_to_absent_id() { + server_response res; + res.send(mk(1, 100)); // never registered + res.add_waiting_task_id(2); + res.remove_waiting_task_id(2); + res.send(mk(2, 200)); // registered then removed + res.add_waiting_task_id(3); + res.send(mk(3, 300)); + auto got = res.recv_with_timeout({3}, 1); + check(got != nullptr && payload_of(got) == 300, "send to absent id is dropped, live id still delivered"); +} + +// FIFO order per reader is what scanning the shared vector from the front used to give +static void t_fifo_order() { + server_response res; + res.add_waiting_task_id(7); + for (int i = 0; i < 32; i++) { res.send(mk(7, i, i == 31)); } + bool ok = true; + for (int i = 0; i < 32; i++) { + auto r = res.recv_with_timeout({7}, 1); + if (payload_of(r) != i) { ok = false; break; } + } + check(ok, "FIFO order preserved for a single reader"); +} + +// two independent readers must not see each other's results +static void t_reader_isolation() { + server_response res; + res.add_waiting_task_ids({10, 11}); + res.add_waiting_task_ids({20, 21}); + res.send(mk(20, 2000)); + res.send(mk(10, 1000)); + res.send(mk(21, 2100)); + res.send(mk(11, 1100)); + + std::vector a, b; + for (int i = 0; i < 2; i++) { a.push_back(payload_of(res.recv_with_timeout({10, 11}, 1))); } + for (int i = 0; i < 2; i++) { b.push_back(payload_of(res.recv_with_timeout({20, 21}, 1))); } + const bool ok = a.size() == 2 && b.size() == 2 && + (a[0] == 1000 && a[1] == 1100) && (b[0] == 2000 && b[1] == 2100); + char d[128]; + snprintf(d, sizeof(d), "a=[%d,%d] b=[%d,%d]", a[0], a[1], b[0], b[1]); + check(ok, "two readers do not steal each other's results", d); + check(res.recv_with_timeout({10, 11}, 1) == nullptr, "reader A drained, no extra result"); +} + +// per-id removal must drop only that id's results +static void t_partial_removal() { + server_response res; + res.add_waiting_task_ids({30, 31}); + res.send(mk(30, 3000)); + res.send(mk(31, 3100)); + res.remove_waiting_task_id(30); + auto r = res.recv_with_timeout({31}, 1); + check(r != nullptr && payload_of(r) == 3100, "removing one id keeps the sibling's result"); + check(res.recv_with_timeout({31}, 1) == nullptr, "the removed id's result is gone"); +} + +// bulk removal then a late send: nothing delivered, nothing leaked, no use after free +static void t_bulk_removal() { + server_response res; + res.add_waiting_task_ids({40, 41, 42}); + res.send(mk(40, 4000)); + res.remove_waiting_task_ids({40, 41, 42}); + res.send(mk(41, 4100)); + check(res.recv_with_timeout({40, 41, 42}, 1) == nullptr, "bulk removal drops queued and late results"); +} + +// broadcast: one copy per registered id, id overridden +static void t_broadcast() { + server_response res; + res.add_waiting_task_ids({50, 51}); + res.add_waiting_task_id(60); + res.broadcast(mk(-1, 9999)); + auto a1 = res.recv_with_timeout({50, 51}, 1); + auto a2 = res.recv_with_timeout({50, 51}, 1); + auto a3 = res.recv_with_timeout({50, 51}, 1); + auto b1 = res.recv_with_timeout({60}, 1); + auto b2 = res.recv_with_timeout({60}, 1); + const bool ok = a1 && a2 && !a3 && b1 && !b2 && + payload_of(a1) == 9999 && payload_of(b1) == 9999; + check(ok, "broadcast delivers one copy per registered id"); + const bool ids_ok = a1 && a2 && (a1->id == 50 || a1->id == 51) && (a2->id == 50 || a2->id == 51) && a1->id != a2->id && b1 && b1->id == 60; + check(ids_ok, "broadcast overrides the result id per target"); +} + +// lost wakeup: park a reader on ids that do not exist yet, then create them and send. +// Both arms must deliver; the head's condition_gone poll bounds the delay. +static void t_late_registration() { + server_response res; + std::atomic got{-2}; + std::atomic done{false}; + std::thread th([&] { + for (int i = 0; i < 60; i++) { + auto r = res.recv_with_timeout({70}, 1); + if (r) { got.store(payload_of(r)); break; } + } + done.store(true); + }); + std::this_thread::sleep_for(ms(300)); + res.add_waiting_task_id(70); + res.send(mk(70, 7000)); + const auto t0 = std::chrono::steady_clock::now(); + while (!done.load() && std::chrono::steady_clock::now() - t0 < std::chrono::seconds(10)) { + std::this_thread::sleep_for(ms(10)); + } + const auto waited = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); + th.join(); + char d[64]; snprintf(d, sizeof(d), "%lldms", (long long) waited); + check(got.load() == 7000, "id registered after the reader parked is still served", d); +} + +// a reader whose ids were dropped parks itself and nothing else. +// The parent commit of the head asserted here, which is GGML_ABORT on the HTTP thread. +static void t_parked_reader_does_not_abort() { + static server_response res; // static: the parked thread outlives this function + std::atomic returned{false}; + std::thread parked([&] { + auto r = res.recv(std::unordered_set{4242}); + (void) r; + returned.store(true); + }); + std::this_thread::sleep_for(ms(2500)); + check(!returned.load(), "recv() on dropped ids parks instead of returning garbage"); + + // and the rest of the queue keeps working while that thread is parked + res.add_waiting_task_id(80); + res.send(mk(80, 8000)); + auto r = res.recv_with_timeout({80}, 2); + check(r != nullptr && payload_of(r) == 8000, "other readers unaffected by a parked reader"); + parked.detach(); +} + +// recv_with_timeout honours its timeout when the waiter exists but is empty +static void t_timeout_honoured() { + server_response res; + res.add_waiting_task_id(90); + const auto t0 = std::chrono::steady_clock::now(); + auto r = res.recv_with_timeout({90}, 1); + const auto waited = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); + char d[64]; snprintf(d, sizeof(d), "%lldms", (long long) waited); + check(r == nullptr && waited >= 900 && waited < 5000, "recv_with_timeout honours its timeout", d); +} + +// concurrent churn: producers, consumers, registration and teardown all at once. +// This is the case ThreadSanitizer is pointed at. +static void t_stress() { + server_response res; + const int n_readers = 16; + const int n_msgs = 200; + std::atomic received{0}; + std::atomic go{false}; + + std::vector readers; + for (int r = 0; r < n_readers; r++) { + readers.emplace_back([&, r] { + while (!go.load()) { std::this_thread::yield(); } + const int base_id = 1000 + r * 10; + std::unordered_set ids{base_id, base_id + 1}; + res.add_waiting_task_ids(ids); + int seen = 0; + while (seen < n_msgs) { + auto p = res.recv_with_timeout(ids, 1); + if (!p) { break; } + seen++; + received.fetch_add(1); + } + res.remove_waiting_task_ids(ids); + }); + } + + std::vector writers; + for (int w = 0; w < 4; w++) { + writers.emplace_back([&, w] { + while (!go.load()) { std::this_thread::yield(); } + for (int i = w; i < n_msgs * n_readers; i += 4) { + const int r = (i / n_msgs) % n_readers; + res.send(mk(1000 + r * 10 + (i % 2), i, false)); + if ((i & 63) == 0) { std::this_thread::sleep_for(ms(1)); } + } + }); + } + + // a churn thread that registers and drops ids nobody waits for + std::thread churn([&] { + while (!go.load()) { std::this_thread::yield(); } + for (int i = 0; i < 2000; i++) { + res.add_waiting_task_id(500000 + i); + res.send(mk(500000 + i, i, false)); + res.remove_waiting_task_id(500000 + i); + } + }); + + go.store(true); + for (auto & t : writers) { t.join(); } + for (auto & t : readers) { t.join(); } + churn.join(); + char d[64]; snprintf(d, sizeof(d), "received=%d", received.load()); + check(received.load() > 0, "concurrent send/recv/register/remove churn survives", d); +} + + +// A single timed receive that starts before the ids exist must still return a result that +// arrives during the call. The shared queue used to notify one condition on every send, so a +// parked receiver woke; per waiter queues have to notify registration explicitly or the caller +// sleeps out its whole timeout and reports a spurious nullptr. +static void t_single_timed_recv_before_registration() { + server_response res; + std::atomic started{false}; + std::thread producer([&] { + while (!started.load()) { std::this_thread::yield(); } + std::this_thread::sleep_for(ms(200)); + res.add_waiting_task_id(300); + res.send(mk(300, 3000)); + }); + started.store(true); + const auto t0 = std::chrono::steady_clock::now(); + auto r = res.recv_with_timeout({300}, 5); // ONE call, not a retry loop + const auto waited = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); + producer.join(); + char d[80]; snprintf(d, sizeof(d), "%lldms, %s", (long long) waited, r ? "got result" : "nullptr"); + check(r != nullptr && payload_of(r) == 3000 && waited < 4000, + "one timed recv sees a result that arrives while it waits", d); +} + +// terminate() has to be honoured even when the caller is parked on ids that are not registered. +// recv_with_timeout() promises std::terminate() there, because the caller is HTTP code that +// cannot return. Run in a child: the correct outcome is that the child dies. +static void t_terminate_while_parked_on_absent_ids() { +#ifdef _WIN32 + // needs fork(): the correct outcome is that the caller terminates, which cannot be asserted + // in-process. The behaviour itself is not platform specific. + printf("%-46s SKIP (needs fork())\n", "terminate() is honoured while parked on absent ids"); +#else + fflush(stdout); + pid_t pid = fork(); + if (pid == 0) { + auto * res = new server_response(); + std::thread killer([res] { + std::this_thread::sleep_for(ms(300)); + res->terminate(); + }); + auto r = res->recv_with_timeout({9999}, 5); + killer.join(); + // reaching here at all means terminate() was ignored + _Exit(r == nullptr ? 20 : 21); + } + int status = 0; + waitpid(pid, &status, 0); + const bool died = WIFSIGNALED(status); + char d[96]; + if (died) { snprintf(d, sizeof(d), "child died on signal %d", WTERMSIG(status)); } + else { snprintf(d, sizeof(d), "child returned %d, terminate() ignored", WEXITSTATUS(status)); } + check(died, "terminate() is honoured while parked on absent ids", d); +#endif +} + +// A result for a sibling id must not be handed to a caller that did not ask for it. +static void t_subset_recv_is_filtered() { + server_response res; + res.add_waiting_task_ids({200, 201}); + res.send(mk(201, 2010)); + auto r = res.recv_with_timeout({200}, 1); + char d[64]; snprintf(d, sizeof(d), "id=%d", r ? r->id : -1); + check(r == nullptr, "recv() does not return a result for an id it was not asked for", d); + // and the sibling's result is still there for the caller that does ask + auto r2 = res.recv_with_timeout({200, 201}, 1); + check(r2 != nullptr && r2->id == 201, "the sibling's result is still delivered to its own reader"); +} + + +// Two readers taking disjoint subsets of one registration share a waiter, so waking only one of +// them can wake the wrong one: it finds nothing matching, sleeps again, and the reader whose +// result is actually queued sits there until its timeout. The shared condition used to +// notify_all(), so every subset receiver got a look. +static void t_subset_receivers_are_all_woken() { + server_response res; + res.add_waiting_task_ids({400, 401}); + + std::atomic parked{0}; + std::atomic got_a{-1}; + std::vector others; + + // four readers waiting on the sibling id park first, so a single notify picks one of them + for (int i = 0; i < 4; i++) { + others.emplace_back([&] { + parked.fetch_add(1); + auto r = res.recv_with_timeout({401}, 3); + (void) r; + }); + } + while (parked.load() < 4) { std::this_thread::yield(); } + std::this_thread::sleep_for(ms(200)); + + std::thread reader_a([&] { + auto r = res.recv_with_timeout({400}, 3); + got_a.store(payload_of(r)); + }); + std::this_thread::sleep_for(ms(200)); + + const auto t0 = std::chrono::steady_clock::now(); + res.send(mk(400, 4000)); + reader_a.join(); + const auto waited = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); + for (auto & t : others) { t.join(); } + + char d[80]; snprintf(d, sizeof(d), "%lldms, payload=%d", (long long) waited, got_a.load()); + check(got_a.load() == 4000 && waited < 2500, + "a subset receiver is woken even when siblings wait too", d); +} + + +// Ids registered by separate calls belong to separate waiters, so no single condition covers a +// receive that names both. The receiver must be woken by a result for either of them, whichever +// waiter the lookup happened to pick, so both directions are driven. +static void t_ids_spanning_two_waiters_one(int base_id, int send_to, const char * label) { + server_response res; + res.add_waiting_task_id(base_id); // two separate registrations, so two waiters + res.add_waiting_task_id(base_id + 1); + + std::atomic got{-1}; + std::thread reader([&] { + auto r = res.recv_with_timeout({base_id, base_id + 1}, 3); + got.store(payload_of(r)); + }); + std::this_thread::sleep_for(ms(300)); + + const auto t0 = std::chrono::steady_clock::now(); + res.send(mk(send_to, 5000 + send_to)); + reader.join(); + const auto waited = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); + + char d[96]; snprintf(d, sizeof(d), "%lldms, payload=%d", (long long) waited, got.load()); + check(got.load() == 5000 + send_to && waited < 2500, label, d); +} + +static void t_ids_spanning_two_waiters() { + t_ids_spanning_two_waiters_one(500, 500, "a receive over two waiters is woken by the first id"); + t_ids_spanning_two_waiters_one(600, 601, "a receive over two waiters is woken by the second id"); +} + + +// Results must come back in arrival order even when the ids live in different waiters. The +// shared vector scanned from the front, so it did. Both orders are driven, because which waiter +// the lookup reaches first depends on the set's iteration order. +static void t_fifo_across_waiters_one(int first, int second, const char * label) { + server_response res; + res.add_waiting_task_id(first); // separate registrations, so separate waiters + res.add_waiting_task_id(second); + + res.send(mk(first, 7000 + first)); + res.send(mk(second, 7000 + second)); + + auto r1 = res.recv_with_timeout({first, second}, 1); + auto r2 = res.recv_with_timeout({first, second}, 1); + + char d[96]; + snprintf(d, sizeof(d), "got %d then %d, wanted %d then %d", + r1 ? r1->id : -1, r2 ? r2->id : -1, first, second); + check(r1 != nullptr && r2 != nullptr && r1->id == first && r2->id == second, label, d); +} + +static void t_fifo_across_waiters() { + t_fifo_across_waiters_one(700, 701, "arrival order kept across waiters, low id first"); + t_fifo_across_waiters_one(711, 710, "arrival order kept across waiters, high id first"); +} + +// A reader already parked on a waiter has to be woken when that waiter is discarded, or it will +// wait out its deadline on a condition nothing will ever fire again while its id is re-registered +// and served on a brand new waiter. +static void t_waiter_replaced_under_a_parked_reader() { + server_response res; + res.add_waiting_task_id(800); + + std::atomic got{-1}; + std::thread reader([&] { + auto r = res.recv_with_timeout({800}, 3); + got.store(payload_of(r)); + }); + std::this_thread::sleep_for(ms(300)); // let the reader select the current waiter and park + + const auto t0 = std::chrono::steady_clock::now(); + res.remove_waiting_task_id(800); // discards the waiter the reader is parked on + res.add_waiting_task_id(800); // a brand new waiter + res.send(mk(800, 8800)); + reader.join(); + const auto waited = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); + + char d[80]; snprintf(d, sizeof(d), "%lldms, payload=%d", (long long) waited, got.load()); + check(got.load() == 8800 && waited < 2500, + "a parked reader is woken when its waiter is replaced", d); +} + + +// A receive can name an id that is not registered yet. One waiter's condition does not cover +// such a call, because the missing id may be registered on a different waiter while it waits. +static void t_partially_registered_receive() { + server_response res; + res.add_waiting_task_id(900); // 901 does not exist yet + + std::atomic got{-1}; + std::thread reader([&] { + auto r = res.recv_with_timeout({900, 901}, 3); + got.store(payload_of(r)); + }); + std::this_thread::sleep_for(ms(300)); // let the reader park on whatever it picked + + const auto t0 = std::chrono::steady_clock::now(); + res.add_waiting_task_id(901); // a separate waiter + res.send(mk(901, 9010)); + reader.join(); + const auto waited = std::chrono::duration_cast(std::chrono::steady_clock::now() - t0).count(); + + char d[80]; snprintf(d, sizeof(d), "%lldms, payload=%d", (long long) waited, got.load()); + check(got.load() == 9010 && waited < 2500, + "a receive naming an id registered later is still woken", d); +} + +static long rss_kb() { + // Linux only; returns -1 elsewhere, and only the optional "leak" mode uses it + FILE * f = fopen("/proc/self/status", "r"); + if (!f) { return -1; } + char line[256]; + long v = -1; + while (fgets(line, sizeof(line), f)) { + if (strncmp(line, "VmRSS:", 6) == 0) { sscanf(line + 6, "%ld", &v); break; } + } + fclose(f); + return v; +} + +// Isolated measurement of what the result queue itself retains. One reader lifecycle per +// iteration: register two ids, receive one result, leave one result queued (which is what a +// disconnect during generation does), then tear the reader down the way stop() does. +static void t_leak(long n) { + server_response res; + const long rss0 = rss_kb(); + for (long i = 0; i < n; i++) { + const int a = 100000 + (int) (i * 2); + const int b = a + 1; + res.add_waiting_task_ids({a, b}); + res.send(mk(a, 1)); + res.send(mk(b, 2)); // left unconsumed on purpose + auto got = res.recv_with_timeout({a, b}, 1); + (void) got; + res.remove_waiting_task_ids({a, b}); // exactly what server_response_reader::stop() does + } + const long rss1 = rss_kb(); + printf("LEAK n=%ld rss_start=%ld kB rss_end=%ld kB growth=%ld kB (%.3f kB per reader)\n", + n, rss0, rss1, rss1 - rss0, (double) (rss1 - rss0) / (double) n); +} + +int main(int argc, char ** argv) { + if (argc > 1 && strcmp(argv[1], "leak") == 0) { + t_leak(argc > 2 ? atol(argv[2]) : 200000); + fflush(stdout); + _Exit(0); + } + t_send_to_absent_id(); + t_fifo_order(); + t_reader_isolation(); + t_partial_removal(); + t_bulk_removal(); + t_broadcast(); + t_late_registration(); + t_timeout_honoured(); + t_stress(); + t_single_timed_recv_before_registration(); + t_terminate_while_parked_on_absent_ids(); + t_subset_recv_is_filtered(); + t_subset_receivers_are_all_woken(); + t_ids_spanning_two_waiters(); + t_fifo_across_waiters(); + t_waiter_replaced_under_a_parked_reader(); + t_partially_registered_receive(); + t_parked_reader_does_not_abort(); + + printf("\nRESULT queue failures=%d\n", g_fail); + fflush(stdout); + _Exit(g_fail == 0 ? 0 : 1); // a thread is parked in recv() by design +}