From fee0642df64deb11805604ea7596d24a40a0abe4 Mon Sep 17 00:00:00 2001 From: Xingrui Yi Date: Wed, 31 Dec 2025 15:32:12 +0800 Subject: [PATCH 1/8] [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; From 6691a9856acf51c4952fd1f8d3dcc54bbee14bcf Mon Sep 17 00:00:00 2001 From: Xingrui Yi Date: Wed, 31 Dec 2025 17:07:46 +0800 Subject: [PATCH 2/8] [Store]: add unit test for tiered backend Signed-off-by: Xingrui Yi --- mooncake-store/tests/CMakeLists.txt | 1 + mooncake-store/tests/tiered_backend_test.cpp | 1162 ++++++++++++++++++ 2 files changed, 1163 insertions(+) create mode 100644 mooncake-store/tests/tiered_backend_test.cpp diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index c0bbf636..9ba66e9f 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -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) diff --git a/mooncake-store/tests/tiered_backend_test.cpp b/mooncake-store/tests/tiered_backend_test.cpp new file mode 100644 index 00000000..0416551e --- /dev/null +++ b/mooncake-store/tests/tiered_backend_test.cpp @@ -0,0 +1,1162 @@ +#include +#include +#include +#include + +#include "tiered_cache/tiered_backend.h" + +// Helper function to parse JSON string using thread-safe CharReaderBuilder +static bool parseJsonString(const std::string& json_str, Json::Value& value, + std::string* error_msg = nullptr) { + Json::CharReaderBuilder builder; + std::unique_ptr reader(builder.newCharReader()); + std::string errs; + + bool success = reader->parse( + json_str.data(), json_str.data() + json_str.size(), &value, &errs); + if (!success && error_msg) { + *error_msg = errs; + } + return success; +} + +namespace mooncake { + +// Test data size constants +static constexpr size_t SMALL_DATA_SIZE = 4 * 1024; // 4KB +static constexpr size_t MEDIUM_DATA_SIZE = 256 * 1024; // 256KB +static constexpr size_t LARGE_DATA_SIZE = 10 * 1024 * 1024; // 10MB + +// Test capacity constants +static constexpr size_t BASIC_CAPACITY = 1024 * 1024 * 1024; // 1GB +static constexpr size_t HIGH_PRIORITY_CAPACITY = 512 * 1024 * 1024; // 512MB +static constexpr size_t LOW_PRIORITY_CAPACITY = 1024 * 1024 * 1024; // 1GB +static constexpr size_t SMALL_CAPACITY = 1 * 1024 * 1024; // 1MB + +class TieredBackendTest : public ::testing::Test { + protected: + void SetUp() override { + // glog is already initialized by gtest_main + } + + void TearDown() override {} + + // Helper: Create test buffer with specified size + std::unique_ptr CreateTestBuffer(size_t size) { + auto buffer = std::make_unique(size); + // Fill with test pattern + for (size_t i = 0; i < size; ++i) { + buffer[i] = static_cast(i % 256); + } + return buffer; + } + + // Helper: Get tier ID by priority from tier views + std::optional GetTierIdByPriority(const TieredBackend& backend, + int priority) { + auto tier_views = backend.GetTierViews(); + for (const auto& view : tier_views) { + if (view.priority == priority) { + return view.id; + } + } + return std::nullopt; + } + + // Helper: Verify tier usage + void VerifyTierUsage(const TieredBackend& backend, UUID tier_id, + size_t expected_usage) { + auto tier_views = backend.GetTierViews(); + for (const auto& view : tier_views) { + if (view.id == tier_id) { + EXPECT_EQ(view.usage, expected_usage); + return; + } + } + FAIL() << "Tier not found with id: " << tier_id.first << "-" + << tier_id.second; + } + + // Helper: Allocate and write (combined operation) + tl::expected AllocateAndWrite( + TieredBackend& backend, size_t size, const char* data, + std::optional preferred_tier = std::nullopt) { + auto alloc_result = backend.Allocate(size, preferred_tier); + if (!alloc_result.has_value()) { + return alloc_result; + } + + AllocationHandle handle = alloc_result.value(); + + // Prepare DataSource for write + DataSource source; + auto buffer = std::make_unique(size); + std::memcpy(buffer.get(), data, size); + source.buffer = std::make_unique(buffer.get(), size); + source.type = MemoryType::DRAM; + + auto write_result = backend.Write(source, handle); + if (!write_result.has_value()) { + return tl::make_unexpected(write_result.error()); + } + + return handle; + } +}; + +// Test basic DRAM tier initialization +TEST_F(TieredBackendTest, BasicDRAMTierInit) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "tags": ["fast", "local"], + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + bool result = backend.Init(config, nullptr, nullptr); + + EXPECT_TRUE(result); + + // Verify tier was created + auto tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views.size(), 1); + + if (!tier_views.empty()) { + EXPECT_EQ(tier_views[0].type, MemoryType::DRAM); + EXPECT_EQ(tier_views[0].capacity, 1073741824); + EXPECT_EQ(tier_views[0].priority, 10); + EXPECT_EQ(tier_views[0].tags.size(), 2); + } +} + +// Test DRAM tier with NUMA node +TEST_F(TieredBackendTest, DRAMTierWithNUMA) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 536870912, + "priority": 20, + "numa_node": 0, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + bool result = backend.Init(config, nullptr, nullptr); + + EXPECT_TRUE(result); + + auto tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views.size(), 1); +} + +// Test multiple DRAM tiers +TEST_F(TieredBackendTest, MultipleDRAMTiers) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 536870912, + "priority": 100, + "allocator_type": "OFFSET" + }, + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 50, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + bool result = backend.Init(config, nullptr, nullptr); + + EXPECT_TRUE(result); + + auto tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views.size(), 2); +} + +// Test missing required fields +TEST_F(TieredBackendTest, MissingTypeField) { + std::string json_config_str = R"({ + "tiers": [ + { + "capacity": 1073741824, + "priority": 10 + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + bool result = backend.Init(config, nullptr, nullptr); + + // Should still succeed but with no tiers created + EXPECT_TRUE(result); + + auto tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views.size(), 0); +} + +// Test missing capacity field +TEST_F(TieredBackendTest, MissingCapacityField) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "priority": 10 + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + bool result = backend.Init(config, nullptr, nullptr); + + // Should still succeed but with no tiers created + EXPECT_TRUE(result); + + auto tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views.size(), 0); +} + +// Test invalid capacity (zero) +TEST_F(TieredBackendTest, InvalidCapacity) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 0, + "priority": 10 + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + bool result = backend.Init(config, nullptr, nullptr); + + // Should still succeed but with no tiers created + EXPECT_TRUE(result); + + auto tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views.size(), 0); +} + +// Test unsupported tier type +TEST_F(TieredBackendTest, UnsupportedTierType) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "NVME", + "capacity": 1073741824, + "priority": 10 + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + bool result = backend.Init(config, nullptr, nullptr); + + // Should succeed but with no tiers created (NVME not supported yet) + EXPECT_TRUE(result); + + auto tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views.size(), 0); +} + +// Test default allocator type +TEST_F(TieredBackendTest, DefaultAllocatorType) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10 + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + bool result = backend.Init(config, nullptr, nullptr); + + EXPECT_TRUE(result); + + auto tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views.size(), 1); +} + +// Test unknown allocator type +TEST_F(TieredBackendTest, UnknownAllocatorType) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "UNKNOWN_TYPE" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + bool result = backend.Init(config, nullptr, nullptr); + + // Should succeed and use default OFFSET allocator + EXPECT_TRUE(result); + + auto tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views.size(), 1); +} + +// namespace mooncake + +// ============================================================================ +// Allocate API Tests +// ============================================================================ + +// Test basic allocation +TEST_F(TieredBackendTest, AllocateBasic) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Allocate space + auto result = backend.Allocate(SMALL_DATA_SIZE); + ASSERT_TRUE(result.has_value()) << "Allocation should succeed"; + + AllocationHandle handle = result.value(); + EXPECT_TRUE(handle) << "Handle should be valid"; + + // Verify tier usage increased + auto tier_views = backend.GetTierViews(); + ASSERT_EQ(tier_views.size(), 1); + EXPECT_GT(tier_views[0].usage, 0) + << "Tier usage should increase after allocation"; +} + +// Test allocation failure due to insufficient space +TEST_F(TieredBackendTest, AllocateInsufficientSpace) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1048576, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Try to allocate more than capacity + auto result = backend.Allocate(LARGE_DATA_SIZE); + EXPECT_FALSE(result.has_value()) + << "Allocation should fail when space insufficient"; + EXPECT_EQ(result.error(), ErrorCode::NO_AVAILABLE_HANDLE); +} + +// Test allocation on specified tier +TEST_F(TieredBackendTest, AllocateOnSpecifiedTier) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 536870912, + "priority": 100, + "allocator_type": "OFFSET" + }, + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 50, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Get low priority tier ID + auto low_priority_tier_id = GetTierIdByPriority(backend, 50); + ASSERT_TRUE(low_priority_tier_id.has_value()); + + // Allocate on low priority tier + auto result = + backend.Allocate(SMALL_DATA_SIZE, low_priority_tier_id.value()); + ASSERT_TRUE(result.has_value()); + + AllocationHandle handle = result.value(); + EXPECT_EQ(handle->loc.tier_id, low_priority_tier_id.value()); +} + +// Test automatic resource release when handle goes out of scope +TEST_F(TieredBackendTest, AllocateAutoRelease) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + auto tier_views = backend.GetTierViews(); + ASSERT_EQ(tier_views.size(), 1); + size_t initial_usage = tier_views[0].usage; + + // Allocate in inner scope + { + auto result = backend.Allocate(MEDIUM_DATA_SIZE); + ASSERT_TRUE(result.has_value()); + + // Verify usage increased + tier_views = backend.GetTierViews(); + size_t usage_after_alloc = tier_views[0].usage; + EXPECT_GT(usage_after_alloc, initial_usage); + } + // Handle goes out of scope, should auto-release + + // Verify usage returned to initial value + tier_views = backend.GetTierViews(); + EXPECT_EQ(tier_views[0].usage, initial_usage) + << "Usage should return to initial value after auto-release"; +} + +// ============================================================================ +// Write API Tests +// ============================================================================ + +// Test basic write operation +TEST_F(TieredBackendTest, WriteBasic) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Allocate space + auto alloc_result = backend.Allocate(SMALL_DATA_SIZE); + ASSERT_TRUE(alloc_result.has_value()); + AllocationHandle handle = alloc_result.value(); + + // Prepare test data + auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); + DataSource source; + source.buffer = + std::make_unique(test_buffer.get(), SMALL_DATA_SIZE); + source.type = MemoryType::DRAM; + + // Write data + auto write_result = backend.Write(source, handle); + EXPECT_TRUE(write_result.has_value()) << "Write should succeed"; +} + +// Test write with invalid handle +TEST_F(TieredBackendTest, WriteInvalidHandle) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Create invalid (null) handle + AllocationHandle invalid_handle; + + // Prepare test data + auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); + DataSource source; + source.buffer = + std::make_unique(test_buffer.get(), SMALL_DATA_SIZE); + source.type = MemoryType::DRAM; + + // Try to write with invalid handle + auto write_result = backend.Write(source, invalid_handle); + EXPECT_FALSE(write_result.has_value()) + << "Write with invalid handle should fail"; + EXPECT_EQ(write_result.error(), ErrorCode::INVALID_PARAMS); +} + +// ============================================================================ +// Commit API Tests +// ============================================================================ + +// Test basic commit operation +TEST_F(TieredBackendTest, CommitBasic) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Allocate and write + auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle_result = + AllocateAndWrite(backend, SMALL_DATA_SIZE, test_buffer.get()); + ASSERT_TRUE(handle_result.has_value()); + AllocationHandle handle = handle_result.value(); + + // Commit + auto commit_result = backend.Commit("test_key", handle); + EXPECT_TRUE(commit_result.has_value()) << "Commit should succeed"; + + // Verify we can get it back + auto get_result = backend.Get("test_key"); + EXPECT_TRUE(get_result.has_value()) + << "Should be able to get committed data"; +} + +// Test commit replaces old data on same tier +TEST_F(TieredBackendTest, CommitReplace) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // First commit + auto test_buffer1 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle1_result = + AllocateAndWrite(backend, SMALL_DATA_SIZE, test_buffer1.get()); + ASSERT_TRUE(handle1_result.has_value()); + auto commit1_result = backend.Commit("key1", handle1_result.value()); + ASSERT_TRUE(commit1_result.has_value()); + + UUID first_tier_id = handle1_result.value()->loc.tier_id; + + // Second commit with same key + auto test_buffer2 = CreateTestBuffer(MEDIUM_DATA_SIZE); + auto handle2_result = + AllocateAndWrite(backend, MEDIUM_DATA_SIZE, test_buffer2.get()); + ASSERT_TRUE(handle2_result.has_value()); + auto commit2_result = backend.Commit("key1", handle2_result.value()); + ASSERT_TRUE(commit2_result.has_value()); + + // Get should return the new handle + auto get_result = backend.Get("key1"); + ASSERT_TRUE(get_result.has_value()); + EXPECT_EQ(get_result.value()->loc.tier_id, first_tier_id); +} + +// Test commit with invalid handle +TEST_F(TieredBackendTest, CommitInvalidHandle) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Create invalid handle + AllocationHandle invalid_handle; + + // Try to commit with invalid handle + auto commit_result = backend.Commit("test_key", invalid_handle); + EXPECT_FALSE(commit_result.has_value()) + << "Commit with invalid handle should fail"; + EXPECT_EQ(commit_result.error(), ErrorCode::INVALID_PARAMS); +} + +// ============================================================================ +// Get API Tests +// ============================================================================ + +// Test basic get operation +TEST_F(TieredBackendTest, GetBasic) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Allocate, write and commit + auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle_result = + AllocateAndWrite(backend, SMALL_DATA_SIZE, test_buffer.get()); + ASSERT_TRUE(handle_result.has_value()); + auto commit_result = backend.Commit("test_key", handle_result.value()); + ASSERT_TRUE(commit_result.has_value()); + + UUID tier_id = handle_result.value()->loc.tier_id; + + // Get the data + auto get_result = backend.Get("test_key"); + ASSERT_TRUE(get_result.has_value()) << "Get should succeed"; + EXPECT_EQ(get_result.value()->loc.tier_id, tier_id); +} + +// Test get with non-existent key +TEST_F(TieredBackendTest, GetNonExistentKey) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Try to get non-existent key + auto get_result = backend.Get("non_existent_key"); + EXPECT_FALSE(get_result.has_value()) + << "Get should fail for non-existent key"; + EXPECT_EQ(get_result.error(), ErrorCode::INVALID_KEY); +} + +// Test get from specified tier +TEST_F(TieredBackendTest, GetFromSpecifiedTier) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 536870912, + "priority": 100, + "allocator_type": "OFFSET" + }, + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 50, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + auto high_tier_id = GetTierIdByPriority(backend, 100); + auto low_tier_id = GetTierIdByPriority(backend, 50); + ASSERT_TRUE(high_tier_id.has_value()); + ASSERT_TRUE(low_tier_id.has_value()); + + // Commit to high priority tier + auto test_buffer1 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle1 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer1.get(), high_tier_id.value()); + ASSERT_TRUE(handle1.has_value()); + backend.Commit("multi_tier_key", handle1.value()); + + // Commit to low priority tier with same key + auto test_buffer2 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle2 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer2.get(), low_tier_id.value()); + ASSERT_TRUE(handle2.has_value()); + backend.Commit("multi_tier_key", handle2.value()); + + // Get from specific tier + auto get_result = backend.Get("multi_tier_key", low_tier_id.value()); + ASSERT_TRUE(get_result.has_value()); + EXPECT_EQ(get_result.value()->loc.tier_id, low_tier_id.value()); +} + +// Test get returns highest priority tier when tier not specified +TEST_F(TieredBackendTest, GetHighestPriority) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 536870912, + "priority": 100, + "allocator_type": "OFFSET" + }, + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 50, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + auto high_tier_id = GetTierIdByPriority(backend, 100); + auto low_tier_id = GetTierIdByPriority(backend, 50); + ASSERT_TRUE(high_tier_id.has_value()); + ASSERT_TRUE(low_tier_id.has_value()); + + // Commit to both tiers + auto test_buffer1 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle1 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer1.get(), high_tier_id.value()); + ASSERT_TRUE(handle1.has_value()); + backend.Commit("priority_key", handle1.value()); + + auto test_buffer2 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle2 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer2.get(), low_tier_id.value()); + ASSERT_TRUE(handle2.has_value()); + backend.Commit("priority_key", handle2.value()); + + // Get without specifying tier should return highest priority + auto get_result = backend.Get("priority_key"); + ASSERT_TRUE(get_result.has_value()); + EXPECT_EQ(get_result.value()->loc.tier_id, high_tier_id.value()) + << "Should return highest priority tier"; +} + +// Test get from tier that doesn't have the key +TEST_F(TieredBackendTest, GetTierNotFound) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 536870912, + "priority": 100, + "allocator_type": "OFFSET" + }, + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 50, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + auto high_tier_id = GetTierIdByPriority(backend, 100); + auto low_tier_id = GetTierIdByPriority(backend, 50); + ASSERT_TRUE(high_tier_id.has_value()); + ASSERT_TRUE(low_tier_id.has_value()); + + // Commit only to high priority tier + auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle = AllocateAndWrite(backend, SMALL_DATA_SIZE, test_buffer.get(), + high_tier_id.value()); + ASSERT_TRUE(handle.has_value()); + backend.Commit("tier_specific_key", handle.value()); + + // Try to get from low priority tier where it doesn't exist + auto get_result = backend.Get("tier_specific_key", low_tier_id.value()); + EXPECT_FALSE(get_result.has_value()); + EXPECT_EQ(get_result.error(), ErrorCode::TIER_NOT_FOUND); +} + +// ============================================================================ +// Delete API Tests +// ============================================================================ + +// Test delete single replica +TEST_F(TieredBackendTest, DeleteSingleReplica) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 536870912, + "priority": 100, + "allocator_type": "OFFSET" + }, + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 50, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + auto high_tier_id = GetTierIdByPriority(backend, 100); + auto low_tier_id = GetTierIdByPriority(backend, 50); + ASSERT_TRUE(high_tier_id.has_value()); + ASSERT_TRUE(low_tier_id.has_value()); + + // Commit to both tiers + auto test_buffer1 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle1 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer1.get(), high_tier_id.value()); + ASSERT_TRUE(handle1.has_value()); + backend.Commit("delete_test_key", handle1.value()); + + auto test_buffer2 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle2 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer2.get(), low_tier_id.value()); + ASSERT_TRUE(handle2.has_value()); + backend.Commit("delete_test_key", handle2.value()); + + // Delete from high priority tier only + auto delete_result = + backend.Delete("delete_test_key", high_tier_id.value()); + EXPECT_TRUE(delete_result.has_value()) << "Delete should succeed"; + + // Should not be able to get from high tier + auto get_high = backend.Get("delete_test_key", high_tier_id.value()); + EXPECT_FALSE(get_high.has_value()); + + // Should still be able to get from low tier + auto get_low = backend.Get("delete_test_key", low_tier_id.value()); + EXPECT_TRUE(get_low.has_value()); +} + +// Test delete all replicas +TEST_F(TieredBackendTest, DeleteAllReplicas) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 536870912, + "priority": 100, + "allocator_type": "OFFSET" + }, + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 50, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + auto high_tier_id = GetTierIdByPriority(backend, 100); + auto low_tier_id = GetTierIdByPriority(backend, 50); + ASSERT_TRUE(high_tier_id.has_value()); + ASSERT_TRUE(low_tier_id.has_value()); + + // Commit to both tiers + auto test_buffer1 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle1 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer1.get(), high_tier_id.value()); + ASSERT_TRUE(handle1.has_value()); + backend.Commit("delete_all_key", handle1.value()); + + auto test_buffer2 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle2 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer2.get(), low_tier_id.value()); + ASSERT_TRUE(handle2.has_value()); + backend.Commit("delete_all_key", handle2.value()); + + // Delete all replicas (no tier_id specified) + auto delete_result = backend.Delete("delete_all_key"); + EXPECT_TRUE(delete_result.has_value()) << "Delete all should succeed"; + + // Should not be able to get from any tier + auto get_result = backend.Get("delete_all_key"); + EXPECT_FALSE(get_result.has_value()); + EXPECT_EQ(get_result.error(), ErrorCode::INVALID_KEY); +} + +// Test delete non-existent key +TEST_F(TieredBackendTest, DeleteNonExistentKey) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + // Try to delete non-existent key + auto delete_result = backend.Delete("non_existent_key"); + EXPECT_FALSE(delete_result.has_value()) + << "Delete should fail for non-existent key"; + EXPECT_EQ(delete_result.error(), ErrorCode::INVALID_KEY); +} + +// ============================================================================ +// Integration Tests +// ============================================================================ + +// Test complete data lifecycle +TEST_F(TieredBackendTest, CompleteDataLifecycle) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + const std::string key = "lifecycle_key"; + + // 1. Allocate + auto alloc_result = backend.Allocate(SMALL_DATA_SIZE); + ASSERT_TRUE(alloc_result.has_value()); + AllocationHandle handle = alloc_result.value(); + + // 2. Write + auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); + DataSource source; + source.buffer = + std::make_unique(test_buffer.get(), SMALL_DATA_SIZE); + source.type = MemoryType::DRAM; + auto write_result = backend.Write(source, handle); + ASSERT_TRUE(write_result.has_value()); + + // 3. Commit + auto commit_result = backend.Commit(key, handle); + ASSERT_TRUE(commit_result.has_value()); + + // 4. Get and verify + auto get_result = backend.Get(key); + ASSERT_TRUE(get_result.has_value()); + EXPECT_EQ(get_result.value()->loc.tier_id, handle->loc.tier_id); + + // 5. Delete + auto delete_result = backend.Delete(key); + ASSERT_TRUE(delete_result.has_value()); + + // 6. Verify deleted + auto get_after_delete = backend.Get(key); + EXPECT_FALSE(get_after_delete.has_value()); + EXPECT_EQ(get_after_delete.error(), ErrorCode::INVALID_KEY); +} + +// Test multi-tier data management +TEST_F(TieredBackendTest, MultiTierDataManagement) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 536870912, + "priority": 100, + "allocator_type": "OFFSET" + }, + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 50, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + auto high_tier_id = GetTierIdByPriority(backend, 100); + auto low_tier_id = GetTierIdByPriority(backend, 50); + ASSERT_TRUE(high_tier_id.has_value()); + ASSERT_TRUE(low_tier_id.has_value()); + + const std::string key = "multi_tier_key"; + + // 1. Commit to high priority tier + auto test_buffer1 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle1 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer1.get(), high_tier_id.value()); + ASSERT_TRUE(handle1.has_value()); + backend.Commit(key, handle1.value()); + + // 2. Commit to low priority tier + auto test_buffer2 = CreateTestBuffer(SMALL_DATA_SIZE); + auto handle2 = AllocateAndWrite(backend, SMALL_DATA_SIZE, + test_buffer2.get(), low_tier_id.value()); + ASSERT_TRUE(handle2.has_value()); + backend.Commit(key, handle2.value()); + + // 3. Get without specifying tier should return high priority + auto get_high = backend.Get(key); + ASSERT_TRUE(get_high.has_value()); + EXPECT_EQ(get_high.value()->loc.tier_id, high_tier_id.value()); + + // 4. Get from low tier explicitly + auto get_low = backend.Get(key, low_tier_id.value()); + ASSERT_TRUE(get_low.has_value()); + EXPECT_EQ(get_low.value()->loc.tier_id, low_tier_id.value()); + + // 5. Delete high priority replica + auto delete_high = backend.Delete(key, high_tier_id.value()); + ASSERT_TRUE(delete_high.has_value()); + + // 6. Get without tier should now return low priority + auto get_after_delete = backend.Get(key); + ASSERT_TRUE(get_after_delete.has_value()); + EXPECT_EQ(get_after_delete.value()->loc.tier_id, low_tier_id.value()); +} + +// Test concurrent allocations +TEST_F(TieredBackendTest, ConcurrentAllocations) { + std::string json_config_str = R"({ + "tiers": [ + { + "type": "DRAM", + "capacity": 1073741824, + "priority": 10, + "allocator_type": "OFFSET" + } + ] + })"; + Json::Value config; + ASSERT_TRUE(parseJsonString(json_config_str, config)); + + TieredBackend backend; + ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + + const int num_allocations = 10; + std::vector handles; + + // Perform multiple allocations + for (int i = 0; i < num_allocations; ++i) { + auto result = backend.Allocate(SMALL_DATA_SIZE); + ASSERT_TRUE(result.has_value()) << "Allocation " << i << " failed"; + handles.push_back(result.value()); + } + + // Verify all handles are valid + EXPECT_EQ(handles.size(), num_allocations); + for (const auto& handle : handles) { + EXPECT_TRUE(handle); + } + + // Verify total usage + auto tier_views = backend.GetTierViews(); + ASSERT_EQ(tier_views.size(), 1); + EXPECT_GT(tier_views[0].usage, 0); +} + +} // namespace mooncake From 925fb18b1121411cadfeef7fdb3e3b3f5a58c76c Mon Sep 17 00:00:00 2001 From: Xingrui Yi Date: Wed, 31 Dec 2025 17:17:15 +0800 Subject: [PATCH 3/8] [Store] Optimize TempDRAMBuffer with RAII memory management Signed-off-by: Xingrui Yi --- mooncake-store/include/tiered_cache/cache_tier.h | 14 ++++++++++---- mooncake-store/src/tiered_cache/data_copier.cpp | 4 +++- mooncake-store/tests/tiered_backend_test.cpp | 15 ++++++++------- 3 files changed, 21 insertions(+), 12 deletions(-) diff --git a/mooncake-store/include/tiered_cache/cache_tier.h b/mooncake-store/include/tiered_cache/cache_tier.h index 9079fdb0..51fe0fbd 100644 --- a/mooncake-store/include/tiered_cache/cache_tier.h +++ b/mooncake-store/include/tiered_cache/cache_tier.h @@ -64,16 +64,22 @@ class DRAMBuffer : public BufferBase { /** * @class TempDRAMBuffer - * @brief Wrapper for temporary DRAM buffers + * @brief Wrapper for temporary DRAM buffers with RAII memory management */ class TempDRAMBuffer : public BufferBase { public: - TempDRAMBuffer(char* ptr, size_t size) : ptr_(ptr), size_(size) {} - uint64_t data() const override { return reinterpret_cast(ptr_); } + // Constructor that takes ownership of the buffer + explicit TempDRAMBuffer(std::unique_ptr buffer, size_t size) + : buffer_(std::move(buffer)), size_(size) {} + + uint64_t data() const override { + return reinterpret_cast(buffer_.get()); + } std::size_t size() const override { return size_; } private: - char* ptr_; + std::unique_ptr + buffer_; // Owns the memory, auto-releases on destruction size_t size_; }; diff --git a/mooncake-store/src/tiered_cache/data_copier.cpp b/mooncake-store/src/tiered_cache/data_copier.cpp index 14bbb24a..e5e0d59b 100644 --- a/mooncake-store/src/tiered_cache/data_copier.cpp +++ b/mooncake-store/src/tiered_cache/data_copier.cpp @@ -106,8 +106,10 @@ tl::expected DataCopier::Copy(const DataSource& src, // Step A: Source -> 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( - temp_dram_buffer.get(), buffer_size); + std::move(temp_dram_buffer), buffer_size); temp_dram.type = MemoryType::DRAM; if (!to_dram_copier(src, temp_dram)) { diff --git a/mooncake-store/tests/tiered_backend_test.cpp b/mooncake-store/tests/tiered_backend_test.cpp index 0416551e..08ef5197 100644 --- a/mooncake-store/tests/tiered_backend_test.cpp +++ b/mooncake-store/tests/tiered_backend_test.cpp @@ -92,7 +92,8 @@ class TieredBackendTest : public ::testing::Test { DataSource source; auto buffer = std::make_unique(size); std::memcpy(buffer.get(), data, size); - source.buffer = std::make_unique(buffer.get(), size); + source.buffer = + std::make_unique(std::move(buffer), size); source.type = MemoryType::DRAM; auto write_result = backend.Write(source, handle); @@ -504,8 +505,8 @@ TEST_F(TieredBackendTest, WriteBasic) { // Prepare test data auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); DataSource source; - source.buffer = - std::make_unique(test_buffer.get(), SMALL_DATA_SIZE); + source.buffer = std::make_unique(std::move(test_buffer), + SMALL_DATA_SIZE); source.type = MemoryType::DRAM; // Write data @@ -537,8 +538,8 @@ TEST_F(TieredBackendTest, WriteInvalidHandle) { // Prepare test data auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); DataSource source; - source.buffer = - std::make_unique(test_buffer.get(), SMALL_DATA_SIZE); + source.buffer = std::make_unique(std::move(test_buffer), + SMALL_DATA_SIZE); source.type = MemoryType::DRAM; // Try to write with invalid handle @@ -1029,8 +1030,8 @@ TEST_F(TieredBackendTest, CompleteDataLifecycle) { // 2. Write auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); DataSource source; - source.buffer = - std::make_unique(test_buffer.get(), SMALL_DATA_SIZE); + source.buffer = std::make_unique(std::move(test_buffer), + SMALL_DATA_SIZE); source.type = MemoryType::DRAM; auto write_result = backend.Write(source, handle); ASSERT_TRUE(write_result.has_value()); From b40bb82830480653db363e3d8b95bf317ff55013 Mon Sep 17 00:00:00 2001 From: Xingrui Yi Date: Tue, 6 Jan 2026 15:24:12 +0800 Subject: [PATCH 4/8] [Store] optimize cache tier and backend api Signed-off-by: Xingrui Yi --- .../include/tiered_cache/cache_tier.h | 16 ++++-- .../include/tiered_cache/dram_tier.h | 9 ++-- .../include/tiered_cache/tiered_backend.h | 3 +- .../src/tiered_cache/data_copier.cpp | 28 ++++++++++ mooncake-store/src/tiered_cache/dram_tier.cpp | 53 +++++++++++-------- .../src/tiered_cache/tiered_backend.cpp | 41 +++++++------- mooncake-store/tests/tiered_backend_test.cpp | 24 ++------- 7 files changed, 105 insertions(+), 69 deletions(-) diff --git a/mooncake-store/include/tiered_cache/cache_tier.h b/mooncake-store/include/tiered_cache/cache_tier.h index 51fe0fbd..0c71d0bb 100644 --- a/mooncake-store/include/tiered_cache/cache_tier.h +++ b/mooncake-store/include/tiered_cache/cache_tier.h @@ -4,9 +4,11 @@ #include #include #include +#include #include "allocator.h" #include "transfer_engine.h" +#include "types.h" namespace mooncake { @@ -104,24 +106,28 @@ class CacheTier { /** * @brief Initializes the cache tier. + * @return tl::expected indicating success or error code. */ - virtual bool Init(TieredBackend* backend, TransferEngine* engine) = 0; + virtual tl::expected 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 indicating success or error code. */ - virtual bool Allocate(size_t size, DataSource& data) = 0; + virtual tl::expected 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 indicating success or error code. */ - virtual bool Free(DataSource data) = 0; + virtual tl::expected Free(DataSource data) = 0; // --- Accessors & Metadata --- virtual UUID GetTierId() const = 0; diff --git a/mooncake-store/include/tiered_cache/dram_tier.h b/mooncake-store/include/tiered_cache/dram_tier.h index 5cea52c3..ab5b1cf5 100644 --- a/mooncake-store/include/tiered_cache/dram_tier.h +++ b/mooncake-store/include/tiered_cache/dram_tier.h @@ -20,9 +20,11 @@ class DramCacheTier : public CacheTier { 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; + tl::expected Init(TieredBackend* backend, + TransferEngine* engine) override; + tl::expected Allocate(size_t size, + DataSource& data) override; + tl::expected Free(DataSource data) override; UUID GetTierId() const override { return tier_id_; } size_t GetCapacity() const override { return capacity_; } @@ -33,7 +35,6 @@ class DramCacheTier : public CacheTier { private: UUID tier_id_; size_t capacity_; - size_t current_usage_; std::vector tags_; std::optional numa_node_; BufferAllocatorType allocator_type_; diff --git a/mooncake-store/include/tiered_cache/tiered_backend.h b/mooncake-store/include/tiered_cache/tiered_backend.h index fada1d5a..2f4a287d 100644 --- a/mooncake-store/include/tiered_cache/tiered_backend.h +++ b/mooncake-store/include/tiered_cache/tiered_backend.h @@ -59,7 +59,8 @@ struct AllocationEntry { TieredBackend* backend; TieredLocation loc; - AllocationEntry(TieredBackend* b, TieredLocation&& l) : backend(b), loc(std::move(l)) {} + AllocationEntry(TieredBackend* b, TieredLocation&& l) + : backend(b), loc(std::move(l)) {} AllocationEntry(const AllocationEntry&) = delete; AllocationEntry& operator=(const AllocationEntry&) = delete; diff --git a/mooncake-store/src/tiered_cache/data_copier.cpp b/mooncake-store/src/tiered_cache/data_copier.cpp index e5e0d59b..8e9847b0 100644 --- a/mooncake-store/src/tiered_cache/data_copier.cpp +++ b/mooncake-store/src/tiered_cache/data_copier.cpp @@ -11,9 +11,37 @@ namespace mooncake { // DRAM <-> DRAM tl::expected 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(src.buffer->data()); void* dest_ptr = reinterpret_cast(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{}; + } + + // 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{}; } diff --git a/mooncake-store/src/tiered_cache/dram_tier.cpp b/mooncake-store/src/tiered_cache/dram_tier.cpp index 38769474..f6010352 100644 --- a/mooncake-store/src/tiered_cache/dram_tier.cpp +++ b/mooncake-store/src/tiered_cache/dram_tier.cpp @@ -14,7 +14,6 @@ DramCacheTier::DramCacheTier(UUID tier_id, size_t capacity, BufferAllocatorType allocator_type) : tier_id_(tier_id), capacity_(capacity), - current_usage_(0), tags_(tags), numa_node_(numa_node), allocator_type_(allocator_type), @@ -22,6 +21,17 @@ DramCacheTier::DramCacheTier(UUID tier_id, size_t capacity, engine_(nullptr) {} DramCacheTier::~DramCacheTier() { + // Check if there are still allocated buffers 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. This may indicate a memory leak."; + } + } + allocator_.reset(); if (engine_ != nullptr && memory_buffer_ != nullptr) { @@ -34,7 +44,8 @@ DramCacheTier::~DramCacheTier() { } } -bool DramCacheTier::Init(TieredBackend* backend, TransferEngine* engine) { +tl::expected DramCacheTier::Init(TieredBackend* backend, + TransferEngine* engine) { int node = -1; std::string location; @@ -45,19 +56,19 @@ bool DramCacheTier::Init(TieredBackend* backend, TransferEngine* engine) { if (numa_node_.has_value()) { if (numa_available() < 0) { LOG(ERROR) << "NUMA not available on this system."; - return false; + return tl::unexpected(ErrorCode::INTERNAL_ERROR); } node = numa_node_.value(); if (node < 0 || node > numa_max_node()) { LOG(ERROR) << "Invalid NUMA node " << node; - return false; + return tl::unexpected(ErrorCode::INVALID_PARAMS); } 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; + return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } memory_buffer_ = std::unique_ptr( mem_ptr, [](char* p) { numa_free(p, 0); }); @@ -71,7 +82,7 @@ bool DramCacheTier::Init(TieredBackend* backend, TransferEngine* engine) { LOG(ERROR) << "Failed to allocate " << capacity_ << " bytes for DramCacheTier " << tier_id_ << ": " << e.what(); - return false; + return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } LOG(INFO) << "Allocated " << capacity_ << " bytes for DramCacheTier " << tier_id_; @@ -83,14 +94,14 @@ bool DramCacheTier::Init(TieredBackend* backend, TransferEngine* engine) { if (numa_node_.has_value()) { location = "cpu:" + std::to_string(node); } else { - location = "*"; + 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 false; + return tl::unexpected(ErrorCode::INTERNAL_ERROR); } else { LOG(INFO) << "registered memory with TransferEngine for DramCacheTier " @@ -118,47 +129,47 @@ bool DramCacheTier::Init(TieredBackend* backend, TransferEngine* engine) { if (engine_) { engine_->unregisterLocalMemory(mem_ptr); } - return false; + 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 true; + return tl::expected{}; } -size_t DramCacheTier::GetUsage() const { return current_usage_; } +size_t DramCacheTier::GetUsage() const { + return allocator_ ? allocator_->size() : 0; +} -bool DramCacheTier::Allocate(size_t size, DataSource& data_source) { +tl::expected DramCacheTier::Allocate(size_t size, + DataSource& data_source) { if (!allocator_) { LOG(ERROR) << "Allocator not initialized for DramCacheTier " << tier_id_; - return false; + 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 false; + return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } 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; + return tl::expected{}; } -bool DramCacheTier::Free(DataSource data_source) { - if (data_source.buffer) { - current_usage_ -= data_source.buffer->size(); - } else { +tl::expected 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 true; + return tl::expected{}; } } // 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 09a62b17..06f5bc56 100644 --- a/mooncake-store/src/tiered_cache/tiered_backend.cpp +++ b/mooncake-store/src/tiered_cache/tiered_backend.cpp @@ -42,18 +42,16 @@ bool TieredBackend::Init(Json::Value root, TransferEngine* engine, for (const auto& tier_config : root["tiers"]) { // Parse required fields if (!tier_config.isMember("type")) { - LOG(ERROR) << "Tier config missing required field 'type', skipping"; - continue; + LOG(ERROR) << "Tier config missing required field 'type'"; + return false; } if (!tier_config.isMember("capacity")) { - LOG(ERROR) - << "Tier config missing required field 'capacity', skipping"; - continue; + LOG(ERROR) << "Tier config missing required field 'capacity'"; + return false; } if (!tier_config.isMember("priority")) { - LOG(ERROR) - << "Tier config missing required field 'priority', skipping"; - continue; + LOG(ERROR) << "Tier config missing required field 'priority'"; + return false; } std::string type = tier_config["type"].asString(); @@ -62,9 +60,8 @@ bool TieredBackend::Init(Json::Value root, TransferEngine* engine, // Validate capacity if (capacity == 0) { - LOG(ERROR) << "Invalid capacity (0) for tier type " << type - << ", skipping"; - continue; + LOG(ERROR) << "Invalid capacity (0) for tier type " << type; + return false; } // Parse tags @@ -114,8 +111,10 @@ bool TieredBackend::Init(Json::Value root, TransferEngine* engine, 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; + auto init_result = tier->Init(this, engine); + if (!init_result) { + LOG(ERROR) << "Failed to initialize DRAM tier: id=" << id + << ", error=" << init_result.error(); return false; } @@ -123,8 +122,8 @@ bool TieredBackend::Init(Json::Value root, TransferEngine* engine, tier_info_[id] = {priority, tags}; LOG(INFO) << "Successfully initialized DRAM tier: id=" << id; } else { - LOG(WARNING) << "Unsupported tier type '" << type << "', skipping"; - continue; + LOG(ERROR) << "Unsupported tier type '" << type << "'"; + return false; } } @@ -153,7 +152,8 @@ 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)) { + auto alloc_result = it->second->Allocate(size, out_loc->data); + if (alloc_result) { out_loc->tier_id = *preferred_tier; return true; } @@ -168,7 +168,8 @@ 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)) { + auto alloc_result = tier->Allocate(size, out_loc->data); + if (alloc_result) { out_loc->tier_id = tier_id; return true; } @@ -179,7 +180,11 @@ bool TieredBackend::AllocateInternalRaw(size_t size, void TieredBackend::FreeInternal(TieredLocation&& loc) { auto it = tiers_.find(loc.tier_id); if (it != tiers_.end()) { - it->second->Free(std::move(loc.data)); + auto free_result = it->second->Free(std::move(loc.data)); + if (!free_result) { + LOG(WARNING) << "Failed to free data from tier " << loc.tier_id + << ", error=" << free_result.error(); + } } } diff --git a/mooncake-store/tests/tiered_backend_test.cpp b/mooncake-store/tests/tiered_backend_test.cpp index 08ef5197..e567e42b 100644 --- a/mooncake-store/tests/tiered_backend_test.cpp +++ b/mooncake-store/tests/tiered_backend_test.cpp @@ -209,11 +209,7 @@ TEST_F(TieredBackendTest, MissingTypeField) { TieredBackend backend; bool result = backend.Init(config, nullptr, nullptr); - // Should still succeed but with no tiers created - EXPECT_TRUE(result); - - auto tier_views = backend.GetTierViews(); - EXPECT_EQ(tier_views.size(), 0); + EXPECT_FALSE(result); } // Test missing capacity field @@ -232,11 +228,7 @@ TEST_F(TieredBackendTest, MissingCapacityField) { TieredBackend backend; bool result = backend.Init(config, nullptr, nullptr); - // Should still succeed but with no tiers created - EXPECT_TRUE(result); - - auto tier_views = backend.GetTierViews(); - EXPECT_EQ(tier_views.size(), 0); + EXPECT_FALSE(result); } // Test invalid capacity (zero) @@ -256,11 +248,7 @@ TEST_F(TieredBackendTest, InvalidCapacity) { TieredBackend backend; bool result = backend.Init(config, nullptr, nullptr); - // Should still succeed but with no tiers created - EXPECT_TRUE(result); - - auto tier_views = backend.GetTierViews(); - EXPECT_EQ(tier_views.size(), 0); + EXPECT_FALSE(result); } // Test unsupported tier type @@ -280,11 +268,7 @@ TEST_F(TieredBackendTest, UnsupportedTierType) { TieredBackend backend; bool result = backend.Init(config, nullptr, nullptr); - // Should succeed but with no tiers created (NVME not supported yet) - EXPECT_TRUE(result); - - auto tier_views = backend.GetTierViews(); - EXPECT_EQ(tier_views.size(), 0); + EXPECT_FALSE(result); } // Test default allocator type From 81d673979c107cb6b9e7611302560717b5ca3bc2 Mon Sep 17 00:00:00 2001 From: Xingrui Yi Date: Tue, 6 Jan 2026 17:11:59 +0800 Subject: [PATCH 5/8] [Store] change uuid to CacheTier* in TieredLocation Signed-off-by: Xingrui Yi --- .../include/tiered_cache/tiered_backend.h | 2 +- .../src/tiered_cache/tiered_backend.cpp | 30 ++++++++++--------- mooncake-store/tests/tiered_backend_test.cpp | 24 ++++++++------- 3 files changed, 30 insertions(+), 26 deletions(-) diff --git a/mooncake-store/include/tiered_cache/tiered_backend.h b/mooncake-store/include/tiered_cache/tiered_backend.h index 2f4a287d..e3ac50e4 100644 --- a/mooncake-store/include/tiered_cache/tiered_backend.h +++ b/mooncake-store/include/tiered_cache/tiered_backend.h @@ -23,7 +23,7 @@ class TieredBackend; // Forward declaration * storage. */ struct TieredLocation { - UUID tier_id; + CacheTier* tier; struct DataSource data; }; diff --git a/mooncake-store/src/tiered_cache/tiered_backend.cpp b/mooncake-store/src/tiered_cache/tiered_backend.cpp index 06f5bc56..5d03f56d 100644 --- a/mooncake-store/src/tiered_cache/tiered_backend.cpp +++ b/mooncake-store/src/tiered_cache/tiered_backend.cpp @@ -154,7 +154,7 @@ bool TieredBackend::AllocateInternalRaw(size_t size, if (it != tiers_.end()) { auto alloc_result = it->second->Allocate(size, out_loc->data); if (alloc_result) { - out_loc->tier_id = *preferred_tier; + out_loc->tier = it->second.get(); return true; } } @@ -170,7 +170,7 @@ bool TieredBackend::AllocateInternalRaw(size_t size, auto& tier = it->second; auto alloc_result = tier->Allocate(size, out_loc->data); if (alloc_result) { - out_loc->tier_id = tier_id; + out_loc->tier = tier.get(); return true; } } @@ -178,11 +178,11 @@ bool TieredBackend::AllocateInternalRaw(size_t size, } void TieredBackend::FreeInternal(TieredLocation&& loc) { - auto it = tiers_.find(loc.tier_id); - if (it != tiers_.end()) { - auto free_result = it->second->Free(std::move(loc.data)); + if (loc.tier) { + auto free_result = loc.tier->Free(std::move(loc.data)); if (!free_result) { - LOG(WARNING) << "Failed to free data from tier " << loc.tier_id + LOG(WARNING) << "Failed to free data from tier " + << loc.tier->GetTierId() << ", error=" << free_result.error(); } } @@ -208,9 +208,8 @@ tl::expected 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); } @@ -222,7 +221,8 @@ tl::expected 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 @@ -260,9 +260,10 @@ tl::expected TieredBackend::Commit(const std::string& key, { std::unique_lock 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; @@ -270,7 +271,7 @@ tl::expected 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& a, const std::pair& b) { @@ -355,11 +356,12 @@ tl::expected 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()); } diff --git a/mooncake-store/tests/tiered_backend_test.cpp b/mooncake-store/tests/tiered_backend_test.cpp index e567e42b..76603437 100644 --- a/mooncake-store/tests/tiered_backend_test.cpp +++ b/mooncake-store/tests/tiered_backend_test.cpp @@ -416,7 +416,7 @@ TEST_F(TieredBackendTest, AllocateOnSpecifiedTier) { ASSERT_TRUE(result.has_value()); AllocationHandle handle = result.value(); - EXPECT_EQ(handle->loc.tier_id, low_priority_tier_id.value()); + EXPECT_EQ(handle->loc.tier->GetTierId(), low_priority_tier_id.value()); } // Test automatic resource release when handle goes out of scope @@ -598,7 +598,7 @@ TEST_F(TieredBackendTest, CommitReplace) { auto commit1_result = backend.Commit("key1", handle1_result.value()); ASSERT_TRUE(commit1_result.has_value()); - UUID first_tier_id = handle1_result.value()->loc.tier_id; + UUID first_tier_id = handle1_result.value()->loc.tier->GetTierId(); // Second commit with same key auto test_buffer2 = CreateTestBuffer(MEDIUM_DATA_SIZE); @@ -611,7 +611,7 @@ TEST_F(TieredBackendTest, CommitReplace) { // Get should return the new handle auto get_result = backend.Get("key1"); ASSERT_TRUE(get_result.has_value()); - EXPECT_EQ(get_result.value()->loc.tier_id, first_tier_id); + EXPECT_EQ(get_result.value()->loc.tier->GetTierId(), first_tier_id); } // Test commit with invalid handle @@ -672,12 +672,12 @@ TEST_F(TieredBackendTest, GetBasic) { auto commit_result = backend.Commit("test_key", handle_result.value()); ASSERT_TRUE(commit_result.has_value()); - UUID tier_id = handle_result.value()->loc.tier_id; + UUID tier_id = handle_result.value()->loc.tier->GetTierId(); // Get the data auto get_result = backend.Get("test_key"); ASSERT_TRUE(get_result.has_value()) << "Get should succeed"; - EXPECT_EQ(get_result.value()->loc.tier_id, tier_id); + EXPECT_EQ(get_result.value()->loc.tier->GetTierId(), tier_id); } // Test get with non-existent key @@ -751,7 +751,7 @@ TEST_F(TieredBackendTest, GetFromSpecifiedTier) { // Get from specific tier auto get_result = backend.Get("multi_tier_key", low_tier_id.value()); ASSERT_TRUE(get_result.has_value()); - EXPECT_EQ(get_result.value()->loc.tier_id, low_tier_id.value()); + EXPECT_EQ(get_result.value()->loc.tier->GetTierId(), low_tier_id.value()); } // Test get returns highest priority tier when tier not specified @@ -799,7 +799,7 @@ TEST_F(TieredBackendTest, GetHighestPriority) { // Get without specifying tier should return highest priority auto get_result = backend.Get("priority_key"); ASSERT_TRUE(get_result.has_value()); - EXPECT_EQ(get_result.value()->loc.tier_id, high_tier_id.value()) + EXPECT_EQ(get_result.value()->loc.tier->GetTierId(), high_tier_id.value()) << "Should return highest priority tier"; } @@ -1027,7 +1027,8 @@ TEST_F(TieredBackendTest, CompleteDataLifecycle) { // 4. Get and verify auto get_result = backend.Get(key); ASSERT_TRUE(get_result.has_value()); - EXPECT_EQ(get_result.value()->loc.tier_id, handle->loc.tier_id); + EXPECT_EQ(get_result.value()->loc.tier->GetTierId(), + handle->loc.tier->GetTierId()); // 5. Delete auto delete_result = backend.Delete(key); @@ -1087,12 +1088,12 @@ TEST_F(TieredBackendTest, MultiTierDataManagement) { // 3. Get without specifying tier should return high priority auto get_high = backend.Get(key); ASSERT_TRUE(get_high.has_value()); - EXPECT_EQ(get_high.value()->loc.tier_id, high_tier_id.value()); + EXPECT_EQ(get_high.value()->loc.tier->GetTierId(), high_tier_id.value()); // 4. Get from low tier explicitly auto get_low = backend.Get(key, low_tier_id.value()); ASSERT_TRUE(get_low.has_value()); - EXPECT_EQ(get_low.value()->loc.tier_id, low_tier_id.value()); + EXPECT_EQ(get_low.value()->loc.tier->GetTierId(), low_tier_id.value()); // 5. Delete high priority replica auto delete_high = backend.Delete(key, high_tier_id.value()); @@ -1101,7 +1102,8 @@ TEST_F(TieredBackendTest, MultiTierDataManagement) { // 6. Get without tier should now return low priority auto get_after_delete = backend.Get(key); ASSERT_TRUE(get_after_delete.has_value()); - EXPECT_EQ(get_after_delete.value()->loc.tier_id, low_tier_id.value()); + EXPECT_EQ(get_after_delete.value()->loc.tier->GetTierId(), + low_tier_id.value()); } // Test concurrent allocations From 80f5659348e45433c44102547a46b150b5264791 Mon Sep 17 00:00:00 2001 From: Xingrui Yi Date: Thu, 8 Jan 2026 11:44:11 +0800 Subject: [PATCH 6/8] [Store] Add wait handle logic in dram iter exit Signed-off-by: Xingrui Yi --- mooncake-store/src/tiered_cache/dram_tier.cpp | 34 ++++++++++++++++--- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/mooncake-store/src/tiered_cache/dram_tier.cpp b/mooncake-store/src/tiered_cache/dram_tier.cpp index f6010352..4f71afe2 100644 --- a/mooncake-store/src/tiered_cache/dram_tier.cpp +++ b/mooncake-store/src/tiered_cache/dram_tier.cpp @@ -1,5 +1,7 @@ #include #include +#include +#include #include "tiered_cache/dram_tier.h" #include "tiered_cache/tiered_backend.h" @@ -21,14 +23,36 @@ DramCacheTier::DramCacheTier(UUID tier_id, size_t capacity, engine_(nullptr) {} DramCacheTier::~DramCacheTier() { - // Check if there are still allocated buffers before destroying allocator + // 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. This may indicate a memory leak."; + 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; + } } } From 90db6b61992a290c504b2f114da8ae2ffdae597a Mon Sep 17 00:00:00 2001 From: Xingrui Yi Date: Thu, 8 Jan 2026 11:49:20 +0800 Subject: [PATCH 7/8] [Store] update tiered backend init return code Signed-off-by: Xingrui Yi --- .../include/tiered_cache/tiered_backend.h | 4 +- .../src/tiered_cache/tiered_backend.cpp | 23 +++--- mooncake-store/tests/tiered_backend_test.cpp | 76 +++++++++---------- 3 files changed, 52 insertions(+), 51 deletions(-) diff --git a/mooncake-store/include/tiered_cache/tiered_backend.h b/mooncake-store/include/tiered_cache/tiered_backend.h index e3ac50e4..378db5da 100644 --- a/mooncake-store/include/tiered_cache/tiered_backend.h +++ b/mooncake-store/include/tiered_cache/tiered_backend.h @@ -92,8 +92,8 @@ class TieredBackend { TieredBackend(); ~TieredBackend() = default; - bool Init(Json::Value root, TransferEngine* engine, - MetadataSyncCallback sync_callback); + tl::expected Init(Json::Value root, TransferEngine* engine, + MetadataSyncCallback sync_callback); // --- Client-Centric Operations --- // All the following operations are designed for Client-Centric, Client diff --git a/mooncake-store/src/tiered_cache/tiered_backend.cpp b/mooncake-store/src/tiered_cache/tiered_backend.cpp index 5d03f56d..42b4e198 100644 --- a/mooncake-store/src/tiered_cache/tiered_backend.cpp +++ b/mooncake-store/src/tiered_cache/tiered_backend.cpp @@ -19,15 +19,16 @@ AllocationEntry::~AllocationEntry() { TieredBackend::TieredBackend() = default; -bool TieredBackend::Init(Json::Value root, TransferEngine* engine, - MetadataSyncCallback sync_callback) { +tl::expected 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 @@ -36,22 +37,22 @@ 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"]) { // Parse required fields if (!tier_config.isMember("type")) { LOG(ERROR) << "Tier config missing required field 'type'"; - return false; + return tl::unexpected(ErrorCode::INVALID_PARAMS); } if (!tier_config.isMember("capacity")) { LOG(ERROR) << "Tier config missing required field 'capacity'"; - return false; + return tl::unexpected(ErrorCode::INVALID_PARAMS); } if (!tier_config.isMember("priority")) { LOG(ERROR) << "Tier config missing required field 'priority'"; - return false; + return tl::unexpected(ErrorCode::INVALID_PARAMS); } std::string type = tier_config["type"].asString(); @@ -61,7 +62,7 @@ bool TieredBackend::Init(Json::Value root, TransferEngine* engine, // Validate capacity if (capacity == 0) { LOG(ERROR) << "Invalid capacity (0) for tier type " << type; - return false; + return tl::unexpected(ErrorCode::INVALID_PARAMS); } // Parse tags @@ -115,7 +116,7 @@ bool TieredBackend::Init(Json::Value root, TransferEngine* engine, if (!init_result) { LOG(ERROR) << "Failed to initialize DRAM tier: id=" << id << ", error=" << init_result.error(); - return false; + return tl::unexpected(init_result.error()); } tiers_[id] = std::move(tier); @@ -123,13 +124,13 @@ bool TieredBackend::Init(Json::Value root, TransferEngine* engine, LOG(INFO) << "Successfully initialized DRAM tier: id=" << id; } else { LOG(ERROR) << "Unsupported tier type '" << type << "'"; - return false; + return tl::unexpected(ErrorCode::INVALID_PARAMS); } } LOG(INFO) << "TieredBackend initialized successfully with " << tier_info_.size() << " tiers."; - return true; + return tl::expected{}; } std::vector TieredBackend::GetSortedTiers() const { diff --git a/mooncake-store/tests/tiered_backend_test.cpp b/mooncake-store/tests/tiered_backend_test.cpp index 76603437..414094a2 100644 --- a/mooncake-store/tests/tiered_backend_test.cpp +++ b/mooncake-store/tests/tiered_backend_test.cpp @@ -122,9 +122,9 @@ TEST_F(TieredBackendTest, BasicDRAMTierInit) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - bool result = backend.Init(config, nullptr, nullptr); + auto result = backend.Init(config, nullptr, nullptr); - EXPECT_TRUE(result); + EXPECT_TRUE(result.has_value()); // Verify tier was created auto tier_views = backend.GetTierViews(); @@ -155,9 +155,9 @@ TEST_F(TieredBackendTest, DRAMTierWithNUMA) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - bool result = backend.Init(config, nullptr, nullptr); + auto result = backend.Init(config, nullptr, nullptr); - EXPECT_TRUE(result); + EXPECT_TRUE(result.has_value()); auto tier_views = backend.GetTierViews(); EXPECT_EQ(tier_views.size(), 1); @@ -185,9 +185,9 @@ TEST_F(TieredBackendTest, MultipleDRAMTiers) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - bool result = backend.Init(config, nullptr, nullptr); + auto result = backend.Init(config, nullptr, nullptr); - EXPECT_TRUE(result); + EXPECT_TRUE(result.has_value()); auto tier_views = backend.GetTierViews(); EXPECT_EQ(tier_views.size(), 2); @@ -207,9 +207,9 @@ TEST_F(TieredBackendTest, MissingTypeField) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - bool result = backend.Init(config, nullptr, nullptr); + auto result = backend.Init(config, nullptr, nullptr); - EXPECT_FALSE(result); + EXPECT_FALSE(result.has_value()); } // Test missing capacity field @@ -226,9 +226,9 @@ TEST_F(TieredBackendTest, MissingCapacityField) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - bool result = backend.Init(config, nullptr, nullptr); + auto result = backend.Init(config, nullptr, nullptr); - EXPECT_FALSE(result); + EXPECT_FALSE(result.has_value()); } // Test invalid capacity (zero) @@ -246,9 +246,9 @@ TEST_F(TieredBackendTest, InvalidCapacity) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - bool result = backend.Init(config, nullptr, nullptr); + auto result = backend.Init(config, nullptr, nullptr); - EXPECT_FALSE(result); + EXPECT_FALSE(result.has_value()); } // Test unsupported tier type @@ -266,9 +266,9 @@ TEST_F(TieredBackendTest, UnsupportedTierType) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - bool result = backend.Init(config, nullptr, nullptr); + auto result = backend.Init(config, nullptr, nullptr); - EXPECT_FALSE(result); + EXPECT_FALSE(result.has_value()); } // Test default allocator type @@ -286,9 +286,9 @@ TEST_F(TieredBackendTest, DefaultAllocatorType) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - bool result = backend.Init(config, nullptr, nullptr); + auto result = backend.Init(config, nullptr, nullptr); - EXPECT_TRUE(result); + EXPECT_TRUE(result.has_value()); auto tier_views = backend.GetTierViews(); EXPECT_EQ(tier_views.size(), 1); @@ -310,10 +310,10 @@ TEST_F(TieredBackendTest, UnknownAllocatorType) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - bool result = backend.Init(config, nullptr, nullptr); + auto result = backend.Init(config, nullptr, nullptr); // Should succeed and use default OFFSET allocator - EXPECT_TRUE(result); + EXPECT_TRUE(result.has_value()); auto tier_views = backend.GetTierViews(); EXPECT_EQ(tier_views.size(), 1); @@ -341,7 +341,7 @@ TEST_F(TieredBackendTest, AllocateBasic) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Allocate space auto result = backend.Allocate(SMALL_DATA_SIZE); @@ -373,7 +373,7 @@ TEST_F(TieredBackendTest, AllocateInsufficientSpace) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Try to allocate more than capacity auto result = backend.Allocate(LARGE_DATA_SIZE); @@ -404,7 +404,7 @@ TEST_F(TieredBackendTest, AllocateOnSpecifiedTier) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Get low priority tier ID auto low_priority_tier_id = GetTierIdByPriority(backend, 50); @@ -435,7 +435,7 @@ TEST_F(TieredBackendTest, AllocateAutoRelease) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); auto tier_views = backend.GetTierViews(); ASSERT_EQ(tier_views.size(), 1); @@ -479,7 +479,7 @@ TEST_F(TieredBackendTest, WriteBasic) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Allocate space auto alloc_result = backend.Allocate(SMALL_DATA_SIZE); @@ -514,7 +514,7 @@ TEST_F(TieredBackendTest, WriteInvalidHandle) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Create invalid (null) handle AllocationHandle invalid_handle; @@ -553,7 +553,7 @@ TEST_F(TieredBackendTest, CommitBasic) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Allocate and write auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); @@ -588,7 +588,7 @@ TEST_F(TieredBackendTest, CommitReplace) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // First commit auto test_buffer1 = CreateTestBuffer(SMALL_DATA_SIZE); @@ -630,7 +630,7 @@ TEST_F(TieredBackendTest, CommitInvalidHandle) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Create invalid handle AllocationHandle invalid_handle; @@ -662,7 +662,7 @@ TEST_F(TieredBackendTest, GetBasic) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Allocate, write and commit auto test_buffer = CreateTestBuffer(SMALL_DATA_SIZE); @@ -696,7 +696,7 @@ TEST_F(TieredBackendTest, GetNonExistentKey) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Try to get non-existent key auto get_result = backend.Get("non_existent_key"); @@ -727,7 +727,7 @@ TEST_F(TieredBackendTest, GetFromSpecifiedTier) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); auto high_tier_id = GetTierIdByPriority(backend, 100); auto low_tier_id = GetTierIdByPriority(backend, 50); @@ -776,7 +776,7 @@ TEST_F(TieredBackendTest, GetHighestPriority) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); auto high_tier_id = GetTierIdByPriority(backend, 100); auto low_tier_id = GetTierIdByPriority(backend, 50); @@ -825,7 +825,7 @@ TEST_F(TieredBackendTest, GetTierNotFound) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); auto high_tier_id = GetTierIdByPriority(backend, 100); auto low_tier_id = GetTierIdByPriority(backend, 50); @@ -871,7 +871,7 @@ TEST_F(TieredBackendTest, DeleteSingleReplica) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); auto high_tier_id = GetTierIdByPriority(backend, 100); auto low_tier_id = GetTierIdByPriority(backend, 50); @@ -927,7 +927,7 @@ TEST_F(TieredBackendTest, DeleteAllReplicas) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); auto high_tier_id = GetTierIdByPriority(backend, 100); auto low_tier_id = GetTierIdByPriority(backend, 50); @@ -973,7 +973,7 @@ TEST_F(TieredBackendTest, DeleteNonExistentKey) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); // Try to delete non-existent key auto delete_result = backend.Delete("non_existent_key"); @@ -1002,7 +1002,7 @@ TEST_F(TieredBackendTest, CompleteDataLifecycle) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); const std::string key = "lifecycle_key"; @@ -1062,7 +1062,7 @@ TEST_F(TieredBackendTest, MultiTierDataManagement) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); auto high_tier_id = GetTierIdByPriority(backend, 100); auto low_tier_id = GetTierIdByPriority(backend, 50); @@ -1122,7 +1122,7 @@ TEST_F(TieredBackendTest, ConcurrentAllocations) { ASSERT_TRUE(parseJsonString(json_config_str, config)); TieredBackend backend; - ASSERT_TRUE(backend.Init(config, nullptr, nullptr)); + ASSERT_TRUE(backend.Init(config, nullptr, nullptr).has_value()); const int num_allocations = 10; std::vector handles; From 751d5f6122af3ee8f7bbd91e8caf0e9358a50b13 Mon Sep 17 00:00:00 2001 From: Xingrui Yi Date: Thu, 8 Jan 2026 11:54:26 +0800 Subject: [PATCH 8/8] [Store] tiered backend remove FreeInternal func Signed-off-by: Xingrui Yi --- .../include/tiered_cache/tiered_backend.h | 3 --- .../src/tiered_cache/tiered_backend.cpp | 22 +++++-------------- 2 files changed, 5 insertions(+), 20 deletions(-) diff --git a/mooncake-store/include/tiered_cache/tiered_backend.h b/mooncake-store/include/tiered_cache/tiered_backend.h index 378db5da..de1182ab 100644 --- a/mooncake-store/include/tiered_cache/tiered_backend.h +++ b/mooncake-store/include/tiered_cache/tiered_backend.h @@ -162,9 +162,6 @@ class TieredBackend { const CacheTier* GetTier(UUID tier_id) const; const DataCopier& GetDataCopier() const; - // Internal API called by AllocationEntry destructor - void FreeInternal(TieredLocation&& loc); - private: struct TierInfo { int priority; diff --git a/mooncake-store/src/tiered_cache/tiered_backend.cpp b/mooncake-store/src/tiered_cache/tiered_backend.cpp index 42b4e198..c3dd2828 100644 --- a/mooncake-store/src/tiered_cache/tiered_backend.cpp +++ b/mooncake-store/src/tiered_cache/tiered_backend.cpp @@ -10,10 +10,9 @@ namespace mooncake { AllocationEntry::~AllocationEntry() { - if (backend) { - // When ref count drops to 0, call back to backend to free physical - // resource. - backend->FreeInternal(std::move(loc)); + if (backend && loc.tier) { + // When ref count drops to 0, free physical resource directly. + loc.tier->Free(std::move(loc.data)); } } @@ -178,24 +177,13 @@ bool TieredBackend::AllocateInternalRaw(size_t size, return false; } -void TieredBackend::FreeInternal(TieredLocation&& loc) { - if (loc.tier) { - auto free_result = loc.tier->Free(std::move(loc.data)); - if (!free_result) { - LOG(WARNING) << "Failed to free data from tier " - << loc.tier->GetTierId() - << ", error=" << free_result.error(); - } - } -} - tl::expected TieredBackend::Allocate( size_t size, std::optional 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. + // destructor triggers Free. return std::make_shared(this, std::move(loc)); } LOG(ERROR) << "Failed to allocate " << size << " bytes"; @@ -444,7 +432,7 @@ tl::expected 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{}; }