mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

[Store] feat: Implement a unified storage interface to simplify integration and extension (#1185)

This commit is contained in:
zhangzuo21 2025-12-17 22:33:05 +08:00 committed by GitHub
parent 31c45e47f5
commit 926c63aec8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 1300 additions and 422 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

86
mooncake-store/include/file_storage.h
View File

@ -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 <typename T>
static T GetEnvOr(const char* name, T default_value);
};
class BucketIterator {
public:
BucketIterator(std::shared_ptr<BucketStorageBackend> storage_backend,
int64_t limit);
tl::expected<void, ErrorCode> HandleNext(
const std::function<
ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas,
const std::vector<int64_t>& buckets)>& handler);
tl::expected<bool, ErrorCode> HasNext();
private:
std::shared_ptr<BucketStorageBackend> storage_backend_;
int64_t limit_;
mutable Mutex mutex_;
int64_t GUARDED_BY(mutex_) next_bucket_ = -1;
};
class FileStorage {
public:
FileStorage(std::shared_ptr<Client> client,
@ -119,21 +55,6 @@ class FileStorage {
tl::expected<void, ErrorCode> OffloadObjects(
const std::unordered_map<std::string, int64_t>& 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<void, ErrorCode> indicating operation status.
*/
tl::expected<void, ErrorCode> GroupOffloadingKeysByBucket(
const std::unordered_map<std::string, int64_t>& offloading_objects,
std::vector<std::vector<std::string>>& 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<bool, ErrorCode> IsEnableOffloading();
tl::expected<void, ErrorCode> BatchOffload(
const std::vector<std::string>& keys);
tl::expected<void, ErrorCode> BatchLoad(
const std::unordered_map<std::string, Slice>& batch_object);
@ -165,12 +83,10 @@ class FileStorage {
std::shared_ptr<Client> client_;
std::string local_rpc_addr_;
FileStorageConfig config_;
std::shared_ptr<BucketStorageBackend> storage_backend_;
std::shared_ptr<StorageBackendInterface> storage_backend_;
std::shared_ptr<ClientBufferAllocator> client_buffer_allocator_;
mutable Mutex offloading_mutex_;
std::unordered_map<std::string, int64_t> GUARDED_BY(offloading_mutex_)
ungrouped_offloading_objects_;
bool GUARDED_BY(offloading_mutex_) enable_offloading_;
std::atomic<bool> heartbeat_running_;
std::thread heartbeat_thread_;

225
mooncake-store/include/storage_backend.h
View File

@ -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<void, ErrorCode> Init() = 0;
virtual tl::expected<int64_t, ErrorCode> BatchOffload(
const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
std::function<ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas)>
complete_handler) = 0;
virtual tl::expected<void, ErrorCode> BatchLoad(
const std::unordered_map<std::string, Slice>& batched_slices) = 0;
virtual tl::expected<bool, ErrorCode> IsExist(const std::string& key) = 0;
virtual tl::expected<bool, ErrorCode> IsEnableOffloading() = 0;
virtual tl::expected<void, ErrorCode> ScanMeta(
const std::function<ErrorCode(
const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& 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<int64_t> 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<void, ErrorCode> Init() override;
tl::expected<int64_t, ErrorCode> BatchOffload(
const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
std::function<ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas)>
complete_handler) override;
tl::expected<void, ErrorCode> BatchLoad(
const std::unordered_map<std::string, Slice>& batched_slices) override;
tl::expected<bool, ErrorCode> IsExist(const std::string& key) override;
tl::expected<bool, ErrorCode> IsEnableOffloading() override;
tl::expected<void, ErrorCode> ScanMeta(
const std::function<ErrorCode(
const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas)>& handler) override;
private:
const FilePerKeyConfig file_per_key_config_;
std::atomic<bool> meta_scanned_{false};
std::unique_ptr<StorageBackend> 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<Slice>& 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<std::string, std::vector<Slice>>& batch_object,
std::function<ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& 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<void, ErrorCode> indicating operation status.
*/
tl::expected<void, ErrorCode> BatchQuery(
const std::vector<std::string>& keys,
std::unordered_map<std::string, StorageObjectMetadata>&
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<void, ErrorCode> indicating operation status.
*/
tl::expected<void, ErrorCode> BatchLoad(
const std::unordered_map<std::string, Slice>& batched_slices);
const std::unordered_map<std::string, Slice>& 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<void, ErrorCode> indicating operation status.
*/
tl::expected<void, ErrorCode> Init();
tl::expected<void, ErrorCode> 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<void, ErrorCode> indicating operation status.
*/
tl::expected<bool, ErrorCode> IsExist(const std::string& key);
tl::expected<bool, ErrorCode> 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<void, ErrorCode> indicating operation status.
*/
tl::expected<void, ErrorCode> ScanMeta(
const std::function<ErrorCode(
const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas)>& handler) override;
/**
* @brief Checks whether the backend is allowed to continue offloading.
* @return tl::expected<bool, ErrorCode>
* - On success: true 表示可以继续 offload；false 表示达到上限/不允许继续。
* - On failure: 返回错误码（例如 IO/内部错误）。
*/
tl::expected<bool, ErrorCode> 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<void, ErrorCode> indicating operation status.
*/
tl::expected<void, ErrorCode> AllocateOffloadingBuckets(
const std::unordered_map<std::string, int64_t>& offloading_objects,
std::vector<std::vector<std::string>>& 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<std::unique_ptr<StorageFile>, ErrorCode> OpenFile(
const std::string& path, FileMode mode) const;
tl::expected<void, ErrorCode> GroupOffloadingKeysByBucket(
const std::unordered_map<std::string, int64_t>& offloading_objects,
std::vector<std::vector<std::string>>& buckets_keys);
tl::expected<void, ErrorCode> HandleNext(
const std::function<
ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas)>& handler);
tl::expected<bool, ErrorCode> HasNext();
private:
std::atomic<bool> initialized_{false};
std::optional<BucketIdGenerator> 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<std::string, StorageObjectMetadata> GUARDED_BY(mutex_)
object_bucket_map_;
std::map<int64_t, std::shared_ptr<BucketMetadata>> GUARDED_BY(
mutex_) buckets_;
int64_t GUARDED_BY(mutex_) next_bucket_ = -1;
BucketBackendConfig bucket_backend_config_;
mutable Mutex offloading_mutex_;
std::unordered_map<std::string, int64_t> GUARDED_BY(offloading_mutex_)
ungrouped_offloading_objects_;
};
} // namespace mooncake
tl::expected<std::shared_ptr<StorageBackendInterface>, ErrorCode>
CreateStorageBackend(const FileStorageConfig& config);
} // namespace mooncake

21
mooncake-store/include/utils.h
View File

@ -285,4 +285,25 @@ int getFreeTcpPort();
int64_t time_gen();
// Helper: Get integer from environment variable, fallback to default
template <typename T>
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<T, uint32_t>) {
if (value < 0 || value > UINT32_MAX) throw std::out_of_range("");
}
return static_cast<T>(value);
} catch (...) {
return default_value;
}
}
std::string GetEnvStringOr(const char* name, const std::string& default_value);
} // namespace mooncake

385
mooncake-store/src/file_storage.cpp
View File

@ -4,53 +4,34 @@
#include <unordered_set>
#include <vector>
#include "storage_backend.h"
#include "utils.h"
namespace mooncake {
// Helper: Get integer from environment variable, fallback to default
template <typename T>
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<T, uint32_t>) {
if (value < 0 || value > UINT32_MAX) throw std::out_of_range("");
}
return static_cast<T>(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<int64_t>(
"MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES", config.local_buffer_size);
config.bucket_iterator_keys_limit =
GetEnvOr<int64_t>("MOONCAKE_OFFLOAD_BUCKET_ITERATOR_KEYS_LIMIT",
config.bucket_iterator_keys_limit);
config.bucket_keys_limit = GetEnvOr<int64_t>(
"MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", config.bucket_keys_limit);
config.bucket_size_limit = GetEnvOr<int64_t>(
"MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", config.bucket_size_limit);
config.scanmeta_iterator_keys_limit =
GetEnvOr<int64_t>("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT",
config.scanmeta_iterator_keys_limit);
config.total_keys_limit = GetEnvOr<int64_t>(
"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_(client),
local_rpc_addr_(local_rpc_addr),
config_(config),
storage_backend_(
std::make_shared<BucketStorageBackend>(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<void, ErrorCode> 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<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas,
const std::vector<int64_t>&) {
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<std::string>& keys,
std::vector<StorageObjectMetadata>& 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<void, ErrorCode> FileStorage::BatchGet(
tl::expected<void, ErrorCode> FileStorage::OffloadObjects(
const std::unordered_map<std::string, int64_t>& offloading_objects) {
std::vector<std::vector<std::string>> 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<BucketStorageBackend>(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<std::string> 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<std::string>& keys,
std::vector<StorageObjectMetadata>& 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<std::string, std::vector<Slice>> 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<void, ErrorCode> FileStorage::OffloadObjects(
}
tl::expected<bool, ErrorCode> 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<void, ErrorCode> FileStorage::Heartbeat() {
return {};
}
tl::expected<void, ErrorCode> FileStorage::BatchOffload(
const std::vector<std::string>& keys) {
auto start_time = std::chrono::steady_clock::now();
std::unordered_map<std::string, std::vector<Slice>> 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<std::string>& keys,
std::vector<StorageObjectMetadata>& 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<std::chrono::microseconds>(
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<void, ErrorCode> FileStorage::BatchLoad(
const std::unordered_map<std::string, Slice>& batch_object) {
auto start_time = std::chrono::steady_clock::now();
@ -501,122 +450,4 @@ tl::expected<FileStorage::AllocatedBatch, ErrorCode> FileStorage::AllocateBatch(
return result;
}
tl::expected<void, ErrorCode> FileStorage::GroupOffloadingKeysByBucket(
const std::unordered_map<std::string, int64_t>& offloading_objects,
std::vector<std::vector<std::string>>& 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<std::string> bucket_keys;
std::unordered_map<std::string, int64_t> 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<int64_t>(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<BucketStorageBackend> storage_backend, int64_t limit)
: storage_backend_(storage_backend), limit_(limit) {};
tl::expected<void, ErrorCode> BucketIterator::HandleNext(
const std::function<ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas,
const std::vector<int64_t>& buckets)>&
handler) {
MutexLocker locker(&mutex_);
std::vector<std::string> keys;
std::vector<StorageObjectMetadata> metadatas;
std::vector<int64_t> 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<bool, ErrorCode> BucketIterator::HasNext() {
MutexLocker locker(&mutex_);
return next_bucket_ != 0;
}
} // namespace mooncake

555
mooncake-store/src/storage_backend.cpp
View File

@ -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<bool>("ENABLE_EVICTION", config.enable_eviction);
return config;
}
BucketBackendConfig BucketBackendConfig::FromEnvironment() {
BucketBackendConfig config;
config.bucket_keys_limit = GetEnvOr<int64_t>(
"MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", config.bucket_keys_limit);
config.bucket_size_limit = GetEnvOr<int64_t>(
"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<void, ErrorCode> StorageBackend::Init(uint64_t quota_bytes = 0) {
std::unique_lock<std::shared_mutex> 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<void, ErrorCode> StorageBackendAdaptor::Init() {
std::string storage_root =
file_storage_config_.storage_filepath + file_per_key_config_.fsdir;
storage_backend_ = std::make_unique<StorageBackend>(
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<std::string>{}(key);
// Use low 8 bits to create 2-level directory structure (e.g. "a1/b2")
char dir1 =
static_cast<char>('a' + (hash & 0x0F)); // Lower 4 bits -> 16 dirs
char dir2 = static_cast<char>(
'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<Slice>& 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<const char*>(s.ptr), s.size);
}
return out;
}
tl::expected<int64_t, ErrorCode> StorageBackendAdaptor::BatchOffload(
const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
std::function<ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& 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<StorageObjectMetadata> metadatas;
std::vector<std::string> 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<int64_t>(kv.key.size()),
static_cast<int64_t>(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<bool, ErrorCode> StorageBackendAdaptor::IsExist(
const std::string& key) {
auto path = ResolvePath(key);
namespace fs = std::filesystem;
return fs::exists(path);
}
tl::expected<void, ErrorCode> StorageBackendAdaptor::BatchLoad(
const std::unordered_map<std::string, Slice>& 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<bool, ErrorCode> 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<void, ErrorCode> StorageBackendAdaptor::ScanMeta(
const std::function<
ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& 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<std::string> keys;
std::vector<StorageObjectMetadata> metas;
auto flush = [&]() -> tl::expected<void, ErrorCode> {
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<int64_t>(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<int64_t, ErrorCode> BucketStorageBackend::BatchOffload(
const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
@ -900,6 +1224,14 @@ tl::expected<int64_t, ErrorCode> 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<iovec> iovs;
std::vector<StorageObjectMetadata> metadatas;
@ -1139,7 +1471,8 @@ tl::expected<void, ErrorCode> 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<bool, ErrorCode> BucketStorageBackend::IsExist(
return false;
}
tl::expected<bool, ErrorCode> 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<void, ErrorCode> BucketStorageBackend::ScanMeta(
const std::function<
ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& 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<int64_t, ErrorCode> BucketStorageBackend::BucketScan(
int64_t bucket_id, std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas,
@ -1191,8 +1564,8 @@ tl::expected<int64_t, ErrorCode> BucketStorageBackend::BucketScan(
auto bucket_it = buckets_.lower_bound(bucket_id);
for (; bucket_it != buckets_.end(); ++bucket_it) {
if (static_cast<int64_t>(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<void, ErrorCode> BucketStorageBackend::AllocateOffloadingBuckets(
const std::unordered_map<std::string, int64_t>& offloading_objects,
std::vector<std::vector<std::string>>& 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<void, ErrorCode> BucketStorageBackend::GroupOffloadingKeysByBucket(
const std::unordered_map<std::string, int64_t>& offloading_objects,
std::vector<std::vector<std::string>>& 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<int64_t>(
offloading_objects.size() + ungrouped_offloading_objects.size());
int64_t total_count = residue_count;
auto is_exist_func =
[this](const std::string& key) -> tl::expected<bool, ErrorCode> {
return IsExist(key);
};
while (it != offloading_objects.cend()) {
std::vector<std::string> bucket_keys;
std::unordered_map<std::string, int64_t> 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<int64_t>(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<int64_t>(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<std::shared_ptr<BucketMetadata>, ErrorCode>
BucketStorageBackend::BuildBucket(
int64_t bucket_id,
@ -1463,4 +1945,63 @@ BucketStorageBackend::OpenFile(const std::string& path, FileMode mode) const {
return std::make_unique<PosixFile>(path, fd);
}
tl::expected<void, ErrorCode> BucketStorageBackend::HandleNext(
const std::function<
ErrorCode(const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas)>& handler) {
MutexLocker locker(&iterator_mutex_);
std::vector<std::string> keys;
std::vector<StorageObjectMetadata> metadatas;
std::vector<int64_t> 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<bool, ErrorCode> BucketStorageBackend::HasNext() {
MutexLocker locker(&iterator_mutex_);
return next_bucket_ != 0;
}
tl::expected<std::shared_ptr<StorageBackendInterface>, 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<BucketStorageBackend>(
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<StorageBackendAdaptor>(
config, file_per_key_backend_config);
}
default: {
LOG(FATAL) << "Unsupported backend type";
return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
}
}
}
} // namespace mooncake

5
mooncake-store/src/utils.cpp
View File

@ -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

126
mooncake-store/tests/file_storage_test.cpp
View File

@ -1,6 +1,7 @@
#include <glog/logging.h>
#include <gtest/gtest.h>
#include <filesystem>
#include <thread>
#include "allocator.h"
@ -70,13 +71,22 @@ class FileStorageTest : public ::testing::Test {
FileStorage& fileStorage,
const std::unordered_map<std::string, int64_t>& offloading_objects,
std::vector<std::vector<std::string>>& buckets_keys) {
return fileStorage.GroupOffloadingKeysByBucket(offloading_objects,
buckets_keys);
auto bucket_backend = std::dynamic_pointer_cast<BucketStorageBackend>(
fileStorage.storage_backend_);
if (!bucket_backend) {
return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
}
return bucket_backend->AllocateOffloadingBuckets(offloading_objects,
buckets_keys);
}
std::unordered_map<std::string, int64_t> GetUngroupedOffloadingObjects(
FileStorage& fileStorage) {
return fileStorage.ungrouped_offloading_objects_;
size_t GetUngroupedOffloadingObjectsSize(FileStorage& fileStorage) {
auto bucket_backend = std::dynamic_pointer_cast<BucketStorageBackend>(
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<std::vector<std::string>> 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<std::vector<std::string>> 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<std::vector<std::string>> 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<std::string> keys;
std::vector<int64_t> sizes;
std::unordered_map<std::string, std::string> 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<char*>(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

317
mooncake-store/tests/storage_backend_test.cpp
View File

@ -6,6 +6,7 @@
#include <filesystem>
#include <iostream>
#include <ranges>
#include <ylt/util/tl/expected.hpp>
#include "allocator.h"
#include "utils.h"
@ -42,7 +43,10 @@ class StorageBackendTest : public ::testing::Test {
TEST_F(StorageBackendTest, StorageBackendAll) {
std::shared_ptr<SimpleAllocator> client_buffer_allocator =
std::make_shared<SimpleAllocator>(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<SimpleAllocator> client_buffer_allocator =
std::make_shared<SimpleAllocator>(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<std::string> 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<SimpleAllocator> 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<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas) {
return ErrorCode::OK;
}));
std::unordered_map<std::string, std::string> test_data = {
{"simple-key", "hello world"},
{"key/with/invalid:chars", "value-2"},
};
std::unordered_map<std::string, std::vector<Slice>> batch_object;
std::vector<std::unique_ptr<char[]>> offload_buffers;
for (auto& [key, value] : test_data) {
auto buf = std::make_unique<char[]>(value.size());
std::memcpy(buf.get(), value.data(), value.size());
std::vector<Slice> slices;
slices.emplace_back(
Slice{buf.get(), static_cast<size_t>(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::string>&,
std::vector<StorageObjectMetadata>&) { 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<std::string, Slice> load_slices;
std::vector<std::unique_ptr<char[]>> load_buffers;
for (auto& [key, value] : test_data) {
auto buf = std::make_unique<char[]>(value.size());
load_slices.emplace(
key, Slice{buf.get(), static_cast<size_t>(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<char*>(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<std::string>& keys,
std::vector<StorageObjectMetadata>& metadatas) {
return ErrorCode::OK;
}));
std::unordered_map<std::string, std::vector<Slice>> empty_batch;
auto res = adaptor.BatchOffload(
empty_batch,
[](const std::vector<std::string>&,
std::vector<StorageObjectMetadata>&) { 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::string>&,
std::vector<StorageObjectMetadata>&) { 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<std::string, std::string> test_data = {
{"k1", std::string(128, 'a')},
{"k2", std::string(256, 'b')},
};
std::unordered_map<std::string, std::vector<Slice>> batch_object;
std::vector<std::unique_ptr<char[]>> offload_buffers;
for (auto& [key, value] : test_data) {
auto buf = std::make_unique<char[]>(value.size());
std::memcpy(buf.get(), value.data(), value.size());
std::vector<Slice> slices;
slices.emplace_back(
Slice{buf.get(), static_cast<size_t>(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::string>&,
std::vector<StorageObjectMetadata>&) { 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<std::string> scan_keys;
std::vector<StorageObjectMetadata> scan_metas;
auto scan_result = restart_adaptor.ScanMeta(
[&](const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& 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::string>&,
std::vector<StorageObjectMetadata>&) { 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<std::string, std::string> 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<std::string, std::vector<Slice>> batch_object;
std::vector<std::unique_ptr<char[]>> write_buffers;
write_buffers.reserve(test_data.size());
for (auto& [key, value] : test_data) {
auto buf = std::make_unique<char[]>(value.size());
std::memcpy(buf.get(), value.data(), value.size());
std::vector<Slice> slices;
slices.emplace_back(
Slice{buf.get(), static_cast<size_t>(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::string>&,
std::vector<StorageObjectMetadata>&) { return ErrorCode::OK; });
ASSERT_TRUE(offload_res);
}
{
StorageBackendAdaptor adaptor(cfg, file_per_key_config);
ASSERT_TRUE(adaptor.Init());
std::vector<std::string> scan_keys;
std::vector<StorageObjectMetadata> scan_metas;
auto scan_res =
adaptor.ScanMeta([&](const std::vector<std::string>& keys,
std::vector<StorageObjectMetadata>& 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<std::string, StorageObjectMetadata> 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<int64_t>(value.size()));
}
std::unordered_map<std::string, Slice> load_slices;
std::vector<std::unique_ptr<char[]>> load_buffers;
load_buffers.reserve(test_data.size());
for (auto& [key, value] : test_data) {
auto buf = std::make_unique<char[]>(value.size());
load_slices.emplace(
key, Slice{buf.get(), static_cast<size_t>(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<char*>(it->second.ptr),
it->second.size);
EXPECT_EQ(loaded, value);
}
}
}
} // namespace mooncake::test

2
mooncake-store/tests/utils/common.h
View File

@ -4,7 +4,7 @@
namespace mooncake {
namespace fs = std::filesystem;
inline tl::expected<void, ErrorCode> BatchOffloadUtil(
BucketStorageBackend& storage_backend, std::vector<std::string>& keys,
StorageBackendInterface& storage_backend, std::vector<std::string>& keys,
std::vector<int64_t>& sizes,
std::unordered_map<std::string, std::string>& batched_data,
std::vector<int64_t>& buckets) {