[Store] feat: introduce tiered backend #1212
Conversation
Summary of ChangesHello @YiXR, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request establishes the foundational architecture for a new tiered cache backend within Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This PR introduces the foundational abstract classes and data plane logic for a tiered cache backend. The design is well-decomposed into CacheTier, DataCopier, and TieredBackend, supporting flexible allocation strategies. The code is generally well-structured. However, I've found several critical issues related to resource management and error handling, particularly in the DataCopier and TieredBackend implementations. There are also some parts of the implementation that are incomplete, preventing the new backend from being functional. My review includes suggestions to fix these critical bugs and address the incomplete implementation.
| std::unique_ptr<char[]> temp_dram_buffer(new char[src.size]); | ||
| if (!temp_dram_buffer) { | ||
| LOG(ERROR) << "Failed to allocate temporary DRAM buffer for " | ||
| "fallback copy."; | ||
| return false; | ||
| } |
There was a problem hiding this comment.
The allocation new char[src.size] will throw a std::bad_alloc exception on failure, but the code checks for a nullptr return value as if new(std::nothrow) was used. This check will never be true for a standard new expression, and an unhandled std::bad_alloc exception will terminate the program. To correctly handle allocation failure without exceptions, you should use the nothrow version of new.
| std::unique_ptr<char[]> temp_dram_buffer(new char[src.size]); | |
| if (!temp_dram_buffer) { | |
| LOG(ERROR) << "Failed to allocate temporary DRAM buffer for " | |
| "fallback copy."; | |
| return false; | |
| } | |
| std::unique_ptr<char[]> temp_dram_buffer(new (std::nothrow) char[src.size]); | |
| if (!temp_dram_buffer) { | |
| LOG(ERROR) << "Failed to allocate temporary DRAM buffer for " | |
| "fallback copy."; | |
| return false; | |
| } |
| if (!dest_tier->WriteAt(dest_loc.offset, source)) { | ||
| LOG(ERROR) << "MoveData failed: Write error."; | ||
| // dest_tier->Free(dest_loc.offset); | ||
| return false; |
There was a problem hiding this comment.
In the MoveData function, if WriteAt fails, the space allocated in the destination tier is not freed. The call to dest_tier->Free() is commented out, which will lead to a resource leak. This rollback logic is critical for correctness. Note that the Free method requires both an offset and a size.
if (!dest_tier->WriteAt(dest_loc.offset, source)) {
LOG(ERROR) << "MoveData failed: Write error.";
dest_tier->Free(dest_loc.offset, source.size);
return false;| if (sync_cb) { | ||
| bool sync_success = sync_cb(key, dest_loc); | ||
| if (!sync_success) { | ||
| LOG(ERROR) << "MoveData aborted: Master sync failed."; | ||
| // Rollback: Free the newly allocated data | ||
| // dest_tier->Free(dest_loc.offset); | ||
| // Keep the source data intact | ||
| return false; | ||
| } |
There was a problem hiding this comment.
If the sync_cb callback fails, the logic to roll back the operation by freeing the newly allocated block in the destination tier is commented out. This will cause a resource leak, as the space will be allocated but never used or tracked. The rollback is essential for maintaining a consistent state. The Free method requires both an offset and a size.
if (sync_cb) {
bool sync_success = sync_cb(key, dest_loc);
if (!sync_success) {
LOG(ERROR) << "MoveData aborted: Master sync failed.";
// Rollback: Free the newly allocated data
dest_tier->Free(dest_loc.offset, source.size);
// Keep the source data intact
return false;
}
}| // TODO: Logic to instantiate specific CacheTier types (DRAM/SSD) goes | ||
| // here. For example: std::unique_ptr<CacheTier> tier = | ||
| // CacheTierFactory::Create(tier_config); tier->Init(this, engine); | ||
| // tiers_[id] = std::move(tier); | ||
|
|
||
| // Placeholder for compilation if Factory is not ready | ||
| // tiers_[id] = std::make_unique<DramTier>(); | ||
|
|
||
| tier_info_[id] = {priority, tags}; |
There was a problem hiding this comment.
The Init function parses tier configurations but doesn't instantiate any CacheTier objects, leaving the tiers_ map empty. This makes the TieredBackend non-functional, as core operations like Allocate, Put, and MoveData will fail or do nothing. While the TODO indicates this is known, for this feature to be testable and usable, at least a placeholder or mock tier implementation should be instantiated here based on the configuration.
There was a problem hiding this comment.
will be realized later.
| if (it == tiers_.end()) { | ||
| throw std::runtime_error( | ||
| fmt::format("Commit failed: Invalid tier ID {}", loc.tier_id)); | ||
| } |
There was a problem hiding this comment.
The Commit method throws a std::runtime_error on failure. However, the PR description states that Commit is designed to be infallible and should raise a fatal error on internal failure. Throwing an exception is inconsistent with this "fatal error" contract. To align with the documented behavior, you should consider using LOG(FATAL) instead of throwing an exception, which will terminate the process and clearly signal a critical, unrecoverable state.
if (it == tiers_.end()) {
LOG(FATAL) << "Commit failed: Invalid tier ID " << loc.tier_id;
}
YiXR
force-pushed
the
xinyi/tiered-backend
branch
2 times, most recently
from
c91b75d to
efc0278
Compare
There was a problem hiding this comment.
Pull request overview
This PR introduces foundational abstract base classes for a tiered cache management system in mooncake-store. The architecture enables flexible data placement across different storage tiers (DRAM, SSD, etc.) with support for both client-centric and master-centric allocation strategies. The design emphasizes RAII-based resource management through reference-counted handles and provides a sophisticated data copying framework with automatic fallback mechanisms.
Key Changes
- Implements core tiered cache abstractions:
CacheTierinterface,TieredBackenddata plane, andDataCopierutility - Introduces decomposed allocation workflow (Allocate → Write → Commit) to support flexible placement strategies
- Provides Make-Before-Break data migration protocol ensuring strong consistency during tier transitions
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 12 comments.
Show a summary per file
| File | Description |
|---|---|
| mooncake-store/include/tiered_cache/cache_tier.h | Defines abstract CacheTier interface with decomposed allocation/write/bind operations and MemoryType enum |
| mooncake-store/include/tiered_cache/tiered_backend.h | Declares TieredBackend class with RAII handle-based resource management and data migration support |
| mooncake-store/include/tiered_cache/data_copier.h | Defines DataCopier and DataCopierBuilder for memory-type-aware data copying with DRAM fallback |
| mooncake-store/include/tiered_cache/copier_registry.h | Provides singleton registry for static registration of memory type copy functions |
| mooncake-store/src/tiered_cache/tiered_backend.cpp | Implements backend initialization, allocation, commit, deletion, and migration logic |
| mooncake-store/src/tiered_cache/data_copier.cpp | Implements builder pattern and copy logic with automatic DRAM-based fallback mechanism |
| mooncake-store/src/tiered_cache/copier_registry.cpp | Implements singleton registry and registrar helper for static initialization |
| mooncake-store/src/CMakeLists.txt | Adds three new source files to the build system |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| /** | ||
| * @brief Retrieves data pointer/info for a key. | ||
| * implementation relies on internal index (Key -> Offset). |
There was a problem hiding this comment.
The documentation for the Get method says "Retrieves data pointer/info for a key" but doesn't document what happens when the key is not found. The method should specify the expected behavior for the return value and output parameters (data and size) when the key doesn't exist.
| * implementation relies on internal index (Key -> Offset). | |
| * Implementation relies on internal index (Key -> Offset). | |
| * | |
| * @param key The key to look up. | |
| * @param data [out] On success, set to point to the data for the key; on failure (key not found), set to nullptr. | |
| * @param size [out] On success, set to the size of the data; on failure (key not found), set to 0. | |
| * @return true if the key exists and data/size are set; false if the key does not exist (data is set to nullptr, size to 0). |
| // Handle is NOT added to map. | ||
| // If caller drops the handle, ref count -> 0, resource freed. |
There was a problem hiding this comment.
The comment at lines 141-142 describes unreachable code logic that appears to be leftover from an earlier implementation. Since the function returns false at line 140, these comments about handle behavior are confusing and should be removed as they don't accurately describe the actual code flow.
| // Handle is NOT added to map. | |
| // If caller drops the handle, ref count -> 0, resource freed. |
|
|
||
| if (src_handle->loc.tier_id == dest_tier_id) return true; | ||
|
|
||
| auto src_tier = tiers_[src_handle->loc.tier_id].get(); |
There was a problem hiding this comment.
Accessing tiers_[src_handle->loc.tier_id] without checking if the tier exists could cause undefined behavior. Using operator[] will create a default-constructed entry if the tier_id doesn't exist, leading to a nullptr dereference at .get(). This should use tiers_.find() or tiers_.at() with proper error handling, consistent with the pattern used at lines 105 and 127.
| auto src_tier = tiers_[src_handle->loc.tier_id].get(); | |
| auto tier_it = tiers_.find(src_handle->loc.tier_id); | |
| if (tier_it == tiers_.end()) { | |
| LOG(ERROR) << "MoveData failed: Source tier_id " << src_handle->loc.tier_id << " not found."; | |
| return false; | |
| } | |
| auto src_tier = tier_it->second.get(); |
| LOG(WARNING) << "Commit failed: Key " << key << " already exists."; | ||
| return false; |
There was a problem hiding this comment.
The Commit method returns false when a key already exists, but doesn't provide a way for the caller to distinguish this from other failure modes. According to the design described in the PR, Commit "is designed to be infallible (must succeed). If an internal error occurs (e.g., OOM), it raises a fatal error." However, this implementation treats duplicate keys as a soft error. This inconsistency between the documented behavior and implementation could lead to confusion and resource leaks if callers don't properly handle the returned AllocationHandle.
| LOG(WARNING) << "Commit failed: Key " << key << " already exists."; | |
| return false; | |
| LOG(FATAL) << "Commit failed: Key " << key << " already exists. Commit is designed to be infallible."; | |
| // LOG(FATAL) will terminate the process, so no return needed. |
| /** | ||
| * @brief Reserve Space (Allocation) | ||
| * Finds free space of `size` bytes. Does NOT copy data. | ||
| * * @param size Bytes to allocate. |
There was a problem hiding this comment.
The documentation has inconsistent formatting. Line 81 has a double asterisk "* *" at the start which appears to be a typo or formatting error. The subsequent parameter descriptions should follow consistent formatting.
| * * @param size Bytes to allocate. | |
| * @param size Bytes to allocate. |
| /** | ||
| * @brief Returns a DataSource descriptor for a key's data. | ||
| * Used for data migration (moving data out of this tier). | ||
| */ |
There was a problem hiding this comment.
The AsDataSource method documentation doesn't specify what should be returned when the key doesn't exist. Based on the usage in tiered_backend.cpp line 196, it appears a DataSource with ptr=nullptr might indicate an error, but this contract should be explicitly documented in the interface.
| */ | |
| * @note If the key does not exist, returns a DataSource with ptr == nullptr to indicate an error. |
| auto& tier = tiers_[handle->loc.tier_id]; | ||
| tier->BindKey(key, handle->loc.offset, handle->size); |
There was a problem hiding this comment.
Accessing tiers_[handle->loc.tier_id] without checking if the tier exists first could cause undefined behavior. While tiers_.find() is used in other methods (lines 105, 127), this line directly accesses the map using operator[], which will create a default-constructed entry if the tier_id doesn't exist. This should use tiers_.find() or tiers_.at() with error handling.
| auto& tier = tiers_[handle->loc.tier_id]; | |
| tier->BindKey(key, handle->loc.offset, handle->size); | |
| auto it = tiers_.find(handle->loc.tier_id); | |
| if (it == tiers_.end()) { | |
| LOG(ERROR) << "Commit failed: Tier " << handle->loc.tier_id << " does not exist."; | |
| return false; | |
| } | |
| it->second->BindKey(key, handle->loc.offset, handle->size); |
| /** | ||
| * @typedef AllocationHandle | ||
| * @brief A reference-counted handle to a storage resource. | ||
| * * Lifecycle Semantics: |
There was a problem hiding this comment.
The documentation comment has inconsistent formatting with "* Lifecycle Semantics:" having an extra asterisk. This should be formatted consistently with other documentation comments in the file.
| * * Lifecycle Semantics: | |
| * Lifecycle Semantics: |
| bool TieredBackend::Init(Json::Value root, TransferEngine* engine) { | ||
| // Initialize DataCopier | ||
| try { | ||
| DataCopierBuilder builder; | ||
| data_copier_ = builder.Build(); | ||
| } catch (const std::logic_error& e) { | ||
| LOG(FATAL) << "Failed to build DataCopier: " << e.what(); | ||
| return false; | ||
| } |
There was a problem hiding this comment.
The error handling is inconsistent here. After logging a FATAL error at line 26, the function returns false at line 27, but LOG(FATAL) typically terminates the program. Either remove the return statement since it's unreachable, or change LOG(FATAL) to LOG(ERROR) if the intent is to allow the caller to handle the failure gracefully.
| } | ||
|
|
||
| src_tier->Delete(key); | ||
|
|
There was a problem hiding this comment.
There is a potential race condition in the MoveData operation. After releasing the map_mutex_ at line 220 and before calling src_tier->Delete(key) at line 222, another thread could call Get(key) and receive the new handle, then attempt to use the old tier's data through AsDataSource before it's deleted. The source tier deletion should happen while still holding the lock, or additional synchronization is needed to ensure the old data isn't accessed during deletion.
| } | |
| src_tier->Delete(key); | |
| src_tier->Delete(key); | |
| } |
YiXR
force-pushed
the
xinyi/tiered-backend
branch
10 times, most recently
from
d79029c to
930d91f
Compare
YiXR
changed the title
YiXR
force-pushed
the
xinyi/tiered-backend
branch
from
930d91f to
e1b2f1a
Compare
Signed-off-by: Xingrui Yi <yixingrui@linux.alibaba.com>
YiXR
force-pushed
the
xinyi/tiered-backend
branch
from
e1b2f1a to
92e557e
Compare
1 participant
Description
This PR introduce the abstract base classes for cache tier management.
CacheTier
mooncake-store/include/tiered_cache/cache_tier.hResponsibilities
This is an abstract base class that defines the unified interface a single cache tier (e.g., DRAM, SSD) must implement. It is responsible for managing the data lifecycle on a specific storage medium.
Core Interface:
Allocate: Finds free space ofsizebytes. Does NOT copy data.Free: Releases space at offset. Used when writes fail or explicitly freeing anonymous blocks.DataCopier
mooncake-store/include/tiered_cache/data_copier.hResponsibilities:
Responsible for efficiently copying data between different memory types (
MemoryType). This is a core utility class that decouples the data movement logic.Core Functionality:
DataCopierBuilder, allowing for the registration of direct copy functions between different memory types (e.g., DRAM -> VRAM).Copyinterface. When a direct copy path is unavailable, it automatically employs a fallback mechanism using DRAM as an intermediate buffer (e.g., VRAM -> DRAM -> SSD) (when implementing a new type, compilation requires copy functions between the new type and DRAM). This greatly simplifies the integration of new storage media.TieredBackendandCacheTierhold and use aDataCopierinstance to execute all data copy operations, whether it's writing new data or moving data between tiers.TieredBackend
mooncake-store/include/tiered_cache/tiered_backend.hResponsibilities:
Serves as the Data Plane for the tiered cache. It does not contain any complex caching policies (e.g., LRU) and is solely responsible for faithfully executing data operation instructions.
Core Data Structures
DataSourceA generic descriptor for source data.
TieredLocationDescribes the physical location of a data segment.
AllocationHandleDecomposed APIs
The traditional atomic Put operation is broken down into three granular steps:
AllocateReserves storage space locally without writing data.
Usage: Called by the AllocationStrategy to find free space.
Returns: true if space is reserved; false otherwise.
Write (local operation)Writes data to a specific, reserved location.
Usage: Performs local write operation. If need remote write, please tranfer data into AllocationHandle by using TE.
Returns: true on success. If failed, the caller must free the allocated space.
CommitRegisters the key in the local metadata index.
Usage: Called after data is successfully written.
Behavior: It binds the key to the offset internally. This operation is designed to be infallible (must succeed). If an internal error occurs (e.g., OOM), it raises a fatal error.
DeleteDelete the key in the local metadata index.
If
tier_idis specified, removes only the replica on that tier. If nullopt, removes ALL replicas for this key.Allocation Workflows
This decoupled design supports two distinct allocation patterns required by the new architecture.
Each backend can only use one allocation strategy after mooncake store start up. If the backend wants to use Master-Centric Allocation, it must realize all the func to register it's info and allocator to Master.
Pattern A: Client-Centric Allocation (Local Decision)
In this mode, the Client decides where to place data (e.g., "Try DRAM first, then SSD") and later informs the Master.
Workflow
Decision
The ClientCentricStrategy calls
Allocate(size, preferred_tier=xxx).TieredBackend checks local capacity and returns a
AllocationHandle.Execution
The strategy calls
Write(data, AllocationHandle), this is a local write operation.Please use TE to transfer data if you want to do remote write.
Local Commit
The strategy calls
Commit(key, AllocationHandle).The key is now visible to local
Get()requests.Global Sync
The Client sends an asynchronous RPC
RegisterAllocationto the Master.Master updates its directory: "Key X is located at Client Y, Tier Z".
Pattern B: Master-Centric Allocation (Global Decision)
In this mode, the Master manages all resources globally (Client should still call
MountSegment()to register all the resources to Master) and tells the Client exactly where to write.Client doesn't need to save metadata for this backend anyway.
Workflow:
Request
The Client sends an RPC
PutStartto the Master.Decision
Master checks its global resource view and returns
{tier_id: 1, offset: 0x2000}.Execution
Data can be write by TE into tier directly.
Data Migration & Consistency
Migrates a key between tiers using a "Make-Before-Break" protocol to ensure strong consistency.
If DataSource is Master-Centric Allocation, Master must 'tell' Client where the data is.
Workflow:
Allocate & Write
Data is copied to the destination tier.
Sync Callback
sync_cb is invoked.
Action
The Client sends an RPC to the Master: "Update Key location to Tier New".
Safety
If this RPC fails, CopyData returns false and rolls back the new copy. The old data remains valid.
Commit
Local metadata is updated to point to the new tier.
Delete
Only after the Master is updated and local commit succeeds, the old replica need to be deleted.
If DataSource is Master-Centric Allocation, Master will delete the key by itself.
If DataSource is Client-Centric Allocation, Client must call
Delete(key)afterCopyData.sequenceDiagram participant Client participant TieredBackend participant CacheTier participant DataCopier participant MemoryPool Client->>TieredBackend: Allocate(size, preferred_tier) activate TieredBackend TieredBackend->>TieredBackend: GetSortedTiers() TieredBackend->>CacheTier: Allocate(size) activate CacheTier CacheTier->>MemoryPool: allocate MemoryPool-->>CacheTier: ptr CacheTier-->>TieredBackend: AllocationHandle(tier_id, DataSource) deactivate CacheTier TieredBackend-->>Client: AllocationHandle deactivate TieredBackend Client->>TieredBackend: Write(source_data, handle) activate TieredBackend TieredBackend->>DataCopier: Copy(source, dest_from_handle) activate DataCopier DataCopier->>DataCopier: FindCopier(src_type, dest_type) alt Direct Path Exists DataCopier->>DataCopier: Execute direct copy else No Direct Path note over DataCopier: Both non-DRAM? DataCopier->>TieredBackend: Allocate temp DRAM buffer DataCopier->>DataCopier: Copy src → DRAM DataCopier->>DataCopier: Copy DRAM → dest end DataCopier-->>TieredBackend: success/failure deactivate DataCopier TieredBackend-->>Client: bool deactivate TieredBackend Client->>TieredBackend: Commit(key, handle) activate TieredBackend TieredBackend->>TieredBackend: Update metadata_index<br/>(per-key replicas) TieredBackend-->>Client: bool deactivate TieredBackendsequenceDiagram participant Client participant TieredBackend participant MetadataIndex participant CacheTier1 participant CacheTier2 Client->>TieredBackend: Get(key, optional_tier_id) activate TieredBackend TieredBackend->>MetadataIndex: Lookup replicas for key activate MetadataIndex MetadataIndex-->>TieredBackend: replica list sorted by priority deactivate MetadataIndex alt tier_id specified TieredBackend->>CacheTier1: Retrieve from specific tier else best available TieredBackend->>CacheTier1: Retrieve from highest-priority tier end activate CacheTier1 CacheTier1-->>TieredBackend: AllocationHandle deactivate CacheTier1 TieredBackend-->>Client: AllocationHandle deactivate TieredBackend Client->>TieredBackend: Delete(key, optional_tier_id) activate TieredBackend TieredBackend->>MetadataIndex: Remove replica from key activate MetadataIndex MetadataIndex->>CacheTier2: Free allocation activate CacheTier2 CacheTier2-->>MetadataIndex: success deactivate CacheTier2 MetadataIndex-->>TieredBackend: bool deactivate MetadataIndex TieredBackend-->>Client: bool deactivate TieredBackendType of Change
How Has This Been Tested?
Checklist