126 lines
4.3 KiB
C++
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
|