zbtrs2
/
Mooncake
forked from mooncake-track/Mooncake
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #1314 from openanolis/xinyi/tiered-backend-next 

[Store] Tiered Backend add Dram tier support
This commit is contained in:
Wan 2026-01-08 12:50:29 +08:00 committed by GitHub
parent 8e85cd4881 751d5f6122
commit bd71efe05a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 1636 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

79
mooncake-store/include/tiered_cache/cache_tier.h
View File

@ -4,8 +4,11 @@
#include <vector>
#include <memory>
#include <optional>
#include <ylt/util/tl/expected.hpp>
#include "allocator.h"
#include "transfer_engine.h"
#include "types.h"
namespace mooncake {
@ -28,14 +31,66 @@ static inline std::string MemoryTypeToString(MemoryType type) {
}
}
/**
* @class BufferBase
* @brief Base class for different types of memory buffers
*/
class BufferBase {
public:
virtual ~BufferBase() = default;
virtual uint64_t data() const = 0;
virtual std::size_t size() const = 0;
};
/**
* @class DRAMBuffer
* @brief Wrapper for DRAM AllocatedBuffer
*/
class DRAMBuffer : public BufferBase {
public:
explicit DRAMBuffer(std::unique_ptr<AllocatedBuffer> buffer)
: dram_buffer_(std::move(buffer)) {}
uint64_t data() const override {
return dram_buffer_ ? reinterpret_cast<uint64_t>(dram_buffer_->data())
: 0;
}
std::size_t size() const override {
return dram_buffer_ ? dram_buffer_->size() : 0;
}
private:
std::unique_ptr<AllocatedBuffer> dram_buffer_;
};
/**
* @class TempDRAMBuffer
* @brief Wrapper for temporary DRAM buffers with RAII memory management
*/
class TempDRAMBuffer : public BufferBase {
public:
// Constructor that takes ownership of the buffer
explicit TempDRAMBuffer(std::unique_ptr<char[]> buffer, size_t size)
: buffer_(std::move(buffer)), size_(size) {}
uint64_t data() const override {
return reinterpret_cast<uint64_t>(buffer_.get());
}
std::size_t size() const override { return size_; }
private:
std::unique_ptr<char[]>
buffer_; // Owns the memory, auto-releases on destruction
size_t size_;
};
/**
* @struct DataSource
* @brief Describes a source of data for copy/write operations.
*/
struct DataSource {
uint64_t ptr; // Pointer to data (if in memory) / file descriptor
uint64_t offset; // Offset within the source (for files/SSDs)
size_t size; // Size in bytes
std::unique_ptr<BufferBase> buffer;
MemoryType type; // Source memory type
};
@ -51,27 +106,31 @@ class CacheTier {
/**
* @brief Initializes the cache tier.
* @return tl::expected<void, ErrorCode> indicating success or error code.
*/
virtual bool Init(TieredBackend* backend, TransferEngine* engine) = 0;
virtual tl::expected<void, ErrorCode> Init(TieredBackend* backend,
TransferEngine* engine) = 0;
/**
* @brief Reserve Space (Allocation)
* Finds free space of `size` bytes. Does NOT copy data.
* * @param size Bytes to allocate.
* @param size Bytes to allocate.
* @param data DataSource struct to fill with allocation info.
* @return true if allocation succeeds.
* @return tl::expected<void, ErrorCode> indicating success or error code.
*/
virtual bool Allocate(size_t size, DataSource& data) = 0;
virtual tl::expected<void, ErrorCode> Allocate(size_t size,
DataSource& data) = 0;
/**
* @brief Free Space (Rollback/Cleanup)
* Releases space at offset. Used when writes fail or explicitly freeing
* anonymous blocks.
* @return tl::expected<void, ErrorCode> indicating success or error code.
*/
virtual bool Free(DataSource data) = 0;
virtual tl::expected<void, ErrorCode> Free(DataSource data) = 0;
// --- Accessors & Metadata ---
virtual uint64_t GetTierId() const = 0;
virtual UUID GetTierId() const = 0;
virtual size_t GetCapacity() const = 0;
virtual size_t GetUsage() const = 0;
virtual MemoryType GetMemoryType() const = 0;
@ -83,4 +142,4 @@ class CacheTier {
TieredBackend* backend_ = nullptr;
};
} // namespace mooncake
} // namespace mooncake

47
mooncake-store/include/tiered_cache/dram_tier.h Normal file
View File

@ -0,0 +1,47 @@
#pragma once
#include <string>
#include <vector>
#include <memory>
#include <unordered_map>
#include <optional>
#include "allocator.h"
#include "tiered_cache/cache_tier.h"
#include "transfer_engine.h"
namespace mooncake {
class DramCacheTier : public CacheTier {
public:
DramCacheTier(
UUID tier_id, size_t capacity, const std::vector<std::string>& tags,
std::optional<int> numa_node = std::nullopt,
BufferAllocatorType allocator_type = BufferAllocatorType::OFFSET);
~DramCacheTier() override;
tl::expected<void, ErrorCode> Init(TieredBackend* backend,
TransferEngine* engine) override;
tl::expected<void, ErrorCode> Allocate(size_t size,
DataSource& data) override;
tl::expected<void, ErrorCode> Free(DataSource data) override;
UUID GetTierId() const override { return tier_id_; }
size_t GetCapacity() const override { return capacity_; }
size_t GetUsage() const override;
const std::vector<std::string>& GetTags() const override { return tags_; }
MemoryType GetMemoryType() const override { return MemoryType::DRAM; }
private:
UUID tier_id_;
size_t capacity_;
std::vector<std::string> tags_;
std::optional<int> numa_node_;
BufferAllocatorType allocator_type_;
std::shared_ptr<BufferAllocatorBase> allocator_;
TransferEngine* engine_;
std::unique_ptr<char[], void (*)(char*)> memory_buffer_{
nullptr, [](char* p) { delete[] p; }};
};
} // namespace mooncake

12
mooncake-store/include/tiered_cache/tiered_backend.h
View File

@ -23,7 +23,7 @@ class TieredBackend; // Forward declaration
* storage.
*/
struct TieredLocation {
UUID tier_id;
CacheTier* tier;
struct DataSource data;
};
@ -59,7 +59,8 @@ struct AllocationEntry {
TieredBackend* backend;
TieredLocation loc;
AllocationEntry(TieredBackend* b, TieredLocation l) : backend(b), loc(l) {}
AllocationEntry(TieredBackend* b, TieredLocation&& l)
: backend(b), loc(std::move(l)) {}
AllocationEntry(const AllocationEntry&) = delete;
AllocationEntry& operator=(const AllocationEntry&) = delete;
@ -91,8 +92,8 @@ class TieredBackend {
TieredBackend();
~TieredBackend() = default;
bool Init(Json::Value root, TransferEngine* engine,
MetadataSyncCallback sync_callback);
tl::expected<void, ErrorCode> Init(Json::Value root, TransferEngine* engine,
MetadataSyncCallback sync_callback);
// --- Client-Centric Operations ---
// All the following operations are designed for Client-Centric, Client
@ -161,9 +162,6 @@ class TieredBackend {
const CacheTier* GetTier(UUID tier_id) const;
const DataCopier& GetDataCopier() const;
// Internal API called by AllocationEntry destructor
void FreeInternal(const TieredLocation& loc);
private:
struct TierInfo {
int priority;

1
mooncake-store/src/CMakeLists.txt
View File

@ -33,6 +33,7 @@ set(MOONCAKE_STORE_SOURCES
tiered_cache/copier_registry.cpp
tiered_cache/data_copier.cpp
tiered_cache/tiered_backend.cpp
tiered_cache/dram_tier.cpp
)
set(EXTRA_LIBS "")

57
mooncake-store/src/tiered_cache/data_copier.cpp
View File

@ -2,12 +2,53 @@
#include <memory>
#include <utility>
#include "tiered_cache/data_copier.h"
#include "tiered_cache/cache_tier.h"
#include "tiered_cache/copier_registry.h"
#include "tiered_cache/data_copier.h"
namespace mooncake {
// DRAM <-> DRAM
tl::expected<void, ErrorCode> CopyDramToDram(const DataSource& src,
const DataSource& dest) {
// Validate buffers exist
if (!src.buffer || !dest.buffer) {
LOG(ERROR) << "Invalid buffer: src.buffer="
<< (src.buffer ? "valid" : "null")
<< ", dest.buffer=" << (dest.buffer ? "valid" : "null");
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
const void* src_ptr = reinterpret_cast<const void*>(src.buffer->data());
void* dest_ptr = reinterpret_cast<void*>(dest.buffer->data());
size_t size = src.buffer->size();
// Validate pointers and size
if (!src_ptr || !dest_ptr) {
LOG(ERROR) << "Invalid pointer: src_ptr=" << src_ptr
<< ", dest_ptr=" << dest_ptr;
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
if (size == 0) {
LOG(WARNING) << "Copy with zero size, skipping memcpy";
return tl::expected<void, ErrorCode>{};
}
// Validate dest buffer size
if (dest.buffer->size() < size) {
LOG(ERROR) << "Destination buffer too small: dest_size="
<< dest.buffer->size() << ", required=" << size;
return tl::unexpected(ErrorCode::BUFFER_OVERFLOW);
}
memcpy(dest_ptr, src_ptr, size);
return tl::expected<void, ErrorCode>{};
}
DataCopierBuilder::DataCopierBuilder() {
// Add the default DRAM<->DRAM copier.
copy_matrix_[{MemoryType::DRAM, MemoryType::DRAM}] = CopyDramToDram;
// Process all registrations from the global registry.
const auto& registry = CopierRegistry::GetInstance();
@ -81,8 +122,10 @@ tl::expected<void, ErrorCode> DataCopier::Copy(const DataSource& src,
auto from_dram_copier = FindCopier(MemoryType::DRAM, dest_type);
if (to_dram_copier && from_dram_copier) {
// Create a temporary DRAM buffer for the fallback path
size_t buffer_size = src.buffer->size();
std::unique_ptr<char[]> temp_dram_buffer(
new (std::nothrow) char[src.size]);
new (std::nothrow) char[buffer_size]);
if (!temp_dram_buffer) {
LOG(ERROR) << "Failed to allocate temporary DRAM buffer for "
"fallback copy.";
@ -90,9 +133,13 @@ tl::expected<void, ErrorCode> DataCopier::Copy(const DataSource& src,
}
// Step A: Source -> DRAM
DataSource temp_dram = {
reinterpret_cast<uint64_t>(temp_dram_buffer.get()), 0, src.size,
MemoryType::DRAM};
DataSource temp_dram;
// Transfer ownership to TempDRAMBuffer (it will be released when
// temp_dram goes out of scope)
temp_dram.buffer = std::make_unique<TempDRAMBuffer>(
std::move(temp_dram_buffer), buffer_size);
temp_dram.type = MemoryType::DRAM;
if (!to_dram_copier(src, temp_dram)) {
LOG(ERROR) << "Fallback copy failed at Step A (Source -> DRAM)";
return tl::make_unexpected(ErrorCode::DATA_COPY_FAILED);

199
mooncake-store/src/tiered_cache/dram_tier.cpp Normal file
View File

@ -0,0 +1,199 @@
#include <glog/logging.h>
#include <numa.h>
#include <chrono>
#include <thread>
#include "tiered_cache/dram_tier.h"
#include "tiered_cache/tiered_backend.h"
#include "tiered_cache/copier_registry.h"
#include "transfer_engine.h"
namespace mooncake {
DramCacheTier::DramCacheTier(UUID tier_id, size_t capacity,
const std::vector<std::string>& tags,
std::optional<int> numa_node,
BufferAllocatorType allocator_type)
: tier_id_(tier_id),
capacity_(capacity),
tags_(tags),
numa_node_(numa_node),
allocator_type_(allocator_type),
allocator_(nullptr),
engine_(nullptr) {}
DramCacheTier::~DramCacheTier() {
// Wait for all allocated buffers to be released before destroying allocator
if (allocator_) {
size_t allocated_size = allocator_->size();
if (allocated_size > 0) {
LOG(WARNING) << "DramCacheTier " << tier_id_
<< " is being destroyed with " << allocated_size
<< " bytes still allocated. Waiting for buffers to be "
"released...";
// Wait for buffers to be released
constexpr int kCheckIntervalMs = 100;
int i = 0;
while (true) {
allocated_size = allocator_->size();
if (allocated_size == 0) {
LOG(INFO) << "All buffers released for DramCacheTier "
<< tier_id_;
break;
}
if (i == 10) { // Log every second
LOG(INFO) << "DramCacheTier " << tier_id_ << " waiting for "
<< allocated_size << " bytes to be released...";
i = 0;
}
std::this_thread::sleep_for(
std::chrono::milliseconds(kCheckIntervalMs));
++i;
}
}
}
allocator_.reset();
if (engine_ != nullptr && memory_buffer_ != nullptr) {
LOG(INFO) << "unregistering memory for DramCacheTier " << tier_id_;
int rc = engine_->unregisterLocalMemory(memory_buffer_.get());
if (rc != 0) {
LOG(ERROR) << "Failed to unregister memory for DramCacheTier "
<< tier_id_ << ", engine ret is " << rc;
}
}
}
tl::expected<void, ErrorCode> DramCacheTier::Init(TieredBackend* backend,
TransferEngine* engine) {
int node = -1;
std::string location;
backend_ = backend;
if (engine != nullptr) engine_ = engine;
// Allocate a contiguous memory block.
if (numa_node_.has_value()) {
if (numa_available() < 0) {
LOG(ERROR) << "NUMA not available on this system.";
return tl::unexpected(ErrorCode::INTERNAL_ERROR);
}
node = numa_node_.value();
if (node < 0 || node > numa_max_node()) {
LOG(ERROR) << "Invalid NUMA node " << node;
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
char* mem_ptr = static_cast<char*>(numa_alloc_onnode(capacity_, node));
if (!mem_ptr) {
LOG(ERROR) << "Failed to allocate " << capacity_
<< " bytes from NUMA node " << node
<< " for DramCacheTier " << tier_id_;
return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
}
memory_buffer_ = std::unique_ptr<char[], void (*)(char*)>(
mem_ptr, [](char* p) { numa_free(p, 0); });
LOG(INFO) << "Allocated " << capacity_ << " bytes from NUMA node "
<< node << " for DramCacheTier " << tier_id_;
} else {
try {
memory_buffer_ = std::unique_ptr<char[], void (*)(char*)>(
new char[capacity_], [](char* p) { delete[] p; });
} catch (const std::bad_alloc& e) {
LOG(ERROR) << "Failed to allocate " << capacity_
<< " bytes for DramCacheTier " << tier_id_ << ": "
<< e.what();
return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
}
LOG(INFO) << "Allocated " << capacity_ << " bytes for DramCacheTier "
<< tier_id_;
}
char* mem_ptr = memory_buffer_.get();
// Register this newly allocated memory with the TransferEngine.
if (engine_) {
if (numa_node_.has_value()) {
location = "cpu:" + std::to_string(node);
} else {
location = kWildcardLocation;
}
int rc = engine_->registerLocalMemory(mem_ptr, capacity_, location);
if (rc != 0) {
LOG(ERROR) << "Failed to register memory with TransferEngine for "
"DramCacheTier "
<< tier_id_ << ", engine ret is " << rc;
return tl::unexpected(ErrorCode::INTERNAL_ERROR);
} else {
LOG(INFO)
<< "registered memory with TransferEngine for DramCacheTier "
<< tier_id_ << " at " << static_cast<void*>(mem_ptr);
}
}
// Use the address of this registered block as the base_address for the
// allocator.
const uintptr_t base_address = reinterpret_cast<uintptr_t>(mem_ptr);
std::string segment_name = "dram_tier_" + std::to_string(tier_id_.first) +
"-" + std::to_string(tier_id_.second);
switch (allocator_type_) {
case BufferAllocatorType::OFFSET:
allocator_ = std::make_shared<OffsetBufferAllocator>(
segment_name, base_address, capacity_, segment_name);
break;
case BufferAllocatorType::CACHELIB:
allocator_ = std::make_shared<CachelibBufferAllocator>(
segment_name, base_address, capacity_, segment_name);
break;
default:
LOG(ERROR) << "Unsupported allocator type for DramCacheTier";
if (engine_) {
engine_->unregisterLocalMemory(mem_ptr);
}
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
LOG(INFO) << "DramCacheTier " << tier_id_ << " initialized and registered "
<< capacity_ << " bytes at base address 0x" << std::hex
<< base_address;
return tl::expected<void, ErrorCode>{};
}
size_t DramCacheTier::GetUsage() const {
return allocator_ ? allocator_->size() : 0;
}
tl::expected<void, ErrorCode> DramCacheTier::Allocate(size_t size,
DataSource& data_source) {
if (!allocator_) {
LOG(ERROR) << "Allocator not initialized for DramCacheTier "
<< tier_id_;
return tl::unexpected(ErrorCode::INTERNAL_ERROR);
}
auto alloc_result = allocator_->allocate(size);
if (!alloc_result) {
LOG(ERROR) << "Failed to allocate " << size
<< " bytes from DramCacheTier " << tier_id_;
return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
}
auto dram_buffer_wrapper =
std::make_unique<DRAMBuffer>(std::move(alloc_result));
data_source.buffer = std::move(dram_buffer_wrapper);
data_source.type = MemoryType::DRAM;
return tl::expected<void, ErrorCode>{};
}
tl::expected<void, ErrorCode> DramCacheTier::Free(DataSource data_source) {
if (!data_source.buffer) {
LOG(WARNING) << "Attempting to free null buffer in DramCacheTier "
<< tier_id_;
}
// RAII will handle the deallocation when buffer_handle goes out of scope.
return tl::expected<void, ErrorCode>{};
}
} // namespace mooncake

161
mooncake-store/src/tiered_cache/tiered_backend.cpp
View File

@ -5,28 +5,29 @@
#include "tiered_cache/tiered_backend.h"
#include "tiered_cache/cache_tier.h"
#include "tiered_cache/dram_tier.h"
namespace mooncake {
AllocationEntry::~AllocationEntry() {
if (backend) {
// When ref count drops to 0, call back to backend to free physical
// resource.
backend->FreeInternal(loc);
if (backend && loc.tier) {
// When ref count drops to 0, free physical resource directly.
loc.tier->Free(std::move(loc.data));
}
}
TieredBackend::TieredBackend() = default;
bool TieredBackend::Init(Json::Value root, TransferEngine* engine,
MetadataSyncCallback sync_callback) {
tl::expected<void, ErrorCode> TieredBackend::Init(
Json::Value root, TransferEngine* engine,
MetadataSyncCallback sync_callback) {
// Initialize DataCopier
try {
DataCopierBuilder builder;
data_copier_ = builder.Build();
} catch (const std::logic_error& e) {
LOG(ERROR) << "Failed to build DataCopier: " << e.what();
return false;
return tl::unexpected(ErrorCode::INTERNAL_ERROR);
}
// Register callback for syncing metadata to Master
@ -35,33 +36,100 @@ bool TieredBackend::Init(Json::Value root, TransferEngine* engine,
// Initialize Tiers
if (!root.isMember("tiers")) {
LOG(ERROR) << "Tiered cache config is missing 'tiers' array.";
return false;
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
for (const auto& tier_config : root["tiers"]) {
UUID id = generate_uuid();
// std::string type = tier_config["type"].asString(); // Unused for now
int priority = tier_config["priority"].asInt();
std::vector<std::string> tags;
if (tier_config.isMember("tags")) {
for (const auto& tag : tier_config["tags"])
tags.push_back(tag.asString());
// Parse required fields
if (!tier_config.isMember("type")) {
LOG(ERROR) << "Tier config missing required field 'type'";
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
if (!tier_config.isMember("capacity")) {
LOG(ERROR) << "Tier config missing required field 'capacity'";
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
if (!tier_config.isMember("priority")) {
LOG(ERROR) << "Tier config missing required field 'priority'";
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
// TODO: Logic to instantiate specific CacheTier types (DRAM/SSD) goes
// here. For example: std::unique_ptr<CacheTier> tier =
// CacheTierFactory::Create(tier_config); tier->Init(this, engine);
// tiers_[id] = std::move(tier);
std::string type = tier_config["type"].asString();
size_t capacity = tier_config["capacity"].asUInt64();
int priority = tier_config["priority"].asInt();
// Placeholder for compilation if Factory is not ready
// tiers_[id] = std::make_unique<DramTier>();
// Validate capacity
if (capacity == 0) {
LOG(ERROR) << "Invalid capacity (0) for tier type " << type;
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
tier_info_[id] = {priority, tags};
// Parse tags
std::vector<std::string> tags;
if (tier_config.isMember("tags")) {
for (const auto& tag : tier_config["tags"]) {
tags.push_back(tag.asString());
}
}
// Generate UUID for this tier
UUID id = generate_uuid();
// Instantiate tier based on type
if (type == "DRAM") {
// Parse NUMA node
std::optional<int> numa_node;
if (tier_config.isMember("numa_node")) {
int node = tier_config["numa_node"].asInt();
if (node < 0) {
LOG(WARNING) << "Invalid NUMA node (" << node
<< "), using default allocation";
} else {
numa_node = node;
}
}
// Parse allocator type
BufferAllocatorType allocator_type = BufferAllocatorType::OFFSET;
if (tier_config.isMember("allocator_type")) {
std::string allocator_str =
tier_config["allocator_type"].asString();
if (allocator_str == "OFFSET") {
allocator_type = BufferAllocatorType::OFFSET;
} else if (allocator_str == "CACHELIB") {
allocator_type = BufferAllocatorType::CACHELIB;
} else {
LOG(WARNING) << "Unknown allocator_type '" << allocator_str
<< "', using default OFFSET";
}
}
LOG(INFO) << "Creating DRAM tier: id=" << id
<< ", capacity=" << capacity << ", priority=" << priority
<< ", allocator_type=" << allocator_type
<< (numa_node.has_value()
? ", numa_node=" + std::to_string(*numa_node)
: "");
auto tier = std::make_unique<DramCacheTier>(
id, capacity, tags, numa_node, allocator_type);
auto init_result = tier->Init(this, engine);
if (!init_result) {
LOG(ERROR) << "Failed to initialize DRAM tier: id=" << id
<< ", error=" << init_result.error();
return tl::unexpected(init_result.error());
}
tiers_[id] = std::move(tier);
tier_info_[id] = {priority, tags};
LOG(INFO) << "Successfully initialized DRAM tier: id=" << id;
} else {
LOG(ERROR) << "Unsupported tier type '" << type << "'";
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
}
LOG(INFO) << "TieredBackend initialized successfully with "
<< tier_info_.size() << " tiers.";
return true;
return tl::expected<void, ErrorCode>{};
}
std::vector<UUID> TieredBackend::GetSortedTiers() const {
@ -84,8 +152,9 @@ bool TieredBackend::AllocateInternalRaw(size_t size,
if (preferred_tier.has_value()) {
auto it = tiers_.find(*preferred_tier);
if (it != tiers_.end()) {
if (it->second->Allocate(size, out_loc->data)) {
out_loc->tier_id = *preferred_tier;
auto alloc_result = it->second->Allocate(size, out_loc->data);
if (alloc_result) {
out_loc->tier = it->second.get();
return true;
}
}
@ -99,29 +168,23 @@ bool TieredBackend::AllocateInternalRaw(size_t size,
auto it = tiers_.find(tier_id);
if (it == tiers_.end() || !it->second) continue;
auto& tier = it->second;
if (tier->Allocate(size, out_loc->data)) {
out_loc->tier_id = tier_id;
auto alloc_result = tier->Allocate(size, out_loc->data);
if (alloc_result) {
out_loc->tier = tier.get();
return true;
}
}
return false;
}
void TieredBackend::FreeInternal(const TieredLocation& loc) {
auto it = tiers_.find(loc.tier_id);
if (it != tiers_.end()) {
it->second->Free(loc.data);
}
}
tl::expected<AllocationHandle, ErrorCode> TieredBackend::Allocate(
size_t size, std::optional<UUID> preferred_tier) {
TieredLocation loc;
if (AllocateInternalRaw(size, preferred_tier, &loc)) {
// Create the handle (Ref count = 1).
// If this handle dies without being committed, AllocationEntry
// destructor triggers FreeInternal.
return std::make_shared<AllocationEntry>(this, loc);
// destructor triggers Free.
return std::make_shared<AllocationEntry>(this, std::move(loc));
}
LOG(ERROR) << "Failed to allocate " << size << " bytes";
return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
@ -134,9 +197,8 @@ tl::expected<void, ErrorCode> TieredBackend::Write(const DataSource& source,
LOG(ERROR) << "TieredBackend not initialized";
return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
}
auto it = tiers_.find(handle->loc.tier_id);
if (it == tiers_.end()) {
LOG(ERROR) << "Tier not found: " << handle->loc.tier_id;
if (!handle->loc.tier) {
LOG(ERROR) << "Tier pointer is null";
return tl::make_unexpected(ErrorCode::TIER_NOT_FOUND);
}
@ -148,7 +210,8 @@ tl::expected<void, ErrorCode> TieredBackend::Commit(const std::string& key,
if (!handle) return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
if (metadata_sync_callback_) {
auto result = metadata_sync_callback_(key, handle->loc.tier_id, COMMIT);
auto result =
metadata_sync_callback_(key, handle->loc.tier->GetTierId(), COMMIT);
if (!result.has_value()) {
LOG(ERROR) << "Failed to Commit key " << key
@ -186,9 +249,10 @@ tl::expected<void, ErrorCode> TieredBackend::Commit(const std::string& key,
{
std::unique_lock<std::shared_mutex> entry_lock(entry->mutex);
// Insert or replace the handle for this tier
UUID current_tier_id = handle->loc.tier->GetTierId();
bool found = false;
for (auto& replica : entry->replicas) {
if (replica.first == handle->loc.tier_id) {
if (replica.first == current_tier_id) {
replica.second = handle;
found = true;
break;
@ -196,7 +260,7 @@ tl::expected<void, ErrorCode> TieredBackend::Commit(const std::string& key,
}
if (!found) {
entry->replicas.emplace_back(handle->loc.tier_id, handle);
entry->replicas.emplace_back(current_tier_id, handle);
std::sort(entry->replicas.begin(), entry->replicas.end(),
[this](const std::pair<UUID, AllocationHandle>& a,
const std::pair<UUID, AllocationHandle>& b) {
@ -281,11 +345,12 @@ tl::expected<void, ErrorCode> TieredBackend::Delete(
if (tier_it != entry->replicas.end()) {
if (metadata_sync_callback_) {
auto result = metadata_sync_callback_(
key, tier_it->second->loc.tier_id, DELETE);
key, tier_it->second->loc.tier->GetTierId(),
DELETE);
if (!result.has_value()) {
LOG(ERROR)
<< "Failed to Delete key " << key << " in Tier "
<< tier_it->second->loc.tier_id
<< tier_it->second->loc.tier->GetTierId()
<< " for Master, error_code=" << result.error();
return tl::make_unexpected(result.error());
}
@ -367,7 +432,7 @@ tl::expected<void, ErrorCode> TieredBackend::Delete(
}
// Handles go out of scope here.
// Ref count drops to 0 -> ~AllocationEntry() -> FreeInternal().
// Ref count drops to 0 -> ~AllocationEntry() -> Free().
// This happens concurrently without holding any locks.
return tl::expected<void, ErrorCode>{};
}
@ -375,11 +440,11 @@ tl::expected<void, ErrorCode> TieredBackend::Delete(
tl::expected<void, ErrorCode> TieredBackend::CopyData(const std::string& key,
const DataSource& source,
UUID dest_tier_id) {
if (source.size == 0) {
LOG(ERROR) << "Invalid size: " << source.size;
if (!source.buffer || source.buffer->size() == 0) {
LOG(ERROR) << "Invalid source buffer or size";
return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
}
auto dest_handle = Allocate(source.size, dest_tier_id);
auto dest_handle = Allocate(source.buffer->size(), dest_tier_id);
if (!dest_handle.has_value()) {
LOG(ERROR) << "Failed to allocate memory for key: " << key
<< " in Tier " << dest_tier_id;

1
mooncake-store/tests/CMakeLists.txt
View File

@ -36,6 +36,7 @@ add_store_test(non_ha_reconnect_test non_ha_reconnect_test.cpp)
add_store_test(storage_backend_test storage_backend_test.cpp)
add_store_test(mutex_test mutex_test.cpp)
add_store_test(file_storage_test file_storage_test.cpp)
add_store_test(tiered_backend_test tiered_backend_test.cpp)
add_subdirectory(e2e)
add_executable(high_availability_test high_availability_test.cpp)

1149
mooncake-store/tests/tiered_backend_test.cpp Normal file
View File

File diff suppressed because it is too large Load Diff