From 926c63aec89d0836a16475fcb0086deda7dcd107 Mon Sep 17 00:00:00 2001 From: zhangzuo21 <99539591+zhangzuo21@users.noreply.github.com> Date: Wed, 17 Dec 2025 22:33:05 +0800 Subject: [PATCH] [Store] feat: Implement a unified storage interface to simplify integration and extension (#1185) --- mooncake-store/include/file_storage.h | 86 +-- mooncake-store/include/storage_backend.h | 225 ++++++- mooncake-store/include/utils.h | 21 + mooncake-store/src/file_storage.cpp | 385 ++++-------- mooncake-store/src/storage_backend.cpp | 555 +++++++++++++++++- mooncake-store/src/utils.cpp | 5 + mooncake-store/tests/file_storage_test.cpp | 126 ++-- mooncake-store/tests/storage_backend_test.cpp | 317 +++++++++- mooncake-store/tests/utils/common.h | 2 +- 9 files changed, 1300 insertions(+), 422 deletions(-) diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h index 26ed7dac..dbfd5a76 100644 --- a/mooncake-store/include/file_storage.h +++ b/mooncake-store/include/file_storage.h @@ -6,70 +6,6 @@ namespace mooncake { -struct FileStorageConfig { - // Path where data files are stored on disk - std::string storage_filepath = "/data/file_storage"; - - // Size of the local client-side buffer (used for caching or batching) - int64_t local_buffer_size = 1280 * 1024 * 1024; // ~1.2 GB - - // Limits for scanning and iteration operations - int64_t bucket_iterator_keys_limit = - 20000; // Max number of keys returned per Scan call - int64_t bucket_keys_limit = - 500; // Max number of keys allowed in a single bucket - int64_t bucket_size_limit = - 256 * 1024 * 1024; // Max total size of a single bucket (256 MB) - - // Global limits across all buckets - int64_t total_keys_limit = 10'000'000; // Maximum total number of keys - int64_t total_size_limit = - 2ULL * 1024 * 1024 * 1024 * 1024; // Maximum total storage size (2 TB) - - // Interval between heartbeats sent to the control plane (in seconds) - uint32_t heartbeat_interval_seconds = 10; - - // Validates the configuration for correctness and consistency - bool Validate() const; - - /** - * @brief Creates a config instance by reading values from environment - * variables. - * - * Uses default values if environment variables are not set or invalid. - * This is a static factory method for easy configuration loading. - * - * @return FileStorageConfig with values from env or defaults - */ - static FileStorageConfig FromEnvironment(); - - static std::string GetEnvStringOr(const char* name, - const std::string& default_value); - - template - static T GetEnvOr(const char* name, T default_value); -}; - -class BucketIterator { - public: - BucketIterator(std::shared_ptr storage_backend, - int64_t limit); - - tl::expected HandleNext( - const std::function< - ErrorCode(const std::vector& keys, - std::vector& metadatas, - const std::vector& buckets)>& handler); - - tl::expected HasNext(); - - private: - std::shared_ptr storage_backend_; - int64_t limit_; - mutable Mutex mutex_; - int64_t GUARDED_BY(mutex_) next_bucket_ = -1; -}; - class FileStorage { public: FileStorage(std::shared_ptr client, @@ -119,21 +55,6 @@ class FileStorage { tl::expected OffloadObjects( const std::unordered_map& offloading_objects); - /** - * @brief Groups offloading keys into buckets based on size and existence - * checks. - * @param offloading_objects Input map of object keys and their sizes - * (e.g., byte size). - * @param buckets_keys Output parameter: receives a 2D vector where: - * - Each outer element represents a bucket. - * - Each inner vector contains the newly allocated - * object keys within that bucket. - * @return tl::expected indicating operation status. - */ - tl::expected GroupOffloadingKeysByBucket( - const std::unordered_map& offloading_objects, - std::vector>& buckets_keys); - /** * @brief Performs a heartbeat operation for the FileStorage component. * 1. Sends object status (e.g., access frequency, size) to the master via @@ -146,9 +67,6 @@ class FileStorage { tl::expected IsEnableOffloading(); - tl::expected BatchOffload( - const std::vector& keys); - tl::expected BatchLoad( const std::unordered_map& batch_object); @@ -165,12 +83,10 @@ class FileStorage { std::shared_ptr client_; std::string local_rpc_addr_; FileStorageConfig config_; - std::shared_ptr storage_backend_; + std::shared_ptr storage_backend_; std::shared_ptr client_buffer_allocator_; mutable Mutex offloading_mutex_; - std::unordered_map GUARDED_BY(offloading_mutex_) - ungrouped_offloading_objects_; bool GUARDED_BY(offloading_mutex_) enable_offloading_; std::atomic heartbeat_running_; std::thread heartbeat_thread_; diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index 7244889f..3ec16011 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -41,6 +41,99 @@ struct OffloadMetadata { enum class FileMode { Read, Write }; +enum class StorageBackendType { kFilePerKey, kBucket }; + +static constexpr size_t kKB = 1024; +static constexpr size_t kMB = kKB * 1024; +static constexpr size_t kGB = kMB * 1024; + +struct FilePerKeyConfig { + std::string fsdir = "file_per_key_dir"; // Subdirectory name + + bool enable_eviction = true; // Enable eviction for storage + + bool Validate() const; + + static FilePerKeyConfig FromEnvironment(); +}; + +struct BucketBackendConfig { + int64_t bucket_size_limit = + 256 * kMB; // Max total size of a single bucket (256 MB) + + int64_t bucket_keys_limit = 500; // Max number of keys allowed in a single + // bucket, required by bucket backend only + bool Validate() const; + + static BucketBackendConfig FromEnvironment(); +}; + +struct FileStorageConfig { + // type of the storage backend + StorageBackendType storage_backend_type = StorageBackendType::kBucket; + + // Path where data files are stored on disk + std::string storage_filepath = "/data/file_storage"; + + // Size of the local client-side buffer (used for caching or batching) + int64_t local_buffer_size = 1280 * kMB; // ~1.2 GB + + // Limits for scanning and iteration operations + int64_t scanmeta_iterator_keys_limit = + 20000; // Max number of keys returned per Scan call, required by bucket + // backend only + // Global limits across all buckets + int64_t total_keys_limit = 10'000'000; // Maximum total number of keys + int64_t total_size_limit = + 2ULL * 1024 * 1024 * 1024 * 1024; // Maximum total storage size (2 TB) + + // Interval between heartbeats sent to the control plane (in seconds) + uint32_t heartbeat_interval_seconds = 10; + + // Validates the configuration for correctness and consistency + bool Validate() const; + + bool ValidatePath(std::string path) const; + + /** + * @brief Creates a config instance by reading values from environment + * variables. + * + * Uses default values if environment variables are not set or invalid. + * This is a static factory method for easy configuration loading. + * + * @return FileStorageConfig with values from env or defaults + */ + static FileStorageConfig FromEnvironment(); +}; + +class StorageBackendInterface { + public: + StorageBackendInterface(const FileStorageConfig& file_storage_config); + + virtual tl::expected Init() = 0; + + virtual tl::expected BatchOffload( + const std::unordered_map>& batch_object, + std::function& keys, + std::vector& metadatas)> + complete_handler) = 0; + + virtual tl::expected BatchLoad( + const std::unordered_map& batched_slices) = 0; + + virtual tl::expected IsExist(const std::string& key) = 0; + + virtual tl::expected IsEnableOffloading() = 0; + + virtual tl::expected ScanMeta( + const std::function& keys, + std::vector& metadatas)>& handler) = 0; + + FileStorageConfig file_storage_config_; +}; + /** * @class StorageBackend * @brief Implementation of StorageBackend interface using local filesystem @@ -388,9 +481,67 @@ class BucketIdGenerator { std::atomic current_id_; }; -class BucketStorageBackend { +class StorageBackendAdaptor : public StorageBackendInterface { public: - BucketStorageBackend(const std::string& storage_filepath); + StorageBackendAdaptor(const FileStorageConfig& file_storage_config, + const FilePerKeyConfig& file_per_key_config); + + tl::expected Init() override; + + tl::expected BatchOffload( + const std::unordered_map>& batch_object, + std::function& keys, + std::vector& metadatas)> + complete_handler) override; + + tl::expected BatchLoad( + const std::unordered_map& batched_slices) override; + + tl::expected IsExist(const std::string& key) override; + + tl::expected IsEnableOffloading() override; + + tl::expected ScanMeta( + const std::function& keys, + std::vector& metadatas)>& handler) override; + + private: + const FilePerKeyConfig file_per_key_config_; + + std::atomic meta_scanned_{false}; + + std::unique_ptr storage_backend_; + + std::string SanitizeKey(const std::string& key) const; + + std::string ResolvePath(const std::string& key) const; + + static std::string ConcatSlicesToString(const std::vector& slices); + + mutable Mutex mutex_; + + int64_t total_keys GUARDED_BY(mutex_); + + int64_t total_size GUARDED_BY(mutex_); + + struct KVEntry { + std::string key; // K tensor or its storage identifier + std::string value; // V tensor or its storage block + + KVEntry() = default; + + KVEntry(std::string k, std::string v) + : key(std::move(k)), value(std::move(v)) {} + + YLT_REFL(KVEntry, key, value); + }; +}; + +class BucketStorageBackend : public StorageBackendInterface { + public: + BucketStorageBackend(const FileStorageConfig& file_storage_config_, + const BucketBackendConfig& bucket_backend_config_); /** * @brief Offload objects in batches @@ -404,20 +555,20 @@ class BucketStorageBackend { const std::unordered_map>& batch_object, std::function& keys, std::vector& metadatas)> - complete_handler); + complete_handler) override; /** * @brief Retrieves metadata for multiple objects in a single batch * operation. * @param keys A list of object keys to query metadata for. - * @param batche_object_metadata Output parameter that receives the + * @param batch_object_metadata Output parameter that receives the * retrieved metadata. * @return tl::expected indicating operation status. */ tl::expected BatchQuery( const std::vector& keys, std::unordered_map& - batche_object_metadata); + batch_object_metadata); /** * @brief Loads data for multiple objects in a batch operation. @@ -426,7 +577,7 @@ class BucketStorageBackend { * @return tl::expected indicating operation status. */ tl::expected BatchLoad( - const std::unordered_map& batched_slices); + const std::unordered_map& batched_slices) override; /** * @brief Retrieves the list of object keys belonging to a specific bucket. @@ -442,7 +593,7 @@ class BucketStorageBackend { * @brief Initializes the bucket storage backend. * @return tl::expected indicating operation status. */ - tl::expected Init(); + tl::expected Init() override; /** * @brief Checks whether an object with the specified key exists in the @@ -450,7 +601,41 @@ class BucketStorageBackend { * @param key The unique identifier of the object to check for existence. * @return tl::expected indicating operation status. */ - tl::expected IsExist(const std::string& key); + tl::expected IsExist(const std::string& key) override; + + /** + * @brief Scan existing object metadata from storage and report via handler. + * @param handler Callback invoked with a batch of keys and metadatas. + * @return tl::expected indicating operation status. + */ + tl::expected ScanMeta( + const std::function& keys, + std::vector& metadatas)>& handler) override; + + /** + * @brief Checks whether the backend is allowed to continue offloading. + * @return tl::expected + * - On success: true 表示可以继续 offload；false 表示达到上限/不允许继续。 + * - On failure: 返回错误码（例如 IO/内部错误）。 + */ + tl::expected IsEnableOffloading() override; + + /** + * @brief 根据后端 bucket 限制（keys/size）将 offloading_objects 分桶。 + * @param offloading_objects Input map of object keys and their sizes + * (bytes). + * @param buckets_keys Output: bucketized keys; each inner vector is a + * bucket. + * @return tl::expected indicating operation status. + */ + tl::expected AllocateOffloadingBuckets( + const std::unordered_map& offloading_objects, + std::vector>& buckets_keys); + + void ClearUngroupedOffloadingObjects(); + + size_t UngroupedOffloadingObjectsSize() const; /** * @brief Iterate over the metadata of stored objects starting from a @@ -510,6 +695,17 @@ class BucketStorageBackend { tl::expected, ErrorCode> OpenFile( const std::string& path, FileMode mode) const; + tl::expected GroupOffloadingKeysByBucket( + const std::unordered_map& offloading_objects, + std::vector>& buckets_keys); + + tl::expected HandleNext( + const std::function< + ErrorCode(const std::vector& keys, + std::vector& metadatas)>& handler); + + tl::expected HasNext(); + private: std::atomic initialized_{false}; std::optional bucket_id_generator_; @@ -525,11 +721,22 @@ class BucketStorageBackend { * - total_size_: cumulative data size of all stored objects */ mutable SharedMutex mutex_; + mutable Mutex iterator_mutex_; std::string storage_path_; int64_t total_size_ GUARDED_BY(mutex_) = 0; std::unordered_map GUARDED_BY(mutex_) object_bucket_map_; std::map> GUARDED_BY( mutex_) buckets_; + int64_t GUARDED_BY(mutex_) next_bucket_ = -1; + BucketBackendConfig bucket_backend_config_; + + mutable Mutex offloading_mutex_; + std::unordered_map GUARDED_BY(offloading_mutex_) + ungrouped_offloading_objects_; }; -} // namespace mooncake + +tl::expected, ErrorCode> +CreateStorageBackend(const FileStorageConfig& config); + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/utils.h b/mooncake-store/include/utils.h index 496d4bdd..6b8064cf 100644 --- a/mooncake-store/include/utils.h +++ b/mooncake-store/include/utils.h @@ -285,4 +285,25 @@ int getFreeTcpPort(); int64_t time_gen(); +// Helper: Get integer from environment variable, fallback to default +template +T GetEnvOr(const char* name, T default_value) { + const char* env_val = std::getenv(name); + if (!env_val || std::string(env_val).empty()) { + return default_value; + } + try { + long long value = std::stoll(env_val); + // Check range for unsigned types + if constexpr (std::is_same_v) { + if (value < 0 || value > UINT32_MAX) throw std::out_of_range(""); + } + return static_cast(value); + } catch (...) { + return default_value; + } +} + +std::string GetEnvStringOr(const char* name, const std::string& default_value); + } // namespace mooncake diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 075c5a1f..3e6cbed7 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -4,53 +4,34 @@ #include #include +#include "storage_backend.h" #include "utils.h" namespace mooncake { -// Helper: Get integer from environment variable, fallback to default -template -T FileStorageConfig::GetEnvOr(const char* name, T default_value) { - const char* env_val = std::getenv(name); - if (!env_val || std::string(env_val).empty()) { - return default_value; - } - try { - long long value = std::stoll(env_val); - // Check range for unsigned types - if constexpr (std::is_same_v) { - if (value < 0 || value > UINT32_MAX) throw std::out_of_range(""); - } - return static_cast(value); - } catch (...) { - return default_value; - } -} - -// Helper: Get string from environment variable, fallback to default -std::string FileStorageConfig::GetEnvStringOr( - const char* name, const std::string& default_value) { - const char* env_val = std::getenv(name); - return env_val ? std::string(env_val) : default_value; -} - FileStorageConfig FileStorageConfig::FromEnvironment() { FileStorageConfig config; + auto storage_backend_descriptor = + GetEnvStringOr("MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR", + "bucket_storage_backend"); + + if (storage_backend_descriptor == "bucket_storage_backend") { + config.storage_backend_type = StorageBackendType::kBucket; + } else if (storage_backend_descriptor == "file_per_key_storage_backend") { + config.storage_backend_type = StorageBackendType::kFilePerKey; + } else { + LOG(ERROR) << "Unknown storage backend."; + } + config.storage_filepath = GetEnvStringOr( "MOONCAKE_OFFLOAD_FILE_STORAGE_PATH", config.storage_filepath); config.local_buffer_size = GetEnvOr( "MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES", config.local_buffer_size); - config.bucket_iterator_keys_limit = - GetEnvOr("MOONCAKE_OFFLOAD_BUCKET_ITERATOR_KEYS_LIMIT", - config.bucket_iterator_keys_limit); - - config.bucket_keys_limit = GetEnvOr( - "MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", config.bucket_keys_limit); - - config.bucket_size_limit = GetEnvOr( - "MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", config.bucket_size_limit); + config.scanmeta_iterator_keys_limit = + GetEnvOr("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT", + config.scanmeta_iterator_keys_limit); config.total_keys_limit = GetEnvOr( "MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT", config.total_keys_limit); @@ -65,12 +46,11 @@ FileStorageConfig FileStorageConfig::FromEnvironment() { return config; } -bool FileStorageConfig::Validate() const { - if (storage_filepath.empty()) { +bool FileStorageConfig::ValidatePath(std::string path) const { + if (path.empty()) { LOG(ERROR) << "FileStorageConfig: storage_filepath is invalid"; return false; } - const std::string& path = storage_filepath; namespace fs = std::filesystem; // 1. Must be an absolute path if (!fs::path(path).is_absolute()) { @@ -123,12 +103,12 @@ bool FileStorageConfig::Validate() const { return false; } } - if (bucket_keys_limit <= 0) { - LOG(ERROR) << "FileStorageConfig: bucket_keys_limit must > 0"; - return false; - } - if (bucket_size_limit <= 0) { - LOG(ERROR) << "FileStorageConfig: bucket_size_limit must > 0"; + + return true; +} + +bool FileStorageConfig::Validate() const { + if (!ValidatePath(storage_filepath)) { return false; } if (total_keys_limit <= 0) { @@ -152,13 +132,18 @@ FileStorage::FileStorage(std::shared_ptr client, : client_(client), local_rpc_addr_(local_rpc_addr), config_(config), - storage_backend_( - std::make_shared(config.storage_filepath)), client_buffer_allocator_( ClientBufferAllocator::create(config.local_buffer_size, "")) { if (!config.Validate()) { throw std::invalid_argument("Invalid FileStorage configuration"); } + + auto create_storage_backend_result = CreateStorageBackend(config_); + if (!create_storage_backend_result) { + LOG(ERROR) << "Failed to create storage backend"; + } + + storage_backend_ = create_storage_backend_result.value(); } FileStorage::~FileStorage() { @@ -200,39 +185,26 @@ tl::expected FileStorage::Init() { } } - BucketIterator bucket_iterator(storage_backend_, - config_.bucket_iterator_keys_limit); - while (true) { - auto has_next_res = bucket_iterator.HasNext(); - if (!has_next_res) { - LOG(ERROR) << "Failed to check for next bucket: " - << has_next_res.error(); - return tl::make_unexpected(has_next_res.error()); - } - if (!has_next_res.value()) { - break; - } - auto add_all_object_res = bucket_iterator.HandleNext( - [this](const std::vector& keys, - std::vector& metadatas, - const std::vector&) { - for (auto& metadata : metadatas) { - metadata.transport_endpoint = local_rpc_addr_; - } - auto add_object_result = - client_->NotifyOffloadSuccess(keys, metadatas); - if (!add_object_result) { - LOG(ERROR) << "Failed to add object to master: " - << add_object_result.error(); - return add_object_result.error(); - } - return ErrorCode::OK; - }); - if (!add_all_object_res) { - LOG(ERROR) << "Failed to add all object to master: " - << add_all_object_res.error(); - return add_all_object_res; - } + auto scan_meta_result = storage_backend_->ScanMeta( + [this](const std::vector& keys, + std::vector& metadatas) { + for (auto& metadata : metadatas) { + metadata.transport_endpoint = local_rpc_addr_; + } + auto add_object_result = + client_->NotifyOffloadSuccess(keys, metadatas); + if (!add_object_result) { + LOG(ERROR) << "Failed to add object to master: " + << add_object_result.error(); + return add_object_result.error(); + } + return ErrorCode::OK; + }); + + if (!scan_meta_result) { + LOG(ERROR) << "Failed to scan meta and send to master: " + << scan_meta_result.error(); + return scan_meta_result; } heartbeat_running_.store(true); @@ -286,33 +258,61 @@ tl::expected FileStorage::BatchGet( tl::expected FileStorage::OffloadObjects( const std::unordered_map& offloading_objects) { std::vector> buckets_keys; - auto allocate_objects_result = - GroupOffloadingKeysByBucket(offloading_objects, buckets_keys); - if (!allocate_objects_result) { - LOG(ERROR) << "GroupKeysByBucket failed with error: " - << allocate_objects_result.error(); - return allocate_objects_result; + if (auto bucket_backend = + std::dynamic_pointer_cast(storage_backend_)) { + auto allocate_res = bucket_backend->AllocateOffloadingBuckets( + offloading_objects, buckets_keys); + if (!allocate_res) { + LOG(ERROR) << "AllocateOffloadingBuckets failed with error: " + << allocate_res.error(); + return allocate_res; + } + } else { + std::vector keys; + keys.reserve(offloading_objects.size()); + for (const auto& it : offloading_objects) { + keys.emplace_back(it.first); + } + buckets_keys.emplace_back(std::move(keys)); } - for (const auto& keys : buckets_keys) { - auto enable_offloading_result = IsEnableOffloading(); - if (!enable_offloading_result) { - LOG(ERROR) << "Get is enable offloading failed with error: " - << enable_offloading_result.error(); - return tl::make_unexpected(enable_offloading_result.error()); + + auto complete_handler = + [this](const std::vector& keys, + std::vector& metadatas) -> ErrorCode { + VLOG(1) << "Success to store objects, keys count: " << keys.size(); + for (auto& metadata : metadatas) { + metadata.transport_endpoint = local_rpc_addr_; } - if (!enable_offloading_result.value()) { - LOG(WARNING) << "Unable to be persisted"; - MutexLocker locker(&offloading_mutex_); - ungrouped_offloading_objects_.clear(); - enable_offloading_ = false; - return tl::make_unexpected(ErrorCode::KEYS_ULTRA_LIMIT); - } - auto result = BatchOffload(keys); + auto result = client_->NotifyOffloadSuccess(keys, metadatas); if (!result) { - LOG(ERROR) << "Failed to store objects with error: " + LOG(ERROR) << "NotifyOffloadSuccess failed with error: " << result.error(); - if (result.error() != ErrorCode::INVALID_READ) { - return result; + return result.error(); + } + return ErrorCode::OK; + }; + + for (const auto& keys : buckets_keys) { + std::unordered_map> batch_object; + auto query_result = BatchQuerySegmentSlices(keys, batch_object); + if (!query_result) { + LOG(ERROR) << "BatchQuerySlices failed with error: " + << query_result.error(); + continue; + } + + auto offload_res = + storage_backend_->BatchOffload(batch_object, complete_handler); + if (!offload_res) { + LOG(ERROR) << "Failed to store objects with error: " + << offload_res.error(); + if (offload_res.error() == ErrorCode::KEYS_ULTRA_LIMIT) { + MutexLocker locker(&offloading_mutex_); + enable_offloading_ = false; + return tl::make_unexpected(offload_res.error()); + } + if (offload_res.error() != ErrorCode::INVALID_READ) { + return tl::make_unexpected(offload_res.error()); } } } @@ -320,26 +320,14 @@ tl::expected FileStorage::OffloadObjects( } tl::expected FileStorage::IsEnableOffloading() { - auto store_metadata_result = storage_backend_->GetStoreMetadata(); - if (!store_metadata_result) { - LOG(ERROR) << "Failed to get store metadata: " - << store_metadata_result.error(); - return tl::make_unexpected(store_metadata_result.error()); + auto is_enable_offloading_result = storage_backend_->IsEnableOffloading(); + if (!is_enable_offloading_result) { + LOG(ERROR) << "Failed to get enabling offload: " + << is_enable_offloading_result.error(); + return tl::make_unexpected(is_enable_offloading_result.error()); } - const auto& store_metadata = store_metadata_result.value(); - auto enable_offloading = - store_metadata.total_keys + config_.bucket_keys_limit <= - config_.total_keys_limit && - store_metadata.total_size + config_.bucket_size_limit <= - config_.total_size_limit; - VLOG(1) << (enable_offloading ? "Enable" : "Unable") - << " offloading,total keys: " << store_metadata.total_keys - << ", bucket keys limit: " << config_.bucket_keys_limit - << ", total keys limit: " << config_.total_keys_limit - << ", total size: " << store_metadata.total_size - << ", bucket size limit: " << config_.bucket_size_limit - << ", total size limit: " << config_.total_size_limit; + auto enable_offloading = is_enable_offloading_result.value(); return enable_offloading; } @@ -377,45 +365,6 @@ tl::expected FileStorage::Heartbeat() { return {}; } -tl::expected FileStorage::BatchOffload( - const std::vector& keys) { - auto start_time = std::chrono::steady_clock::now(); - std::unordered_map> batch_object; - auto query_result = BatchQuerySegmentSlices(keys, batch_object); - if (!query_result) { - LOG(ERROR) << "BatchQuerySlices failed with error: " - << query_result.error(); - return tl::make_unexpected(ErrorCode::INVALID_READ); - } - auto result = storage_backend_->BatchOffload( - batch_object, [this](const std::vector& keys, - std::vector& metadatas) { - VLOG(1) << "Success to store objects, keys count: " << keys.size(); - for (auto& metadata : metadatas) { - metadata.transport_endpoint = local_rpc_addr_; - } - auto result = client_->NotifyOffloadSuccess(keys, metadatas); - if (!result) { - LOG(ERROR) << "NotifyOffloadSuccess failed with error: " - << result.error(); - return result.error(); - } - return ErrorCode::OK; - }); - auto end_time = std::chrono::steady_clock::now(); - auto elapsed_time = std::chrono::duration_cast( - end_time - start_time) - .count(); - VLOG(1) << "Time taken for BatchStore: " << elapsed_time - << "us,with keys count: " << keys.size(); - if (!result) { - LOG(ERROR) << "Batch store object failed, err_code = " - << result.error(); - return tl::make_unexpected(result.error()); - } - return {}; -} - tl::expected FileStorage::BatchLoad( const std::unordered_map& batch_object) { auto start_time = std::chrono::steady_clock::now(); @@ -501,122 +450,4 @@ tl::expected FileStorage::AllocateBatch( return result; } -tl::expected FileStorage::GroupOffloadingKeysByBucket( - const std::unordered_map& offloading_objects, - std::vector>& buckets_keys) { - MutexLocker locker(&offloading_mutex_); - auto it = offloading_objects.cbegin(); - int64_t residue_count = - offloading_objects.size() + ungrouped_offloading_objects_.size(); - int64_t total_count = - offloading_objects.size() + ungrouped_offloading_objects_.size(); - while (it != offloading_objects.cend()) { - std::vector bucket_keys; - std::unordered_map bucket_objects; - int64_t bucket_data_size = 0; - // Process previously ungrouped objects first - if (!ungrouped_offloading_objects_.empty()) { - for (const auto& ungrouped_objects_it : - ungrouped_offloading_objects_) { - bucket_data_size += ungrouped_objects_it.second; - bucket_keys.push_back(ungrouped_objects_it.first); - bucket_objects.emplace(ungrouped_objects_it.first, - ungrouped_objects_it.second); - } - VLOG(1) << "Ungrouped offloading objects have been processed and " - "cleared; count=" - << ungrouped_offloading_objects_.size(); - ungrouped_offloading_objects_.clear(); - } - - // Fill the rest of the bucket with new offloading objects - for (int64_t i = static_cast(bucket_keys.size()); - i < config_.bucket_keys_limit; ++i) { - if (it == offloading_objects.cend()) { - // No more objects to add — move current batch to ungrouped pool - for (const auto& bucket_object : bucket_objects) { - ungrouped_offloading_objects_.emplace(bucket_object.first, - bucket_object.second); - } - VLOG(1) << "Add offloading objects to ungrouped pool. " - << "Total ungrouped count: " - << ungrouped_offloading_objects_.size(); - return {}; - } - if (it->second > config_.bucket_size_limit) { - LOG(ERROR) << "Object size exceeds bucket size limit: " - << "key=" << it->first - << ", object_size=" << it->second - << ", limit=" << config_.bucket_size_limit; - ++it; - continue; - } - auto is_exist_result = storage_backend_->IsExist(it->first); - if (!is_exist_result) { - LOG(ERROR) << "Failed to check existence in storage backend: " - << "key=" << it->first - << ", error=" << is_exist_result.error(); - } - if (is_exist_result.value()) { - ++it; - continue; - } - if (bucket_data_size + it->second > config_.bucket_size_limit) { - break; - } - bucket_data_size += it->second; - bucket_keys.push_back(it->first); - bucket_objects.emplace(it->first, it->second); - ++it; - if (bucket_data_size == config_.bucket_size_limit) { - break; - } - } - auto bucket_keys_count = bucket_keys.size(); - // Finalize current bucket - residue_count -= bucket_keys_count; - buckets_keys.push_back(std::move(bucket_keys)); - VLOG(1) << "Group objects with total object count: " << total_count - << ", current bucket object count: " << bucket_keys_count - << ", current bucket data size: " << bucket_data_size - << ", grouped bucket count: " << buckets_keys.size() - << ", residue object count: " << residue_count; - } - return {}; -} - -BucketIterator::BucketIterator( - std::shared_ptr storage_backend, int64_t limit) - : storage_backend_(storage_backend), limit_(limit) {}; - -tl::expected BucketIterator::HandleNext( - const std::function& keys, - std::vector& metadatas, - const std::vector& buckets)>& - handler) { - MutexLocker locker(&mutex_); - std::vector keys; - std::vector metadatas; - std::vector buckets; - auto key_iterator_result = storage_backend_->BucketScan( - next_bucket_, keys, metadatas, buckets, limit_); - if (!key_iterator_result) { - LOG(ERROR) << "Bucket scan failed, error : " - << key_iterator_result.error(); - return tl::make_unexpected(key_iterator_result.error()); - } - auto handle_result = handler(keys, metadatas, buckets); - if (handle_result != ErrorCode::OK) { - LOG(ERROR) << "Key iterator failed, error : " << handle_result; - return tl::make_unexpected(handle_result); - } - next_bucket_ = key_iterator_result.value(); - return {}; -} - -tl::expected BucketIterator::HasNext() { - MutexLocker locker(&mutex_); - return next_bucket_ != 0; -} - } // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 820b108f..b3ad3688 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -20,6 +20,53 @@ namespace mooncake { +bool FilePerKeyConfig::Validate() const { + if (fsdir.empty()) { + LOG(ERROR) << "FilePerKeyConfig: fsdir is invalid"; + return false; + } + return true; +} + +bool BucketBackendConfig::Validate() const { + if (bucket_keys_limit <= 0) { + LOG(ERROR) << "BucketBackendConfig: bucket_keys_limit must > 0"; + return false; + } + if (bucket_size_limit <= 0) { + LOG(ERROR) << "BucketBackendConfig: bucket_size_limit must > 0"; + return false; + } + return true; +} + +FilePerKeyConfig FilePerKeyConfig::FromEnvironment() { + FilePerKeyConfig config; + + config.fsdir = GetEnvStringOr("MOONCAKE_OFFLOAD_FSDIR", config.fsdir); + + config.enable_eviction = + GetEnvOr("ENABLE_EVICTION", config.enable_eviction); + + return config; +} + +BucketBackendConfig BucketBackendConfig::FromEnvironment() { + BucketBackendConfig config; + + config.bucket_keys_limit = GetEnvOr( + "MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", config.bucket_keys_limit); + + config.bucket_size_limit = GetEnvOr( + "MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", config.bucket_size_limit); + + return config; +} + +StorageBackendInterface::StorageBackendInterface( + const FileStorageConfig& config) + : file_storage_config_(config) {} + std::string StorageBackend::GetActualFsdir() const { std::string actual_fsdir = fsdir_; if (actual_fsdir.rfind("moon_", 0) == 0) { @@ -170,8 +217,8 @@ tl::expected StorageBackend::Init(uint64_t quota_bytes = 0) { std::unique_lock lock(space_mutex_); RecalculateAvailableSpace(); - LOG(INFO) << "Init: " << "Quota: " << total_space_ - << ", Used: " << used_space_ + LOG(INFO) << "Init: " + << "Quota: " << total_space_ << ", Used: " << used_space_ << ", Available: " << available_space_; } @@ -866,6 +913,279 @@ void StorageBackend::ReleaseSpace(uint64_t size_to_release) { } } +StorageBackendAdaptor::StorageBackendAdaptor( + const FileStorageConfig& file_storage_config, + const FilePerKeyConfig& file_per_key_config) + : StorageBackendInterface(file_storage_config), + file_per_key_config_(file_per_key_config), + total_keys(0), + total_size(0) {} + +tl::expected StorageBackendAdaptor::Init() { + std::string storage_root = + file_storage_config_.storage_filepath + file_per_key_config_.fsdir; + + storage_backend_ = std::make_unique( + file_storage_config_.storage_filepath, file_per_key_config_.fsdir, + file_per_key_config_.enable_eviction); + auto init_result = storage_backend_->Init(); + if (!init_result) { + LOG(ERROR) << "Failed to init storage backend"; + return init_result; + } + return {}; +} + +std::string StorageBackendAdaptor::SanitizeKey(const std::string& key) const { + // Set of invalid filesystem characters to be replaced + constexpr std::string_view kInvalidChars = "/\\:*?\"<>|"; + std::string sanitized_key; + sanitized_key.reserve(key.size()); + + for (char c : key) { + // Replace invalid characters with underscore + sanitized_key.push_back( + kInvalidChars.find(c) != std::string_view::npos ? '_' : c); + } + return sanitized_key; +} + +std::string StorageBackendAdaptor::ResolvePath(const std::string& key) const { + // Compute hash of the key + size_t hash = std::hash{}(key); + + // Use low 8 bits to create 2-level directory structure (e.g. "a1/b2") + char dir1 = + static_cast('a' + (hash & 0x0F)); // Lower 4 bits -> 16 dirs + char dir2 = static_cast( + 'a' + ((hash >> 4) & 0x0F)); // Next 4 bits -> 16 subdirs + + // Safely construct path using std::filesystem + namespace fs = std::filesystem; + fs::path dir_path = fs::path(std::string(1, dir1)) / std::string(1, dir2); + + // Combine directory path with sanitized filename + fs::path full_path = fs::path(file_storage_config_.storage_filepath) / + file_per_key_config_.fsdir / dir_path / + SanitizeKey(key); + + return full_path.lexically_normal().string(); +} + +std::string StorageBackendAdaptor::ConcatSlicesToString( + const std::vector& slices) { + size_t total = 0; + for (const auto& s : slices) { + if (s.size == 0) continue; + total += s.size; + } + + std::string out; + out.reserve(total); + + for (const auto& s : slices) { + if (s.size == 0) continue; + out.append(reinterpret_cast(s.ptr), s.size); + } + return out; +} + +tl::expected StorageBackendAdaptor::BatchOffload( + const std::unordered_map>& batch_object, + std::function& keys, + std::vector& metadatas)> + complete_handler) { + if (batch_object.empty()) { + LOG(ERROR) << "batch object is empty"; + return tl::make_unexpected(ErrorCode::INVALID_KEY); + } + + auto enable_offloading_res = IsEnableOffloading(); + if (!enable_offloading_res) { + return tl::make_unexpected(enable_offloading_res.error()); + } + std::vector metadatas; + std::vector keys; + metadatas.reserve(batch_object.size()); + keys.reserve(batch_object.size()); + for (auto& object : batch_object) { + KVEntry kv; + kv.key = object.first; + auto value = object.second; + + auto path = ResolvePath(kv.key); + kv.value = ConcatSlicesToString(value); + + std::string kv_buf; + struct_pb::to_pb(kv, kv_buf); + auto store_result = storage_backend_->StoreObject(path, kv_buf); + if (!store_result) { + LOG(ERROR) << "Failed to store object"; + return tl::make_unexpected(store_result.error()); + } + + { + MutexLocker lock(&mutex_); + total_keys++; + total_size += kv_buf.size(); + } + + metadatas.emplace_back( + StorageObjectMetadata{-1, 0, static_cast(kv.key.size()), + static_cast(kv.value.size())}); + keys.emplace_back(kv.key); + } + + if (complete_handler != nullptr) { + auto error_code = complete_handler(keys, metadatas); + if (error_code != ErrorCode::OK) { + LOG(ERROR) << "Complete handler failed: " << error_code; + return tl::make_unexpected(error_code); + } + } + + return 0; +} + +tl::expected StorageBackendAdaptor::IsExist( + const std::string& key) { + auto path = ResolvePath(key); + namespace fs = std::filesystem; + return fs::exists(path); +} + +tl::expected StorageBackendAdaptor::BatchLoad( + const std::unordered_map& batched_slices) { + for (const auto& [key, slice] : batched_slices) { + KVEntry kv; + kv.key = key; + auto path = ResolvePath(kv.key); + + kv.value.resize(slice.size); + + std::string kv_buf; + struct_pb::to_pb(kv, kv_buf); + + auto r = storage_backend_->LoadObject(path, kv_buf, kv_buf.size()); + if (!r) { + LOG(ERROR) << "Failed to load from file"; + return tl::make_unexpected(r.error()); + } + + struct_pb::from_pb(kv, kv_buf); + + if (!kv.value.empty()) { + std::memcpy(slice.ptr, kv.value.data(), kv.value.size()); + } + } + return {}; +} + +tl::expected StorageBackendAdaptor::IsEnableOffloading() { + if (storage_backend_->enable_eviction_) { + return true; + } + + if (!meta_scanned_.load(std::memory_order_acquire)) { + LOG(ERROR) << "Metadata has not been loaded yet"; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + + MutexLocker lock(&mutex_); + + auto is_enable_offloading = + total_keys <= file_storage_config_.total_keys_limit && + total_size <= file_storage_config_.total_size_limit; + + return is_enable_offloading; +} + +tl::expected StorageBackendAdaptor::ScanMeta( + const std::function< + ErrorCode(const std::vector& keys, + std::vector& metadatas)>& handler) { + namespace fs = std::filesystem; + + fs::path root = fs::path(file_storage_config_.storage_filepath) / + file_per_key_config_.fsdir; + if (!fs::exists(root)) { + meta_scanned_.store(true, std::memory_order_acquire); + return {}; + } + + std::vector keys; + std::vector metas; + + auto flush = [&]() -> tl::expected { + if (keys.empty()) return {}; + auto ec = handler(keys, metas); + if (ec != ErrorCode::OK) return tl::make_unexpected(ec); + keys.clear(); + metas.clear(); + return {}; + }; + + MutexLocker lock(&mutex_); + + std::error_code ec_root; + for (auto it1 = fs::directory_iterator(root, ec_root); + !ec_root && it1 != fs::directory_iterator(); it1.increment(ec_root)) { + if (ec_root) break; + if (!it1->is_directory(ec_root) || ec_root) continue; + + const auto& d1 = it1->path(); + + std::error_code ec_d1; + for (auto it2 = fs::directory_iterator(d1, ec_d1); + !ec_d1 && it2 != fs::directory_iterator(); it2.increment(ec_d1)) { + if (ec_d1) break; + if (!it2->is_directory(ec_d1) || ec_d1) continue; + + const auto& leaf = it2->path(); + + std::error_code ec_leaf; + for (auto it = fs::directory_iterator(leaf, ec_leaf); + !ec_leaf && it != fs::directory_iterator(); + it.increment(ec_leaf)) { + if (ec_leaf) break; + const auto& p = it->path(); + if (!it->is_regular_file(ec_leaf) || ec_leaf) continue; + + uintmax_t sz = fs::file_size(p, ec_leaf); + if (ec_leaf) continue; + + std::string buf; + auto r = + storage_backend_->LoadObject(p.string(), buf, (int64_t)sz); + if (!r) continue; + + KVEntry kv; + struct_pb::from_pb(kv, buf); + + total_keys++; + total_size += buf.size(); + + keys.emplace_back(std::move(kv.key)); + metas.emplace_back(StorageObjectMetadata{ + -1, 0, (int64_t)keys.back().size(), + static_cast(kv.value.size())}); + + if ((int64_t)keys.size() >= + file_storage_config_.scanmeta_iterator_keys_limit) { + auto fr = flush(); + if (!fr) return fr; + } + } + + auto fr = flush(); + if (!fr) return fr; + } + } + + meta_scanned_.store(true, std::memory_order_acquire); + return {}; +} + BucketIdGenerator::BucketIdGenerator(int64_t start) { if (start <= 0) { auto cur_time_stamp = time_gen(); @@ -883,8 +1203,12 @@ int64_t BucketIdGenerator::CurrentId() { return current_id_.load(std::memory_order_relaxed); } -BucketStorageBackend::BucketStorageBackend(const std::string& storage_path) - : storage_path_(storage_path) {} +BucketStorageBackend::BucketStorageBackend( + const FileStorageConfig& file_storage_config_, + const BucketBackendConfig& bucket_backend_config_) + : StorageBackendInterface(file_storage_config_), + storage_path_(file_storage_config_.storage_filepath), + bucket_backend_config_(bucket_backend_config_) {} tl::expected BucketStorageBackend::BatchOffload( const std::unordered_map>& batch_object, @@ -900,6 +1224,14 @@ tl::expected BucketStorageBackend::BatchOffload( LOG(ERROR) << "batch object is empty"; return tl::make_unexpected(ErrorCode::INVALID_KEY); } + + auto enable_offloading_res = IsEnableOffloading(); + if (!enable_offloading_res) { + return tl::make_unexpected(enable_offloading_res.error()); + } + if (!enable_offloading_res.value()) { + return tl::make_unexpected(ErrorCode::KEYS_ULTRA_LIMIT); + } auto bucket_id = bucket_id_generator_->NextId(); std::vector iovs; std::vector metadatas; @@ -1139,7 +1471,8 @@ tl::expected BucketStorageBackend::Init() { orphaned_space_freed += file_size; LOG(WARNING) << "Removed orphaned bucket file (no metadata): " << entry.path().string() << " (size: " << file_size - << " bytes, " << "bucket_id: " << bucket_id << ")"; + << " bytes, " + << "bucket_id: " << bucket_id << ")"; } else if (cleanup_ec) { LOG(ERROR) << "Failed to remove orphaned bucket file: " << entry.path().string() @@ -1183,6 +1516,46 @@ tl::expected BucketStorageBackend::IsExist( return false; } +tl::expected BucketStorageBackend::IsEnableOffloading() { + auto store_metadata_result = GetStoreMetadata(); + if (!store_metadata_result) { + LOG(ERROR) << "Failed to get store metadata: " + << store_metadata_result.error(); + return tl::make_unexpected(store_metadata_result.error()); + } + const auto& store_metadata = store_metadata_result.value(); + auto enable_offloading = + store_metadata.total_keys + bucket_backend_config_.bucket_keys_limit <= + file_storage_config_.total_keys_limit && + store_metadata.total_size + bucket_backend_config_.bucket_size_limit <= + file_storage_config_.total_size_limit; + return enable_offloading; +} + +tl::expected BucketStorageBackend::ScanMeta( + const std::function< + ErrorCode(const std::vector& keys, + std::vector& metadatas)>& handler) { + while (true) { + auto has_next_res = HasNext(); + if (!has_next_res) { + LOG(ERROR) << "Failed to check for next bucket: " + << has_next_res.error(); + return tl::make_unexpected(has_next_res.error()); + } + if (!has_next_res.value()) { + break; + } + auto add_all_object_res = HandleNext(handler); + if (!add_all_object_res) { + LOG(ERROR) << "Failed to add all object to master: " + << add_all_object_res.error(); + return add_all_object_res; + } + } + return {}; +} + tl::expected BucketStorageBackend::BucketScan( int64_t bucket_id, std::vector& keys, std::vector& metadatas, @@ -1191,8 +1564,8 @@ tl::expected BucketStorageBackend::BucketScan( auto bucket_it = buckets_.lower_bound(bucket_id); for (; bucket_it != buckets_.end(); ++bucket_it) { if (static_cast(bucket_it->second->keys.size()) > limit) { - LOG(ERROR) << "Bucket key count exceeds limit: " << "bucket_id=" - << bucket_it->first + LOG(ERROR) << "Bucket key count exceeds limit: " + << "bucket_id=" << bucket_it->first << ", current_size=" << bucket_it->second->keys.size() << ", limit=" << limit; return tl::make_unexpected(ErrorCode::KEYS_EXCEED_BUCKET_LIMIT); @@ -1220,6 +1593,115 @@ BucketStorageBackend::GetStoreMetadata() { return metadata; } +tl::expected BucketStorageBackend::AllocateOffloadingBuckets( + const std::unordered_map& offloading_objects, + std::vector>& buckets_keys) { + return GroupOffloadingKeysByBucket(offloading_objects, buckets_keys); +} + +void BucketStorageBackend::ClearUngroupedOffloadingObjects() { + MutexLocker locker(&offloading_mutex_); + ungrouped_offloading_objects_.clear(); +} + +size_t BucketStorageBackend::UngroupedOffloadingObjectsSize() const { + MutexLocker locker(&offloading_mutex_); + return ungrouped_offloading_objects_.size(); +} + +tl::expected BucketStorageBackend::GroupOffloadingKeysByBucket( + const std::unordered_map& offloading_objects, + std::vector>& buckets_keys) { + MutexLocker offloading_locker(&offloading_mutex_); + auto& ungrouped_offloading_objects = ungrouped_offloading_objects_; + auto it = offloading_objects.cbegin(); + int64_t residue_count = static_cast( + offloading_objects.size() + ungrouped_offloading_objects.size()); + int64_t total_count = residue_count; + + auto is_exist_func = + [this](const std::string& key) -> tl::expected { + return IsExist(key); + }; + + while (it != offloading_objects.cend()) { + std::vector bucket_keys; + std::unordered_map bucket_objects; + int64_t bucket_data_size = 0; + + if (!ungrouped_offloading_objects.empty()) { + for (const auto& ungrouped_it : ungrouped_offloading_objects) { + bucket_data_size += ungrouped_it.second; + bucket_keys.push_back(ungrouped_it.first); + bucket_objects.emplace(ungrouped_it.first, ungrouped_it.second); + } + VLOG(1) << "Ungrouped offloading objects have been processed and " + "cleared; count=" + << ungrouped_offloading_objects.size(); + ungrouped_offloading_objects.clear(); + } + + for (int64_t i = static_cast(bucket_keys.size()); + i < bucket_backend_config_.bucket_keys_limit; ++i) { + if (it == offloading_objects.cend()) { + for (const auto& bucket_object : bucket_objects) { + ungrouped_offloading_objects.emplace(bucket_object.first, + bucket_object.second); + } + VLOG(1) << "Add offloading objects to ungrouped pool. " + << "Total ungrouped count: " + << ungrouped_offloading_objects.size(); + return {}; + } + + if (it->second > bucket_backend_config_.bucket_size_limit) { + LOG(ERROR) << "Object size exceeds bucket size limit: " + << "key=" << it->first + << ", object_size=" << it->second << ", limit=" + << bucket_backend_config_.bucket_size_limit; + ++it; + continue; + } + + auto is_exist_result = is_exist_func(it->first); + if (!is_exist_result) { + LOG(ERROR) << "Failed to check existence in storage backend: " + << "key=" << it->first + << ", error=" << is_exist_result.error(); + } + if (is_exist_result && is_exist_result.value()) { + ++it; + continue; + } + + if (bucket_data_size + it->second > + bucket_backend_config_.bucket_size_limit) { + break; + } + + bucket_data_size += it->second; + bucket_keys.push_back(it->first); + bucket_objects.emplace(it->first, it->second); + ++it; + + if (bucket_data_size == bucket_backend_config_.bucket_size_limit) { + break; + } + } + + auto bucket_keys_count = static_cast(bucket_keys.size()); + residue_count -= bucket_keys_count; + buckets_keys.push_back(std::move(bucket_keys)); + VLOG(1) << "Group objects with total object count: " << total_count + << ", current bucket object count: " << bucket_keys_count + << ", current bucket data size: " << bucket_data_size + << ", grouped bucket count: " << buckets_keys.size() + << ", residue object count: " << residue_count; + } + + return {}; +} + tl::expected, ErrorCode> BucketStorageBackend::BuildBucket( int64_t bucket_id, @@ -1463,4 +1945,63 @@ BucketStorageBackend::OpenFile(const std::string& path, FileMode mode) const { return std::make_unique(path, fd); } +tl::expected BucketStorageBackend::HandleNext( + const std::function< + ErrorCode(const std::vector& keys, + std::vector& metadatas)>& handler) { + MutexLocker locker(&iterator_mutex_); + std::vector keys; + std::vector metadatas; + std::vector buckets; + auto key_iterator_result = + BucketScan(next_bucket_, keys, metadatas, buckets, + file_storage_config_.scanmeta_iterator_keys_limit); + if (!key_iterator_result) { + LOG(ERROR) << "Bucket scan failed, error : " + << key_iterator_result.error(); + return tl::make_unexpected(key_iterator_result.error()); + } + auto handle_result = handler(keys, metadatas); + if (handle_result != ErrorCode::OK) { + LOG(ERROR) << "Key iterator failed, error : " << handle_result; + return tl::make_unexpected(handle_result); + } + next_bucket_ = key_iterator_result.value(); + return {}; +} + +tl::expected BucketStorageBackend::HasNext() { + MutexLocker locker(&iterator_mutex_); + return next_bucket_ != 0; +} + +tl::expected, ErrorCode> +CreateStorageBackend(const FileStorageConfig& config) { + switch (config.storage_backend_type) { + case StorageBackendType::kBucket: { + auto bucket_backend_config = BucketBackendConfig::FromEnvironment(); + if (!bucket_backend_config.Validate()) { + throw std::invalid_argument( + "Invalid StorageBackend configuration"); + } + return std::make_shared( + config, bucket_backend_config); + } + case StorageBackendType::kFilePerKey: { + auto file_per_key_backend_config = + FilePerKeyConfig::FromEnvironment(); + if (!file_per_key_backend_config.Validate()) { + throw std::invalid_argument( + "Invalid StorageBackend configuration"); + } + return std::make_shared( + config, file_per_key_backend_config); + } + default: { + LOG(FATAL) << "Unsupported backend type"; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + } +} + } // namespace mooncake diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp index 449bed19..4a5808cb 100644 --- a/mooncake-store/src/utils.cpp +++ b/mooncake-store/src/utils.cpp @@ -173,4 +173,9 @@ int64_t time_gen() { .count(); } +std::string GetEnvStringOr(const char *name, const std::string &default_value) { + const char *env_val = std::getenv(name); + return env_val ? std::string(env_val) : default_value; +} + } // namespace mooncake diff --git a/mooncake-store/tests/file_storage_test.cpp b/mooncake-store/tests/file_storage_test.cpp index cf8e5035..eda6bd4b 100644 --- a/mooncake-store/tests/file_storage_test.cpp +++ b/mooncake-store/tests/file_storage_test.cpp @@ -1,6 +1,7 @@ #include #include +#include #include #include "allocator.h" @@ -70,13 +71,22 @@ class FileStorageTest : public ::testing::Test { FileStorage& fileStorage, const std::unordered_map& offloading_objects, std::vector>& buckets_keys) { - return fileStorage.GroupOffloadingKeysByBucket(offloading_objects, - buckets_keys); + auto bucket_backend = std::dynamic_pointer_cast( + fileStorage.storage_backend_); + if (!bucket_backend) { + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + return bucket_backend->AllocateOffloadingBuckets(offloading_objects, + buckets_keys); } - std::unordered_map GetUngroupedOffloadingObjects( - FileStorage& fileStorage) { - return fileStorage.ungrouped_offloading_objects_; + size_t GetUngroupedOffloadingObjectsSize(FileStorage& fileStorage) { + auto bucket_backend = std::dynamic_pointer_cast( + fileStorage.storage_backend_); + if (!bucket_backend) { + return 0; + } + return bucket_backend->UngroupedOffloadingObjectsSize(); } void TearDown() override { @@ -103,22 +113,25 @@ TEST_F(FileStorageTest, IsEnableOffloading) { auto enable_offloading_result1 = FileStorageIsEnableOffloading(fileStorage1); ASSERT_TRUE(enable_offloading_result1 && enable_offloading_result1.value()); - file_storage_config.bucket_keys_limit = 10; - file_storage_config.total_keys_limit = 91; + // bucket_keys_limit/bucket_size_limit moved to BucketBackendConfig. + // With current semantics, backend prevents offloading once it would exceed + // limits, so we validate IsEnableOffloading directly under tight limits. + + // Case 2: total_keys_limit < bucket_keys_limit => cannot offload + SetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "10"); + file_storage_config.total_keys_limit = 9; FileStorage fileStorage2(nullptr, "localhost:9003", file_storage_config); - keys.clear(); - sizes.clear(); - ASSERT_TRUE(FileStorageBatchOffload(fileStorage2, keys, sizes, batch_data)); auto enable_offloading_result2 = FileStorageIsEnableOffloading(fileStorage2); ASSERT_TRUE(enable_offloading_result2 && !enable_offloading_result2.value()); - file_storage_config.bucket_size_limit = 969; + + // Case 3: total_size_limit < bucket_size_limit => cannot offload + SetEnv("MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", "969"); + file_storage_config.total_keys_limit = 10'000'000; + file_storage_config.total_size_limit = 100; FileStorage fileStorage3(nullptr, "localhost:9003", file_storage_config); - keys.clear(); - sizes.clear(); - ASSERT_TRUE(FileStorageBatchOffload(fileStorage3, keys, sizes, batch_data)); auto enable_offloading_result3 = FileStorageIsEnableOffloading(fileStorage3); ASSERT_TRUE(enable_offloading_result3 && @@ -156,8 +169,8 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_keys_limit) { std::vector> buckets_keys; auto file_storage_config = FileStorageConfig::FromEnvironment(); file_storage_config.storage_filepath = data_path; - file_storage_config.bucket_keys_limit = 10; - file_storage_config.bucket_iterator_keys_limit = 969; + file_storage_config.scanmeta_iterator_keys_limit = 969; + SetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "10"); FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); @@ -165,7 +178,7 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_keys_limit) { for (const auto& bucket_keys : buckets_keys) { ASSERT_EQ(bucket_keys.size(), 10); } - ASSERT_EQ(GetUngroupedOffloadingObjects(fileStorage).size(), 5); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 5); buckets_keys.clear(); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); @@ -173,7 +186,7 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_keys_limit) { for (const auto& bucket_keys : buckets_keys) { ASSERT_EQ(bucket_keys.size(), 10); } - ASSERT_EQ(GetUngroupedOffloadingObjects(fileStorage).size(), 0); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 0); } TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_size_limit) { @@ -184,7 +197,7 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_size_limit) { std::vector> buckets_keys; auto file_storage_config = FileStorageConfig::FromEnvironment(); file_storage_config.storage_filepath = data_path; - file_storage_config.bucket_size_limit = 10; + SetEnv("MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", "10"); FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); @@ -192,7 +205,7 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_size_limit) { for (const auto& bucket_keys : buckets_keys) { ASSERT_EQ(bucket_keys.size(), 10); } - ASSERT_EQ(GetUngroupedOffloadingObjects(fileStorage).size(), 5); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 5); buckets_keys.clear(); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); @@ -200,7 +213,7 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_size_limit) { for (const auto& bucket_keys : buckets_keys) { ASSERT_EQ(bucket_keys.size(), 10); } - ASSERT_EQ(GetUngroupedOffloadingObjects(fileStorage).size(), 0); + ASSERT_EQ(GetUngroupedOffloadingObjectsSize(fileStorage), 0); } TEST_F(FileStorageTest, @@ -212,8 +225,8 @@ TEST_F(FileStorageTest, std::vector> buckets_keys; auto file_storage_config = FileStorageConfig::FromEnvironment(); file_storage_config.storage_filepath = data_path; - file_storage_config.bucket_keys_limit = 9; - file_storage_config.bucket_size_limit = 496; + SetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "9"); + SetEnv("MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", "496"); FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config); ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket( fileStorage, offloading_objects, buckets_keys)); @@ -254,12 +267,13 @@ TEST_F(FileStorageTest, TEST_F(FileStorageTest, DefaultValuesWhenNoEnvSet) { auto config = FileStorageConfig::FromEnvironment(); + auto bucket_backend_config = BucketBackendConfig::FromEnvironment(); EXPECT_EQ(config.storage_filepath, "/data/file_storage"); EXPECT_EQ(config.local_buffer_size, 1280 * 1024 * 1024); - EXPECT_EQ(config.bucket_iterator_keys_limit, 20000); - EXPECT_EQ(config.bucket_keys_limit, 500); - EXPECT_EQ(config.bucket_size_limit, 256 * 1024 * 1024); + EXPECT_EQ(config.scanmeta_iterator_keys_limit, 20000); + EXPECT_EQ(bucket_backend_config.bucket_keys_limit, 500); + EXPECT_EQ(bucket_backend_config.bucket_size_limit, 256 * 1024 * 1024); EXPECT_EQ(config.total_keys_limit, 10'000'000); EXPECT_EQ(config.total_size_limit, 2ULL * 1024 * 1024 * 1024 * 1024); EXPECT_EQ(config.heartbeat_interval_seconds, 10u); @@ -278,9 +292,10 @@ TEST_F(FileStorageTest, ReadInt64FromEnv) { SetEnv("MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT", "5000000"); auto config = FileStorageConfig::FromEnvironment(); + auto bucket_backend_config = BucketBackendConfig::FromEnvironment(); EXPECT_EQ(config.local_buffer_size, 2147483648); - EXPECT_EQ(config.bucket_keys_limit, 1000); + EXPECT_EQ(bucket_backend_config.bucket_keys_limit, 1000); EXPECT_EQ(config.total_keys_limit, 5000000); } @@ -297,8 +312,9 @@ TEST_F(FileStorageTest, InvalidIntValueUsesDefault) { SetEnv("MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS", "-1"); auto config = FileStorageConfig::FromEnvironment(); + auto bucket_backend_config = BucketBackendConfig::FromEnvironment(); - EXPECT_EQ(config.bucket_keys_limit, 500); + EXPECT_EQ(bucket_backend_config.bucket_keys_limit, 500); EXPECT_EQ(config.total_size_limit, 2ULL * 1024 * 1024 * 1024 * 1024); EXPECT_EQ(config.heartbeat_interval_seconds, 10u); } @@ -316,14 +332,13 @@ TEST_F(FileStorageTest, EmptyEnvValueUsesDefault) { SetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", ""); // empty string auto config = FileStorageConfig::FromEnvironment(); - EXPECT_EQ(config.bucket_keys_limit, 500); // fallback + auto bucket_backend_config = BucketBackendConfig::FromEnvironment(); + EXPECT_EQ(bucket_backend_config.bucket_keys_limit, 500); // fallback } TEST_F(FileStorageTest, ValidateSuccessWithValidConfig) { FileStorageConfig config; config.storage_filepath = std::filesystem::current_path().string(); - config.bucket_keys_limit = 100; - config.bucket_size_limit = 100000; config.total_keys_limit = 1000000; config.total_size_limit = 1073741824; // 1GB config.heartbeat_interval_seconds = 5; @@ -357,14 +372,6 @@ TEST_F(FileStorageTest, ValidateFailsOnInvalidLimits) { FileStorageConfig config; config.storage_filepath = "/tmp"; - config.bucket_keys_limit = 0; - EXPECT_FALSE(config.Validate()); - - config.bucket_keys_limit = 1; - config.bucket_size_limit = 0; - EXPECT_FALSE(config.Validate()); - - config.bucket_size_limit = 1; config.total_keys_limit = 0; EXPECT_FALSE(config.Validate()); @@ -377,4 +384,45 @@ TEST_F(FileStorageTest, ValidateFailsOnInvalidLimits) { EXPECT_FALSE(config.Validate()); } +TEST_F(FileStorageTest, BatchLoad_WithStorageBackendAdaptor) { + std::vector keys; + std::vector sizes; + std::unordered_map batch_data; + + auto file_storage_config = FileStorageConfig::FromEnvironment(); + file_storage_config.storage_backend_type = StorageBackendType::kFilePerKey; + file_storage_config.storage_filepath = data_path; + file_storage_config.local_buffer_size = 128 * 1024 * 1024; + FilePerKeyConfig file_per_key_config; + file_per_key_config.fsdir = "FileStorageTestDir"; + + auto total_path = fs::path(data_path) / file_per_key_config.fsdir; + fs::create_directories(total_path); + + FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config); + + auto offload_res = + FileStorageBatchOffload(fileStorage, keys, sizes, batch_data); + ASSERT_TRUE(offload_res) << "FileStorageBatchOffload failed"; + + auto allocate_res = FileStorageAllocateBatch(fileStorage, keys, sizes); + ASSERT_TRUE(allocate_res) << "FileStorageAllocateBatch failed"; + + auto batch = std::move(allocate_res.value()); + + auto load_res = FileStorageBatchLoad(fileStorage, batch.slices); + ASSERT_TRUE(load_res) << "FileStorageBatchLoad failed"; + + for (const auto& it : batch.slices) { + const std::string& key = it.first; + const Slice& slice = it.second; + std::string data(static_cast(slice.ptr), slice.size); + + auto found = batch_data.find(key); + ASSERT_TRUE(found != batch_data.end()) + << "key not found in batch_data: " << key; + EXPECT_EQ(data, found->second); + } +} + } // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp index a4fec33e..d4f1bfd5 100644 --- a/mooncake-store/tests/storage_backend_test.cpp +++ b/mooncake-store/tests/storage_backend_test.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include "allocator.h" #include "utils.h" @@ -42,7 +43,10 @@ class StorageBackendTest : public ::testing::Test { TEST_F(StorageBackendTest, StorageBackendAll) { std::shared_ptr client_buffer_allocator = std::make_shared(128 * 1024 * 1024); - BucketStorageBackend storage_backend(data_path); + FileStorageConfig config; + BucketBackendConfig bucket_config; + config.storage_filepath = data_path; + BucketStorageBackend storage_backend(config, bucket_config); ASSERT_TRUE(storage_backend.Init()); ASSERT_TRUE(fs::directory_iterator(data_path) == fs::directory_iterator{}); @@ -95,7 +99,10 @@ TEST_F(StorageBackendTest, StorageBackendAll) { TEST_F(StorageBackendTest, BucketScan) { std::shared_ptr client_buffer_allocator = std::make_shared(128 * 1024 * 1024); - BucketStorageBackend storage_backend(data_path); + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + BucketStorageBackend storage_backend(config, bucket_config); ASSERT_TRUE(storage_backend.Init()); ASSERT_TRUE(!storage_backend.Init()); std::vector keys; @@ -281,8 +288,11 @@ TEST_F(StorageBackendTest, OrphanedBucketFileCleanup) { } } + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; // Create a valid bucket with data and metadata - BucketStorageBackend storage_backend(data_path); + BucketStorageBackend storage_backend(config, bucket_config); ASSERT_TRUE(storage_backend.Init()); std::shared_ptr client_buffer_allocator = @@ -347,7 +357,7 @@ TEST_F(StorageBackendTest, OrphanedBucketFileCleanup) { // Re-initialize the storage backend (orphan cleanup enabled by default now) // This should trigger orphan cleanup - BucketStorageBackend storage_backend_2(data_path); + BucketStorageBackend storage_backend_2(config, bucket_config); auto init_result = storage_backend_2.Init(); ASSERT_TRUE(init_result); @@ -384,4 +394,303 @@ TEST_F(StorageBackendTest, OrphanedBucketFileCleanup) { ASSERT_TRUE(is_exist.value()); } +TEST_F(StorageBackendTest, AdaptorBatchOffloadAndBatchLoad) { + FileStorageConfig cfg; + + cfg.storage_filepath = data_path; + FilePerKeyConfig file_per_key_config; + file_per_key_config.fsdir = "file_per_key_dir_offload_load"; + file_per_key_config.enable_eviction = false; + + StorageBackendAdaptor adaptor(cfg, file_per_key_config); + ASSERT_TRUE(adaptor.Init()); + ASSERT_TRUE( + adaptor.ScanMeta([](const std::vector& keys, + std::vector& metadatas) { + return ErrorCode::OK; + })); + + std::unordered_map test_data = { + {"simple-key", "hello world"}, + {"key/with/invalid:chars", "value-2"}, + }; + + std::unordered_map> batch_object; + std::vector> offload_buffers; + + for (auto& [key, value] : test_data) { + auto buf = std::make_unique(value.size()); + std::memcpy(buf.get(), value.data(), value.size()); + + std::vector slices; + slices.emplace_back( + Slice{buf.get(), static_cast(value.size())}); + batch_object.emplace(key, std::move(slices)); + + offload_buffers.push_back(std::move(buf)); + } + + auto offload_result = adaptor.BatchOffload( + batch_object, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(offload_result); + + auto exist_simple = adaptor.IsExist("simple-key"); + ASSERT_TRUE(exist_simple); + EXPECT_TRUE(exist_simple.value()); + + auto exist_not = adaptor.IsExist("not-exist-key"); + ASSERT_TRUE(exist_not); + EXPECT_FALSE(exist_not.value()); + + std::unordered_map load_slices; + std::vector> load_buffers; + + for (auto& [key, value] : test_data) { + auto buf = std::make_unique(value.size()); + load_slices.emplace( + key, Slice{buf.get(), static_cast(value.size())}); + load_buffers.push_back(std::move(buf)); + } + + auto load_result = adaptor.BatchLoad(load_slices); + ASSERT_TRUE(load_result); + + for (auto& [key, value] : test_data) { + auto it = load_slices.find(key); + ASSERT_NE(it, load_slices.end()); + + std::string loaded(static_cast(it->second.ptr), it->second.size); + EXPECT_EQ(loaded, value); + } +} + +TEST_F(StorageBackendTest, AdaptorBatchOffloadEmptyShouldFail) { + FileStorageConfig cfg; + cfg.storage_filepath = data_path + "/"; + + FilePerKeyConfig file_per_key_config; + file_per_key_config.fsdir = "file_per_key_dir_offload_empty"; + file_per_key_config.enable_eviction = false; + + StorageBackendAdaptor adaptor(cfg, file_per_key_config); + ASSERT_TRUE(adaptor.Init()); + ASSERT_TRUE( + adaptor.ScanMeta([](const std::vector& keys, + std::vector& metadatas) { + return ErrorCode::OK; + })); + + std::unordered_map> empty_batch; + + auto res = adaptor.BatchOffload( + empty_batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + + EXPECT_FALSE(res); + EXPECT_EQ(res.error(), ErrorCode::INVALID_KEY); +} + +TEST_F(StorageBackendTest, AdaptorScanMetaAndIsEnableOffloading) { + FileStorageConfig cfg; + cfg.storage_filepath = data_path + "/"; + cfg.total_keys_limit = 10; + cfg.total_size_limit = 1024 * 1024; + + FilePerKeyConfig file_per_key_config; + file_per_key_config.fsdir = "file_per_key_dir_is_enable_offloading"; + file_per_key_config.enable_eviction = false; + + StorageBackendAdaptor adaptor(cfg, file_per_key_config); + ASSERT_TRUE(adaptor.Init()); + + auto enable_before = adaptor.IsEnableOffloading(); + EXPECT_FALSE(enable_before); + EXPECT_EQ(enable_before.error(), ErrorCode::INTERNAL_ERROR); + + // New behavior: must call ScanMeta once before BatchOffload when eviction + // is disabled, otherwise meta_scanned_ is false and BatchOffload is + // rejected. + auto scan_init_res = adaptor.ScanMeta( + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(scan_init_res); + + auto enable_empty = adaptor.IsEnableOffloading(); + ASSERT_TRUE(enable_empty); + EXPECT_TRUE(enable_empty.value()); + + std::unordered_map test_data = { + {"k1", std::string(128, 'a')}, + {"k2", std::string(256, 'b')}, + }; + + std::unordered_map> batch_object; + std::vector> offload_buffers; + + for (auto& [key, value] : test_data) { + auto buf = std::make_unique(value.size()); + std::memcpy(buf.get(), value.data(), value.size()); + + std::vector slices; + slices.emplace_back( + Slice{buf.get(), static_cast(value.size())}); + batch_object.emplace(key, std::move(slices)); + + offload_buffers.push_back(std::move(buf)); + } + + auto offload_result = adaptor.BatchOffload( + batch_object, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(offload_result); + + auto enable_after_write = adaptor.IsEnableOffloading(); + ASSERT_TRUE(enable_after_write); + EXPECT_TRUE(enable_after_write.value()); + + // Verify scan results via a fresh adaptor instance to avoid double-counting + // totals if ScanMeta is called again on the same adaptor. + StorageBackendAdaptor restart_adaptor(cfg, file_per_key_config); + ASSERT_TRUE(restart_adaptor.Init()); + + std::vector scan_keys; + std::vector scan_metas; + + auto scan_result = restart_adaptor.ScanMeta( + [&](const std::vector& keys, + std::vector& metas) { + scan_keys.insert(scan_keys.end(), keys.begin(), keys.end()); + scan_metas.insert(scan_metas.end(), metas.begin(), metas.end()); + return ErrorCode::OK; + }); + ASSERT_TRUE(scan_result); + + EXPECT_EQ(scan_keys.size(), test_data.size()); + EXPECT_EQ(scan_metas.size(), test_data.size()); + + auto enable_after = restart_adaptor.IsEnableOffloading(); + ASSERT_TRUE(enable_after); + EXPECT_TRUE(enable_after.value()); + + FileStorageConfig strict_cfg = cfg; + strict_cfg.total_keys_limit = 1; + strict_cfg.total_size_limit = 1; + + StorageBackendAdaptor strict_adaptor(strict_cfg, file_per_key_config); + ASSERT_TRUE(strict_adaptor.Init()); + + auto strict_scan_result = strict_adaptor.ScanMeta( + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(strict_scan_result); + + auto enable_strict = strict_adaptor.IsEnableOffloading(); + ASSERT_TRUE(enable_strict); + EXPECT_FALSE(enable_strict.value()); +} + +TEST_F(StorageBackendTest, AdaptorScanMetaAndBatchLoadAcrossRestart) { + FileStorageConfig cfg; + cfg.storage_filepath = data_path; + cfg.scanmeta_iterator_keys_limit = 16; + cfg.total_keys_limit = 100; + cfg.total_size_limit = 1 << 20; + + FilePerKeyConfig file_per_key_config; + file_per_key_config.fsdir = "file_per_key_dir_batch_load_restart"; + file_per_key_config.enable_eviction = true; + + std::unordered_map test_data = { + {"simple-key", "hello world"}, + {"key/with:illegal*chars?", "value-2"}, + {"another_key", "third-value"}, + }; + + { + StorageBackendAdaptor adaptor(cfg, file_per_key_config); + ASSERT_TRUE(adaptor.Init()); + + std::unordered_map> batch_object; + std::vector> write_buffers; + write_buffers.reserve(test_data.size()); + + for (auto& [key, value] : test_data) { + auto buf = std::make_unique(value.size()); + std::memcpy(buf.get(), value.data(), value.size()); + + std::vector slices; + slices.emplace_back( + Slice{buf.get(), static_cast(value.size())}); + + batch_object.emplace(key, std::move(slices)); + write_buffers.emplace_back(std::move(buf)); + } + + auto offload_res = adaptor.BatchOffload( + batch_object, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(offload_res); + } + + { + StorageBackendAdaptor adaptor(cfg, file_per_key_config); + ASSERT_TRUE(adaptor.Init()); + + std::vector scan_keys; + std::vector scan_metas; + + auto scan_res = + adaptor.ScanMeta([&](const std::vector& keys, + std::vector& metas) { + scan_keys.insert(scan_keys.end(), keys.begin(), keys.end()); + scan_metas.insert(scan_metas.end(), metas.begin(), metas.end()); + return ErrorCode::OK; + }); + ASSERT_TRUE(scan_res); + + ASSERT_EQ(scan_keys.size(), test_data.size()); + ASSERT_EQ(scan_metas.size(), test_data.size()); + + std::unordered_map meta_map; + for (size_t i = 0; i < scan_keys.size(); ++i) { + meta_map.emplace(scan_keys[i], scan_metas[i]); + } + + for (auto& [key, value] : test_data) { + auto it = meta_map.find(key); + ASSERT_NE(it, meta_map.end()) + << "Meta for key " << key << " not found"; + EXPECT_EQ(it->second.data_size, static_cast(value.size())); + } + + std::unordered_map load_slices; + std::vector> load_buffers; + load_buffers.reserve(test_data.size()); + + for (auto& [key, value] : test_data) { + auto buf = std::make_unique(value.size()); + load_slices.emplace( + key, Slice{buf.get(), static_cast(value.size())}); + load_buffers.emplace_back(std::move(buf)); + } + + auto load_res = adaptor.BatchLoad(load_slices); + ASSERT_TRUE(load_res); + + for (auto& [key, value] : test_data) { + auto it = load_slices.find(key); + ASSERT_NE(it, load_slices.end()); + + std::string loaded(static_cast(it->second.ptr), + it->second.size); + EXPECT_EQ(loaded, value); + } + } +} + } // namespace mooncake::test diff --git a/mooncake-store/tests/utils/common.h b/mooncake-store/tests/utils/common.h index d3484669..7f2e857c 100644 --- a/mooncake-store/tests/utils/common.h +++ b/mooncake-store/tests/utils/common.h @@ -4,7 +4,7 @@ namespace mooncake { namespace fs = std::filesystem; inline tl::expected BatchOffloadUtil( - BucketStorageBackend& storage_backend, std::vector& keys, + StorageBackendInterface& storage_backend, std::vector& keys, std::vector& sizes, std::unordered_map& batched_data, std::vector& buckets) {