From 2cf86bf38546760d725f6491051c85aef8d12cf7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=A5=A0=E5=BF=97?= <240190801+zhuxinjie-nz@users.noreply.github.com> Date: Fri, 31 Oct 2025 10:19:05 +0800 Subject: [PATCH] [Store] Add BucketStorageBackend (#968) --- .github/workflows/ci.yml | 13 +- mooncake-store/include/mutex.h | 135 +++- mooncake-store/include/storage_backend.h | 193 +++++- mooncake-store/include/types.h | 5 + mooncake-store/include/utils.h | 2 + mooncake-store/src/storage_backend.cpp | 579 ++++++++++++++++-- mooncake-store/src/utils.cpp | 6 + mooncake-store/tests/CMakeLists.txt | 3 +- mooncake-store/tests/mutex_test.cpp | 176 ++++++ mooncake-store/tests/storage_backend_test.cpp | 313 ++++++++++ 10 files changed, 1378 insertions(+), 47 deletions(-) create mode 100644 mooncake-store/tests/mutex_test.cpp create mode 100644 mooncake-store/tests/storage_backend_test.cpp diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6b877178..b1a7986c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -187,6 +187,8 @@ jobs: sudo rm -rf /usr/share/dotnet sudo rm -rf /opt/ghc sudo rm -rf /opt/hostedtoolcache/CodeQL + sudo rm -rf /usr/local/lib/android + df -h - name: Install CUDA Toolkit uses: Jimver/cuda-toolkit@v0.2.24 @@ -216,6 +218,7 @@ jobs: sudo apt update -y sudo bash -x dependencies.sh -y pip install torch==2.8.0 + df -h shell: bash - name: Build transfer engine only @@ -226,8 +229,9 @@ jobs: export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH cmake .. -DUSE_ETCD=OFF -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs" - make -j + make -j4 sudo make install + df -h shell: bash - name: Configure project with all settings are ON @@ -241,8 +245,9 @@ jobs: - name: Build project with all settings are ON run: | cd build - make -j + make -j4 sudo make install + df -h shell: bash - name: Configure project with unit tests and examples @@ -255,7 +260,7 @@ jobs: - name: Build project with unit tests and examples run: | cd build - make -j + make -j4 sudo make install shell: bash @@ -269,7 +274,7 @@ jobs: - name: Build project run: | cd build - make -j + make -j4 sudo make install shell: bash diff --git a/mooncake-store/include/mutex.h b/mooncake-store/include/mutex.h index 66de7031..a92dae8c 100644 --- a/mooncake-store/include/mutex.h +++ b/mooncake-store/include/mutex.h @@ -2,6 +2,7 @@ #define THREAD_SAFETY_ANALYSIS_MUTEX_H #include +#include // Enable thread safety attributes only with clang. // The attributes can be safely erased when compiling with other compilers. @@ -31,12 +32,21 @@ #define ACQUIRE(...) \ THREAD_ANNOTATION_ATTRIBUTE__(acquire_capability(__VA_ARGS__)) +#define ACQUIRE_SHARED(...) \ + THREAD_ANNOTATION_ATTRIBUTE__(acquire_shared_capability(__VA_ARGS__)) + #define RELEASE(...) \ THREAD_ANNOTATION_ATTRIBUTE__(release_capability(__VA_ARGS__)) +#define RELEASE_SHARED(...) \ + THREAD_ANNOTATION_ATTRIBUTE__(release_shared_capability(__VA_ARGS__)) + #define TRY_ACQUIRE(...) \ THREAD_ANNOTATION_ATTRIBUTE__(try_acquire_capability(__VA_ARGS__)) +#define TRY_ACQUIRE_SHARED(...) \ + THREAD_ANNOTATION_ATTRIBUTE__(try_acquire_shared_capability(__VA_ARGS__)) + #define EXCLUDES(...) THREAD_ANNOTATION_ATTRIBUTE__(locks_excluded(__VA_ARGS__)) #define ASSERT_CAPABILITY(x) THREAD_ANNOTATION_ATTRIBUTE__(assert_capability(x)) @@ -62,19 +72,50 @@ class CAPABILITY("mutex") Mutex { bool try_lock() TRY_ACQUIRE(true) { return mutex_.try_lock(); } // For negative capabilities. - const Mutex &operator!() const { return *this; } + const Mutex& operator!() const { return *this; } +}; + +// Simple shared_mutex implementation using std::shared_mutex for exclusive +// locking only. +class CAPABILITY("shared_mutex") SharedMutex { + private: + std::shared_mutex mutex_; + + public: + // Acquire/lock this mutex exclusively. + void lock() ACQUIRE() { mutex_.lock(); } + + // Acquire/lock this mutex shared. + void lock_shared() ACQUIRE_SHARED() { mutex_.lock_shared(); } + + // Release/unlock the mutex. + void unlock() RELEASE() { mutex_.unlock(); } + + // Release/unlock a shared mutex. + void unlock_shared() RELEASE_SHARED() { mutex_.unlock_shared(); } + + // Try to acquire the mutex. Returns true on success, and false on failure. + bool try_lock() TRY_ACQUIRE(true) { return mutex_.try_lock(); } + + // Try to acquire the mutex for read operations. + bool try_lock_shared() TRY_ACQUIRE_SHARED(true) { + return mutex_.try_lock_shared(); + } + + // For negative capabilities. + const SharedMutex& operator!() const { return *this; } }; // MutexLocker is an RAII class that acquires a mutex in its constructor, and // releases it in its destructor. class SCOPED_CAPABILITY MutexLocker { private: - Mutex *mut; + Mutex* mut; bool locked; public: // Acquire mu, implicitly acquire *this and associate it with mu. - MutexLocker(Mutex *mu) ACQUIRE(mu) : mut(mu), locked(true) { mu->lock(); } + MutexLocker(Mutex* mu) ACQUIRE(mu) : mut(mu), locked(true) { mu->lock(); } // Release *this and all associated mutexes, if they are still held. ~MutexLocker() RELEASE() { @@ -99,4 +140,92 @@ class SCOPED_CAPABILITY MutexLocker { } }; +// Tag types for selecting a constructor. +struct shared_lock_t { +} inline constexpr shared_lock = {}; + +// SharedMutexLocker is an RAII class that acquires a shared mutex in its +// constructor, and releases it in its destructor. +class SCOPED_CAPABILITY SharedMutexLocker { + private: + SharedMutex* mut; + bool is_exclusive; // true if holding the lock in exclusive mode + bool locked; + + public: + // Constructor: Acquire the mutex in exclusive mode + explicit SharedMutexLocker(SharedMutex* mu) ACQUIRE(mu) + : mut(mu), is_exclusive(true), locked(true) { + if (mut) { + mut->lock(); + } + } + + // Constructor: Acquire the mutex in shared mode + SharedMutexLocker(SharedMutex* mu, const shared_lock_t&) ACQUIRE_SHARED(mu) + : mut(mu), is_exclusive(false), locked(true) { + if (mut) { + mut->lock_shared(); + } + } + + // Destructor: Automatically release the mutex + ~SharedMutexLocker() RELEASE() { + if (locked && mut) { + if (is_exclusive) { + mut->unlock(); + } else { + mut->unlock_shared(); + } + } + } + + // Prevent copying and assignment + SharedMutexLocker(const SharedMutexLocker&) = delete; + SharedMutexLocker& operator=(const SharedMutexLocker&) = delete; + + // Acquire the mutex in exclusive mode + void lock() ACQUIRE() { + if (!mut || locked) return; + mut->lock(); + is_exclusive = true; + locked = true; + } + + // Acquire the mutex in shared mode + void lock_shared() ACQUIRE_SHARED() { + if (!mut || locked) return; + mut->lock_shared(); + is_exclusive = false; + locked = true; + } + + // Try to acquire the mutex in exclusive mode; returns true on success + bool try_lock() TRY_ACQUIRE(true) { + if (!mut || locked) return false; + locked = mut->try_lock(); + if (locked) is_exclusive = true; + return locked; + } + + // Try to acquire the mutex in shared mode; returns true on success + bool try_lock_shared() TRY_ACQUIRE_SHARED(true) { + if (!mut || locked) return false; + locked = mut->try_lock_shared(); + if (locked) is_exclusive = false; + return locked; + } + + // Release the mutex according to the current mode (exclusive or shared) + void unlock() RELEASE() { + if (!locked || !mut) return; + if (is_exclusive) { + mut->unlock(); + } else { + mut->unlock_shared(); + } + locked = false; + } +}; + #endif // THREAD_SAFETY_ANALYSIS_MUTEX_H diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index 1f5ea775..b58d60b4 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -1,15 +1,46 @@ #pragma once #include - +#include #include #include #include +#include "mutex.h" #include "types.h" #include "file_interface.h" namespace mooncake { + +struct StorageObjectMetadata { + int64_t bucket_id; + int64_t offset; + int64_t key_size; + int64_t data_size; +}; + +struct BucketObjectMetadata { + int64_t offset; + int64_t key_size; + int64_t data_size; +}; +YLT_REFL(BucketObjectMetadata, offset, key_size, data_size); + +struct BucketMetadata { + int64_t meta_size; + int64_t data_size; + std::unordered_map object_metadata; + std::vector keys; +}; +YLT_REFL(BucketMetadata, data_size, object_metadata, keys); + +struct OffloadMetadata { + int64_t total_keys; + int64_t total_size; +}; + +enum class FileMode { Read, Write }; + /** * @class StorageBackend * @brief Implementation of StorageBackend interface using local filesystem @@ -114,7 +145,7 @@ class StorageBackend { */ tl::expected LoadObject(const std::string& path, std::vector& slices, - size_t length); + int64_t length); /** * @brief Loads an object as a string @@ -124,7 +155,7 @@ class StorageBackend { * @return tl::expected indicating operation status */ tl::expected LoadObject(const std::string& path, - std::string& str, size_t length); + std::string& str, int64_t length); /** * @brief Deletes the physical file associated with the given object key @@ -175,4 +206,160 @@ class StorageBackend { FileMode mode) const; }; +class BucketIdGenerator { + public: + explicit BucketIdGenerator(int64_t start); + + int64_t NextId(); + + int64_t CurrentId(); + static constexpr int64_t INIT_NEW_START_ID = -1; + + private: + static constexpr int SEQUENCE_BITS = 12; + static constexpr int SEQUENCE_ID_SHIFT = 0; + static constexpr int TIMESTAMP_SHIFT = SEQUENCE_BITS; + static constexpr int64_t SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1; + std::atomic current_id_; +}; + +class BucketStorageBackend { + public: + BucketStorageBackend(const std::string& storage_filepath); + + /** + * @brief Offload objects in batches + * @param batch_object A map from object key to a list of data slices to be + * stored. + * @param complete_handler A callback function that is invoked after all + * data is stored successfully. + * @return tl::expected indicating operation status + */ + tl::expected BatchOffload( + const std::unordered_map>& batch_object, + std::function&)> + complete_handler); + + /** + * @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 + * retrieved metadata. + * @return tl::expected indicating operation status. + */ + tl::expected BatchQuery( + const std::vector& keys, + std::unordered_map& + batche_object_metadata); + + /** + * @brief Loads data for multiple objects in a batch operation. + * @param batched_slices A map from object key to a pre-allocated writable + * buffer (Slice). + * @return tl::expected indicating operation status. + */ + tl::expected BatchLoad( + std::unordered_map& batched_slices); + + /** + * @brief Retrieves the list of object keys belonging to a specific bucket. + * @param bucket_id The unique identifier of the bucket to query. + * @param bucket_keys Output parameter that will be populated with the list + * of object keys. + * @return tl::expected indicating operation status. + */ + tl::expected GetBucketKeys( + int64_t bucket_id, std::vector& bucket_keys); + + /** + * @brief Initializes the bucket storage backend. + * @return tl::expected indicating operation status. + */ + tl::expected Init(); + + /** + * @brief Checks whether an object with the specified key exists in the + * storage system. + * @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); + + /** + * @brief Iterate over the metadata of stored objects starting from a + * specified bucket. + * @param bucket_id The ID of the bucket to start scanning from. + * @param objects Output parameter: a map from object key to its metadata. + * @param buckets Output parameter: a list of bucket IDs encountered during + * iteration. + * @param limit Maximum number of objects to return in this iteration. + * @return tl::expected + * - On success: the bucket ID where the next iteration should start (or 0 + * if all data has been scanned). + * - On failure: returns an error code (e.g., BUCKET_NOT_FOUND, IO_ERROR). + */ + tl::expected BucketScan( + int64_t bucket_id, + std::unordered_map& objects, + std::vector& buckets, int64_t limit); + + /** + * @brief Retrieves the global metadata of the store. + * @return On success: `tl::expected` containing a `StoreMetadata` + * object. On failure: an error code. + */ + tl::expected GetStoreMetadata(); + + private: + tl::expected, ErrorCode> BuildBucket( + const std::unordered_map>& batch_object, + std::vector& iovs); + + tl::expected WriteBucket( + int64_t bucket_id, std::shared_ptr bucket_metadata, + std::vector& iovs); + + tl::expected StoreBucketMetadata( + int64_t bucket_id, std::shared_ptr bucket_metadata); + + tl::expected LoadBucketMetadata( + int64_t bucket_id, std::shared_ptr bucket_metadata); + + tl::expected BatchLoadBucket( + int64_t bucket_id, const std::vector& keys, + std::unordered_map& batched_slices); + + tl::expected CreateBucketId(); + + tl::expected GetBucketMetadataPath( + int64_t bucket_id); + + tl::expected GetBucketDataPath(int64_t bucket_id); + + tl::expected, ErrorCode> OpenFile( + const std::string& path, FileMode mode) const; + + private: + std::atomic initialized_{false}; + std::optional bucket_id_generator_; + static constexpr const char* BUCKET_METADATA_FILE_SUFFIX = ".meta"; + /** + * @brief A shared mutex to protect concurrent access to metadata. + * + * This mutex is used to synchronize read/write operations on the following + * metadata members: + * - object_bucket_map_: maps object keys to bucket IDs + * - buckets_: ordered map of bucket ID to bucket metadata + * - total_size_: cumulative data size of all stored objects + */ + mutable SharedMutex 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_; +}; } // namespace mooncake diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h index d9e617df..2364f074 100644 --- a/mooncake-store/include/types.h +++ b/mooncake-store/include/types.h @@ -145,6 +145,11 @@ enum class ErrorCode : int32_t { FILE_INVALID_BUFFER = -1104, ///< File buffer is wrong. FILE_LOCK_FAIL = -1105, ///< File lock operation failed. FILE_INVALID_HANDLE = -1106, ///< Invalid file handle. + + BUCKET_NOT_FOUND = -1200, ///< Bucket not found. + BUCKET_ALREADY_EXISTS = -1201, ///< Bucket already exists. + KEYS_ULTRA_BUCKET_LIMIT = -1202, ///< Keys ultra bucket limit. + UNABLE_OFFLOAD = -1300, ///< The offload functionality is not enabled }; int32_t toInt(ErrorCode errorCode) noexcept; diff --git a/mooncake-store/include/utils.h b/mooncake-store/include/utils.h index 5d981b57..356ff4f7 100644 --- a/mooncake-store/include/utils.h +++ b/mooncake-store/include/utils.h @@ -183,4 +183,6 @@ tl::expected httpGet(const std::string& url); // Network utility: obtain an available TCP port on loopback by binding to 0 int getFreeTcpPort(); +int64_t time_gen(); + } // namespace mooncake diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 75d18292..acbd1fa4 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -2,10 +2,12 @@ #include #include -#include #include #include #include +#include +#include "utils.h" +#include "mutex.h" namespace mooncake { @@ -14,12 +16,12 @@ tl::expected StorageBackend::StoreObject( ResolvePath(path); auto file = create_file(path, FileMode::Write); if (!file) { - LOG(INFO) << "Failed to open file for writing: " << path; + LOG(ERROR) << "Failed to open file for writing: " << path; return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL); } std::vector iovs; - size_t slices_total_size = 0; + int64_t slices_total_size = 0; for (const auto& slice : slices) { iovec io{slice.ptr, slice.size}; iovs.push_back(io); @@ -29,15 +31,15 @@ tl::expected StorageBackend::StoreObject( auto write_result = file->vector_write(iovs.data(), static_cast(iovs.size()), 0); if (!write_result) { - LOG(INFO) << "vector_write failed for: " << path - << ", error: " << write_result.error(); + LOG(ERROR) << "vector_write failed for: " << path + << ", error: " << write_result.error(); return tl::make_unexpected(write_result.error()); } if (*write_result != slices_total_size) { - LOG(INFO) << "Write size mismatch for: " << path - << ", expected: " << slices_total_size - << ", got: " << *write_result; + LOG(ERROR) << "Write size mismatch for: " << path + << ", expected: " << slices_total_size + << ", got: " << *write_result; return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } @@ -54,22 +56,22 @@ tl::expected StorageBackend::StoreObject( ResolvePath(path); auto file = create_file(path, FileMode::Write); if (!file) { - LOG(INFO) << "Failed to open file for writing: " << path; + LOG(ERROR) << "Failed to open file for writing: " << path; return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL); } - size_t file_total_size = data.size(); + int64_t file_total_size = data.size(); auto write_result = file->write(data, file_total_size); if (!write_result) { - LOG(INFO) << "Write failed for: " << path - << ", error: " << write_result.error(); + LOG(ERROR) << "Write failed for: " << path + << ", error: " << write_result.error(); return tl::make_unexpected(write_result.error()); } if (*write_result != file_total_size) { - LOG(INFO) << "Write size mismatch for: " << path - << ", expected: " << file_total_size - << ", got: " << *write_result; + LOG(ERROR) << "Write size mismatch for: " << path + << ", expected: " << file_total_size + << ", got: " << *write_result; return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } @@ -77,20 +79,20 @@ tl::expected StorageBackend::StoreObject( } tl::expected StorageBackend::LoadObject( - const std::string& path, std::vector& slices, size_t length) { + const std::string& path, std::vector& slices, int64_t length) { ResolvePath(path); auto file = create_file(path, FileMode::Read); if (!file) { - LOG(INFO) << "Failed to open file for reading: " << path; + LOG(ERROR) << "Failed to open file for reading: " << path; return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL); } off_t current_offset = 0; - size_t total_bytes_processed = 0; + int64_t total_bytes_processed = 0; std::vector iovs_chunk; off_t chunk_start_offset = 0; - size_t chunk_length = 0; + int64_t chunk_length = 0; auto process_chunk = [&]() -> tl::expected { if (iovs_chunk.empty()) { @@ -101,15 +103,15 @@ tl::expected StorageBackend::LoadObject( iovs_chunk.data(), static_cast(iovs_chunk.size()), chunk_start_offset); if (!read_result) { - LOG(INFO) << "vector_read failed for chunk at offset " - << chunk_start_offset << " for path: " << path - << ", error: " << read_result.error(); + LOG(ERROR) << "vector_read failed for chunk at offset " + << chunk_start_offset << " for path: " << path + << ", error: " << read_result.error(); return tl::make_unexpected(read_result.error()); } if (*read_result != chunk_length) { - LOG(INFO) << "Read size mismatch for chunk in path: " << path - << ", expected: " << chunk_length - << ", got: " << *read_result; + LOG(ERROR) << "Read size mismatch for chunk in path: " << path + << ", expected: " << chunk_length + << ", got: " << *read_result; return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); } @@ -146,9 +148,9 @@ tl::expected StorageBackend::LoadObject( } if (total_bytes_processed != length) { - LOG(INFO) << "Total read size mismatch for: " << path - << ", expected: " << length - << ", got: " << total_bytes_processed; + LOG(ERROR) << "Total read size mismatch for: " << path + << ", expected: " << length + << ", got: " << total_bytes_processed; return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); } @@ -156,23 +158,23 @@ tl::expected StorageBackend::LoadObject( } tl::expected StorageBackend::LoadObject( - const std::string& path, std::string& str, size_t length) { + const std::string& path, std::string& str, int64_t length) { ResolvePath(path); auto file = create_file(path, FileMode::Read); if (!file) { - LOG(INFO) << "Failed to open file for reading: " << path; + LOG(ERROR) << "Failed to open file for reading: " << path; return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL); } auto read_result = file->read(str, length); if (!read_result) { - LOG(INFO) << "read failed for: " << path - << ", error: " << read_result.error(); + LOG(ERROR) << "read failed for: " << path + << ", error: " << read_result.error(); return tl::make_unexpected(read_result.error()); } if (*read_result != length) { - LOG(INFO) << "Read size mismatch for: " << path - << ", expected: " << length << ", got: " << *read_result; + LOG(ERROR) << "Read size mismatch for: " << path + << ", expected: " << length << ", got: " << *read_result; return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); } @@ -267,8 +269,8 @@ void StorageBackend::ResolvePath(const std::string& path) const { fs::path parent_path = full_path.parent_path(); if (!parent_path.empty() && !fs::exists(parent_path)) { if (!fs::create_directories(parent_path, ec) && ec) { - LOG(INFO) << "Failed to create directories: " << parent_path - << ", error: " << ec.message(); + LOG(ERROR) << "Failed to create directories: " << parent_path + << ", error: " << ec.message(); } } } @@ -306,4 +308,509 @@ std::unique_ptr StorageBackend::create_file( return std::make_unique(path, fd); } +BucketIdGenerator::BucketIdGenerator(int64_t start) { + if (start <= 0) { + auto cur_time_stamp = time_gen(); + current_id_ = (cur_time_stamp << TIMESTAMP_SHIFT) | SEQUENCE_ID_SHIFT; + } else { + current_id_ = start; + } +} + +int64_t BucketIdGenerator::NextId() { + return current_id_.fetch_add(1, std::memory_order_relaxed) + 1; +} + +int64_t BucketIdGenerator::CurrentId() { + return current_id_.load(std::memory_order_relaxed); +} + +BucketStorageBackend::BucketStorageBackend(const std::string& storage_path) + : storage_path_(storage_path) {} + +tl::expected BucketStorageBackend::BatchOffload( + const std::unordered_map>& batch_object, + std::function< + ErrorCode(const std::unordered_map&)> + complete_handler) { + if (!initialized_.load(std::memory_order_acquire)) { + LOG(ERROR) + << "Storage backend is not initialized. Call Init() before use."; + return tl::unexpected(ErrorCode::INTERNAL_ERROR); + } + if (batch_object.empty()) { + LOG(ERROR) << "batch object is empty"; + return tl::make_unexpected(ErrorCode::INVALID_KEY); + } + auto bucket_id = bucket_id_generator_->NextId(); + std::vector iovs; + auto build_bucket_result = BuildBucket(batch_object, iovs); + if (!build_bucket_result) { + LOG(ERROR) << "Failed to build bucket with id: " << bucket_id; + return tl::make_unexpected(build_bucket_result.error()); + } + auto bucket = build_bucket_result.value(); + auto write_bucket_result = WriteBucket(bucket_id, bucket, iovs); + if (!write_bucket_result) { + LOG(ERROR) << "Failed to write bucket with id: " << bucket_id; + return tl::make_unexpected(write_bucket_result.error()); + } + if (complete_handler != nullptr) { + auto error_code = complete_handler(bucket->object_metadata); + if (error_code != ErrorCode::OK) { + LOG(ERROR) << "Sync Store object failed,err_code = " << error_code; + return tl::make_unexpected(error_code); + } + } + SharedMutexLocker lock(&mutex_); + total_size_ += bucket->data_size + bucket->meta_size; + for (auto object_metadata_it : bucket->object_metadata) { + object_bucket_map_.emplace( + object_metadata_it.first, + StorageObjectMetadata{bucket_id, object_metadata_it.second.offset, + object_metadata_it.second.key_size, + object_metadata_it.second.data_size}); + } + buckets_.emplace(bucket_id, std::move(bucket)); + return bucket_id; +} + +tl::expected BucketStorageBackend::BatchQuery( + const std::vector& keys, + std::unordered_map& + batch_object_metadata) { + SharedMutexLocker lock(&mutex_, shared_lock); + for (const auto& key : keys) { + auto object_metadata_it = object_bucket_map_.find(key); + if (object_metadata_it != object_bucket_map_.end()) { + batch_object_metadata.emplace(key, object_metadata_it->second); + } else { + LOG(ERROR) << "Key " << key << " does not exist"; + return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + } + } + return {}; +} + +tl::expected BucketStorageBackend::BatchLoad( + std::unordered_map& batch_object) { + std::unordered_map> bucket_key_map; + { + SharedMutexLocker lock(&mutex_, shared_lock); + for (const auto& key_it : batch_object) { + auto object_bucket_it = object_bucket_map_.find(key_it.first); + if (object_bucket_it == object_bucket_map_.end()) { + LOG(ERROR) << "key " << key_it.first << " does not exist"; + return tl::make_unexpected(ErrorCode::INVALID_KEY); + } + auto [bucket_keys_it, _] = + bucket_key_map.try_emplace(object_bucket_it->second.bucket_id); + bucket_keys_it->second.emplace_back(key_it.first); + } + } + for (const auto& bucket_key_it : bucket_key_map) { + auto result = BatchLoadBucket(bucket_key_it.first, bucket_key_it.second, + batch_object); + if (!result) { + LOG(ERROR) << "Failed to load bucket " << bucket_key_it.first; + return result; + } + } + return {}; +} + +tl::expected BucketStorageBackend::GetBucketKeys( + int64_t bucket_id, std::vector& bucket_keys) { + SharedMutexLocker locker(&mutex_, shared_lock); + auto bucket_it = buckets_.find(bucket_id); + if (bucket_it == buckets_.end()) { + return tl::make_unexpected(ErrorCode::BUCKET_NOT_FOUND); + } + for (const auto& key : bucket_it->second->keys) { + bucket_keys.emplace_back(key); + } + return {}; +} + +tl::expected BucketStorageBackend::Init() { + namespace fs = std::filesystem; + try { + if (initialized_.load(std::memory_order_acquire)) { + LOG(ERROR) << "Storage backend already initialized"; + return tl::unexpected(ErrorCode::INTERNAL_ERROR); + } + SharedMutexLocker lock(&mutex_); + object_bucket_map_.clear(); + buckets_.clear(); + total_size_ = 0; + int64_t max_bucket_id = BucketIdGenerator::INIT_NEW_START_ID; + for (const auto& entry : + fs::recursive_directory_iterator(storage_path_)) { + if (entry.is_regular_file() && + entry.path().extension() == BUCKET_METADATA_FILE_SUFFIX) { + auto bucket_id_str = entry.path().stem(); + int64_t bucket_id = std::stoll(bucket_id_str); + auto [metadata_it, success] = buckets_.try_emplace( + bucket_id, std::make_shared()); + if (!success) { + LOG(ERROR) << "Failed to load bucket " << bucket_id_str; + return tl::unexpected(ErrorCode::BUCKET_ALREADY_EXISTS); + } + auto load_bucket_metadata_result = + LoadBucketMetadata(bucket_id, metadata_it->second); + if (!load_bucket_metadata_result) { + LOG(ERROR) + << "Failed to load metadata for bucket: " + << bucket_id_str + << ", will delete the bucket's data and metadata"; + + auto bucket_data_path_res = GetBucketDataPath(bucket_id); + if (bucket_data_path_res) { + fs::remove(bucket_data_path_res.value()); + } + + auto bucket_meta_path_res = + GetBucketMetadataPath(bucket_id); + if (bucket_meta_path_res) { + fs::remove(bucket_meta_path_res.value()); + } + + buckets_.erase(bucket_id); + continue; + } + auto& meta = *(metadata_it->second); + if (meta.data_size == 0 || meta.meta_size == 0 || + meta.object_metadata.empty() || meta.keys.empty()) { + LOG(ERROR) << "Metadata validation failed for bucket: " + << bucket_id_str + << ", will delete the bucket's data and " + "metadata. Detailed values:"; + LOG(ERROR) << " data_size: " << meta.data_size + << " (should not be 0)"; + LOG(ERROR) << " meta_size: " << meta.meta_size + << " (should not be 0)"; + LOG(ERROR) + << " object_metadata.size(): " + << meta.object_metadata.size() << " (empty: " + << (meta.object_metadata.empty() ? "true" : "false") + << ")"; + + LOG(ERROR) + << " keys.size(): " << meta.keys.size() + << " (empty: " << (meta.keys.empty() ? "true" : "false") + << ")"; + auto bucket_data_path_res = GetBucketDataPath(bucket_id); + if (bucket_data_path_res) { + fs::remove(bucket_data_path_res.value()); + } + + auto bucket_meta_path_res = + GetBucketMetadataPath(bucket_id); + if (bucket_meta_path_res) { + fs::remove(bucket_meta_path_res.value()); + } + + buckets_.erase(bucket_id); + continue; + } + if (bucket_id > max_bucket_id) { + max_bucket_id = bucket_id; + } + total_size_ += metadata_it->second->data_size + + metadata_it->second->meta_size; + for (const auto& object_metadata_it : + metadata_it->second->object_metadata) { + object_bucket_map_.emplace( + object_metadata_it.first, + StorageObjectMetadata{ + metadata_it->first, + object_metadata_it.second.offset, + object_metadata_it.second.key_size, + object_metadata_it.second.data_size}); + } + } + } + bucket_id_generator_.emplace(max_bucket_id); + if (max_bucket_id == BucketIdGenerator::INIT_NEW_START_ID) { + LOG(INFO) << "Initialized BucketIdGenerator with fresh start. " + "No existing buckets found; starting from ID: " + << bucket_id_generator_->CurrentId(); + } else { + LOG(INFO) << "Initialized BucketIdGenerator from existing state. " + << "Last used bucket ID was " << max_bucket_id; + } + initialized_.store(true, std::memory_order_release); + } catch (const std::exception& e) { + LOG(ERROR) << "Bucket storage backend initialize error: " << e.what() + << std::endl; + return tl::unexpected(ErrorCode::INTERNAL_ERROR); + } + + return {}; +} + +tl::expected BucketStorageBackend::IsExist( + const std::string& key) { + SharedMutexLocker lock(&mutex_, shared_lock); + auto bucket_id_it = object_bucket_map_.find(key); + if (bucket_id_it != object_bucket_map_.end()) { + return true; + } + return false; +} + +tl::expected BucketStorageBackend::BucketScan( + int64_t bucket_id, + std::unordered_map& objects, + std::vector& buckets, int64_t limit) { + SharedMutexLocker lock(&mutex_, shared_lock); + auto bucket_it = buckets_.lower_bound(bucket_id); + for (; bucket_it != buckets_.end(); ++bucket_it) { + if (bucket_it->second->keys.size() > limit) { + 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_ULTRA_BUCKET_LIMIT); + } + if (bucket_it->second->keys.size() + objects.size() > limit) { + return bucket_it->first; + } + buckets.emplace_back(bucket_it->first); + for (const auto& object_it : bucket_it->second->object_metadata) { + objects.emplace(object_it.first, object_it.second); + } + } + return 0; +} + +tl::expected +BucketStorageBackend::GetStoreMetadata() { + SharedMutexLocker lock(&mutex_, shared_lock); + OffloadMetadata metadata{object_bucket_map_.size(), total_size_}; + return metadata; +} + +tl::expected, ErrorCode> +BucketStorageBackend::BuildBucket( + const std::unordered_map>& batch_object, + std::vector& iovs) { + auto bucket = std::make_shared(); + int64_t storage_offset = 0; + for (const auto& object : batch_object) { + if (object.second.empty()) { + LOG(ERROR) << "Failed to create bucket, object is empty"; + return tl::make_unexpected(ErrorCode::INVALID_KEY); + } + int64_t object_total_size = 0; + iovs.emplace_back( + iovec{const_cast(object.first.data()), object.first.size()}); + for (const auto& slice : object.second) { + object_total_size += slice.size; + iovs.emplace_back(iovec{slice.ptr, slice.size}); + } + bucket->data_size += object_total_size + object.first.size(); + bucket->object_metadata.emplace( + object.first, + BucketObjectMetadata{storage_offset, object.first.size(), + object_total_size}); + bucket->keys.push_back(object.first); + storage_offset += object_total_size + object.first.size(); + } + return bucket; +} + +tl::expected BucketStorageBackend::WriteBucket( + int64_t bucket_id, std::shared_ptr bucket_metadata, + std::vector& iovs) { + auto bucket_data_path_res = GetBucketDataPath(bucket_id); + if (!bucket_data_path_res) { + LOG(ERROR) << "Failed to get bucket data path, bucket_id=" << bucket_id; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + auto bucket_data_path = bucket_data_path_res.value(); + auto open_file_result = OpenFile(bucket_data_path, FileMode::Write); + if (!open_file_result) { + LOG(ERROR) << "Failed to open file for bucket writing: " + << bucket_data_path; + return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL); + } + auto file = std::move(open_file_result.value()); + + auto write_result = file->vector_write(iovs.data(), iovs.size(), 0); + if (!write_result) { + LOG(ERROR) << "vector_write failed for: " << bucket_id + << ", error: " << write_result.error(); + return tl::make_unexpected(write_result.error()); + } + + auto store_bucket_metadata_result = + StoreBucketMetadata(bucket_id, bucket_metadata); + if (!store_bucket_metadata_result) { + LOG(ERROR) << "Failed to store bucket metadata, error: " + << store_bucket_metadata_result.error(); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + return {}; +} + +tl::expected BucketStorageBackend::StoreBucketMetadata( + int64_t id, std::shared_ptr metadata) { + auto meta_path_res = GetBucketMetadataPath(id); + if (!meta_path_res) { + LOG(ERROR) << "Failed to get bucket metadata path, bucket_id=" << id; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + auto meta_path = meta_path_res.value(); + auto open_file_result = OpenFile(meta_path, FileMode::Write); + if (!open_file_result) { + LOG(ERROR) << "Failed to open file for bucket writing: " << meta_path; + return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL); + } + auto file = std::move(open_file_result.value()); + std::string str; + struct_pb::to_pb(*metadata, str); + auto write_result = file->write(str, str.size()); + if (!write_result) { + LOG(ERROR) << "Write failed for: " << meta_path + << ", error: " << write_result.error(); + return tl::make_unexpected(write_result.error()); + } + metadata->meta_size = str.size(); + return {}; +} + +tl::expected BucketStorageBackend::LoadBucketMetadata( + int64_t id, std::shared_ptr metadata) { + auto meta_path_res = GetBucketMetadataPath(id); + if (!meta_path_res) { + LOG(ERROR) << "Failed to get bucket metadata path, bucket_id=" << id; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + auto meta_path = meta_path_res.value(); + auto open_file_result = OpenFile(meta_path, FileMode::Read); + if (!open_file_result) { + LOG(ERROR) << "Failed to open file for reading: " << meta_path; + return tl::make_unexpected(open_file_result.error()); + } + auto file = std::move(open_file_result.value()); + std::string str; + int64_t size = std::filesystem::file_size(meta_path); + auto read_result = file->read(str, size); + if (!read_result) { + LOG(ERROR) << "read failed for: " << meta_path + << ", error: " << read_result.error(); + return tl::make_unexpected(read_result.error()); + } + if (*read_result != size) { + LOG(ERROR) << "Read size mismatch for: " << meta_path + << ", expected: " << size << ", got: " << *read_result; + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + try { + struct_pb::from_pb(*metadata, str); + metadata->meta_size = size; + } catch (const std::exception& e) { + LOG(ERROR) << "Metadata parsing failed with exception: " << e.what(); + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } catch (...) { + LOG(ERROR) << "Metadata parsing failed with unknown exception"; + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + return {}; +} + +tl::expected BucketStorageBackend::BatchLoadBucket( + int64_t bucket_id, const std::vector& keys, + std::unordered_map& batched_slices) { + SharedMutexLocker locker(&mutex_, shared_lock); + auto storage_filepath_res = GetBucketDataPath(bucket_id); + if (!storage_filepath_res) { + LOG(ERROR) << "Failed to get bucket data path, bucket_id=" << bucket_id; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + auto storage_filepath = storage_filepath_res.value(); + auto open_file_result = OpenFile(storage_filepath, FileMode::Read); + if (!open_file_result) { + LOG(ERROR) << "Failed to open file for reading: " << storage_filepath; + return tl::make_unexpected(open_file_result.error()); + } + auto file = std::move(open_file_result.value()); + for (const auto& key : keys) { + int64_t offset; + auto slice = batched_slices[key]; + auto bucket = buckets_.find(bucket_id); + if (bucket == buckets_.end()) { + LOG(ERROR) << "Bucket not found with id: " << bucket_id; + return tl::make_unexpected(ErrorCode::BUCKET_NOT_FOUND); + } + auto object_metadata = buckets_[bucket_id]->object_metadata.find(key); + if (object_metadata == buckets_[bucket_id]->object_metadata.end()) { + LOG(ERROR) << "Object metadata not found for key '" << key + << "' in bucket " << bucket_id; + return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + } + if (object_metadata->second.data_size != slice.size) { + LOG(ERROR) << "Read size mismatch for: " << storage_filepath + << ", expected: " << object_metadata->second.data_size + << ", got: " << slice.size; + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + offset = object_metadata->second.offset; + std::vector iovs; + iovs.emplace_back(iovec{slice.ptr, slice.size}); + auto read_result = file->vector_read( + iovs.data(), static_cast(iovs.size()), offset + key.size()); + if (!read_result) { + LOG(ERROR) << "vector_read failed for: " << storage_filepath + << ", error: " << read_result.error(); + return tl::make_unexpected(read_result.error()); + } + if (*read_result != slice.size) { + LOG(ERROR) << "Read size mismatch for: " << storage_filepath + << ", expected: " << slice.size + << ", got: " << *read_result; + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + } + return {}; +} + +tl::expected BucketStorageBackend::GetBucketDataPath( + int64_t bucket_id) { + std::string sep = + storage_path_.empty() || storage_path_.back() == '/' ? "" : "/"; + return storage_path_ + sep + std::to_string(bucket_id); +} + +tl::expected +BucketStorageBackend::GetBucketMetadataPath(int64_t bucket_id) { + auto bucket_data_path_res = GetBucketDataPath(bucket_id); + if (!bucket_data_path_res) { + LOG(ERROR) << "Failed to get bucket data path, bucket_id=" << bucket_id; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + return bucket_data_path_res.value() + ".meta"; +} + +tl::expected, ErrorCode> +BucketStorageBackend::OpenFile(const std::string& path, FileMode mode) const { + int flags = O_CLOEXEC; + int access_mode = 0; + switch (mode) { + case FileMode::Read: + access_mode = O_RDONLY; + break; + case FileMode::Write: + access_mode = O_WRONLY | O_CREAT | O_TRUNC; + break; + } + + int fd = open(path.c_str(), flags | access_mode, 0644); + if (fd < 0) { + return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL); + } + return std::make_unique(path, fd); +} + } // namespace mooncake diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp index 8923ffd6..449bed19 100644 --- a/mooncake-store/src/utils.cpp +++ b/mooncake-store/src/utils.cpp @@ -167,4 +167,10 @@ int getFreeTcpPort() { return port; } +int64_t time_gen() { + return std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); +} + } // namespace mooncake diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 864d109d..28bd63c9 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -30,7 +30,8 @@ add_store_test(pybind_client_test pybind_client_test.cpp) add_store_test(client_metrics_test client_metrics_test.cpp) add_store_test(serializer_test serializer_test.cpp) add_store_test(non_ha_reconnect_test non_ha_reconnect_test.cpp) - +add_store_test(storage_backend_test storage_backend_test.cpp) +add_store_test(mutex_test mutex_test.cpp) add_subdirectory(e2e) add_executable(high_availability_test high_availability_test.cpp) diff --git a/mooncake-store/tests/mutex_test.cpp b/mooncake-store/tests/mutex_test.cpp new file mode 100644 index 00000000..144bcb27 --- /dev/null +++ b/mooncake-store/tests/mutex_test.cpp @@ -0,0 +1,176 @@ +#include "mutex.h" + +#include +#include + +#include +#include +#include +#include + +namespace mooncake::test { + +class SharedMutexTest : public ::testing::Test { + protected: + void SetUp() override { + google::InitGoogleLogging("SharedMutexTest"); + FLAGS_logtostderr = true; + } + void TearDown() override { google::ShutdownGoogleLogging(); } +}; + +TEST(SharedMutexTest, CanLockExclusive) { + SharedMutex mtx; + EXPECT_NO_THROW({ + mtx.lock(); + mtx.unlock(); + }); +} + +TEST(SharedMutexTest, CanLockShared) { + SharedMutex mtx; + EXPECT_NO_THROW({ + mtx.lock_shared(); + mtx.unlock_shared(); + }); +} + +TEST(SharedMutexTest, ExclusiveAccessIsMutuallyExclusive) { + SharedMutex mtx; + std::atomic counter{0}; + std::vector threads; + + for (int i = 0; i < 5; ++i) { + threads.emplace_back([&mtx, &counter]() { + mtx.lock(); + int val = counter.load(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + counter.store(val + 1); + mtx.unlock(); + }); + } + + for (auto& t : threads) t.join(); + EXPECT_EQ(counter.load(), 5); +} + +TEST(SharedMutexTest, SharedAccessIsConcurrent) { + SharedMutex mtx; + std::atomic active_readers{0}; + std::atomic peak_readers{0}; + std::vector threads; + + for (int i = 0; i < 5; ++i) { + threads.emplace_back([&mtx, &active_readers, &peak_readers]() { + mtx.lock_shared(); + int current = ++active_readers; + if (current > peak_readers) peak_readers = current; + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + --active_readers; + mtx.unlock_shared(); + }); + } + + for (auto& t : threads) t.join(); + + EXPECT_EQ(active_readers.load(), 0); // All readers have exited + EXPECT_GE(peak_readers.load(), 2); // At least two readers ran concurrently +} + +TEST(SharedMutexTest, WriterBlocksReaders) { + SharedMutex mtx; + std::atomic reader_started{false}; + std::atomic writer_proceeded{false}; + + std::thread reader([&]() { + mtx.lock_shared(); + reader_started = true; + while (!writer_proceeded) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + mtx.unlock_shared(); + }); + + // Give reader time to start + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + + auto try_write = + mtx.try_lock(); // Should fail because reader holds shared lock + EXPECT_FALSE(try_write); + + // Let reader finish + writer_proceeded = true; + reader.join(); + + // Now writer should be able to acquire the lock + mtx.lock(); + EXPECT_TRUE(true); // No deadlock occurred + mtx.unlock(); +} + +TEST(SharedMutexTest, LocksOnConstructionExclusive) { + SharedMutex mtx; + { + SharedMutexLocker locker(&mtx); + // Lock should still be held before destruction + EXPECT_FALSE(mtx.try_lock()); // Cannot acquire another exclusive lock + EXPECT_FALSE(mtx.try_lock_shared()); // Shared lock may also be blocked + // (implementation-defined) + } // Destructor automatically unlocks + EXPECT_TRUE(mtx.try_lock()); // After destruction, lock should be available +} + +TEST(SharedMutexTest, LocksOnConstructionShared) { + SharedMutex mtx; + { + SharedMutexLocker locker(&mtx, shared_lock); + EXPECT_TRUE( + mtx.try_lock_shared()); // Multiple shared locks should be allowed + // Note: This test does not attempt recursive locking (UB), just checks + // concurrent shared access. + } + EXPECT_TRUE( + mtx.try_lock_shared()); // Should still be available after unlock +} + +TEST(SharedMutexTest, ManualLockUnlock) { + SharedMutex mtx; + SharedMutexLocker locker(nullptr); // Initialize without a mutex + EXPECT_NO_THROW(locker.unlock()); // Unlocking a null locker should be safe + + SharedMutexLocker temp(&mtx); + temp.unlock(); // Manually release the lock + EXPECT_TRUE(mtx.try_lock()); // Now we should be able to acquire it +} + +TEST(SharedMutexTest, TryLockSuccess) { + SharedMutex mtx; + SharedMutexLocker locker(&mtx); + locker.unlock(); // Ensure it's released first + + bool result = locker.try_lock(); + EXPECT_TRUE(result); + EXPECT_FALSE(locker.try_lock()); // Should not allow re-locking +} + +TEST(SharedMutexTest, TryLockSharedSuccess) { + SharedMutex mtx; + SharedMutexLocker locker(&mtx, shared_lock); + locker.unlock(); + + bool result = locker.try_lock_shared(); + EXPECT_TRUE(result); + EXPECT_FALSE(locker.try_lock_shared()); // Should not allow re-locking +} + +TEST(SharedMutexTest, HandlesNullptrSafely) { + SharedMutexLocker locker(nullptr); + EXPECT_NO_THROW(locker.lock()); + EXPECT_NO_THROW(locker.lock_shared()); + EXPECT_NO_THROW(locker.try_lock()); + EXPECT_NO_THROW(locker.try_lock_shared()); + EXPECT_NO_THROW(locker.unlock()); + // Should not crash under any operation +} + +} // namespace mooncake::test \ No newline at end of file diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp new file mode 100644 index 00000000..766477d9 --- /dev/null +++ b/mooncake-store/tests/storage_backend_test.cpp @@ -0,0 +1,313 @@ +#include +#include +#include "storage_backend.h" +#include "allocator.h" +#include "utils.h" +#include +#include +#include +#include + +namespace fs = std::filesystem; +namespace mooncake::test { +class StorageBackendTest : public ::testing::Test { + protected: + void SetUp() override { + google::InitGoogleLogging("StorageBackendTest"); + FLAGS_logtostderr = true; + } + + void TearDown() override { + google::ShutdownGoogleLogging(); + std::string data_path = + std::filesystem::current_path().string() + "/data"; + LOG(INFO) << "Clear test data..."; + for (const auto& entry : fs::directory_iterator(data_path)) { + if (entry.is_regular_file()) { + fs::remove(entry.path()); + } + } + } +}; + +tl::expected BatchOffload( + std::vector& keys, + std::unordered_map& batch_data, + std::shared_ptr& client_buffer_allocator, + BucketStorageBackend& storage_backend, std::vector& buckets) { + size_t bucket_sz = 10; + size_t batch_sz = 10; + size_t data_sz = 10; + for (size_t i = 0; i < bucket_sz; i++) { + std::unordered_map> batched_slices; + for (size_t j = 0; j < batch_sz; j++) { + std::string key = + "test_key_i_" + std::to_string(i) + "_j_" + std::to_string(j); + size_t data_size = 0; + std::string all_data; + std::vector slices; + for (size_t k = 0; k < data_sz; k++) { + std::string data = "test_data_i_" + std::to_string(i) + "_j_" + + std::to_string(j) + "_k_" + + std::to_string(k); + all_data += data; + data_size += data.size(); + void* buffer = client_buffer_allocator->allocate(data.size()); + memcpy(buffer, data.data(), data.size()); + slices.emplace_back(Slice{buffer, data.size()}); + } + batched_slices.emplace(key, slices); + batch_data.emplace(key, all_data); + keys.emplace_back(key); + } + auto batch_store_object_one_result = storage_backend.BatchOffload( + batched_slices, + [&](const std::unordered_map& + keys) { + if (keys.size() != batched_slices.size()) { + return ErrorCode::INVALID_KEY; + } + for (const auto& key : keys) { + if (batched_slices.find(key.first) == + batched_slices.end()) { + return ErrorCode::INVALID_KEY; + } + } + return ErrorCode::OK; + }); + if (!batch_store_object_one_result) { + return tl::make_unexpected(batch_store_object_one_result.error()); + } + buckets.emplace_back(batch_store_object_one_result.value()); + } + return {}; +} + +TEST_F(StorageBackendTest, StorageBackendAll) { + std::string data_path = std::filesystem::current_path().string() + "/data"; + fs::create_directories(data_path); + std::shared_ptr client_buffer_allocator = + std::make_shared(128 * 1024 * 1024); + BucketStorageBackend storage_backend(data_path); + for (const auto& entry : fs::directory_iterator(data_path)) { + if (entry.is_regular_file()) { + fs::remove(entry.path()); + } + } + ASSERT_TRUE(storage_backend.Init()); + ASSERT_TRUE(fs::directory_iterator(data_path) == fs::directory_iterator{}); + ASSERT_TRUE(!storage_backend.Init()); + std::unordered_map test_data; + std::vector keys; + std::vector buckets; + auto test_batch_store_object_result = BatchOffload( + keys, test_data, client_buffer_allocator, storage_backend, buckets); + ASSERT_TRUE(test_batch_store_object_result); + + std::unordered_map + batche_object_metadata; + auto batch_query_object_result_two = + storage_backend.BatchQuery(keys, batche_object_metadata); + ASSERT_TRUE(batch_query_object_result_two); + ASSERT_EQ(batche_object_metadata.size(), test_data.size()); + for (const auto& keys_it : test_data) { + auto metadata = batche_object_metadata[keys_it.first]; + ASSERT_EQ(keys_it.second.size(), metadata.data_size); + } + + std::unordered_map batche_object; + + for (auto test_data_it : test_data) { + void* buffer = + client_buffer_allocator->allocate(test_data_it.second.size()); + batche_object.emplace(test_data_it.first, + Slice{buffer, test_data_it.second.size()}); + } + + auto batch_load_object_result = storage_backend.BatchLoad(batche_object); + ASSERT_TRUE(batch_load_object_result); + ASSERT_EQ(batche_object.size(), test_data.size()); + for (const auto& test_data_it : test_data) { + auto is_exist_object_result = + storage_backend.IsExist(test_data_it.first); + ASSERT_TRUE(is_exist_object_result); + ASSERT_TRUE(is_exist_object_result.value()); + auto object_it = batche_object.find(test_data_it.first); + ASSERT_TRUE(object_it != batche_object.end()); + char* buf = new char[object_it->second.size + 1]; + buf[object_it->second.size] = '\0'; + memcpy(buf, object_it->second.ptr, object_it->second.size); + auto data = std::string(buf); + ASSERT_EQ(data, test_data_it.second); + delete[] buf; + } +} +TEST_F(StorageBackendTest, BucketScan) { + std::string data_path = std::filesystem::current_path().string() + "/data"; + fs::create_directories(data_path); + std::shared_ptr client_buffer_allocator = + std::make_shared(128 * 1024 * 1024); + BucketStorageBackend storage_backend(data_path); + for (const auto& entry : fs::directory_iterator(data_path)) { + if (entry.is_regular_file()) { + fs::remove(entry.path()); + } + } + ASSERT_TRUE(storage_backend.Init()); + ASSERT_TRUE(!storage_backend.Init()); + std::unordered_map test_data; + std::vector keys; + std::vector buckets; + auto test_batch_store_object_result = BatchOffload( + keys, test_data, client_buffer_allocator, storage_backend, buckets); + ASSERT_TRUE(test_batch_store_object_result); + std::unordered_map objects; + std::vector scan_buckets; + auto res = storage_backend.BucketScan(0, objects, scan_buckets, 10); + ASSERT_TRUE(res); + ASSERT_EQ(res.value(), buckets.at(1)); + for (const auto& object : objects) { + ASSERT_EQ(object.second.data_size, test_data.at(object.first).size()); + ASSERT_EQ(object.second.key_size, object.first.size()); + } + ASSERT_EQ(scan_buckets.size(), 1); + ASSERT_EQ(scan_buckets.at(0), buckets.at(0)); + objects.clear(); + scan_buckets.clear(); + res = storage_backend.BucketScan(0, objects, scan_buckets, 45); + ASSERT_TRUE(res); + ASSERT_EQ(res.value(), buckets.at(4)); + ASSERT_EQ(scan_buckets.size(), 4); + for (int i = 0; i < 4; i++) { + ASSERT_EQ(scan_buckets.at(i), buckets.at(i)); + } + + objects.clear(); + scan_buckets.clear(); + res = storage_backend.BucketScan(buckets.at(4), objects, scan_buckets, 45); + ASSERT_TRUE(res); + ASSERT_EQ(res.value(), buckets.at(8)); + ASSERT_EQ(scan_buckets.size(), 4); + for (int i = 0; i < 4; i++) { + ASSERT_EQ(scan_buckets.at(i), buckets.at(i + 4)); + } + + objects.clear(); + scan_buckets.clear(); + res = storage_backend.BucketScan(buckets.at(9), objects, scan_buckets, 45); + ASSERT_TRUE(res); + ASSERT_EQ(res.value(), 0); + ASSERT_EQ(scan_buckets.size(), 1); + ASSERT_EQ(scan_buckets.at(0), buckets.at(9)); + + objects.clear(); + scan_buckets.clear(); + res = storage_backend.BucketScan(buckets.at(9) + 10, objects, scan_buckets, + 45); + ASSERT_TRUE(res); + ASSERT_EQ(res.value(), 0); + ASSERT_EQ(scan_buckets.size(), 0); + + objects.clear(); + scan_buckets.clear(); + res = storage_backend.BucketScan(0, objects, scan_buckets, 8); + ASSERT_TRUE(!res); + ASSERT_EQ(res.error(), ErrorCode::KEYS_ULTRA_BUCKET_LIMIT); + ASSERT_EQ(scan_buckets.size(), 0); + ASSERT_EQ(objects.size(), 0); +} + +TEST_F(StorageBackendTest, InitializeWithValidStart) { + BucketIdGenerator gen(100); + EXPECT_EQ(gen.CurrentId(), 100); + EXPECT_EQ(gen.NextId(), 101); + EXPECT_EQ(gen.NextId(), 102); +} + +TEST_F(StorageBackendTest, InitializeWithInvalidStart_UseTimestampFallback) { + auto time = time_gen(); + int64_t expected = (time << 12) | 0; + BucketIdGenerator gen( + BucketIdGenerator::INIT_NEW_START_ID); // invalid start + LOG(INFO) << "expected is: " << expected << " gen is: " << gen.CurrentId(); + EXPECT_TRUE(expected <= gen.CurrentId()); +} + +TEST_F(StorageBackendTest, NextIdReturnsNewValue) { + BucketIdGenerator gen(10); + + EXPECT_EQ(gen.NextId(), 11); // Returns the new value: old + 1 = 11 + EXPECT_EQ(gen.NextId(), 12); + EXPECT_EQ(gen.CurrentId(), 12); +} + +TEST_F(StorageBackendTest, IdsAreMonotonicallyIncreasing) { + BucketIdGenerator gen(100); + + int64_t id1 = gen.NextId(); // 101 + int64_t id2 = gen.NextId(); // 102 + int64_t id3 = gen.NextId(); // 103 + + EXPECT_LT(id1, id2); + EXPECT_LT(id2, id3); + EXPECT_EQ(id1 + 1, id2); + EXPECT_EQ(id2 + 1, id3); +} + +TEST_F(StorageBackendTest, Concurrency_UniquenessAndNoDuplicates) { + const int num_threads = 4; + const int iterations_per_thread = 1000; + + BucketIdGenerator gen(1); + std::vector threads; + std::vector all_ids; + std::mutex mutex; + + auto worker = [&gen, &all_ids, &mutex] { + for (int i = 0; i < iterations_per_thread; ++i) { + int64_t id = gen.NextId(); // Returns the next ID (new value) + { + std::lock_guard lock(mutex); + all_ids.push_back(id); + } + } + }; + + for (int i = 0; i < num_threads; ++i) { + threads.emplace_back(worker); + } + + for (auto& t : threads) { + t.join(); + } + + // Check uniqueness: no duplicate IDs should exist + std::set unique_ids(all_ids.begin(), all_ids.end()); + EXPECT_EQ(unique_ids.size(), all_ids.size()) + << "Duplicate IDs detected in concurrent execution!"; +} + +TEST_F(StorageBackendTest, CurrentIdReturnsLatestValue) { + BucketIdGenerator gen(50); + + EXPECT_EQ(gen.CurrentId(), 50); + EXPECT_EQ(gen.NextId(), 51); + EXPECT_EQ(gen.CurrentId(), 51); + EXPECT_EQ(gen.NextId(), 52); + EXPECT_EQ(gen.CurrentId(), 52); +} + +TEST_F(StorageBackendTest, LargeNumberOfIds_NoOverflowInLifetime) { + BucketIdGenerator gen(1000); + int64_t last_id = 1000; + + for (int i = 0; i < 100000; ++i) { + int64_t id = gen.NextId(); + EXPECT_EQ(id, last_id + 1); // Each ID increments by exactly 1 + last_id = id; + } + + EXPECT_GE(last_id, 101000); // Should have increased by at least 100,000 +} + +} // namespace mooncake::test \ No newline at end of file