From fee0642df64deb11805604ea7596d24a40a0abe4 Mon Sep 17 00:00:00 2001 From: Xingrui Yi Date: Wed, 31 Dec 2025 15:32:12 +0800 Subject: [PATCH] [Store]: add dram tier support Signed-off-by: Xingrui Yi --- .../include/tiered_cache/cache_tier.h | 57 +++++- .../include/tiered_cache/dram_tier.h | 46 +++++ .../include/tiered_cache/tiered_backend.h | 4 +- mooncake-store/src/CMakeLists.txt | 1 + .../src/tiered_cache/data_copier.cpp | 27 ++- mooncake-store/src/tiered_cache/dram_tier.cpp | 164 ++++++++++++++++++ .../src/tiered_cache/tiered_backend.cpp | 111 +++++++++--- 7 files changed, 377 insertions(+), 33 deletions(-) create mode 100644 mooncake-store/include/tiered_cache/dram_tier.h create mode 100644 mooncake-store/src/tiered_cache/dram_tier.cpp diff --git a/mooncake-store/include/tiered_cache/cache_tier.h b/mooncake-store/include/tiered_cache/cache_tier.h index 1c8021cd..9079fdb0 100644 --- a/mooncake-store/include/tiered_cache/cache_tier.h +++ b/mooncake-store/include/tiered_cache/cache_tier.h @@ -5,6 +5,7 @@ #include #include +#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 buffer) + : dram_buffer_(std::move(buffer)) {} + + uint64_t data() const override { + return dram_buffer_ ? reinterpret_cast(dram_buffer_->data()) + : 0; + } + + std::size_t size() const override { + return dram_buffer_ ? dram_buffer_->size() : 0; + } + + private: + std::unique_ptr 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(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 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 \ No newline at end of file diff --git a/mooncake-store/include/tiered_cache/dram_tier.h b/mooncake-store/include/tiered_cache/dram_tier.h new file mode 100644 index 00000000..5cea52c3 --- /dev/null +++ b/mooncake-store/include/tiered_cache/dram_tier.h @@ -0,0 +1,46 @@ +#pragma once + +#include +#include +#include +#include +#include + +#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& tags, + std::optional 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& 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 tags_; + std::optional numa_node_; + BufferAllocatorType allocator_type_; + std::shared_ptr allocator_; + TransferEngine* engine_; + std::unique_ptr memory_buffer_{ + nullptr, [](char* p) { delete[] p; }}; +}; + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/tiered_cache/tiered_backend.h b/mooncake-store/include/tiered_cache/tiered_backend.h index 694a3896..fada1d5a 100644 --- a/mooncake-store/include/tiered_cache/tiered_backend.h +++ b/mooncake-store/include/tiered_cache/tiered_backend.h @@ -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 { diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index cd593878..5ae716eb 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -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 "") diff --git a/mooncake-store/src/tiered_cache/data_copier.cpp b/mooncake-store/src/tiered_cache/data_copier.cpp index 79a12877..14bbb24a 100644 --- a/mooncake-store/src/tiered_cache/data_copier.cpp +++ b/mooncake-store/src/tiered_cache/data_copier.cpp @@ -2,12 +2,25 @@ #include #include -#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 CopyDramToDram(const DataSource& src, + const DataSource& dest) { + const void* src_ptr = reinterpret_cast(src.buffer->data()); + void* dest_ptr = reinterpret_cast(dest.buffer->data()); + size_t size = src.buffer->size(); + memcpy(dest_ptr, src_ptr, size); + return tl::expected{}; +} + 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 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 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 DataCopier::Copy(const DataSource& src, } // Step A: Source -> DRAM - DataSource temp_dram = { - reinterpret_cast(temp_dram_buffer.get()), 0, src.size, - MemoryType::DRAM}; + DataSource temp_dram; + temp_dram.buffer = std::make_unique( + 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); diff --git a/mooncake-store/src/tiered_cache/dram_tier.cpp b/mooncake-store/src/tiered_cache/dram_tier.cpp new file mode 100644 index 00000000..38769474 --- /dev/null +++ b/mooncake-store/src/tiered_cache/dram_tier.cpp @@ -0,0 +1,164 @@ +#include +#include + +#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& tags, + std::optional 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(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( + 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( + 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(mem_ptr); + } + } + + // Use the address of this registered block as the base_address for the + // allocator. + const uintptr_t base_address = reinterpret_cast(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( + segment_name, base_address, capacity_, segment_name); + break; + case BufferAllocatorType::CACHELIB: + allocator_ = std::make_shared( + 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(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 \ No newline at end of file diff --git a/mooncake-store/src/tiered_cache/tiered_backend.cpp b/mooncake-store/src/tiered_cache/tiered_backend.cpp index d209e28d..09a62b17 100644 --- a/mooncake-store/src/tiered_cache/tiered_backend.cpp +++ b/mooncake-store/src/tiered_cache/tiered_backend.cpp @@ -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 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 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(); + // 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 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 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( + 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 TieredBackend::Allocate( // Create the handle (Ref count = 1). // If this handle dies without being committed, AllocationEntry // destructor triggers FreeInternal. - return std::make_shared(this, loc); + return std::make_shared(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 TieredBackend::Delete( tl::expected 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;