Mooncake/mooncake-store/include/replica.h

549 lines
18 KiB
C++

#pragma once
#include <glog/logging.h>
#include <boost/functional/hash.hpp>
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_set>
#include <variant>
#include <vector>
#include <unordered_map>
#include <optional>
#include <string_view>
#include <ostream>
#include "types.h"
#include "allocator.h"
#include "master_metric_manager.h"
namespace mooncake {
/**
* @brief Globally unique replica identification.
*/
using ReplicaID = uint64_t;
/**
* @brief Type of buffer allocator used in the system
*/
enum class ReplicaType {
MEMORY, // Memory replica
DISK, // Disk replica
LOCAL_DISK // Local disk replica
};
/**
* @brief Stream operator for ReplicaType
*/
inline std::ostream& operator<<(std::ostream& os,
const ReplicaType& replicaType) noexcept {
static const std::unordered_map<ReplicaType, std::string_view>
replica_type_strings{{ReplicaType::MEMORY, "MEMORY"},
{ReplicaType::DISK, "DISK"},
{ReplicaType::LOCAL_DISK, "LOCAL_DISK"}};
os << (replica_type_strings.count(replicaType)
? replica_type_strings.at(replicaType)
: "UNKNOWN");
return os;
}
/**
* @brief Status of a replica in the system
*/
enum class ReplicaStatus {
UNDEFINED = 0, // Uninitialized
INITIALIZED, // Space allocated, waiting for write
PROCESSING, // Write in progress
COMPLETE, // Write complete, replica is available
REMOVED, // Replica has been removed
FAILED, // Failed state (can be used for reassignment)
};
/**
* @brief Stream operator for ReplicaStatus
*/
inline std::ostream& operator<<(std::ostream& os,
const ReplicaStatus& status) noexcept {
static const std::unordered_map<ReplicaStatus, std::string_view>
status_strings{{ReplicaStatus::UNDEFINED, "UNDEFINED"},
{ReplicaStatus::INITIALIZED, "INITIALIZED"},
{ReplicaStatus::PROCESSING, "PROCESSING"},
{ReplicaStatus::COMPLETE, "COMPLETE"},
{ReplicaStatus::REMOVED, "REMOVED"},
{ReplicaStatus::FAILED, "FAILED"}};
os << (status_strings.count(status) ? status_strings.at(status)
: "UNKNOWN");
return os;
}
/**
* @brief Configuration for replica management
*/
struct ReplicateConfig {
size_t replica_num{1};
bool with_soft_pin{false};
bool with_hard_pin{false}; // Hard pin: object cannot be evicted
std::vector<std::string>
preferred_segments{}; // Preferred segments for allocation
std::string preferred_segment{}; // Deprecated: Single preferred segment
// for backward compatibility
bool prefer_alloc_in_same_node{false};
ObjectDataType data_type{ObjectDataType::UNKNOWN};
friend std::ostream& operator<<(std::ostream& os,
const ReplicateConfig& config) noexcept {
os << "ReplicateConfig: { replica_num: " << config.replica_num
<< ", with_soft_pin: " << config.with_soft_pin
<< ", with_hard_pin: " << config.with_hard_pin
<< ", preferred_segments: [";
for (size_t i = 0; i < config.preferred_segments.size(); ++i) {
os << config.preferred_segments[i];
if (i < config.preferred_segments.size() - 1) os << ", ";
}
os << "]";
if (!config.preferred_segment.empty()) {
os << ", preferred_segment (deprecated): "
<< config.preferred_segment;
}
os << ", prefer_alloc_in_same_node: "
<< config.prefer_alloc_in_same_node
<< ", data_type: " << config.data_type << " }";
return os;
}
};
struct MemoryReplicaData {
std::unique_ptr<AllocatedBuffer> buffer;
};
struct DiskReplicaData {
std::string file_path;
uint64_t object_size = 0;
};
struct LocalDiskReplicaData {
UUID client_id;
uint64_t object_size = 0;
std::string transport_endpoint;
};
struct MemoryDescriptor {
AllocatedBuffer::Descriptor buffer_descriptor;
YLT_REFL(MemoryDescriptor, buffer_descriptor);
};
struct DiskDescriptor {
std::string file_path{};
uint64_t object_size = 0;
YLT_REFL(DiskDescriptor, file_path, object_size);
};
struct LocalDiskDescriptor {
UUID client_id;
uint64_t object_size = 0;
std::string transport_endpoint;
YLT_REFL(LocalDiskDescriptor, client_id, object_size, transport_endpoint);
};
class Replica {
public:
struct Descriptor;
// memory replica constructor
Replica(std::unique_ptr<AllocatedBuffer> buffer, ReplicaStatus status)
: id_(next_id_.fetch_add(1)),
data_(MemoryReplicaData{std::move(buffer)}),
status_(status),
refcnt_(0) {}
// disk replica constructor
Replica(std::string file_path, uint64_t object_size, ReplicaStatus status)
: id_(next_id_.fetch_add(1)),
data_(DiskReplicaData{std::move(file_path), object_size}),
status_(status),
refcnt_(0) {
// Automatic update allocated_file_size via RAII
MasterMetricManager::instance().inc_allocated_file_size(object_size);
}
// local disk replica constructor
Replica(UUID client_id, uint64_t object_size,
std::string transport_endpoint, ReplicaStatus status)
: id_(next_id_.fetch_add(1)),
data_(LocalDiskReplicaData{client_id, object_size,
std::move(transport_endpoint)}),
status_(status),
refcnt_(0) {
MasterMetricManager::instance().inc_allocated_file_size(object_size);
}
~Replica() {
if (status_ == ReplicaStatus::UNDEFINED) return;
if (is_disk_replica()) {
MasterMetricManager::instance().dec_allocated_file_size(
std::get<DiskReplicaData>(data_).object_size);
} else if (is_local_disk_replica()) {
MasterMetricManager::instance().dec_allocated_file_size(
std::get<LocalDiskReplicaData>(data_).object_size);
}
}
// Copy-construction is not allowed.
Replica(const Replica&) = delete;
Replica& operator=(const Replica&) = delete;
// Move-construction is allowed.
Replica(Replica&& src) noexcept
: id_(src.id_),
data_(std::move(src.data_)),
status_(src.status_),
refcnt_(src.refcnt_.exchange(0)) {
// Mark the source as moved-from so its destructor doesn't
// double-decrement metrics.
src.status_ = ReplicaStatus::UNDEFINED;
}
Replica& operator=(Replica&& src) noexcept {
if (this == &src) {
// Same object, skip moving.
return *this;
}
// Decrement metric for the current object before overwriting.
if (status_ != ReplicaStatus::UNDEFINED) {
if (is_disk_replica()) {
MasterMetricManager::instance().dec_allocated_file_size(
std::get<DiskReplicaData>(data_).object_size);
} else if (is_local_disk_replica()) {
MasterMetricManager::instance().dec_allocated_file_size(
std::get<LocalDiskReplicaData>(data_).object_size);
}
}
id_ = src.id_;
data_ = std::move(src.data_);
status_ = src.status_;
refcnt_.store(src.refcnt_.exchange(0));
// Mark src as moved-from.
src.status_ = ReplicaStatus::UNDEFINED;
return *this;
}
[[nodiscard]] Descriptor get_descriptor() const;
[[nodiscard]] ReplicaID id() const { return id_; }
[[nodiscard]] ReplicaStatus status() const { return status_; }
[[nodiscard]] bool is_completed() const {
return status_ == ReplicaStatus::COMPLETE;
}
[[nodiscard]] static bool fn_is_completed(const Replica& replica) {
return replica.is_completed();
}
[[nodiscard]] bool is_processing() const {
return status_ == ReplicaStatus::PROCESSING;
}
[[nodiscard]] static bool fn_is_processing(const Replica& replica) {
return replica.is_processing();
}
[[nodiscard]] bool is_busy() const { return refcnt_.load() > 0; }
[[nodiscard]] static bool fn_is_busy(const Replica& replica) {
return replica.is_busy();
}
[[nodiscard]] ReplicaType type() const {
return std::visit(ReplicaTypeVisitor{}, data_);
}
[[nodiscard]] bool is_memory_replica() const {
return std::holds_alternative<MemoryReplicaData>(data_);
}
[[nodiscard]] static bool fn_is_memory_replica(const Replica& replica) {
return replica.is_memory_replica();
}
[[nodiscard]] bool is_disk_replica() const {
return std::holds_alternative<DiskReplicaData>(data_);
}
[[nodiscard]] static bool fn_is_disk_replica(const Replica& replica) {
return replica.is_disk_replica();
}
[[nodiscard]] bool is_local_disk_replica() const {
return std::holds_alternative<LocalDiskReplicaData>(data_);
}
[[nodiscard]] static bool fn_is_local_disk_replica(const Replica& replica) {
return replica.is_local_disk_replica();
}
[[nodiscard]] bool has_invalid_mem_handle() const {
if (is_memory_replica()) {
const auto& mem_data = std::get<MemoryReplicaData>(data_);
return !mem_data.buffer->isAllocatorValid();
}
return false; // DiskReplicaData does not have handles
}
/**
* @brief Check if a local_disk replica's owner client is still alive.
* Used by CleanupStaleHandles to remove replicas belonging to expired
* clients. For non-local_disk replicas, always returns false.
* @param alive_clients Set of currently alive client IDs.
* @return true if this is a local_disk replica whose client is not alive.
*/
[[nodiscard]] bool has_stale_local_disk_client(
const std::unordered_set<UUID, boost::hash<UUID>>& alive_clients)
const {
auto client_id = get_local_disk_client_id();
if (client_id.has_value()) {
return alive_clients.find(client_id.value()) == alive_clients.end();
}
return false;
}
/**
* @brief Get the client_id for local_disk replicas.
* @return The client_id if this is a local_disk replica, std::nullopt
* otherwise.
*/
[[nodiscard]] std::optional<UUID> get_local_disk_client_id() const {
if (is_local_disk_replica()) {
const auto& disk_data = std::get<LocalDiskReplicaData>(data_);
return disk_data.client_id;
}
return std::nullopt;
}
[[nodiscard]] size_t get_memory_buffer_size() const {
if (is_memory_replica()) {
const auto& mem_data = std::get<MemoryReplicaData>(data_);
return mem_data.buffer->size();
} else {
LOG(ERROR) << "Invalid replica type: " << type();
return 0;
}
}
[[nodiscard]] std::vector<std::optional<std::string>> get_segment_names()
const;
void mark_complete() {
if (status_ == ReplicaStatus::PROCESSING) {
status_ = ReplicaStatus::COMPLETE;
} else if (status_ == ReplicaStatus::COMPLETE) {
LOG(WARNING) << "Replica already marked as complete";
} else {
LOG(ERROR) << "Invalid replica status: " << status_;
}
}
void mark_processing() {
if (status_ == ReplicaStatus::COMPLETE) {
status_ = ReplicaStatus::PROCESSING;
} else {
LOG(ERROR) << "Cannot mark_processing from status: " << status_;
}
}
void inc_refcnt() { refcnt_.fetch_add(1); }
void dec_refcnt() { refcnt_.fetch_sub(1); }
uint32_t get_refcnt() const { return refcnt_.load(); }
friend std::ostream& operator<<(std::ostream& os, const Replica& replica);
struct ReplicaTypeVisitor {
ReplicaType operator()(const MemoryReplicaData&) const {
return ReplicaType::MEMORY;
}
ReplicaType operator()(const DiskReplicaData&) const {
return ReplicaType::DISK;
}
ReplicaType operator()(const LocalDiskReplicaData&) const {
return ReplicaType::LOCAL_DISK;
}
};
struct Descriptor {
ReplicaID id;
std::variant<MemoryDescriptor, DiskDescriptor, LocalDiskDescriptor>
descriptor_variant;
ReplicaStatus status;
YLT_REFL(Descriptor, id, descriptor_variant, status);
// Helper functions
bool is_memory_replica() noexcept {
return std::holds_alternative<MemoryDescriptor>(descriptor_variant);
}
bool is_memory_replica() const noexcept {
return std::holds_alternative<MemoryDescriptor>(descriptor_variant);
}
bool is_disk_replica() noexcept {
return std::holds_alternative<DiskDescriptor>(descriptor_variant);
}
bool is_disk_replica() const noexcept {
return std::holds_alternative<DiskDescriptor>(descriptor_variant);
}
bool is_local_disk_replica() noexcept {
return std::holds_alternative<LocalDiskDescriptor>(
descriptor_variant);
}
bool is_local_disk_replica() const noexcept {
return std::holds_alternative<LocalDiskDescriptor>(
descriptor_variant);
}
MemoryDescriptor& get_memory_descriptor() {
if (auto* desc =
std::get_if<MemoryDescriptor>(&descriptor_variant)) {
return *desc;
}
throw std::runtime_error("Expected MemoryDescriptor");
}
DiskDescriptor& get_disk_descriptor() {
if (auto* desc = std::get_if<DiskDescriptor>(&descriptor_variant)) {
return *desc;
}
throw std::runtime_error("Expected DiskDescriptor");
}
LocalDiskDescriptor& get_local_disk_descriptor() {
if (auto* desc =
std::get_if<LocalDiskDescriptor>(&descriptor_variant)) {
return *desc;
}
throw std::runtime_error("Expected LocalDiskDescriptor");
}
const MemoryDescriptor& get_memory_descriptor() const {
if (auto* desc =
std::get_if<MemoryDescriptor>(&descriptor_variant)) {
return *desc;
}
throw std::runtime_error("Expected MemoryDescriptor");
}
const DiskDescriptor& get_disk_descriptor() const {
if (auto* desc = std::get_if<DiskDescriptor>(&descriptor_variant)) {
return *desc;
}
throw std::runtime_error("Expected DiskDescriptor");
}
const LocalDiskDescriptor& get_local_disk_descriptor() const {
if (auto* desc =
std::get_if<LocalDiskDescriptor>(&descriptor_variant)) {
return *desc;
}
throw std::runtime_error("Expected LocalDiskDescriptor");
}
};
private:
inline static std::atomic<ReplicaID> next_id_{1};
ReplicaID id_;
std::variant<MemoryReplicaData, DiskReplicaData, LocalDiskReplicaData>
data_;
ReplicaStatus status_{ReplicaStatus::UNDEFINED};
friend class Serializer<Replica>;
friend class MasterService; // For MetadataSerializer to access next_id_
std::atomic<uint32_t> refcnt_{0};
};
inline Replica::Descriptor Replica::get_descriptor() const {
Replica::Descriptor desc;
desc.id = id_;
desc.status = status_;
if (is_memory_replica()) {
const auto& mem_data = std::get<MemoryReplicaData>(data_);
MemoryDescriptor mem_desc;
if (mem_data.buffer) {
mem_desc.buffer_descriptor = mem_data.buffer->get_descriptor();
} else {
mem_desc.buffer_descriptor.size_ = 0;
mem_desc.buffer_descriptor.buffer_address_ = 0;
mem_desc.buffer_descriptor.transport_endpoint_ = "";
LOG(ERROR) << "Trying to get invalid memory replica descriptor";
}
desc.descriptor_variant = std::move(mem_desc);
} else if (is_disk_replica()) {
const auto& disk_data = std::get<DiskReplicaData>(data_);
DiskDescriptor disk_desc;
disk_desc.file_path = disk_data.file_path;
disk_desc.object_size = disk_data.object_size;
desc.descriptor_variant = std::move(disk_desc);
} else if (is_local_disk_replica()) {
const auto& disk_data = std::get<LocalDiskReplicaData>(data_);
LocalDiskDescriptor local_disk_desc;
local_disk_desc.client_id = disk_data.client_id;
local_disk_desc.object_size = disk_data.object_size;
local_disk_desc.transport_endpoint = disk_data.transport_endpoint;
desc.descriptor_variant = std::move(local_disk_desc);
}
return desc;
}
inline std::vector<std::optional<std::string>> Replica::get_segment_names()
const {
if (is_memory_replica()) {
const auto& mem_data = std::get<MemoryReplicaData>(data_);
std::vector<std::optional<std::string>> segment_names;
if (mem_data.buffer && mem_data.buffer->isAllocatorValid()) {
segment_names.push_back(mem_data.buffer->getSegmentName());
} else {
segment_names.push_back(std::nullopt);
}
return segment_names;
}
return std::vector<std::optional<std::string>>();
}
inline std::ostream& operator<<(std::ostream& os, const Replica& replica) {
os << "Replica: { id: " << replica.id_ << ", status: " << replica.status_
<< ", ";
if (replica.is_memory_replica()) {
const auto& mem_data = std::get<MemoryReplicaData>(replica.data_);
os << "type: MEMORY, buffers: [";
if (mem_data.buffer) {
os << *mem_data.buffer;
}
os << "]";
} else if (replica.is_disk_replica()) {
const auto& disk_data = std::get<DiskReplicaData>(replica.data_);
os << "type: DISK, file_path: " << disk_data.file_path
<< ", object_size: " << disk_data.object_size;
}
os << ", refcnt: " << replica.refcnt_.load() << " }";
return os;
}
} // namespace mooncake