diff --git a/mooncake-store/include/tiered_cache/cache_tier.h b/mooncake-store/include/tiered_cache/cache_tier.h new file mode 100644 index 000000000..f571e9ba0 --- /dev/null +++ b/mooncake-store/include/tiered_cache/cache_tier.h @@ -0,0 +1,126 @@ +#pragma once + +#include +#include +#include +#include "transfer_engine.h" + +namespace mooncake { +struct DataSource; +enum class MemoryType; +} // namespace mooncake + +namespace mooncake { + +class TieredBackend; + +/** + * @enum MemoryType + * @brief Defines the physical storage medium type for a cache tier. + */ +enum class MemoryType { DRAM, UNKNOWN }; + +static inline std::string MemoryTypeToString(MemoryType type) { + switch (type) { + case MemoryType::DRAM: + return "DRAM"; + default: + return "UNKNOWN"; + } +} + +/** + * @struct DataSource + * @brief Describes a source of data for a copy operation. + * + * This struct is used as a generic descriptor for a block of memory, allowing + * data to be described abstractly regardless of its physical location. + */ +struct DataSource { + const void* + ptr; // Pointer to the data. Its interpretation depends on the `type`. + size_t size; // Size of the data in bytes. + MemoryType type; // The memory type where the data resides. +}; + +/** + * @class CacheTier + * @brief Abstract base class for a single tier in the tiered cache system. + * + * This class defines the common interface that all storage media (DRAM, VRAM, + * SSD, etc.) must implement. The interface is designed to be simple and focuses + * on the essential operations of a storage layer, leaving complex eviction and + * promotion logic to the TieredBackend and CacheScheduler. + */ +class CacheTier { + public: + virtual ~CacheTier() = default; + + /** + * @brief Initializes the cache tier. + * @param backend A pointer to the parent TieredBackend for coordination. + * @param te A pointer to the active TransferEngine, for tiers that need it. + * @return True on success, false otherwise. + */ + virtual bool Init(TieredBackend* backend, TransferEngine* engine) = 0; + + /** + * @brief Retrieves a pointer to the data for a given key. + * @param key The key to look up. + * @param data [out] A reference to a void pointer that will be set to the + * data's location. + * @param size [out] A reference that will be set to the data's size. + * @return True if the key is found, false otherwise. + */ + virtual bool Get(const std::string& key, void*& data, size_t& size) = 0; + + /** + * @brief Puts data into the tier from a generic data source. + * This is the sole method for writing data. The implementation must always + * allocate its own memory and copy the data from the source, using the + * backend's DataCopier. + * @param key The key for the data. + * @param source The descriptor for the source data (pointer, size, type). + * @return True on success, false otherwise. + */ + virtual bool Put(const std::string& key, const DataSource& source) = 0; + + /** + * @brief Deletes a key and its associated data from the tier. + * @param key The key to delete. + * @return True if the key was found and deleted, false otherwise. + */ + virtual bool Delete(const std::string& key) = 0; + + /** + * @brief Checks if the tier contains a given key. + * @param key The key to check. + * @return True if the key exists in this tier, false otherwise. + */ + virtual bool Contains(const std::string& key) const = 0; + + /** + * @brief Returns a DataSource descriptor for a key's data within this tier. + * This is used to describe the data as a source for a subsequent copy + * operation when it needs to be moved to another tier. + * @param key The key to describe. + * @return A DataSource object. If the key is not found, the `ptr` member + * will be null. + */ + virtual DataSource AsDataSource(const std::string& key) = 0; + + // --- Accessors for tier properties --- + + virtual uint64_t GetTierId() const = 0; + virtual size_t GetCapacity() const = 0; + virtual size_t GetUsage() const = 0; + virtual const std::vector& GetTags() const = 0; + virtual MemoryType GetMemoryType() const = 0; + + protected: + // A pointer to the parent backend, allowing tiers to access shared services + // like the DataCopier. + TieredBackend* backend_ = nullptr; +}; + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/tiered_cache/copier_registry.h b/mooncake-store/include/tiered_cache/copier_registry.h new file mode 100644 index 000000000..80c518ae7 --- /dev/null +++ b/mooncake-store/include/tiered_cache/copier_registry.h @@ -0,0 +1,86 @@ +#pragma once + +#include "tiered_cache/cache_tier.h" +#include "tiered_cache/data_copier.h" +#include +#include +#include +#include + +namespace mooncake { + +// Forward declaration from data_copier.h to avoid circular dependency +class DataCopierBuilder; + +// Holds the registration information for a memory type. +struct MemoryTypeRegistration { + MemoryType type; + CopyFunction to_dram_func; + CopyFunction from_dram_func; +}; + +// Holds the registration for an optimized direct path. +struct DirectPathRegistration { + MemoryType src_type; + MemoryType dest_type; + CopyFunction func; +}; + +/** + * @brief A singleton registry for data copier functions. + * + * Modules can register their copy functions here during static initialization. + * The DataCopierBuilder will then use this registry to construct a DataCopier. + */ +class CopierRegistry { + public: + /** + * @brief Get the singleton instance of the registry. + */ + static CopierRegistry& GetInstance(); + + /** + * @brief Registers the to/from DRAM copy functions for a memory type. + */ + void RegisterMemoryType(MemoryType type, CopyFunction to_dram, + CopyFunction from_dram); + + /** + * @brief Registers an optional, optimized direct copy path. + */ + void RegisterDirectPath(MemoryType src, MemoryType dest, CopyFunction func); + + // These methods are used by the DataCopierBuilder to collect all + // registrations. + const std::vector& GetMemoryTypeRegistrations() + const; + const std::vector& GetDirectPathRegistrations() + const; + + private: + friend class DataCopierBuilder; + + CopierRegistry() = default; + ~CopierRegistry() = default; + CopierRegistry(const CopierRegistry&) = delete; + CopierRegistry& operator=(const CopierRegistry&) = delete; + + std::vector memory_type_regs_; + std::vector direct_path_regs_; +}; + +/** + * @brief A helper class to automatically register copiers at static + * initialization time. + * + * To register a new memory type, simply declare a static instance of this class + * in the corresponding .cpp file, providing the type and its to/from DRAM + * copiers. + */ +class CopierRegistrar { + public: + CopierRegistrar(MemoryType type, CopyFunction to_dram, + CopyFunction from_dram); +}; + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/tiered_cache/data_copier.h b/mooncake-store/include/tiered_cache/data_copier.h new file mode 100644 index 000000000..7c6415fb0 --- /dev/null +++ b/mooncake-store/include/tiered_cache/data_copier.h @@ -0,0 +1,90 @@ +#pragma once + +#include "tiered_cache/cache_tier.h" +#include +#include +#include +#include +#include +#include + +namespace mooncake { + +using CopyFunction = std::function; + +class DataCopier; + +/** + * @brief A helper class to build a valid DataCopier. + * + * This builder enforces the rule that for any new memory type added, + * its copy functions to and from DRAM *must* be provided via the + * CopierRegistry. + */ +class DataCopierBuilder { + public: + /** + * @brief Constructs a builder. It automatically pulls all existing + * registrations from the global CopierRegistry. + */ + DataCopierBuilder(); + + /** + * @brief (Optional) Registers a highly optimized direct copy path. + * This will be used instead of the DRAM fallback. Can be used for testing + * or for paths that are not self-registered. + * @return A reference to the builder for chaining. + */ + DataCopierBuilder& AddDirectPath(MemoryType src_type, MemoryType dest_type, + CopyFunction func); + + /** + * @brief Builds the final, immutable DataCopier object. + * It verifies that all memory types defined in the MemoryType enum + * have been registered via the registry before creating the object. + * @return A unique_ptr to the new DataCopier. + * @throws std::logic_error if a required to/from DRAM copier is missing. + */ + std::unique_ptr Build() const; + + private: + std::map, CopyFunction> copy_matrix_; +}; + +/** + * @brief A central utility for copying data between different memory types. + * It supports a fallback mechanism via DRAM for any copy paths that are not + * explicitly registered as a direct path. + */ +class DataCopier { + public: + // The constructor is private. Use DataCopierBuilder to create an instance. + ~DataCopier() = default; + DataCopier(const DataCopier&) = delete; + DataCopier& operator=(const DataCopier&) = delete; + + /** + * @brief Executes a copy from a source to a destination. + * It first attempts to find a direct copy function (e.g., VRAM -> VRAM). + * If not found, it automatically falls back to a two-step copy via a + * temporary DRAM buffer (e.g., VRAM -> DRAM -> SSD). + * @param src The data source descriptor. + * @param dest_type The memory type of the destination. + * @param dest_ptr A pointer to the destination (memory address, handle, + * etc.). + * @return True if the copy was successful, false otherwise. + */ + bool Copy(const DataSource& src, MemoryType dest_type, + void* dest_ptr) const; + + private: + friend class DataCopierBuilder; // Allow builder to access the constructor. + DataCopier( + std::map, CopyFunction> copy_matrix); + + CopyFunction FindCopier(MemoryType src_type, MemoryType dest_type) const; + const std::map, CopyFunction> + copy_matrix_; +}; + +} // 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 new file mode 100644 index 000000000..a500073e2 --- /dev/null +++ b/mooncake-store/include/tiered_cache/tiered_backend.h @@ -0,0 +1,75 @@ +#pragma once + +#include "tiered_cache/cache_tier.h" +#include "tiered_cache/data_copier.h" +#include +#include +#include +#include +#include +#include +#include + +namespace mooncake { + +/** + * @struct TierView + * @brief A snapshot of a CacheTier's status for the upper layer (e.g., Worker). + */ +struct TierView { + uint64_t id; + MemoryType type; + size_t capacity; + size_t usage; + int priority; + std::vector tags; +}; + +/** + * @class TieredBackend + * @brief A pure data plane for the tiered caching system. + */ +class TieredBackend { + public: + TieredBackend(); + ~TieredBackend() = default; + + bool Init(Json::Value root, TransferEngine* engine); + bool Get(const std::string& key, void*& data, size_t& size); + bool Put(const std::string& key, uint64_t target_tier_id, + const DataSource& source); + bool Delete(const std::string& key); + bool MoveData(const std::string& key, uint64_t src_tier_id, + uint64_t dest_tier_id); + + std::optional FindKey(const std::string& key) const; + std::vector GetTierViews() const; + const CacheTier* GetTier(uint64_t tier_id) const; + const DataCopier& GetDataCopier() const; + + private: + /** + * @struct TierInfo + * @brief Internal struct to hold static configuration for each tier. + */ + struct TierInfo { + int priority; + std::vector tags; + }; + + bool DeleteFromTier(const std::string& key, uint64_t tier_id); + + // Map from tier ID to the actual CacheTier instance. + std::unordered_map> tiers_; + + // Map from tier ID to its static configuration info. + std::unordered_map tier_info_; + + // A fast lookup map from a key to the ID of the tier that holds it. + std::unordered_map key_to_tier_map_; + mutable std::shared_mutex map_mutex_; // Protects key_to_tier_map_ + + std::unique_ptr data_copier_; +}; + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 385c84045..b835396ea 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -24,6 +24,9 @@ set(MOONCAKE_STORE_SOURCES client_buffer.cpp pybind_client.cpp http_metadata_server.cpp + tiered_cache/copier_registry.cpp + tiered_cache/data_copier.cpp + tiered_cache/tiered_backend.cpp ) set(EXTRA_LIBS "") diff --git a/mooncake-store/src/tiered_cache/copier_registry.cpp b/mooncake-store/src/tiered_cache/copier_registry.cpp new file mode 100644 index 000000000..d6554bdfd --- /dev/null +++ b/mooncake-store/src/tiered_cache/copier_registry.cpp @@ -0,0 +1,41 @@ +#include "tiered_cache/copier_registry.h" +#include "tiered_cache/data_copier.h" +#include + +namespace mooncake { + +CopierRegistry& CopierRegistry::GetInstance() { + static CopierRegistry instance; + return instance; +} + +void CopierRegistry::RegisterMemoryType(MemoryType type, CopyFunction to_dram, + CopyFunction from_dram) { + memory_type_regs_.push_back( + {type, std::move(to_dram), std::move(from_dram)}); +} + +void CopierRegistry::RegisterDirectPath(MemoryType src, MemoryType dest, + CopyFunction func) { + direct_path_regs_.push_back({src, dest, std::move(func)}); +} + +const std::vector& +CopierRegistry::GetMemoryTypeRegistrations() const { + return memory_type_regs_; +} + +const std::vector& +CopierRegistry::GetDirectPathRegistrations() const { + return direct_path_regs_; +} + +CopierRegistrar::CopierRegistrar(MemoryType type, CopyFunction to_dram, + CopyFunction from_dram) { + // When a static CopierRegistrar object is created, it registers the memory + // type. + CopierRegistry::GetInstance().RegisterMemoryType(type, std::move(to_dram), + std::move(from_dram)); +} + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/src/tiered_cache/data_copier.cpp b/mooncake-store/src/tiered_cache/data_copier.cpp new file mode 100644 index 000000000..8e6a6204f --- /dev/null +++ b/mooncake-store/src/tiered_cache/data_copier.cpp @@ -0,0 +1,114 @@ +#include "tiered_cache/data_copier.h" +#include "tiered_cache/copier_registry.h" +#include +#include +#include + +namespace mooncake { + +DataCopierBuilder::DataCopierBuilder() { + // Process all registrations from the global registry. + const auto& registry = CopierRegistry::GetInstance(); + + for (const auto& reg : registry.GetMemoryTypeRegistrations()) { + copy_matrix_[{reg.type, MemoryType::DRAM}] = reg.to_dram_func; + copy_matrix_[{MemoryType::DRAM, reg.type}] = reg.from_dram_func; + } + for (const auto& reg : registry.GetDirectPathRegistrations()) { + copy_matrix_[{reg.src_type, reg.dest_type}] = reg.func; + } +} + +DataCopierBuilder& DataCopierBuilder::AddDirectPath(MemoryType src_type, + MemoryType dest_type, + CopyFunction func) { + copy_matrix_[{src_type, dest_type}] = std::move(func); + return *this; +} + +std::unique_ptr DataCopierBuilder::Build() const { + const auto& registry = CopierRegistry::GetInstance(); + for (const auto& reg : registry.GetMemoryTypeRegistrations()) { + if (reg.type == MemoryType::DRAM) { + continue; + } + if (copy_matrix_.find({reg.type, MemoryType::DRAM}) == + copy_matrix_.end()) { + throw std::logic_error( + "DataCopierBuilder Error: Missing copy function for type " + + MemoryTypeToString(reg.type) + " TO DRAM."); + } + if (copy_matrix_.find({MemoryType::DRAM, reg.type}) == + copy_matrix_.end()) { + throw std::logic_error( + "DataCopierBuilder Error: Missing copy function for DRAM TO " + "type " + + MemoryTypeToString(reg.type) + "."); + } + } + + return std::unique_ptr(new DataCopier(copy_matrix_)); +} + +DataCopier::DataCopier( + std::map, CopyFunction> copy_matrix) + : copy_matrix_(std::move(copy_matrix)) {} + +CopyFunction DataCopier::FindCopier(MemoryType src_type, + MemoryType dest_type) const { + auto it = copy_matrix_.find({src_type, dest_type}); + return (it != copy_matrix_.end()) ? it->second : nullptr; +} + +bool DataCopier::Copy(const DataSource& src, MemoryType dest_type, + void* dest_ptr) const { + // Try to find a direct copy function. + if (auto direct_copier = FindCopier(src.type, dest_type)) { + VLOG(1) << "Using direct copier for " << MemoryTypeToString(src.type) + << " -> " << MemoryTypeToString(dest_type); + return direct_copier(src, dest_ptr); + } + + // If no direct copier, try fallback via DRAM. + if (src.type != MemoryType::DRAM && dest_type != MemoryType::DRAM) { + VLOG(1) << "No direct copier. Attempting fallback via DRAM for " + << MemoryTypeToString(src.type) << " -> " + << MemoryTypeToString(dest_type); + + auto to_dram_copier = FindCopier(src.type, MemoryType::DRAM); + auto from_dram_copier = FindCopier(MemoryType::DRAM, dest_type); + + if (to_dram_copier && from_dram_copier) { + std::unique_ptr temp_dram_buffer(new char[src.size]); + if (!temp_dram_buffer) { + LOG(ERROR) << "Failed to allocate temporary DRAM buffer for " + "fallback copy."; + return false; + } + + // Step A: Source -> DRAM + if (!to_dram_copier(src, temp_dram_buffer.get())) { + LOG(ERROR) << "Fallback copy failed at Step A (Source -> DRAM)"; + return false; + } + + // Step B: DRAM -> Destination + DataSource temp_dram_source = {temp_dram_buffer.get(), src.size, + MemoryType::DRAM}; + if (!from_dram_copier(temp_dram_source, dest_ptr)) { + LOG(ERROR) + << "Fallback copy failed at Step B (DRAM -> Destination)"; + return false; + } + return true; + } + } + + LOG(ERROR) << "No copier registered for transfer from memory type " + << MemoryTypeToString(src.type) << " to " + << MemoryTypeToString(dest_type) + << ", and fallback path is not available."; + return false; +} + +} // 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 new file mode 100644 index 000000000..df3285785 --- /dev/null +++ b/mooncake-store/src/tiered_cache/tiered_backend.cpp @@ -0,0 +1,203 @@ +#include "tiered_cache/tiered_backend.h" +#include "tiered_cache/cache_tier.h" + +#include +#include +#include +#include +#include +#include + +namespace mooncake { + +TieredBackend::TieredBackend() = default; + +bool TieredBackend::Init(Json::Value root, TransferEngine* engine) { + // Initialize the DataCopier + try { + DataCopierBuilder builder; + data_copier_ = builder.Build(); + } catch (const std::logic_error& e) { + LOG(FATAL) << "Failed to build DataCopier: " << e.what(); + return false; + } + + // Create CacheTier instances and store their static info + if (!root.isMember("tiers")) { + LOG(ERROR) << "Tiered cache config is missing 'tiers' array."; + return false; + } + + for (const auto& tier_config : root["tiers"]) { + uint64_t id = tier_config["id"].asUInt(); + std::string type = tier_config["type"].asString(); + // size_t capacity = tier_config["capacity"].asInt(); + int priority = tier_config["priority"].asInt(); + std::vector tags; + if (tier_config.isMember("tags") && tier_config["tags"].isArray()) { + for (const auto& tag_node : tier_config["tags"]) { + tags.push_back(tag_node.asString()); + } + } + std::unique_ptr tier; + + // TODO: add specific cache tier types init logic here + /* + if (!tier->Init(this, engine)) { + LOG(ERROR) << "Failed to initialize tier " << id; + return false; + } + */ + + tiers_[id] = std::move(tier); + tier_info_[id] = {priority, tags}; + } + + LOG(INFO) << "TieredBackend initialized successfully with " << tiers_.size() + << " tiers."; + return true; +} + +std::vector TieredBackend::GetTierViews() const { + std::vector views; + views.reserve(tiers_.size()); + for (const auto& [id, tier] : tiers_) { + const auto& info = tier_info_.at(id); + views.push_back({tier->GetTierId(), tier->GetMemoryType(), + tier->GetCapacity(), tier->GetUsage(), info.priority, + info.tags}); + } + return views; +} + +bool TieredBackend::Get(const std::string& key, void*& data, size_t& size) { + auto maybe_tier_id = FindKey(key); + if (!maybe_tier_id) { + return false; + } + return tiers_.at(*maybe_tier_id)->Get(key, data, size); +} + +bool TieredBackend::Put(const std::string& key, uint64_t target_tier_id, + const DataSource& source) { + auto it = tiers_.find(target_tier_id); + if (it == tiers_.end()) { + LOG(ERROR) << "Put failed: Invalid target tier ID " << target_tier_id; + return false; + } + auto& target_tier = it->second; + + if (target_tier->Put(key, source)) { + std::unique_lock lock(map_mutex_); + key_to_tier_map_[key] = target_tier_id; + return true; + } + + return false; +} + +bool TieredBackend::Delete(const std::string& key) { + std::optional tier_id_opt; + { + std::unique_lock lock(map_mutex_); + auto it = key_to_tier_map_.find(key); + if (it != key_to_tier_map_.end()) { + tier_id_opt = it->second; + key_to_tier_map_.erase(it); + } + } + + if (tier_id_opt) { + return DeleteFromTier(key, *tier_id_opt); + } + + return false; +} + +bool TieredBackend::MoveData(const std::string& key, uint64_t src_tier_id, + uint64_t dest_tier_id) { + VLOG(1) << "Moving key '" << key << "' from tier " << src_tier_id << " to " + << dest_tier_id; + + auto src_it = tiers_.find(src_tier_id); + auto dest_it = tiers_.find(dest_tier_id); + if (src_it == tiers_.end() || dest_it == tiers_.end()) { + LOG(ERROR) << "MoveData failed: Invalid tier ID. Source: " + << src_tier_id << ", Dest: " << dest_tier_id; + return false; + } + auto& src_tier = src_it->second; + auto& dest_tier = dest_it->second; + + std::unique_lock lock(map_mutex_); + + auto key_it = key_to_tier_map_.find(key); + if (key_it == key_to_tier_map_.end() || key_it->second != src_tier_id) { + LOG(WARNING) << "MoveData failed: Key '" << key + << "' is not in the expected source tier " << src_tier_id; + return false; + } + + DataSource source = src_tier->AsDataSource(key); + if (source.ptr == nullptr) { + LOG(WARNING) << "Key '" << key << "' disappeared from tier " + << src_tier_id << " during move operation."; + return false; + } + + if (!dest_tier->Put(key, source)) { + LOG(WARNING) << "Could not move key '" << key + << "' to destination tier " << dest_tier_id + << " (likely full)."; + return false; + } + + if (!src_tier->Delete(key)) { + LOG(ERROR) << "CRITICAL INCONSISTENCY: Moved key '" << key + << "' to tier " << dest_tier_id + << " but failed to delete from source " << src_tier_id + << ". Attempting rollback."; + // Attempt to roll back by deleting the key from the destination tier. + if (!dest_tier->Delete(key)) { + LOG(FATAL) << "Rollback failed. Data for key '" << key + << "' is now duplicated in tiers " << src_tier_id + << " and " << dest_tier_id + << ". Manual intervention required."; + } + // Even if rollback succeeds, the original move operation failed. + return false; + } + + key_it->second = dest_tier_id; + + return true; +} + +bool TieredBackend::DeleteFromTier(const std::string& key, uint64_t tier_id) { + auto it = tiers_.find(tier_id); + if (it == tiers_.end()) { + return false; + } + return it->second->Delete(key); +} + +std::optional TieredBackend::FindKey(const std::string& key) const { + std::shared_lock lock(map_mutex_); + auto it = key_to_tier_map_.find(key); + if (it != key_to_tier_map_.end()) { + return it->second; + } + return std::nullopt; +} + +const CacheTier* TieredBackend::GetTier(uint64_t tier_id) const { + auto it = tiers_.find(tier_id); + if (it != tiers_.end()) { + return it->second.get(); + } + return nullptr; +} + +const DataCopier& TieredBackend::GetDataCopier() const { return *data_copier_; } + +} // namespace mooncake \ No newline at end of file