Files
clice/src/server/worker_pool.h
2026-03-22 23:37:08 +08:00

126 lines
4.3 KiB
C++

#pragma once
#include <chrono>
#include <cstdint>
#include <functional>
#include <list>
#include <memory>
#include "eventide/async/async.h"
#include "eventide/ipc/peer.h"
#include "server/protocol.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
namespace clice {
/// Default timeout for IPC requests to worker processes.
constexpr inline auto kWorkerRequestTimeout = std::chrono::milliseconds(30000);
namespace et = eventide;
using et::ipc::RequestResult;
struct WorkerPoolOptions {
std::string self_path;
std::uint32_t stateless_count = 2;
std::uint32_t stateful_count = 2;
std::uint64_t worker_memory_limit = 4ULL * 1024 * 1024 * 1024; // 4GB default
};
class WorkerPool {
public:
WorkerPool(et::event_loop& loop) : loop(loop) {}
/// Spawn all worker processes. Returns false on failure.
bool start(const WorkerPoolOptions& options);
/// Gracefully stop all workers.
et::task<> stop();
/// Send a request to a stateful worker with path_id affinity routing.
template <typename Params>
RequestResult<Params> send_stateful(std::uint32_t path_id,
const Params& params,
et::ipc::request_options opts = {});
/// Send a request to a stateless worker with round-robin dispatch.
template <typename Params>
RequestResult<Params> send_stateless(const Params& params, et::ipc::request_options opts = {});
/// Send a notification to the stateful worker owning path_id (if any).
template <typename Params>
void notify_stateful(std::uint32_t path_id, const Params& params);
/// Remove path_id from ownership tracking (e.g. when the master learns a
/// document was evicted).
void remove_owner(std::uint32_t path_id);
/// Callback invoked when a stateful worker sends an EvictedParams notification.
/// The master should translate the path to a path_id and call remove_owner().
std::function<void(const std::string& path)> on_evicted;
private:
struct WorkerProcess {
et::process proc;
std::unique_ptr<et::ipc::BincodePeer> peer;
std::size_t owned_documents = 0;
};
et::event_loop& loop;
llvm::SmallVector<WorkerProcess> stateless_workers;
llvm::SmallVector<WorkerProcess> stateful_workers;
std::size_t next_stateless = 0;
// Stateful worker routing: path_id -> worker index with LRU tracking
llvm::DenseMap<std::uint32_t, std::size_t> owner;
std::list<std::uint32_t> owner_lru;
llvm::DenseMap<std::uint32_t, std::list<std::uint32_t>::iterator> owner_lru_index;
std::size_t assign_worker(std::uint32_t path_id);
void clear_owner(std::size_t worker_index);
std::size_t pick_least_loaded();
bool spawn_worker(const std::string& self_path, bool stateful, std::uint64_t memory_limit);
};
// --- Template implementations ---------------------------------------------------
template <typename Params>
RequestResult<Params> WorkerPool::send_stateful(std::uint32_t path_id,
const Params& params,
et::ipc::request_options opts) {
if(stateful_workers.empty()) {
co_return et::outcome_error(et::ipc::Error{"No stateful workers available"});
}
if(!opts.timeout.has_value()) {
opts.timeout = kWorkerRequestTimeout;
}
auto idx = assign_worker(path_id);
co_return co_await stateful_workers[idx].peer->send_request(params, opts);
}
template <typename Params>
RequestResult<Params> WorkerPool::send_stateless(const Params& params,
et::ipc::request_options opts) {
if(stateless_workers.empty()) {
co_return et::outcome_error(et::ipc::Error{"No stateless workers available"});
}
if(!opts.timeout.has_value()) {
opts.timeout = kWorkerRequestTimeout;
}
auto idx = next_stateless;
next_stateless = (next_stateless + 1) % stateless_workers.size();
co_return co_await stateless_workers[idx].peer->send_request(params, opts);
}
template <typename Params>
void WorkerPool::notify_stateful(std::uint32_t path_id, const Params& params) {
auto it = owner.find(path_id);
if(it == owner.end())
return;
stateful_workers[it->second].peer->send_notification(params);
}
} // namespace clice