forked from mooncake-track/Mooncake
[Store] Add BucketStorageBackend (#968)
This commit is contained in:
楠志
GitHub
parent
c31bb270ae
commit
2cf86bf385
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
#define THREAD_SAFETY_ANALYSIS_MUTEX_H
|
||||
|
||||
#include <mutex>
|
||||
#include <shared_mutex>
|
||||
|
||||
// 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
|
||||
|
|
|
|||
|
|
@ -1,15 +1,46 @@
|
|||
#pragma once
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <filesystem>
|
||||
#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<std::string, BucketObjectMetadata> object_metadata;
|
||||
std::vector<std::string> 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<void, ErrorCode> LoadObject(const std::string& path,
|
||||
std::vector<Slice>& slices,
|
||||
size_t length);
|
||||
int64_t length);
|
||||
|
||||
/**
|
||||
* @brief Loads an object as a string
|
||||
|
|
@ -124,7 +155,7 @@ class StorageBackend {
|
|||
* @return tl::expected<void, ErrorCode> indicating operation status
|
||||
*/
|
||||
tl::expected<void, ErrorCode> 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<int64_t> 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<void, ErrorCode> indicating operation status
|
||||
*/
|
||||
tl::expected<int64_t, ErrorCode> BatchOffload(
|
||||
const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
|
||||
std::function<ErrorCode(
|
||||
const std::unordered_map<std::string, BucketObjectMetadata>&)>
|
||||
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<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);
|
||||
|
||||
/**
|
||||
* @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<void, ErrorCode> indicating operation status.
|
||||
*/
|
||||
tl::expected<void, ErrorCode> BatchLoad(
|
||||
std::unordered_map<std::string, Slice>& 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<void, ErrorCode> indicating operation status.
|
||||
*/
|
||||
tl::expected<void, ErrorCode> GetBucketKeys(
|
||||
int64_t bucket_id, std::vector<std::string>& bucket_keys);
|
||||
|
||||
/**
|
||||
* @brief Initializes the bucket storage backend.
|
||||
* @return tl::expected<void, ErrorCode> indicating operation status.
|
||||
*/
|
||||
tl::expected<void, ErrorCode> 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<void, ErrorCode> indicating operation status.
|
||||
*/
|
||||
tl::expected<bool, ErrorCode> 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<int64_t, ErrorCode>
|
||||
* - 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<int64_t, ErrorCode> BucketScan(
|
||||
int64_t bucket_id,
|
||||
std::unordered_map<std::string, BucketObjectMetadata>& objects,
|
||||
std::vector<int64_t>& 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<OffloadMetadata, ErrorCode> GetStoreMetadata();
|
||||
|
||||
private:
|
||||
tl::expected<std::shared_ptr<BucketMetadata>, ErrorCode> BuildBucket(
|
||||
const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
|
||||
std::vector<iovec>& iovs);
|
||||
|
||||
tl::expected<void, ErrorCode> WriteBucket(
|
||||
int64_t bucket_id, std::shared_ptr<BucketMetadata> bucket_metadata,
|
||||
std::vector<iovec>& iovs);
|
||||
|
||||
tl::expected<void, ErrorCode> StoreBucketMetadata(
|
||||
int64_t bucket_id, std::shared_ptr<BucketMetadata> bucket_metadata);
|
||||
|
||||
tl::expected<void, ErrorCode> LoadBucketMetadata(
|
||||
int64_t bucket_id, std::shared_ptr<BucketMetadata> bucket_metadata);
|
||||
|
||||
tl::expected<void, ErrorCode> BatchLoadBucket(
|
||||
int64_t bucket_id, const std::vector<std::string>& keys,
|
||||
std::unordered_map<std::string, Slice>& batched_slices);
|
||||
|
||||
tl::expected<int64_t, ErrorCode> CreateBucketId();
|
||||
|
||||
tl::expected<std::string, ErrorCode> GetBucketMetadataPath(
|
||||
int64_t bucket_id);
|
||||
|
||||
tl::expected<std::string, ErrorCode> GetBucketDataPath(int64_t bucket_id);
|
||||
|
||||
tl::expected<std::unique_ptr<StorageFile>, ErrorCode> OpenFile(
|
||||
const std::string& path, FileMode mode) const;
|
||||
|
||||
private:
|
||||
std::atomic<bool> initialized_{false};
|
||||
std::optional<BucketIdGenerator> 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<std::string, StorageObjectMetadata> GUARDED_BY(mutex_)
|
||||
object_bucket_map_;
|
||||
std::map<int64_t, std::shared_ptr<BucketMetadata>> GUARDED_BY(
|
||||
mutex_) buckets_;
|
||||
};
|
||||
} // namespace mooncake
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -183,4 +183,6 @@ tl::expected<std::string, int> 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
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@
|
|||
|
||||
#include <fcntl.h>
|
||||
#include <unistd.h>
|
||||
#include <sys/uio.h>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <regex>
|
||||
#include <ylt/struct_pb.hpp>
|
||||
#include "utils.h"
|
||||
#include "mutex.h"
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
|
|
@ -14,12 +16,12 @@ tl::expected<void, ErrorCode> 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<iovec> 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<void, ErrorCode> StorageBackend::StoreObject(
|
|||
auto write_result =
|
||||
file->vector_write(iovs.data(), static_cast<int>(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<void, ErrorCode> 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<void, ErrorCode> StorageBackend::StoreObject(
|
|||
}
|
||||
|
||||
tl::expected<void, ErrorCode> StorageBackend::LoadObject(
|
||||
const std::string& path, std::vector<Slice>& slices, size_t length) {
|
||||
const std::string& path, std::vector<Slice>& 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<iovec> iovs_chunk;
|
||||
off_t chunk_start_offset = 0;
|
||||
size_t chunk_length = 0;
|
||||
int64_t chunk_length = 0;
|
||||
|
||||
auto process_chunk = [&]() -> tl::expected<void, ErrorCode> {
|
||||
if (iovs_chunk.empty()) {
|
||||
|
|
@ -101,15 +103,15 @@ tl::expected<void, ErrorCode> StorageBackend::LoadObject(
|
|||
iovs_chunk.data(), static_cast<int>(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<void, ErrorCode> 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<void, ErrorCode> StorageBackend::LoadObject(
|
|||
}
|
||||
|
||||
tl::expected<void, ErrorCode> 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<StorageFile> StorageBackend::create_file(
|
|||
return std::make_unique<PosixFile>(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<int64_t, ErrorCode> BucketStorageBackend::BatchOffload(
|
||||
const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
|
||||
std::function<
|
||||
ErrorCode(const std::unordered_map<std::string, BucketObjectMetadata>&)>
|
||||
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<iovec> 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<void, ErrorCode> BucketStorageBackend::BatchQuery(
|
||||
const std::vector<std::string>& keys,
|
||||
std::unordered_map<std::string, StorageObjectMetadata>&
|
||||
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<void, ErrorCode> BucketStorageBackend::BatchLoad(
|
||||
std::unordered_map<std::string, Slice>& batch_object) {
|
||||
std::unordered_map<int64_t, std::vector<std::string>> 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<void, ErrorCode> BucketStorageBackend::GetBucketKeys(
|
||||
int64_t bucket_id, std::vector<std::string>& 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<void, ErrorCode> 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<BucketMetadata>());
|
||||
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<bool, ErrorCode> 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<int64_t, ErrorCode> BucketStorageBackend::BucketScan(
|
||||
int64_t bucket_id,
|
||||
std::unordered_map<std::string, BucketObjectMetadata>& objects,
|
||||
std::vector<int64_t>& 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<OffloadMetadata, ErrorCode>
|
||||
BucketStorageBackend::GetStoreMetadata() {
|
||||
SharedMutexLocker lock(&mutex_, shared_lock);
|
||||
OffloadMetadata metadata{object_bucket_map_.size(), total_size_};
|
||||
return metadata;
|
||||
}
|
||||
|
||||
tl::expected<std::shared_ptr<BucketMetadata>, ErrorCode>
|
||||
BucketStorageBackend::BuildBucket(
|
||||
const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
|
||||
std::vector<iovec>& iovs) {
|
||||
auto bucket = std::make_shared<BucketMetadata>();
|
||||
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<char*>(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<void, ErrorCode> BucketStorageBackend::WriteBucket(
|
||||
int64_t bucket_id, std::shared_ptr<BucketMetadata> bucket_metadata,
|
||||
std::vector<iovec>& 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<void, ErrorCode> BucketStorageBackend::StoreBucketMetadata(
|
||||
int64_t id, std::shared_ptr<BucketMetadata> 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<void, ErrorCode> BucketStorageBackend::LoadBucketMetadata(
|
||||
int64_t id, std::shared_ptr<BucketMetadata> 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<void, ErrorCode> BucketStorageBackend::BatchLoadBucket(
|
||||
int64_t bucket_id, const std::vector<std::string>& keys,
|
||||
std::unordered_map<std::string, Slice>& 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<iovec> iovs;
|
||||
iovs.emplace_back(iovec{slice.ptr, slice.size});
|
||||
auto read_result = file->vector_read(
|
||||
iovs.data(), static_cast<int>(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<std::string, ErrorCode> 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<std::string, ErrorCode>
|
||||
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<std::unique_ptr<StorageFile>, 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<PosixFile>(path, fd);
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
|
|||
|
|
@ -167,4 +167,10 @@ int getFreeTcpPort() {
|
|||
return port;
|
||||
}
|
||||
|
||||
int64_t time_gen() {
|
||||
return std::chrono::duration_cast<std::chrono::seconds>(
|
||||
std::chrono::system_clock::now().time_since_epoch())
|
||||
.count();
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,176 @@
|
|||
#include "mutex.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
#include <chrono>
|
||||
#include <stdexcept>
|
||||
|
||||
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<int> counter{0};
|
||||
std::vector<std::thread> 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<int> active_readers{0};
|
||||
std::atomic<int> peak_readers{0};
|
||||
std::vector<std::thread> 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<bool> reader_started{false};
|
||||
std::atomic<bool> 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
|
||||
|
|
@ -0,0 +1,313 @@
|
|||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include "storage_backend.h"
|
||||
#include "allocator.h"
|
||||
#include "utils.h"
|
||||
#include <ylt/struct_pb.hpp>
|
||||
#include <filesystem>
|
||||
#include <iostream>
|
||||
#include <ranges>
|
||||
|
||||
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<void, ErrorCode> BatchOffload(
|
||||
std::vector<std::string>& keys,
|
||||
std::unordered_map<std::string, std::string>& batch_data,
|
||||
std::shared_ptr<SimpleAllocator>& client_buffer_allocator,
|
||||
BucketStorageBackend& storage_backend, std::vector<int64_t>& 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<std::string, std::vector<Slice>> 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<Slice> 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<std::string, BucketObjectMetadata>&
|
||||
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<SimpleAllocator> client_buffer_allocator =
|
||||
std::make_shared<SimpleAllocator>(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<std::string, std::string> test_data;
|
||||
std::vector<std::string> keys;
|
||||
std::vector<int64_t> 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<std::string, StorageObjectMetadata>
|
||||
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<std::string, Slice> 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<SimpleAllocator> client_buffer_allocator =
|
||||
std::make_shared<SimpleAllocator>(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<std::string, std::string> test_data;
|
||||
std::vector<std::string> keys;
|
||||
std::vector<int64_t> 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<std::string, BucketObjectMetadata> objects;
|
||||
std::vector<int64_t> 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<std::thread> threads;
|
||||
std::vector<int64_t> 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<std::mutex> 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<int64_t> 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
|
||||
Loading…
Reference in New Issue