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

[Store]: add dram tier support 

Signed-off-by: Xingrui Yi <yixingrui@linux.alibaba.com>
This commit is contained in:
Xingrui Yi 2025-12-31 15:32:12 +08:00
parent 6fb4a84e6e
commit fee0642df6
7 changed files with 377 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -5,6 +5,7 @@
#include <memory>
#include <optional>
#include "allocator.h"
#include "transfer_engine.h"
namespace mooncake {
@ -28,14 +29,60 @@ 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
*/
class TempDRAMBuffer : public BufferBase {
public:
TempDRAMBuffer(char* ptr, size_t size) : ptr_(ptr), size_(size) {}
uint64_t data() const override { return reinterpret_cast<uint64_t>(ptr_); }
std::size_t size() const override { return size_; }
private:
char* ptr_;
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
};
@ -71,7 +118,7 @@ class CacheTier {
virtual bool 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 +130,4 @@ class CacheTier {
TieredBackend* backend_ = nullptr;
};
} // namespace mooncake
} // namespace mooncake

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

@ -0,0 +1,46 @@
#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;
bool Init(TieredBackend* backend, TransferEngine* engine) override;
bool Allocate(size_t size, DataSource& data) override;
bool 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_;
size_t current_usage_;
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

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

@ -59,7 +59,7 @@ 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;
@ -162,7 +162,7 @@ class TieredBackend {
const DataCopier& GetDataCopier() const;
// Internal API called by AllocationEntry destructor
void FreeInternal(const TieredLocation& loc);
void FreeInternal(TieredLocation&& loc);
private:
struct TierInfo {

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

@ -29,6 +29,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 "")

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

@ -2,12 +2,25 @@
#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) {
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();
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 +94,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 +105,11 @@ 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;
temp_dram.buffer = std::make_unique<TempDRAMBuffer>(
temp_dram_buffer.get(), 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);

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

@ -0,0 +1,164 @@
#include <glog/logging.h>
#include <numa.h>
#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),
current_usage_(0),
tags_(tags),
numa_node_(numa_node),
allocator_type_(allocator_type),
allocator_(nullptr),
engine_(nullptr) {}
DramCacheTier::~DramCacheTier() {
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;
}
}
}
bool 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 false;
}
node = numa_node_.value();
if (node < 0 || node > numa_max_node()) {
LOG(ERROR) << "Invalid NUMA node " << node;
return false;
}
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 false;
}
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 false;
}
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 = "*";
}
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 false;
} 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 false;
}
LOG(INFO) << "DramCacheTier " << tier_id_ << " initialized and registered "
<< capacity_ << " bytes at base address 0x" << std::hex
<< base_address;
return true;
}
size_t DramCacheTier::GetUsage() const { return current_usage_; }
bool DramCacheTier::Allocate(size_t size, DataSource& data_source) {
if (!allocator_) {
LOG(ERROR) << "Allocator not initialized for DramCacheTier "
<< tier_id_;
return false;
}
auto alloc_result = allocator_->allocate(size);
if (!alloc_result) {
LOG(ERROR) << "Failed to allocate " << size
<< " bytes from DramCacheTier " << tier_id_;
return false;
}
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;
current_usage_ += data_source.buffer->size();
return true;
}
bool DramCacheTier::Free(DataSource data_source) {
if (data_source.buffer) {
current_usage_ -= data_source.buffer->size();
} else {
LOG(WARNING) << "Attempting to free null buffer in DramCacheTier "
<< tier_id_;
}
// RAII will handle the deallocation when buffer_handle goes out of scope.
return true;
}
} // namespace mooncake

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

@ -5,6 +5,7 @@
#include "tiered_cache/tiered_backend.h"
#include "tiered_cache/cache_tier.h"
#include "tiered_cache/dram_tier.h"
namespace mooncake {
@ -12,7 +13,7 @@ AllocationEntry::~AllocationEntry() {
if (backend) {
// When ref count drops to 0, call back to backend to free physical
// resource.
backend->FreeInternal(loc);
backend->FreeInternal(std::move(loc));
}
}
@ -39,24 +40,92 @@ bool TieredBackend::Init(Json::Value root, TransferEngine* engine,
}
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', skipping";
continue;
}
if (!tier_config.isMember("capacity")) {
LOG(ERROR)
<< "Tier config missing required field 'capacity', skipping";
continue;
}
if (!tier_config.isMember("priority")) {
LOG(ERROR)
<< "Tier config missing required field 'priority', skipping";
continue;
}
// 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
<< ", skipping";
continue;
}
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);
if (!tier->Init(this, engine)) {
LOG(ERROR) << "Failed to initialize DRAM tier: id=" << id;
return false;
}
tiers_[id] = std::move(tier);
tier_info_[id] = {priority, tags};
LOG(INFO) << "Successfully initialized DRAM tier: id=" << id;
} else {
LOG(WARNING) << "Unsupported tier type '" << type << "', skipping";
continue;
}
}
LOG(INFO) << "TieredBackend initialized successfully with "
@ -107,10 +176,10 @@ bool TieredBackend::AllocateInternalRaw(size_t size,
return false;
}
void TieredBackend::FreeInternal(const TieredLocation& loc) {
void TieredBackend::FreeInternal(TieredLocation&& loc) {
auto it = tiers_.find(loc.tier_id);
if (it != tiers_.end()) {
it->second->Free(loc.data);
it->second->Free(std::move(loc.data));
}
}
@ -121,7 +190,7 @@ tl::expected<AllocationHandle, ErrorCode> TieredBackend::Allocate(
// Create the handle (Ref count = 1).
// If this handle dies without being committed, AllocationEntry
// destructor triggers FreeInternal.
return std::make_shared<AllocationEntry>(this, loc);
return std::make_shared<AllocationEntry>(this, std::move(loc));
}
LOG(ERROR) << "Failed to allocate " << size << " bytes";
return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
@ -375,11 +444,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;