server-queue.cpp

#include "server-task.h"
#include "server-queue.h"

#include "log.h"

#include <algorithm>
#include <chrono>
#include <thread>

#define QUE_INF(fmt, ...) LOG_INF("que  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
#define QUE_WRN(fmt, ...) LOG_WRN("que  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
#define QUE_ERR(fmt, ...) LOG_ERR("que  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
#define QUE_DBG(fmt, ...) LOG_DBG("que  %12.*s: " fmt, 12, __func__, __VA_ARGS__)

#define RES_INF(fmt, ...) LOG_INF("res  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
#define RES_WRN(fmt, ...) LOG_WRN("res  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
#define RES_ERR(fmt, ...) LOG_ERR("res  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
#define RES_DBG(fmt, ...) LOG_DBG("res  %12.*s: " fmt, 12, __func__, __VA_ARGS__)

//
// server_queue
//

static bool task_resets_idle_timer(server_task_type type) {
    return type != SERVER_TASK_TYPE_METRICS;
}

int server_queue::post(server_task && task, bool front) {
    std::unique_lock<std::mutex> lock(mutex_tasks);
    GGML_ASSERT(task.id != -1);
    // if this is cancel task make sure to clean up pending tasks
    if (task.type == SERVER_TASK_TYPE_CANCEL) {
        cleanup_pending_task(task.id_target);
    }
    const int  task_id     = task.id;
    const bool reset_timer = task_resets_idle_timer(task.type);
    QUE_DBG("new task, id = %d, front = %d\n", task_id, front);
    if (front) {
        queue_tasks.push_front(std::move(task));
    } else {
        queue_tasks.push_back(std::move(task));
    }
    if (reset_timer) {
        time_last_task = ggml_time_ms();
    }
    condition_tasks.notify_one();
    return task_id;
}

int server_queue::post(std::vector<server_task> && tasks, bool front) {
    std::unique_lock<std::mutex> lock(mutex_tasks);
    bool reset_timer = false;
    for (auto & task : tasks) {
        if (task.id == -1) {
            task.id = id++;
        }
        // if this is cancel task make sure to clean up pending tasks
        if (task.type == SERVER_TASK_TYPE_CANCEL) {
            cleanup_pending_task(task.id_target);
        }
        reset_timer |= task_resets_idle_timer(task.type);
        QUE_DBG("new task, id = %d/%d, front = %d\n", task.id, (int) tasks.size(), front);
        if (front) {
            queue_tasks.push_front(std::move(task));
        } else {
            queue_tasks.push_back(std::move(task));
        }
    }
    if (reset_timer) {
        time_last_task = ggml_time_ms();
    }
    condition_tasks.notify_one();
    return 0;
}

void server_queue::defer(server_task && task) {
    std::unique_lock<std::mutex> lock(mutex_tasks);
    QUE_DBG("defer task, id = %d\n", task.id);
    queue_tasks_deferred.push_back(std::move(task));
    time_last_task = ggml_time_ms();
    condition_tasks.notify_one();
}

int server_queue::get_new_id() {
    std::unique_lock<std::mutex> lock(mutex_tasks);
    int new_id = id++;
    return new_id;
}

void server_queue::pop_deferred_task(int id_slot) {
    std::unique_lock<std::mutex> lock(mutex_tasks);
    if (!queue_tasks_deferred.empty()) {
        // try to find a task that uses the specified slot
        bool found = false;
        for (auto it = queue_tasks_deferred.begin(); it != queue_tasks_deferred.end(); ++it) {
            if (it->id_slot == id_slot) {
                QUE_DBG("pop deferred task (use slot %d), id_task = %d\n", id_slot, it->id);
                queue_tasks.emplace_front(std::move(*it));
                queue_tasks_deferred.erase(it);
                found = true;
                break;
            }
        }
        // if not tasks found using the slot, just pop the first deferred task (default behavior)
        if (!found) {
            QUE_DBG("pop deferred task, id_task = %d\n", queue_tasks_deferred.front().id);
            queue_tasks.emplace_front(std::move(queue_tasks_deferred.front()));
            queue_tasks_deferred.pop_front();
        }
    }
    time_last_task = ggml_time_ms();
    condition_tasks.notify_one();
}

void server_queue::wait_until_no_sleep() {
    std::unique_lock<std::mutex> lock(mutex_tasks);
    if (!sleeping) {
        return;
    } else {
        if (!req_stop_sleeping) {
            QUE_DBG("%s", "requesting to stop sleeping\n");
            req_stop_sleeping = true;
            condition_tasks.notify_one(); // only main thread is waiting on this
        }
        QUE_DBG("%s", "waiting until no sleep\n");
        condition_tasks.wait(lock, [&]{
            return !sleeping;
        });
    }
}

void server_queue::terminate() {
    std::unique_lock<std::mutex> lock(mutex_tasks);
    running = false;
    condition_tasks.notify_all();
}

bool server_queue::process_new_tasks(bool is_yielding) {
    while (true) {
        std::unique_lock<std::mutex> lock(mutex_tasks);
        if (!running) {
            QUE_DBG("%s", "terminate\n");
            return true;
        }
        if (queue_tasks.empty()) {
            return false;
        }
        server_task task = std::move(queue_tasks.front());
        queue_tasks.pop_front();
        lock.unlock();

        QUE_DBG("processing task, id = %d\n", task.id);
        if (!callback_new_task(std::move(task), is_yielding)) {
            // set it aside, do not put it back in the queue, else we offer it again in a loop
            GGML_ASSERT(is_yielding && "a task can only be declined while yielding");
            QUE_DBG("task declined, id = %d\n", task.id);
            lock.lock();
            queue_tasks_unhandled.push_back(std::move(task));
        }
    }
}

void server_queue::worker_loop() {
    while (true) {
        {
            std::unique_lock<std::mutex> lock(mutex_tasks);
            // wait on busy instead of yielding - busy stays set even when the yield already ended
            worker.cv.wait(lock, [&]{
                return worker.stop || worker.busy;
            });
            if (worker.stop) {
                return;
            }
        }

        // process tasks while the yield is active
        while (true) {
            bool terminated = false;
            try {
                // note: do not hold any lock here, the callback may post new tasks
                terminated = process_new_tasks(true);
            } catch (...) {
                std::unique_lock<std::mutex> lock(mutex_tasks);
                worker.exception = std::current_exception();
                break;
            }

            std::unique_lock<std::mutex> lock(mutex_tasks);
            if (terminated || worker.stop || !worker.yielding) {
                break;
            }
            if (!queue_tasks.empty()) {
                continue; // a new task arrived in the meantime
            }
            condition_tasks.wait(lock, [&]{
                return worker.stop || !running || !worker.yielding || !queue_tasks.empty();
            });
        }

        // signal to yield_to_queue() that no more tasks will be processed
        {
            std::unique_lock<std::mutex> lock(mutex_tasks);
            worker.busy = false;
        }
        condition_tasks.notify_all();
    }
}

void server_queue::worker_stop() {
    if (!worker.thread.joinable()) {
        return;
    }
    {
        std::unique_lock<std::mutex> lock(mutex_tasks);
        worker.stop = true;
    }
    worker.cv.notify_one();
    condition_tasks.notify_all();
    worker.thread.join();
}

void server_queue::yield_to_queue(std::function<void()> && work) {
    GGML_ASSERT(worker.thread.joinable() && "yield_to_queue() requires start_loop() to be running");

    QUE_DBG("%s", "yielding to queue\n");

    {
        std::unique_lock<std::mutex> lock(mutex_tasks);
        GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested");
        worker.busy     = true;
        worker.yielding = true;
    }
    worker.cv.notify_one();

    // run the work on the current thread, so that all ggml compute stays on the same thread
    std::exception_ptr exception;
    try {
        work();
    } catch (...) {
        exception = std::current_exception();
    }

    {
        std::unique_lock<std::mutex> lock(mutex_tasks);

        // the yield is over, wait for the worker to finish its current task
        worker.yielding = false;
        condition_tasks.notify_all();
        condition_tasks.wait(lock, [&]{
            return !worker.busy;
        });

        // put the declined tasks back, keeping their order
        while (!queue_tasks_unhandled.empty()) {
            queue_tasks.push_front(std::move(queue_tasks_unhandled.back()));
            queue_tasks_unhandled.pop_back();
        }

        // make sure to avoid idle timeout here
        time_last_task = ggml_time_ms();

        // an exception from work() takes precedence over the one from the worker
        if (!exception) {
            std::swap(exception, worker.exception);
        } else {
            worker.exception = nullptr;
        }
    }

    QUE_DBG("%s", "done yielding to queue\n");

    // note: rethrow only after the declined tasks are back in the queue, so they are not lost
    if (exception) {
        std::rethrow_exception(exception);
    }
}

void server_queue::start_loop(int64_t idle_sleep_ms) {
    running = true;
    time_last_task = ggml_time_ms();

    // spawn the worker thread used by yield_to_queue()
    GGML_ASSERT(!worker.thread.joinable() && "start_loop() is already running");
    worker.stop     = false;
    worker.busy     = false;
    worker.yielding = false;
    worker.thread = std::thread([this]() { worker_loop(); });

    constexpr auto max_wait_time = std::chrono::seconds(1);
    auto should_sleep = [&]() -> bool {
        // caller must hold mutex_tasks
        if (idle_sleep_ms < 0) {
            return false;
        }
        int64_t now = ggml_time_ms();
        return (now - time_last_task) >= idle_sleep_ms;
    };

    while (true) {
        QUE_DBG("%s", "processing new tasks\n");
        if (process_new_tasks(false)) {
            break; // terminate
        }

        // all tasks in the current loop is processed, slots data is now ready
        QUE_DBG("%s", "update slots\n");

        // this will run the main inference process for all slots
        const int64_t t_update_slots = ggml_time_ms();
        callback_update_slots();
        {
            // update_slots() may take a while to finish, we need to make sure it's not counted as idle
            // shift instead of reset, so that non-task_resets_idle_timer tasks do not delay the sleep
            std::unique_lock<std::mutex> lock(mutex_tasks);
            const int64_t now = ggml_time_ms();
            time_last_task = std::min(now, time_last_task + (now - t_update_slots));
        }

        QUE_DBG("%s", "waiting for new tasks\n");
        while (true) {
            std::unique_lock<std::mutex> lock(mutex_tasks);
            if (!running || !queue_tasks.empty()) {
                break; // go back to process new tasks or terminate
            }

            // no tasks, check for sleeping state
            if (should_sleep()) {
                QUE_INF("%s", "entering sleeping state\n");
                sleeping = true;
                // Call order cb0 -> cb1 -> cb{N}
                for (auto & cb : callback_sleeping_state) {
                    cb(true);
                }
                req_stop_sleeping = false;
                // wait until we are requested to exit sleeping state
                condition_tasks.wait(lock, [&]{
                    return (!running || req_stop_sleeping);
                });
                if (!running) { // may changed during sleep
                    break; // terminate
                }
                QUE_INF("%s", "exiting sleeping state\n");
                req_stop_sleeping = false;
                // Call order cb{N} -> cb1 -> cb0
                for (size_t i = callback_sleeping_state.size(); i > 0; i--) {
                    callback_sleeping_state[i - 1](false);
                }
                sleeping = false;
                time_last_task = ggml_time_ms();
                condition_tasks.notify_all(); // notify wait_until_no_sleep()
                break; // process new tasks
            } else {
                // wait for new tasks or timeout for checking sleeping condition
                bool res = condition_tasks.wait_for(lock, max_wait_time, [&]{
                    return (!queue_tasks.empty() || !running);
                });
                if (res) {
                    break; // new task arrived or terminate
                }
                // otherwise, loop again to check sleeping condition
            }
        }
    }

    worker_stop();
}

void server_queue::cleanup_pending_task(int id_target) {
    // no need lock because this is called exclusively by post()
    auto rm_func = [id_target](const server_task & task) {
        return task.id == id_target;
    };
    queue_tasks.erase(
        std::remove_if(queue_tasks.begin(),           queue_tasks.end(),           rm_func),
        queue_tasks.end());
    queue_tasks_deferred.erase(
        std::remove_if(queue_tasks_deferred.begin(),  queue_tasks_deferred.end(),  rm_func),
        queue_tasks_deferred.end());
    // a task declined while yielding is not in queue_tasks yet, but it can still be cancelled
    queue_tasks_unhandled.erase(
        std::remove_if(queue_tasks_unhandled.begin(), queue_tasks_unhandled.end(), rm_func),
        queue_tasks_unhandled.end());
}

//
// server_response
//

void server_response::add_waiting_task_id(int id_task) {
    std::unique_lock<std::mutex> lock(mutex_results);

    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<waiter>());

    // 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<int> & id_tasks) {
    std::unique_lock<std::mutex> lock(mutex_results);

    // one waiter for the whole set: these ids belong to one reader
    auto w = std::make_shared<waiter>();

    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.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) {
    std::unique_lock<std::mutex> lock(mutex_results);

    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;
        }),
        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<int> & id_tasks) {
    std::unique_lock<std::mutex> lock(mutex_results);

    std::vector<waiter_ptr> 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<int> & id_tasks) const {
    for (const auto & id_task : id_tasks) {
        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<int> & 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<int> & 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<pending>::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<pending>::iterator       best_it;
    uint64_t                            best_seq = 0;
    std::vector<const waiter *>         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<int> & id_tasks) {
    std::unique_lock<std::mutex> lock(mutex_results);

    while (true) {
        if (!running) {
            RES_DBG("%s : queue result stop\n", "recv");
            std::terminate(); // we cannot return here since the caller is HTTP code
        }

        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<int> & id_tasks, int timeout) {
    std::unique_lock<std::mutex> lock(mutex_results);

    // 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);

    while (true) {
        if (!running) {
            RES_DBG("%s : queue result stop\n", __func__);
            std::terminate(); // we cannot return here since the caller is HTTP code
        }

        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;
        }
    }

    // should never reach here
}

server_task_result_ptr server_response::recv(int id_task) {
    std::unordered_set<int> id_tasks = {id_task};
    return recv(id_tasks);
}

void server_response::send(server_task_result_ptr && result) {
    RES_DBG("sending result for task id = %d\n", result->id);

    std::unique_lock<std::mutex> lock(mutex_results);

    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<std::mutex> lock(mutex_results);
    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
        w->results.push_back(pending{next_seq++, std::move(res_copy)});
        w->cv.notify_all();
    }

    if (n_split_readers > 0) {
        condition_gone.notify_all();
    }
}

void server_response::terminate() {
    std::unique_lock<std::mutex> lock(mutex_results);
    running = false;
    for (const auto & [id_task, w] : waiting) {
        (void) id_task;
        w->cv.notify_all();
    }
    condition_gone.notify_all();
}

//
// server_response_reader
//

void server_response_reader::post_task(server_task && task, bool front) {
    GGML_ASSERT(id_tasks.empty() && "post_task() can only be called once per reader");
    GGML_ASSERT(!task.is_parent() && "not supported, use post_tasks() instead");
    task.index = 0;
    id_tasks.insert(task.id);
    states.push_back(task.create_state());
    queue_results.add_waiting_task_id(task.id);
    queue_tasks.post(std::move(task), front);
}

void server_response_reader::post_tasks(std::vector<server_task> && tasks, bool front) {
    GGML_ASSERT(id_tasks.empty() && "post_tasks() can only be called once per reader");
    id_tasks = server_task::get_list_id(tasks);
    states.reserve(tasks.size());
    size_t index = 0;
    for (auto & task : tasks) {
        task.index = index++;
        states.push_back(task.create_state());
        // for child tasks
        for (auto & child_task : task.child_tasks) {
            child_task.index = index++;
            states.push_back(child_task.create_state());
        }
    }
    GGML_ASSERT(states.size() == id_tasks.size());
    queue_results.add_waiting_task_ids(id_tasks);
    queue_tasks.post(std::move(tasks), front);
}

bool server_response_reader::has_next() const {
    return !cancelled && received_count < id_tasks.size();
}

// return nullptr if should_stop() is true before receiving a result
// note: if one error is received, it will stop further processing and return error result
server_task_result_ptr server_response_reader::next(const std::function<bool()> & should_stop) {
    while (true) {
        server_task_result_ptr result = queue_results.recv_with_timeout(id_tasks, polling_interval_seconds);
        if (result == nullptr) {
            // timeout, check stop condition
            if (should_stop()) {
                return nullptr;
            }
        } else {
            if (result->is_error()) {
                stop(); // cancel remaining tasks
                SRV_DBG("%s", "received error result, stopping further processing\n");
                return result;
            }
            if (!states.empty()) {
                // update the generation state if needed
                const size_t idx = result->index;
                GGML_ASSERT(idx < states.size());
                result->update(states[idx]);
            }
            if (result->is_stop()) {
                received_count++;
            }
            return result;
        }
    }

    // should not reach here
}

server_response_reader::batch_response server_response_reader::wait_for_all(const std::function<bool()> & should_stop) {
    batch_response batch_res;
    batch_res.results.clear();
    batch_res.results.resize(id_tasks.size());
    while (has_next()) {
        auto res = next(should_stop);
        if (res == nullptr) {
            batch_res.is_terminated = true;
            return batch_res;
        }
        if (res->is_error()) {
            batch_res.error = std::move(res);
            return batch_res;
        }
        const size_t idx = res->index;
        GGML_ASSERT(idx < batch_res.results.size() && "index out of range");
        GGML_ASSERT(batch_res.results[idx] == nullptr && "duplicate result received");
        batch_res.results[idx] = std::move(res);
    }
    return batch_res;
}

void server_response_reader::stop() {
    queue_results.remove_waiting_task_ids(id_tasks);
    if (has_next() && !cancelled) {
        // if tasks is not finished yet, cancel them
        cancelled = true;
        std::vector<server_task> cancel_tasks;
        cancel_tasks.reserve(id_tasks.size());
        for (const auto & id_task : id_tasks) {
            SRV_WRN("cancel task, id_task = %d\n", id_task);
            server_task task(SERVER_TASK_TYPE_CANCEL);
            task.id_target = id_task;
            queue_results.remove_waiting_task_id(id_task);
            cancel_tasks.push_back(std::move(task));
        }
        // push to beginning of the queue, so it has highest priority
        queue_tasks.post(std::move(cancel_tasks), true);
    } else {
        SRV_DBG("%s", "all tasks already finished, no need to cancel\n");
    }
}