chenyu520
/
Mooncake
forked from mooncake-track/Mooncake
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Store] High Availability V2: Client Failover (#501) 

Major changes include: Client failover. Refactor SegmentManager. On the client side, limit the segment name to be equal to the localhost name. Add clientctl for manual e2e tests.
This commit is contained in:
ykwd 2025-06-20 16:59:55 +08:00 committed by GitHub
parent 01df366801
commit 91c5778166
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
30 changed files with 2267 additions and 565 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mooncake-integration/store/store_py.cpp
View File

@ -197,8 +197,7 @@ int DistributedObjectStore::setup(const std::string &local_hostname,
return 1;
}
segment_ptr_.reset(ptr);
error_code = client_->MountSegment(this->local_hostname, segment_ptr_.get(),
global_segment_size);
error_code = client_->MountSegment(segment_ptr_.get(), global_segment_size);
if (error_code != ErrorCode::OK) {
LOG(ERROR) << "Failed to mount segment: " << toString(error_code);
return 1;

114
mooncake-store/include/allocation_strategy.h
View File

@ -21,7 +21,8 @@ class AllocationStrategy {
/**
* @brief Given all mounted BufferAllocators and required object size,
* the strategy can freely choose a suitable BufferAllocator.
* @param allocators Container of mounted allocators, key is segment_name,
* @param allocators Container of mounted allocators
* @param allocators_by_name Container of mounted allocators, key is segment_name,
* value is the corresponding allocator
* @param objectSize Size of object to be allocated
* @param config Replica configuration
@ -29,8 +30,9 @@ class AllocationStrategy {
* or no suitable allocator is found
*/
virtual std::unique_ptr<AllocatedBuffer> Allocate(
const std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>&
allocators,
const std::vector<std::shared_ptr<BufferAllocator>>& allocators,
const std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>&
allocators_by_name,
size_t objectSize, const ReplicateConfig& config) = 0;
};
@ -46,22 +48,23 @@ class RandomAllocationStrategy : public AllocationStrategy {
RandomAllocationStrategy() : rng_(std::random_device{}()) {}
std::unique_ptr<AllocatedBuffer> Allocate(
const std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>&
allocators,
const std::vector<std::shared_ptr<BufferAllocator>>& allocators,
const std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>&
allocators_by_name,
size_t objectSize, const ReplicateConfig& config) override {
// Fast path: single allocator case
if (allocators.size() == 1) {
return allocators.begin()->second->allocate(objectSize);
return allocators[0]->allocate(objectSize);
}
// Try preferred segment first if specified
if (auto preferred_buffer =
TryPreferredAllocation(allocators, objectSize, config)) {
TryPreferredAllocate(allocators_by_name, objectSize, config)) {
return preferred_buffer;
}
// Fall back to random allocation among all eligible allocators
return RandomAllocateFromEligible(allocators, objectSize);
return TryRandomAllocate(allocators, objectSize);
}
private:
@ -73,8 +76,8 @@ class RandomAllocationStrategy : public AllocationStrategy {
* @brief Attempts allocation from preferred segment if available and
* eligible
*/
std::unique_ptr<AllocatedBuffer> TryPreferredAllocation(
const std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>&
std::unique_ptr<AllocatedBuffer> TryPreferredAllocate(
const std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>&
allocators,
size_t objectSize, const ReplicateConfig& config) {
if (config.preferred_segment.empty()) {
@ -86,97 +89,48 @@ class RandomAllocationStrategy : public AllocationStrategy {
return nullptr;
}
auto& preferred_allocator = preferred_it->second;
if (MayHasSufficientSpace(preferred_allocator, objectSize)) {
return preferred_allocator->allocate(objectSize);
}
return nullptr;
}
/**
* @brief Performs random allocation from eligible allocators with retry
* logic
*/
std::unique_ptr<AllocatedBuffer> RandomAllocateFromEligible(
const std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>&
allocators,
size_t objectSize) {
auto eligible = CollectEligibleAllocators(allocators, objectSize);
if (eligible.empty()) {
return nullptr;
}
return TryAllocateWithRetry(eligible, objectSize);
}
/**
* @brief Collects all allocators with sufficient available space
*/
std::vector<std::shared_ptr<BufferAllocator>> CollectEligibleAllocators(
const std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>&
allocators,
size_t objectSize) {
std::vector<std::shared_ptr<BufferAllocator>> eligible;
eligible.reserve(allocators.size());
for (const auto& [segment_name, allocator] : allocators) {
if (MayHasSufficientSpace(allocator, objectSize)) {
eligible.push_back(allocator);
auto& preferred_allocators = preferred_it->second;
for (auto& allocator : preferred_allocators) {
auto buffer = allocator->allocate(objectSize);
if (buffer != nullptr) {
return buffer;
}
}
return eligible;
return nullptr;
}
/**
* @brief Attempts allocation with random selection and retry logic
*/
std::unique_ptr<AllocatedBuffer> TryAllocateWithRetry(
std::vector<std::shared_ptr<BufferAllocator>>& eligible,
std::unique_ptr<AllocatedBuffer> TryRandomAllocate(
const std::vector<std::shared_ptr<BufferAllocator>>& allocators,
size_t objectSize) {
const size_t max_tries = std::min(kMaxRetryLimit, eligible.size());
const size_t max_tries = std::min(kMaxRetryLimit, allocators.size());
std::vector<size_t> allocator_indices(allocators.size());
std::iota(allocator_indices.begin(), allocator_indices.end(), 0);
for (size_t try_count = 0; try_count < max_tries; ++try_count) {
// Randomly select an allocator
std::uniform_int_distribution<size_t> dist(0, eligible.size() - 1);
const size_t random_index = dist(rng_);
std::uniform_int_distribution<size_t> dist(
0, allocator_indices.size() - 1);
const size_t random_index = allocator_indices[dist(rng_)];
auto& allocator = eligible[random_index];
auto& allocator = allocators[random_index];
if (auto buffer = allocator->allocate(objectSize)) {
return buffer;
}
// Remove failed allocator and continue with remaining ones
RemoveAllocatorAtIndex(eligible, random_index);
if (random_index + 1 != allocator_indices.size()) {
std::swap(allocator_indices[random_index],
allocator_indices[allocator_indices.size() - 1]);
}
allocator_indices.pop_back();
}
return nullptr;
}
/**
* @brief Checks if allocator has sufficient available space
*/
static bool MayHasSufficientSpace(
const std::shared_ptr<BufferAllocator>& allocator,
size_t required_size) {
const size_t capacity = allocator->capacity();
const size_t used = allocator->size();
const size_t available = capacity > used ? (capacity - used) : 0;
return available >= required_size;
}
/**
* @brief Efficiently removes allocator at given index using swap-and-pop
*/
static void RemoveAllocatorAtIndex(
std::vector<std::shared_ptr<BufferAllocator>>& allocators,
size_t index) {
if (index + 1 != allocators.size()) {
std::swap(allocators[index], allocators.back());
}
allocators.pop_back();
}
};
} // namespace mooncake

23
mooncake-store/include/client.h
View File

@ -5,6 +5,7 @@
#include <optional>
#include <string>
#include <vector>
#include <boost/functional/hash.hpp>
#include "master_client.h"
#include "rpc_service.h"
@ -145,21 +146,19 @@ class Client {
/**
* @brief Registers a memory segment to master for allocation
* @param segment_name Unique identifier for the segment
* @param buffer Memory buffer to register
* @param size Size of the buffer in bytes
* @return ErrorCode indicating success/failure
*/
ErrorCode MountSegment(const std::string& segment_name, const void* buffer,
size_t size);
ErrorCode MountSegment(const void* buffer, size_t size);
/**
* @brief Unregisters a memory segment from master
* @param segment_name Name of the segment to unregister
* @param addr Memory address to unregister
* @param buffer Memory buffer to unregister
* @param size Size of the buffer in bytes
* @return ErrorCode indicating success/failure
*/
ErrorCode UnmountSegment(const std::string& segment_name, void* addr);
ErrorCode UnmountSegment(const void* buffer, size_t size);
/**
* @brief Registers memory buffer with TransferEngine for data transfer
@ -233,14 +232,9 @@ class Client {
MasterClient master_client_;
std::unique_ptr<TransferSubmitter> transfer_submitter_;
// Client local segments
struct Segment{
void* buffer;
size_t size;
};
// Mutex to protect mounted_segments_
std::mutex mounted_segments_mutex_;
std::unordered_map<std::string, Segment> mounted_segments_;
std::unordered_map<UUID, Segment, boost::hash<UUID>> mounted_segments_;
// Configuration
const std::string local_hostname_;
@ -250,7 +244,10 @@ class Client {
MasterViewHelper master_view_helper_;
std::thread ping_thread_;
std::atomic<bool> ping_running_{false};
void PingThreadFunc(int current_version);
void PingThreadFunc();
// Client identification
UUID client_id_;
};
} // namespace mooncake

2
mooncake-store/include/ha_helper.h
View File

@ -80,6 +80,7 @@ class MasterServiceSupervisor {
bool enable_metric_reporting, int metrics_port,
int64_t default_kv_lease_ttl, double eviction_ratio,
double eviction_high_watermark_ratio,
int64_t client_live_ttl_sec,
const std::string& etcd_endpoints = "0.0.0.0:2379",
const std::string& local_hostname = "0.0.0.0:50051");
int Start();
@ -95,6 +96,7 @@ class MasterServiceSupervisor {
int64_t default_kv_lease_ttl_;
double eviction_ratio_;
double eviction_high_watermark_ratio_;
int64_t client_live_ttl_sec_;
// coro_rpc server thread
std::thread server_thread_;

29
mooncake-store/include/master_client.h
View File

@ -132,29 +132,42 @@ class MasterClient {
/**
* @brief Registers a segment to master for allocation
* @param segment_name hostname:port of the segment
* @param buffer Buffer address of the segment
* @param size Size of the segment in bytes
* @param segment Segment to register
* @param client_id The uuid of the client
* @return ErrorCode indicating success/failure
*/
[[nodiscard]] MountSegmentResponse MountSegment(
const std::string& segment_name, const void* buffer, size_t size);
const Segment& segment, const UUID& client_id);
/**
* @brief Re-mount segments, invoked when the client is the first time to
* connect to the master or the client Ping TTL is expired and need
* to remount. This function is idempotent. Client should retry if the
* return code is not ErrorCode::OK.
* @param segments Segments to remount
* @param client_id The uuid of the client
* @return ErrorCode indicating success/failure
*/
[[nodiscard]] ReMountSegmentResponse ReMountSegment(
const std::vector<Segment>& segments, const UUID& client_id);
/**
* @brief Unregisters a memory segment from master
* @param segment_name Name which is used to register the segment
* @param segment_id ID of the segment to unmount
* @param client_id The uuid of the client
* @return ErrorCode indicating success/failure
*/
[[nodiscard]] UnmountSegmentResponse UnmountSegment(
const std::string& segment_name);
const UUID& segment_id, const UUID& client_id);
/**
* @brief Pings master to check its availability
* @param No parameters
* @param client_id The uuid of the client
* @return current master view version
* @return client status from the master
* @return ErrorCode indicating success/failure
*/
[[nodiscard]] PingResponse Ping();
[[nodiscard]] PingResponse Ping(const UUID& client_id);
private:
coro_rpc_client client_;

25
mooncake-store/include/master_metric_manager.h
View File

@ -34,6 +34,11 @@ class MasterMetricManager {
void observe_value_size(int64_t size);
int64_t get_key_count();
// Cluster Metrics
void inc_active_clients(int64_t val = 1);
void dec_active_clients(int64_t val = 1);
int64_t get_active_clients();
// Operation Statistics (Counters)
void inc_put_start_requests(int64_t val = 1);
void inc_put_start_failures(int64_t val = 1);
@ -53,7 +58,11 @@ class MasterMetricManager {
void inc_mount_segment_failures(int64_t val = 1);
void inc_unmount_segment_requests(int64_t val = 1);
void inc_unmount_segment_failures(int64_t val = 1);
void inc_remount_segment_requests(int64_t val = 1);
void inc_remount_segment_failures(int64_t val = 1);
void inc_ping_requests(int64_t val = 1);
void inc_ping_failures(int64_t val = 1);
// Operation Statistics Getters
int64_t get_put_start_requests();
@ -74,7 +83,10 @@ class MasterMetricManager {
int64_t get_mount_segment_failures();
int64_t get_unmount_segment_requests();
int64_t get_unmount_segment_failures();
int64_t get_remount_segment_requests();
int64_t get_remount_segment_failures();
int64_t get_ping_requests();
int64_t get_ping_failures();
// Eviction Metrics
void inc_eviction_success(int64_t key_count, int64_t size);
@ -99,6 +111,9 @@ class MasterMetricManager {
*/
std::string get_summary_string();
// --- Setters ---
void set_enable_ha(bool enable_ha);
private:
// --- Private Constructor & Destructor ---
MasterMetricManager();
@ -114,6 +129,9 @@ class MasterMetricManager {
ylt::metric::gauge_t key_count_;
ylt::metric::histogram_t value_size_distribution_;
// Cluster Metrics
ylt::metric::gauge_t active_clients_;
// Operation Statistics
ylt::metric::counter_t put_start_requests_;
ylt::metric::counter_t put_start_failures_;
@ -133,13 +151,20 @@ class MasterMetricManager {
ylt::metric::counter_t mount_segment_failures_;
ylt::metric::counter_t unmount_segment_requests_;
ylt::metric::counter_t unmount_segment_failures_;
ylt::metric::counter_t remount_segment_requests_;
ylt::metric::counter_t remount_segment_failures_;
ylt::metric::counter_t ping_requests_;
ylt::metric::counter_t ping_failures_;
// Eviction Metrics
ylt::metric::counter_t eviction_success_;
ylt::metric::counter_t eviction_attempts_;
ylt::metric::counter_t evicted_key_count_;
ylt::metric::counter_t evicted_size_;
// Some metrics are used only in HA mode. Use a flag to control the output
// content.
bool enable_ha_{false};
};
} // namespace mooncake

135
mooncake-store/include/master_service.h
View File

@ -2,6 +2,7 @@
#include <atomic>
#include <boost/lockfree/queue.hpp>
#include <boost/functional/hash.hpp>
#include <chrono>
#include <cstdint>
#include <memory>
@ -9,6 +10,7 @@
#include <string>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <optional>
@ -16,6 +18,7 @@
#include "eviction_strategy.h"
#include "allocator.h"
#include "types.h"
#include "segment.h"
namespace mooncake {
@ -38,48 +41,13 @@ struct GCTask {
}
};
class BufferAllocatorManager {
public:
BufferAllocatorManager() = default;
~BufferAllocatorManager() = default;
/**
* @brief Register a new buffer for allocation
* @return ErrorCode::OK on success, ErrorCode::INVALID_PARAMS if segment
* exists
*/
ErrorCode AddSegment(const std::string& segment_name, uint64_t base,
uint64_t size);
/**
* @brief Unregister a buffer
* @return ErrorCode::OK on success, ErrorCode::INVALID_PARAMS if segment
* not found
*/
ErrorCode RemoveSegment(const std::string& segment_name);
/**
* @brief Get the map of buffer allocators
* @note Caller must hold the mutex while accessing the map
*/
const std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>&
GetAllocators() const {
return buf_allocators_;
}
/**
* @brief Get the mutex for thread-safe access
*/
std::shared_mutex& GetMutex() { return allocator_mutex_; }
private:
// Protects the buffer allocator map (BufferAllocator is thread-safe by
// itself)
mutable std::shared_mutex allocator_mutex_;
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
buf_allocators_;
};
/*
* @brief MasterService is the main class for the master server.
* Lock order: To avoid deadlocks, the following lock order should be followed:
* 1. client_mutex_
* 2. metadata_shards_[shard_idx_].mutex
* 3. segment_mutex_
*/
class MasterService {
private:
// Comparator for GC tasks priority queue
@ -93,23 +61,44 @@ class MasterService {
MasterService(bool enable_gc = true,
uint64_t default_kv_lease_ttl = DEFAULT_DEFAULT_KV_LEASE_TTL,
double eviction_ratio = DEFAULT_EVICTION_RATIO,
double eviction_high_watermark_ratio = DEFAULT_EVICTION_HIGH_WATERMARK_RATIO);
double eviction_high_watermark_ratio = DEFAULT_EVICTION_HIGH_WATERMARK_RATIO,
ViewVersionId view_version = 0,
int64_t client_live_ttl_sec = DEFAULT_CLIENT_LIVE_TTL_SEC,
bool enable_ha = false);
~MasterService();
/**
* @brief Mount a memory segment for buffer allocation
* @return ErrorCode::OK on success, ErrorCode::INVALID_PARAMS if segment
* exists or params invalid, ErrorCode::INTERNAL_ERROR if allocation fails
* @brief Mount a memory segment for buffer allocation. This function is
* idempotent.
* @return ErrorCode::OK on success,
* ErrorCode::INVALID_PARAMS on invalid parameters,
* ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS if the segment cannot
* be mounted temporarily,
* ErrorCode::INTERNAL_ERROR on internal errors.
*/
ErrorCode MountSegment(uint64_t buffer, uint64_t size,
const std::string& segment_name);
ErrorCode MountSegment(const Segment& segment, const UUID& client_id);
/**
* @brief Unmount a memory segment
* @return ErrorCode::OK on success, ErrorCode::INVALID_PARAMS if segment
* not found
* @brief Re-mount segments, invoked when the client is the first time to
* connect to the master or the client Ping TTL is expired and need
* to remount. This function is idempotent. Client should retry if the
* return code is not ErrorCode::OK.
* @return ErrorCode::OK means either all segments are remounted successfully
* or the fail is not solvable by a new remount request.
* ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS if the segment cannot
* be mounted temporarily.
* ErrorCode::INTERNAL_ERROR if something temporary error happens.
*/
ErrorCode UnmountSegment(const std::string& segment_name);
ErrorCode ReMountSegment(const std::vector<Segment>& segments,
const UUID& client_id);
/**
* @brief Unmount a memory segment. This function is idempotent.
* @return ErrorCode::OK on success,
* ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS if the segment is
* currently unmounting.
*/
ErrorCode UnmountSegment(const UUID& segment_id, const UUID& client_id);
/**
* @brief Check if an object exists
@ -240,6 +229,17 @@ class MasterService {
*/
size_t GetKeyCount() const;
/**
* @brief Heartbeat from client
* @param client_id The uuid of the client
* @param[out] view_version The view version of the master
* @param[out] client_status The status of the client from the master
* @return ErrorCode::OK on success, ErrorCode::INTERNAL_ERROR if the client
* ping queue is full
*/
ErrorCode Ping(const UUID& client_id, ViewVersionId& view_version,
ClientStatus& client_status);
private:
// GC thread function
void GCThreadFunc();
@ -247,6 +247,9 @@ class MasterService {
// Check all shards and try to evict some keys
void BatchEvict(double eviction_ratio);
// Clear invalid handles in all shards
void ClearInvalidHandles();
// Internal data structures
struct ObjectMetadata {
std::vector<Replica> replicas;
@ -284,11 +287,10 @@ class MasterService {
}
};
// Buffer allocator management
std::shared_ptr<BufferAllocatorManager> buffer_allocator_manager_;
// Segment management
SegmentManager segment_manager_;
std::shared_ptr<AllocationStrategy> allocation_strategy_;
static constexpr size_t kNumShards = 1024; // Number of metadata shards
// Sharded metadata maps and their mutexes
@ -373,6 +375,29 @@ class MasterService {
};
friend class MetadataAccessor;
ViewVersionId view_version_;
// Client related members
mutable std::shared_mutex client_mutex_;
std::unordered_set<UUID, boost::hash<UUID>> ok_client_; // client with ok status
void ClientMonitorFunc();
std::thread client_monitor_thread_;
std::atomic<bool> client_monitor_running_{false};
static constexpr uint64_t kClientMonitorSleepMs =
1000; // 1000 ms sleep between client monitor checks
// boost lockfree queue requires trivial assignment operator
struct PodUUID {
uint64_t first;
uint64_t second;
};
static constexpr size_t kClientPingQueueSize =
128 * 1024; // Size of the client ping queue
boost::lockfree::queue<PodUUID> client_ping_queue_{kClientPingQueueSize};
const int64_t client_live_ttl_sec_;
// if high availability features enabled
const bool enable_ha_;
};
} // namespace mooncake

81
mooncake-store/include/rpc_service.h
View File

@ -80,6 +80,11 @@ struct MountSegmentResponse {
};
YLT_REFL(MountSegmentResponse, error_code)
struct ReMountSegmentResponse {
ErrorCode error_code = ErrorCode::OK;
};
YLT_REFL(ReMountSegmentResponse, error_code)
struct UnmountSegmentResponse {
ErrorCode error_code = ErrorCode::OK;
};
@ -87,6 +92,7 @@ YLT_REFL(UnmountSegmentResponse, error_code)
struct PingResponse {
ViewVersionId view_version = 0;
ClientStatus client_status = ClientStatus::UNDEFINED;
ErrorCode error_code = ErrorCode::OK;
};
YLT_REFL(PingResponse, view_version, error_code)
@ -95,21 +101,27 @@ constexpr uint64_t kMetricReportIntervalSeconds = 10;
class WrappedMasterService {
public:
WrappedMasterService(bool enable_gc, uint64_t default_kv_lease_ttl,
bool enable_metric_reporting = true,
uint16_t http_port = 9003,
double eviction_ratio = DEFAULT_EVICTION_RATIO,
double eviction_high_watermark_ratio =
DEFAULT_EVICTION_HIGH_WATERMARK_RATIO,
ViewVersionId view_version = 0)
WrappedMasterService(
bool enable_gc, uint64_t default_kv_lease_ttl,
bool enable_metric_reporting = true, uint16_t http_port = 9003,
double eviction_ratio = DEFAULT_EVICTION_RATIO,
double eviction_high_watermark_ratio =
DEFAULT_EVICTION_HIGH_WATERMARK_RATIO,
ViewVersionId view_version = 0,
int64_t client_live_ttl_sec = DEFAULT_CLIENT_LIVE_TTL_SEC,
bool enable_ha = false)
: master_service_(enable_gc, default_kv_lease_ttl, eviction_ratio,
eviction_high_watermark_ratio),
eviction_high_watermark_ratio, view_version,
client_live_ttl_sec, enable_ha),
http_server_(4, http_port),
metric_report_running_(enable_metric_reporting),
view_version_(view_version) {
// Initialize HTTP server for metrics
init_http_server();
// Set the config for metric reporting
MasterMetricManager::instance().set_enable_ha(enable_ha);
// Start metric reporting thread if enabled
if (enable_metric_reporting) {
metric_report_thread_ = std::thread([this]() {
@ -450,40 +462,57 @@ class WrappedMasterService {
return response;
}
MountSegmentResponse MountSegment(uint64_t buffer, uint64_t size,
const std::string& segment_name) {
MountSegmentResponse MountSegment(const Segment& segment, const UUID& client_id) {
ScopedVLogTimer timer(1, "MountSegment");
timer.LogRequest("buffer=", buffer, ", size=", size,
", segment_name=", segment_name);
timer.LogRequest("base=", segment.base, ", size=", segment.size,
", segment_name=", segment.name, ", id=", segment.id);
// Increment request metric
MasterMetricManager::instance().inc_mount_segment_requests();
MountSegmentResponse response;
response.error_code =
master_service_.MountSegment(buffer, size, segment_name);
response.error_code = master_service_.MountSegment(segment, client_id);
// Track failures if needed
if (response.error_code != ErrorCode::OK) {
MasterMetricManager::instance().inc_mount_segment_failures();
} else {
// Update total capacity on successful mount
MasterMetricManager::instance().inc_total_capacity(size);
}
timer.LogResponseJson(response);
return response;
}
UnmountSegmentResponse UnmountSegment(const std::string& segment_name) {
ReMountSegmentResponse ReMountSegment(const std::vector<Segment>& segments,
const UUID& client_id) {
ScopedVLogTimer timer(1, "ReMountSegment");
timer.LogRequest("segments_count=", segments.size(),
", client_id=", client_id);
// Increment request metric
MasterMetricManager::instance().inc_remount_segment_requests();
ReMountSegmentResponse response;
response.error_code =
master_service_.ReMountSegment(segments, client_id);
// Track failures if needed
if (response.error_code != ErrorCode::OK) {
MasterMetricManager::instance().inc_remount_segment_failures();
}
timer.LogResponseJson(response);
return response;
}
UnmountSegmentResponse UnmountSegment(const UUID& segment_id, const UUID& client_id) {
ScopedVLogTimer timer(1, "UnmountSegment");
timer.LogRequest("segment_name=", segment_name);
timer.LogRequest("segment_id=", segment_id);
// Increment request metric
MasterMetricManager::instance().inc_unmount_segment_requests();
UnmountSegmentResponse response;
response.error_code = master_service_.UnmountSegment(segment_name);
response.error_code = master_service_.UnmountSegment(segment_id, client_id);
// Track failures if needed
if (response.error_code != ErrorCode::OK) {
@ -494,13 +523,19 @@ class WrappedMasterService {
return response;
}
PingResponse Ping() {
PingResponse Ping(const UUID& client_id) {
ScopedVLogTimer timer(1, "Ping");
timer.LogRequest("action=ping");
timer.LogRequest("client_id=", client_id);
MasterMetricManager::instance().inc_ping_requests();
PingResponse response(view_version_, ErrorCode::OK);
PingResponse response;
response.error_code = master_service_.Ping(
client_id, response.view_version, response.client_status);
if (response.error_code != ErrorCode::OK) {
MasterMetricManager::instance().inc_ping_failures();
}
timer.LogResponseJson(response);
return response;

182
mooncake-store/include/segment.h Normal file
View File

@ -0,0 +1,182 @@
#pragma once
#include <boost/functional/hash.hpp>
#include <ostream>
#include <shared_mutex>
#include <string>
#include <string_view>
#include <unordered_map>
#include <vector>
#include "allocation_strategy.h"
#include "allocator.h"
#include "types.h"
namespace mooncake {
/**
* @brief Status of a mounted segment in master
*/
enum class SegmentStatus {
UNDEFINED = 0, // Uninitialized
OK, // Segment is mounted and available for allocation
UNMOUNTING, // Segment is under unmounting
};
/**
* @brief Stream operator for SegmentStatus
*/
inline std::ostream& operator<<(std::ostream& os,
const SegmentStatus& status) noexcept {
static const std::unordered_map<SegmentStatus, std::string_view>
status_strings{{SegmentStatus::UNDEFINED, "UNDEFINED"},
{SegmentStatus::OK, "OK"},
{SegmentStatus::UNMOUNTING, "UNMOUNTING"}};
os << (status_strings.count(status) ? status_strings.at(status)
: "UNKNOWN");
return os;
}
struct MountedSegment {
Segment segment;
SegmentStatus status;
std::shared_ptr<BufferAllocator> buf_allocator;
};
// Forward declarations
class SegmentManager;
/**
* @brief RAII-style access to segment mutex for thread-safe segment operations
*/
class ScopedSegmentAccess {
public:
/**
* @brief Acquires a lock on the segment mutex
* @param mutex Reference to the segment mutex
*/
explicit ScopedSegmentAccess(SegmentManager* segment_manager,
std::shared_mutex& mutex)
: segment_manager_(segment_manager), lock_(mutex) {}
/**
* @brief Mount a segment
*/
ErrorCode MountSegment(const Segment& segment, const UUID& client_id);
/**
* @brief Re-mount a segment. To avoid infinite remount trying, only the
* errors that may be solved by subsequent remount tryings are considered as
* errors. When encounters unsolvable errors, the segment will not be mounted
* while the return value will be OK.
*/
ErrorCode ReMountSegment(const std::vector<Segment>& segments,
const UUID& client_id);
/**
* @brief Prepare to unmount a segment by deleting its allocator
*/
ErrorCode PrepareUnmountSegment(const UUID& segment_id,
size_t& metrics_dec_capacity);
/**
* @brief Deleting the segment to complete the unmounting operation
*/
ErrorCode CommitUnmountSegment(const UUID& segment_id,
const UUID& client_id,
const size_t& metrics_dec_capacity);
/**
* @brief Get all the segments of a client
*/
ErrorCode GetClientSegments(const UUID& client_id,
std::vector<Segment>& segments) const;
/**
* @brief Get the names of all the segments
*/
ErrorCode GetAllSegments(std::vector<std::string>& all_segments);
/**
* @brief Get the segment by name. If there are multiple segments with the
* same name, return the first one.
*/
ErrorCode QuerySegments(const std::string& segment, size_t& used,
size_t& capacity);
private:
SegmentManager* segment_manager_;
std::unique_lock<std::shared_mutex> lock_;
};
/**
* @brief RAII-style access to allocators for thread-safe allocator usage
*/
class ScopedAllocatorAccess {
public:
explicit ScopedAllocatorAccess(
std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocator>>>&
allocators_by_name,
std::vector<std::shared_ptr<BufferAllocator>>& allocators,
std::shared_mutex& mutex)
: allocators_by_name_(allocators_by_name),
allocators_(allocators),
lock_(mutex) {}
const std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocator>>>&
getAllocatorsByName() {
return allocators_by_name_;
}
const std::vector<std::shared_ptr<BufferAllocator>>& getAllocators() {
return allocators_;
}
private:
const std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocator>>>&
allocators_by_name_; // segment name -> allocators
const std::vector<std::shared_ptr<BufferAllocator>>& allocators_;
std::shared_lock<std::shared_mutex> lock_;
};
class SegmentManager {
public:
/**
* @brief Get RAII-style access to segment management operations
* @return ScopedSegmentAccess object that holds the lock
*/
ScopedSegmentAccess getSegmentAccess() {
return ScopedSegmentAccess(this, segment_mutex_);
}
/**
* @brief Get RAII-style access to use allocators
* @return ScopedAllocatorAccess object that holds the lock
*/
ScopedAllocatorAccess getAllocatorAccess() {
return ScopedAllocatorAccess(allocators_by_name_, allocators_,
segment_mutex_);
}
private:
mutable std::shared_mutex segment_mutex_;
std::shared_ptr<AllocationStrategy> allocation_strategy_;
// Each allocator is put into both of allocators_by_name_ and allocators_.
// These two containers only contain allocators whose segment status is OK.
std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocator>>>
allocators_by_name_; // segment name -> allocators
std::vector<std::shared_ptr<BufferAllocator>> allocators_; // allocators
std::unordered_map<UUID, MountedSegment, boost::hash<UUID>>
mounted_segments_; // segment_id -> mounted segment
std::unordered_map<UUID, std::vector<UUID>, boost::hash<UUID>>
client_segments_; // client_id -> segment_ids
friend class ScopedSegmentAccess;
friend class SegmentTest; // for unit tests
};
} // namespace mooncake

66
mooncake-store/include/types.h
View File

@ -28,6 +28,7 @@ static constexpr uint64_t DEFAULT_DEFAULT_KV_LEASE_TTL =
static constexpr double DEFAULT_EVICTION_RATIO = 0.1;
static constexpr double DEFAULT_EVICTION_HIGH_WATERMARK_RATIO = 1.0;
static constexpr int64_t ETCD_MASTER_VIEW_LEASE_TTL = 5; // in seconds
static constexpr int64_t DEFAULT_CLIENT_LIVE_TTL_SEC = 10; // in seconds
// Forward declarations
class BufferAllocator;
@ -55,6 +56,15 @@ using ViewVersionId = int64_t;
using EtcdLeaseId = int64_t;
#endif
using UUID = std::pair<uint64_t, uint64_t>;
inline std::ostream& operator<<(std::ostream& os, const UUID& uuid) noexcept {
os << uuid.first << "-" << uuid.second;
return os;
}
UUID generate_uuid();
/**
* @brief Error codes for various operations in the system
*/
@ -67,7 +77,8 @@ enum class ErrorCode : int32_t {
// Segment selection errors (Range: -100 to -199)
SHARD_INDEX_OUT_OF_RANGE = -100, ///< Shard index is out of bounds.
AVAILABLE_SEGMENT_EMPTY = -101, ///< No available segments found.
SEGMENT_NOT_FOUND = -101, ///< No available segments found.
SEGMENT_ALREADY_EXISTS = -102, ///< Segment already exists.
// Handle selection errors (Range: -200 to -299)
NO_AVAILABLE_HANDLE = -200, ///< No available handles.
@ -99,11 +110,15 @@ enum class ErrorCode : int32_t {
// RPC errors (Range: -900 to -999)
RPC_FAIL = -900, ///< RPC operation failed.
// ETCD errors (Range: -1000 to -1099)
ETCD_OPERATION_ERROR = -1000, ///< etcd operation failed.
ETCD_KEY_NOT_EXIST = -1001, ///< key not found in etcd.
// High availability errors (Range: -1000 to -1099)
ETCD_OPERATION_ERROR = -1000, ///< etcd operation failed.
ETCD_KEY_NOT_EXIST = -1001, ///< key not found in etcd.
ETCD_TRANSACTION_FAIL = -1002, ///< etcd transaction failed.
ETCD_CTX_CANCELLED = -1003, ///< etcd context cancelled.
ETCD_CTX_CANCELLED = -1003, ///< etcd context cancelled.
UNAVAILABLE_IN_CURRENT_STATUS =
-1010, ///< Request cannot be done in current status.
UNAVAILABLE_IN_CURRENT_MODE =
-1011, ///< Request cannot be done in current mode.
};
int32_t toInt(ErrorCode errorCode) noexcept;
@ -345,4 +360,45 @@ const static uint64_t kMinSliceSize = facebook::cachelib::Slab::kMinAllocSize;
const static uint64_t kMaxSliceSize =
facebook::cachelib::Slab::kSize - 16; // should be lower than limit
/**
* @brief Represents a contiguous memory region
*/
struct Segment {
UUID id{0, 0};
std::string name{}; // The name of the segment, also might be the hostname
// of the server that owns the segment
uintptr_t base{0};
size_t size{0};
Segment() = default;
Segment(const UUID& id, const std::string& name, uintptr_t base,
size_t size)
: id(id), name(name), base(base), size(size) {}
};
YLT_REFL(Segment, id, name, base, size);
/**
* @brief Client status from the master's perspective
*/
enum class ClientStatus {
UNDEFINED = 0, // Uninitialized
OK, // Client is alive, no need to remount for now
NEED_REMOUNT, // Ping ttl expired, or the first time connect to master, so
// need to remount
};
/**
* @brief Stream operator for ClientStatus
*/
inline std::ostream& operator<<(std::ostream& os,
const ClientStatus& status) noexcept {
static const std::unordered_map<ClientStatus, std::string_view>
status_strings{{ClientStatus::UNDEFINED, "UNDEFINED"},
{ClientStatus::OK, "OK"},
{ClientStatus::NEED_REMOUNT, "NEED_REMOUNT"}};
os << (status_strings.count(status) ? status_strings.at(status)
: "UNKNOWN");
return os;
}
} // namespace mooncake

1
mooncake-store/src/CMakeLists.txt
View File

@ -11,6 +11,7 @@ set(MOONCAKE_STORE_SOURCES
master_metric_manager.cpp
etcd_helper.cpp
ha_helper.cpp
segment.cpp
transfer_task.cpp
etcd_helper.cpp
ha_helper.cpp

197
mooncake-store/src/client.cpp
View File

@ -26,24 +26,35 @@ namespace mooncake {
Client::Client(const std::string& local_hostname,
const std::string& metadata_connstring)
: local_hostname_(local_hostname),
metadata_connstring_(metadata_connstring) {}
metadata_connstring_(metadata_connstring) {
client_id_ = generate_uuid();
LOG(INFO) << "client_id=" << client_id_;
}
Client::~Client() {
// No need for mutex here since the client is being destroyed(protected by
// shared_ptr)
// Make a copy of mounted_segments_ to avoid modifying while iterating
std::unordered_map<std::string, Segment> segments_to_unmount =
mounted_segments_;
std::vector<Segment> segments_to_unmount;
{
std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
segments_to_unmount.reserve(mounted_segments_.size());
for (auto& entry : mounted_segments_) {
segments_to_unmount.push_back(entry.second);
}
}
for (auto& entry : segments_to_unmount) {
auto err_code = UnmountSegment(entry.first, entry.second.buffer);
for (auto& segment : segments_to_unmount) {
auto err_code = UnmountSegment(reinterpret_cast<void*>(segment.base),
segment.size);
if (err_code != ErrorCode::OK) {
LOG(ERROR) << "Failed to unmount segment: " << toString(err_code);
}
}
// Clear any remaining segments
mounted_segments_.clear();
{
std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
mounted_segments_.clear();
}
// Stop ping thread only after no need to contact master anymore
if (ping_running_) {
@ -138,7 +149,7 @@ ErrorCode Client::ConnectToMaster(const std::string& master_server_entry) {
// if needed
ping_running_ = true;
ping_thread_ =
std::thread(&Client::PingThreadFunc, this, master_version);
std::thread(&Client::PingThreadFunc, this);
return ErrorCode::OK;
} else {
@ -201,7 +212,7 @@ std::optional<std::shared_ptr<Client>> Client::Create(
// Initialize transfer engine
err = client->InitTransferEngine(local_hostname, metadata_connstring,
protocol, protocol_args);
protocol, protocol_args);
if (err != ErrorCode::OK) {
LOG(ERROR) << "Failed to initialize transfer engine";
return std::nullopt;
@ -533,8 +544,7 @@ ErrorCode Client::Remove(const ObjectKey& key) {
long Client::RemoveAll() { return master_client_.RemoveAll().removed_count; }
ErrorCode Client::MountSegment(const std::string& segment_name,
const void* buffer, size_t size) {
ErrorCode Client::MountSegment(const void* buffer, size_t size) {
if (buffer == nullptr || size == 0 ||
reinterpret_cast<uintptr_t>(buffer) % facebook::cachelib::Slab::kSize ||
size % facebook::cachelib::Slab::kSize) {
@ -543,11 +553,19 @@ ErrorCode Client::MountSegment(const std::string& segment_name,
return ErrorCode::INVALID_PARAMS;
}
{
std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
if (mounted_segments_.find(segment_name) != mounted_segments_.end()) {
LOG(ERROR) << "segment_already_exists segment_name="
<< segment_name;
std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
// Check if the segment overlaps with any existing segment
for (auto& it : mounted_segments_) {
auto& mtseg = it.second;
uintptr_t l1 = reinterpret_cast<uintptr_t>(mtseg.base);
uintptr_t r1 = reinterpret_cast<uintptr_t>(mtseg.size) + l1;
uintptr_t l2 = reinterpret_cast<uintptr_t>(buffer);
uintptr_t r2 = reinterpret_cast<uintptr_t>(size) + l2;
if (std::max(l1, l2) < std::min(r1, r2)) {
LOG(ERROR) << "segment_overlaps base1=" << mtseg.base
<< " size1=" << mtseg.size << " base2=" << buffer
<< " size2=" << size;
return ErrorCode::INVALID_PARAMS;
}
}
@ -555,52 +573,66 @@ ErrorCode Client::MountSegment(const std::string& segment_name,
int rc = transfer_engine_.registerLocalMemory(
(void*)buffer, size, kWildcardLocation, true, true);
if (rc != 0) {
LOG(ERROR) << "register_local_memory_failed segment_name="
<< segment_name;
LOG(ERROR) << "register_local_memory_failed base=" << buffer
<< " size=" << size << ", error=" << rc;
return ErrorCode::INVALID_PARAMS;
}
Segment segment(generate_uuid(), local_hostname_,
reinterpret_cast<uintptr_t>(buffer), size);
ErrorCode err =
master_client_.MountSegment(segment, client_id_).error_code;
if (err != ErrorCode::OK) {
LOG(ERROR) << "mount_segment_to_master_failed base=" << buffer
<< " size=" << size << ", error=" << err;
return err;
}
mounted_segments_[segment.id] = segment;
return ErrorCode::OK;
}
ErrorCode Client::UnmountSegment(const void* buffer, size_t size) {
std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
auto segment = mounted_segments_.end();
for (auto it = mounted_segments_.begin(); it != mounted_segments_.end();
++it) {
if (it->second.base == reinterpret_cast<uintptr_t>(buffer) &&
it->second.size == size) {
segment = it;
break;
}
}
if (segment == mounted_segments_.end()) {
LOG(ERROR) << "segment_not_found base=" << buffer << " size=" << size;
return ErrorCode::INVALID_PARAMS;
}
ErrorCode err =
master_client_.MountSegment(segment_name, buffer, size).error_code;
if (err != ErrorCode::OK) {
return err;
}
{
std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
mounted_segments_[segment_name] = {(void*)buffer, size};
}
return ErrorCode::OK;
}
ErrorCode Client::UnmountSegment(const std::string& segment_name, void* addr) {
void* segment_addr = nullptr;
{
std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
auto it = mounted_segments_.find(segment_name);
if (it == mounted_segments_.end() || it->second.buffer != addr) {
LOG(ERROR) << "segment_not_found segment_name=" << segment_name;
return ErrorCode::INVALID_PARAMS;
}
segment_addr = it->second.buffer;
// Remove from map first to prevent any further access to this segment
mounted_segments_.erase(it);
}
ErrorCode err = master_client_.UnmountSegment(segment_name).error_code;
master_client_.UnmountSegment(segment->second.id, client_id_)
.error_code;
if (err != ErrorCode::OK) {
LOG(ERROR) << "Failed to unmount segment from master: "
<< toString(err);
return err;
}
int rc = transfer_engine_.unregisterLocalMemory(segment_addr);
int rc = transfer_engine_.unregisterLocalMemory(
reinterpret_cast<void*>(segment->second.base));
if (rc != 0) {
LOG(ERROR) << "Failed to unregister transfer buffer with transfer "
"engine ret is "
<< rc;
return ErrorCode::INVALID_PARAMS;
if (rc != ERR_ADDRESS_NOT_REGISTERED) {
return ErrorCode::INTERNAL_ERROR;
}
// Otherwise, the segment is already unregistered from transfer engine,
// we can continue
}
mounted_segments_.erase(segment);
return ErrorCode::OK;
}
@ -668,7 +700,7 @@ ErrorCode Client::TransferRead(
return TransferData(handles, slices, TransferRequest::READ);
}
void Client::PingThreadFunc(int current_version) {
void Client::PingThreadFunc() {
// How many failed pings before getting latest master view from etcd
const int max_ping_fail_count = 3;
// How long to wait for next ping after success
@ -677,46 +709,44 @@ void Client::PingThreadFunc(int current_version) {
const int fail_ping_interval_ms = 1000;
// Increment after a ping failure, reset after a ping success
int ping_fail_count = 0;
// Set to true when there is a view change.
// When set true, will try to remount periodically.
bool need_remount = false;
auto remount_segment = [this]() {
// This lock must be held until the remount rpc is finished,
// otherwise there will be corner cases, e.g., a segment is unmounted
// successfully first, and then remounted again in this thread.
std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
std::vector<Segment> segments;
for (auto it : mounted_segments_) {
auto& name = it.first;
auto& segment = it.second;
auto err =
master_client_.MountSegment(name, segment.buffer, segment.size)
.error_code;
// If err is INVALID_PARAMS, it means the segment is already
// mounted, or cannot be mounted with current parameters. Either
// way, there is nothing we can do for this segment.
if (err != ErrorCode::OK && err != ErrorCode::INVALID_PARAMS) {
LOG(ERROR) << "Failed to remount segment " << name << ": "
<< toString(err);
return err;
}
segments.push_back(segment);
}
ErrorCode err =
master_client_.ReMountSegment(segments, client_id_).error_code;
if (err != ErrorCode::OK) {
LOG(ERROR) << "Failed to remount segments: " << err;
}
return ErrorCode::OK;
};
// Use another thread to remount segments to avoid blocking the ping thread
std::future<void> remount_segment_future;
while (ping_running_) {
auto ping_result = master_client_.Ping();
// Join the remount segment thread if it is ready
if (remount_segment_future.valid() &&
remount_segment_future.wait_for(std::chrono::seconds(0)) ==
std::future_status::ready) {
remount_segment_future = std::future<void>();
}
// Ping master
auto ping_result = master_client_.Ping(client_id_);
if (ping_result.error_code == ErrorCode::OK) {
// Reset ping failure count
ping_fail_count = 0;
if (ping_result.view_version > current_version) {
// There is an unknown view change, we need to update
// local view version and remount segments.
LOG(ERROR) << "Master view version has changed, need to "
"remount segments";
current_version = ping_result.view_version;
need_remount = true;
}
// Only try to remount if the ping succeeds and need_remount is true
if (need_remount && remount_segment() == ErrorCode::OK) {
LOG(INFO) << "Successfully remounted all segments";
need_remount = false;
if (ping_result.client_status == ClientStatus::NEED_REMOUNT &&
!remount_segment_future.valid()) {
// Ensure at most one remount segment thread is running
remount_segment_future =
std::async(std::launch::async, remount_segment);
}
std::this_thread::sleep_for(
std::chrono::milliseconds(success_ping_interval_ms));
@ -757,11 +787,10 @@ void Client::PingThreadFunc(int current_version) {
LOG(INFO) << "Reconnected to master " << master_address;
ping_fail_count = 0;
if (next_version > current_version) {
// Master view has changed
current_version = next_version;
need_remount = true;
}
}
// Explicitly wait for the remount segment thread to finish
if (remount_segment_future.valid()) {
remount_segment_future.wait();
}
}

8
mooncake-store/src/ha_helper.cpp
View File

@ -93,8 +93,8 @@ MasterServiceSupervisor::MasterServiceSupervisor(
int port, int server_thread_num, bool enable_gc,
bool enable_metric_reporting, int metrics_port,
int64_t default_kv_lease_ttl, double eviction_ratio,
double eviction_high_watermark_ratio, const std::string& etcd_endpoints,
const std::string& local_hostname)
double eviction_high_watermark_ratio, int64_t client_live_ttl_sec,
const std::string& etcd_endpoints, const std::string& local_hostname)
: port_(port),
server_thread_num_(server_thread_num),
enable_gc_(enable_gc),
@ -103,6 +103,7 @@ MasterServiceSupervisor::MasterServiceSupervisor(
default_kv_lease_ttl_(default_kv_lease_ttl),
eviction_ratio_(eviction_ratio),
eviction_high_watermark_ratio_(eviction_high_watermark_ratio),
client_live_ttl_sec_(client_live_ttl_sec),
etcd_endpoints_(etcd_endpoints),
local_hostname_(local_hostname) {}
@ -136,10 +137,11 @@ int MasterServiceSupervisor::Start() {
std::this_thread::sleep_for(std::chrono::seconds(waiting_time));
LOG(INFO) << "Starting master service...";
bool enable_ha = true;
mooncake::WrappedMasterService wrapped_master_service(
enable_gc_, default_kv_lease_ttl_, enable_metric_reporting_,
metrics_port_, eviction_ratio_, eviction_high_watermark_ratio_,
version);
version, client_live_ttl_sec_, enable_ha);
mooncake::RegisterRpcService(server, wrapped_master_service);
// Metric reporting is now handled by WrappedMasterService.

13
mooncake-store/src/master.cpp
View File

@ -39,6 +39,9 @@ DEFINE_string(
"Endpoints of ETCD server, separated by semicolon, required in HA mode");
DEFINE_string(local_hostname, "",
"Local host address (IP:Port), required in HA mode");
DEFINE_int64(client_ttl, mooncake::DEFAULT_CLIENT_LIVE_TTL_SEC,
"How long a client is considered alive after the last ping, only "
"used in HA mode");
int main(int argc, char* argv[]) {
easylog::set_min_severity(easylog::Severity::WARN);
@ -56,7 +59,8 @@ int main(int argc, char* argv[]) {
<< FLAGS_eviction_high_watermark_ratio
<< ", enable_ha=" << FLAGS_enable_ha
<< ", etcd_endpoints=" << FLAGS_etcd_endpoints
<< ", local_hostname=" << FLAGS_local_hostname;
<< ", local_hostname=" << FLAGS_local_hostname
<< ", client_ttl=" << FLAGS_client_ttl;
int server_thread_num =
std::min(FLAGS_max_threads,
@ -85,8 +89,8 @@ int main(int argc, char* argv[]) {
FLAGS_port, server_thread_num, FLAGS_enable_gc,
FLAGS_enable_metric_reporting, FLAGS_metrics_port,
FLAGS_default_kv_lease_ttl, FLAGS_eviction_ratio,
FLAGS_eviction_high_watermark_ratio, FLAGS_etcd_endpoints,
FLAGS_local_hostname);
FLAGS_eviction_high_watermark_ratio, FLAGS_client_ttl,
FLAGS_etcd_endpoints, FLAGS_local_hostname);
return supervisor.Start();
} else {
@ -96,7 +100,8 @@ int main(int argc, char* argv[]) {
mooncake::WrappedMasterService wrapped_master_service(
FLAGS_enable_gc, FLAGS_default_kv_lease_ttl,
FLAGS_enable_metric_reporting, FLAGS_metrics_port,
FLAGS_eviction_ratio, FLAGS_eviction_high_watermark_ratio, version);
FLAGS_eviction_ratio, FLAGS_eviction_high_watermark_ratio, version,
FLAGS_client_ttl, FLAGS_enable_ha);
mooncake::RegisterRpcService(server, wrapped_master_service);
return server.start();

58
mooncake-store/src/master_client.cpp
View File

@ -331,20 +331,19 @@ RemoveAllResponse MasterClient::RemoveAll() {
return result.value();
}
MountSegmentResponse MasterClient::MountSegment(const std::string& segment_name,
const void* buffer,
size_t size) {
MountSegmentResponse MasterClient::MountSegment(const Segment& segment,
const UUID& client_id) {
ScopedVLogTimer timer(1, "MasterClient::MountSegment");
timer.LogRequest("segment_name=", segment_name, ", buffer=", buffer,
", size=", size);
timer.LogRequest("base=", segment.base, ", size=", segment.size,
", name=", segment.name, ", id=", segment.id,
", client_id=", client_id);
std::optional<MountSegmentResponse> result =
syncAwait([&]() -> coro::Lazy<std::optional<MountSegmentResponse>> {
Lazy<async_rpc_result<MountSegmentResponse>> handler =
co_await client_
.send_request<&WrappedMasterService::MountSegment>(
reinterpret_cast<uint64_t>(buffer),
static_cast<uint64_t>(size), segment_name);
segment, client_id);
async_rpc_result<MountSegmentResponse> result = co_await handler;
if (!result) {
co_return std::nullopt;
@ -361,14 +360,41 @@ MountSegmentResponse MasterClient::MountSegment(const std::string& segment_name,
return result.value();
}
UnmountSegmentResponse MasterClient::UnmountSegment(
const std::string& segment_name) {
ReMountSegmentResponse MasterClient::ReMountSegment(
const std::vector<Segment>& segments, const UUID& client_id) {
ScopedVLogTimer timer(1, "MasterClient::ReMountSegment");
timer.LogRequest("segments_num=", segments.size(), ", client_id=", client_id);
std::optional<ReMountSegmentResponse> result =
syncAwait([&]() -> coro::Lazy<std::optional<ReMountSegmentResponse>> {
Lazy<async_rpc_result<ReMountSegmentResponse>> handler =
co_await client_
.send_request<&WrappedMasterService::ReMountSegment>(
segments, client_id);
async_rpc_result<ReMountSegmentResponse> result = co_await handler;
if (!result) {
co_return std::nullopt;
}
co_return result->result();
}());
if (!result) {
LOG(ERROR) << "Failed to remount segment due to rpc error";
auto response = ReMountSegmentResponse{ErrorCode::RPC_FAIL};
timer.LogResponseJson(response);
return response;
}
timer.LogResponseJson(result.value());
return result.value();
}
UnmountSegmentResponse MasterClient::UnmountSegment(const UUID& segment_id,
const UUID& client_id) {
ScopedVLogTimer timer(1, "MasterClient::UnmountSegment");
timer.LogRequest("segment_name=", segment_name);
timer.LogRequest("segment_id=", segment_id, ", client_id=", client_id);
auto request_result =
client_.send_request<&WrappedMasterService::UnmountSegment>(
segment_name);
client_.send_request<&WrappedMasterService::UnmountSegment>(segment_id,
client_id);
std::optional<UnmountSegmentResponse> result = coro::syncAwait(
[&]() -> coro::Lazy<std::optional<UnmountSegmentResponse>> {
auto result = co_await co_await request_result;
@ -388,12 +414,12 @@ UnmountSegmentResponse MasterClient::UnmountSegment(
return result.value();
}
PingResponse MasterClient::Ping() {
PingResponse MasterClient::Ping(const UUID& client_id) {
ScopedVLogTimer timer(1, "MasterClient::Ping");
timer.LogRequest("action=ping");
timer.LogRequest("client_id=", client_id);
auto request_result =
client_.send_request<&WrappedMasterService::Ping>();
client_.send_request<&WrappedMasterService::Ping>(client_id);
std::optional<PingResponse> result =
coro::syncAwait([&]() -> coro::Lazy<std::optional<PingResponse>> {
auto result = co_await co_await request_result;
@ -405,7 +431,7 @@ PingResponse MasterClient::Ping() {
}());
if (!result) {
auto response = PingResponse{0, ErrorCode::RPC_FAIL};
auto response = PingResponse{0, ClientStatus::UNDEFINED, ErrorCode::RPC_FAIL};
timer.LogResponseJson(response);
return response;
}

74
mooncake-store/src/master_metric_manager.cpp
View File

@ -27,6 +27,10 @@ MasterMetricManager::MasterMetricManager()
"Distribution of object value sizes",
{4096, 65536, 262144, 1048576, 4194304,
16777216, 67108864}),
// Initialize cluster metrics
active_clients_("master_active_clients",
"Total number of active clients"),
// Initialize Request Counters
put_start_requests_("master_put_start_requests_total",
"Total number of PutStart requests received"),
@ -71,8 +75,16 @@ MasterMetricManager::MasterMetricManager()
unmount_segment_failures_(
"master_unmount_segment_failures_total",
"Total number of failed UnmountSegment requests"),
remount_segment_requests_(
"master_remount_segment_requests_total",
"Total number of RemountSegment requests received"),
remount_segment_failures_(
"master_remount_segment_failures_total",
"Total number of failed RemountSegment requests"),
ping_requests_("master_ping_requests_total",
"Total number of ping requests received"),
ping_failures_("master_ping_failures_total",
"Total number of failed ping requests"),
// Initialize Eviction Counters
eviction_success_("master_successful_evictions_total",
@ -130,6 +142,19 @@ int64_t MasterMetricManager::get_key_count() {
return key_count_.value();
}
// Cluster Metrics
void MasterMetricManager::inc_active_clients(int64_t val) {
active_clients_.inc(val);
}
void MasterMetricManager::dec_active_clients(int64_t val) {
active_clients_.dec(val);
}
int64_t MasterMetricManager::get_active_clients() {
return active_clients_.value();
}
// Operation Statistics (Counters)
void MasterMetricManager::inc_exist_key_requests(int64_t val) {
exist_key_requests_.inc(val);
@ -185,9 +210,18 @@ void MasterMetricManager::inc_unmount_segment_requests(int64_t val) {
void MasterMetricManager::inc_unmount_segment_failures(int64_t val) {
unmount_segment_failures_.inc(val);
}
void MasterMetricManager::inc_remount_segment_requests(int64_t val) {
remount_segment_requests_.inc(val);
}
void MasterMetricManager::inc_remount_segment_failures(int64_t val) {
remount_segment_failures_.inc(val);
}
void MasterMetricManager::inc_ping_requests(int64_t val) {
ping_requests_.inc(val);
}
void MasterMetricManager::inc_ping_failures(int64_t val) {
ping_failures_.inc(val);
}
int64_t MasterMetricManager::get_put_start_requests() {
return put_start_requests_.value();
@ -261,10 +295,22 @@ int64_t MasterMetricManager::get_unmount_segment_failures() {
return unmount_segment_failures_.value();
}
int64_t MasterMetricManager::get_remount_segment_requests() {
return remount_segment_requests_.value();
}
int64_t MasterMetricManager::get_remount_segment_failures() {
return remount_segment_failures_.value();
}
int64_t MasterMetricManager::get_ping_requests() {
return ping_requests_.value();
}
int64_t MasterMetricManager::get_ping_failures() {
return ping_failures_.value();
}
// Eviction Metrics
void MasterMetricManager::inc_eviction_success(int64_t key_count, int64_t size) {
evicted_key_count_.inc(key_count);
@ -293,6 +339,11 @@ int64_t MasterMetricManager::get_evicted_size() {
return evicted_size_.value();
}
// --- Setters ---
void MasterMetricManager::set_enable_ha(bool enable_ha) {
enable_ha_ = enable_ha;
}
// --- Serialization ---
std::string MasterMetricManager::serialize_metrics() {
// Note: Following Prometheus style, metrics with value 0 that haven't
@ -311,6 +362,9 @@ std::string MasterMetricManager::serialize_metrics() {
serialize_metric(allocated_size_);
serialize_metric(total_capacity_);
serialize_metric(key_count_);
if (enable_ha_) {
serialize_metric(active_clients_);
}
// Serialize Histogram
serialize_metric(value_size_distribution_);
@ -334,7 +388,12 @@ std::string MasterMetricManager::serialize_metrics() {
serialize_metric(mount_segment_failures_);
serialize_metric(unmount_segment_requests_);
serialize_metric(unmount_segment_failures_);
serialize_metric(ping_requests_);
serialize_metric(remount_segment_requests_);
serialize_metric(remount_segment_failures_);
if (enable_ha_) {
serialize_metric(ping_requests_);
serialize_metric(ping_failures_);
}
// Serialize Eviction Counters
serialize_metric(eviction_success_);
@ -369,6 +428,7 @@ std::string MasterMetricManager::get_summary_string() {
int64_t allocated = allocated_size_.value();
int64_t capacity = total_capacity_.value();
int64_t keys = key_count_.value();
int64_t active_clients = active_clients_.value();
// Request counters
int64_t exist_keys = exist_key_requests_.value();
@ -383,7 +443,6 @@ std::string MasterMetricManager::get_summary_string() {
int64_t remove_fails = remove_failures_.value();
int64_t remove_all = remove_all_requests_.value();
int64_t remove_all_fails = remove_all_failures_.value();
int64_t pings = ping_requests_.value();
// Eviction counters
int64_t eviction_success = eviction_success_.value();
@ -391,6 +450,10 @@ std::string MasterMetricManager::get_summary_string() {
int64_t evicted_key_count = evicted_key_count_.value();
int64_t evicted_size = evicted_size_.value();
// Ping counters
int64_t ping = ping_requests_.value();
int64_t ping_fails = ping_failures_.value();
// --- Format the summary string ---
ss << "Storage: " << format_bytes(allocated) << " / "
<< format_bytes(capacity);
@ -399,6 +462,9 @@ std::string MasterMetricManager::get_summary_string() {
<< ((double) allocated / (double)capacity * 100.0) << "%)";
}
ss << " | Keys: " << keys;
if (enable_ha_) {
ss << " | Clients: " << active_clients;
}
// Request summary - focus on the most important metrics
ss << " | Requests (Success/Total): ";
@ -409,7 +475,9 @@ std::string MasterMetricManager::get_summary_string() {
ss << "Exist=" << exist_keys - exist_key_fails << "/" << exist_keys << ", ";
ss << "Del=" << removes - remove_fails << "/" << removes << ", ";
ss << "DelAll=" << remove_all - remove_all_fails << "/" << remove_all << ", ";
ss << "Ping=" << pings;
if (enable_ha_) {
ss << "Ping=" << ping - ping_fails << "/" << ping << ", ";
}
// Eviction summary
ss << " | Eviction: "

409
mooncake-store/src/master_service.cpp
View File

@ -10,71 +10,18 @@
namespace mooncake {
ErrorCode BufferAllocatorManager::AddSegment(const std::string& segment_name,
uint64_t base, uint64_t size) {
// Check if parameters are valid before allocating memory.
if (base == 0 || size == 0 ||
reinterpret_cast<uintptr_t>(base) % facebook::cachelib::Slab::kSize ||
size % facebook::cachelib::Slab::kSize) {
LOG(ERROR) << "base_address=" << base << " or size=" << size
<< " is not aligned to " << facebook::cachelib::Slab::kSize;
return ErrorCode::INVALID_PARAMS;
}
std::unique_lock<std::shared_mutex> lock(allocator_mutex_);
// Check if segment already exists
if (buf_allocators_.find(segment_name) != buf_allocators_.end()) {
LOG(WARNING) << "segment_name=" << segment_name
<< ", error=segment_already_exists";
return ErrorCode::INVALID_PARAMS;
}
std::shared_ptr<BufferAllocator> allocator;
try {
// SlabAllocator may throw an exception if the size or base is invalid
// for the slab allocator.
allocator = std::make_shared<BufferAllocator>(segment_name, base, size);
if (!allocator) {
LOG(ERROR) << "segment_name=" << segment_name
<< ", error=failed_to_create_allocator";
return ErrorCode::INVALID_PARAMS;
}
} catch (...) {
LOG(ERROR) << "segment_name=" << segment_name
<< ", error=unknown_exception_during_allocator_creation";
return ErrorCode::INVALID_PARAMS;
}
buf_allocators_[segment_name] = std::move(allocator);
return ErrorCode::OK;
}
ErrorCode BufferAllocatorManager::RemoveSegment(
const std::string& segment_name) {
std::unique_lock<std::shared_mutex> lock(allocator_mutex_);
auto it = buf_allocators_.find(segment_name);
if (it == buf_allocators_.end()) {
LOG(WARNING) << "segment_name=" << segment_name
<< ", error=segment_not_found";
return ErrorCode::INVALID_PARAMS;
}
MasterMetricManager::instance().dec_total_capacity(it->second->capacity());
buf_allocators_.erase(it);
return ErrorCode::OK;
}
MasterService::MasterService(bool enable_gc, uint64_t default_kv_lease_ttl,
double eviction_ratio,
double eviction_high_watermark_ratio)
: buffer_allocator_manager_(std::make_shared<BufferAllocatorManager>()),
allocation_strategy_(std::make_shared<RandomAllocationStrategy>()),
double eviction_high_watermark_ratio,
ViewVersionId view_version,
int64_t client_live_ttl_sec, bool enable_ha)
: allocation_strategy_(std::make_shared<RandomAllocationStrategy>()),
enable_gc_(enable_gc),
default_kv_lease_ttl_(default_kv_lease_ttl),
eviction_ratio_(eviction_ratio),
eviction_high_watermark_ratio_(eviction_high_watermark_ratio) {
eviction_high_watermark_ratio_(eviction_high_watermark_ratio),
client_live_ttl_sec_(client_live_ttl_sec),
enable_ha_(enable_ha) {
if (eviction_ratio_ < 0.0 || eviction_ratio_ > 1.0) {
LOG(ERROR) << "Eviction ratio must be between 0.0 and 1.0, "
<< "current value: " << eviction_ratio_;
@ -90,14 +37,24 @@ MasterService::MasterService(bool enable_gc, uint64_t default_kv_lease_ttl,
gc_running_ = true;
gc_thread_ = std::thread(&MasterService::GCThreadFunc, this);
VLOG(1) << "action=start_gc_thread";
if (enable_ha) {
client_monitor_running_ = true;
client_monitor_thread_ = std::thread(&MasterService::ClientMonitorFunc, this);
VLOG(1) << "action=start_client_monitor_thread";
}
}
MasterService::~MasterService() {
// Stop and join the GC thread
// Stop and join the threads
gc_running_ = false;
client_monitor_running_ = false;
if (gc_thread_.joinable()) {
gc_thread_.join();
}
if (client_monitor_thread_.joinable()) {
client_monitor_thread_.join();
}
// Clean up any remaining GC tasks
GCTask* task = nullptr;
@ -108,23 +65,91 @@ MasterService::~MasterService() {
}
}
ErrorCode MasterService::MountSegment(uint64_t buffer, uint64_t size,
const std::string& segment_name) {
if (buffer == 0 || size == 0) {
LOG(ERROR) << "buffer=" << buffer << ", size=" << size
<< ", error=invalid_buffer_params";
return ErrorCode::INVALID_PARAMS;
ErrorCode MasterService::MountSegment(const Segment& segment,
const UUID& client_id) {
ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
if (enable_ha_) {
// Tell the client monitor thread to start timing for this client. To
// avoid the following undesired situations, this message must be sent
// after locking the segment mutex and before the mounting operation
// completes:
// 1. Sending the message before the lock: the client expires and
// unmouting invokes before this mounting are completed, which prevents
// this segment being able to be unmounted forever;
// 2. Sending the message after mounting the segment: After mounting
// this segment, when trying to push id to the queue, the queue is
// already full. However, at this point, the message must be sent,
// otherwise this client cannot be monitored and expired.
PodUUID pod_client_id;
pod_client_id.first = client_id.first;
pod_client_id.second = client_id.second;
if (!client_ping_queue_.push(pod_client_id)) {
LOG(ERROR) << "segment_name=" << segment.name
<< ", error=client_ping_queue_full";
return ErrorCode::INTERNAL_ERROR;
}
}
return buffer_allocator_manager_->AddSegment(segment_name, buffer, size);
auto err = segment_access.MountSegment(segment, client_id);
if (err == ErrorCode::SEGMENT_ALREADY_EXISTS) {
// Return OK because this is an idempotent operation
return ErrorCode::OK;
} else {
return err;
}
}
ErrorCode MasterService::UnmountSegment(const std::string& segment_name) {
// 1. Remove the segment from the allocator
auto ret = buffer_allocator_manager_->RemoveSegment(segment_name);
if (ret != ErrorCode::OK) return ret;
ErrorCode MasterService::ReMountSegment(const std::vector<Segment>& segments,
const UUID& client_id) {
if (!enable_ha_) {
LOG(ERROR) << "ReMountSegment is only available in HA mode";
return ErrorCode::UNAVAILABLE_IN_CURRENT_MODE;
}
// 2. Remove the metadata of the related objects
std::unique_lock<std::shared_mutex> lock(client_mutex_);
if (ok_client_.contains(client_id)) {
LOG(WARNING) << "client_id=" << client_id
<< ", warn=client_already_remounted";
// Return OK because this is an idempotent operation
return ErrorCode::OK;
}
ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
// Tell the client monitor thread to start timing for this client. To
// avoid the following undesired situations, this message must be sent
// after locking the segment mutex or client mutex and before the remounting
// operation completes:
// 1. Sending the message before the lock: the client expires and
// unmouting invokes before this remounting are completed, which prevents
// this segment being able to be unmounted forever;
// 2. Sending the message after remounting the segments: After remounting
// these segments, when trying to push id to the queue, the queue is
// already full. However, at this point, the message must be sent,
// otherwise this client cannot be monitored and expired.
PodUUID pod_client_id;
pod_client_id.first = client_id.first;
pod_client_id.second = client_id.second;
if (!client_ping_queue_.push(pod_client_id)) {
LOG(ERROR) << "client_id=" << client_id
<< ", error=client_ping_queue_full";
return ErrorCode::INTERNAL_ERROR;
}
ErrorCode err = segment_access.ReMountSegment(segments, client_id);
if (err != ErrorCode::OK) {
return err;
}
// Change the client status to OK
ok_client_.insert(client_id);
MasterMetricManager::instance().inc_active_clients();
return ErrorCode::OK;
}
void MasterService::ClearInvalidHandles() {
for (auto& shard : metadata_shards_) {
std::unique_lock lock(shard.mutex);
auto it = shard.metadata.begin();
@ -146,8 +171,35 @@ ErrorCode MasterService::UnmountSegment(const std::string& segment_name) {
}
}
}
}
return ErrorCode::OK;
ErrorCode MasterService::UnmountSegment(const UUID& segment_id,
const UUID& client_id) {
size_t metrics_dec_capacity = 0; // to update the metrics
// 1. Prepare to unmount the segment by deleting its allocator
{
ScopedSegmentAccess segment_access =
segment_manager_.getSegmentAccess();
ErrorCode err = segment_access.PrepareUnmountSegment(
segment_id, metrics_dec_capacity);
if (err == ErrorCode::SEGMENT_NOT_FOUND) {
// Return OK because this is an idempotent operation
return ErrorCode::OK;
}
if (err != ErrorCode::OK) {
return err;
}
} // Release the segment mutex before long-running step 2 and avoid
// deadlocks
// 2. Remove the metadata of the related objects
ClearInvalidHandles();
// 3. Commit the unmount operation
ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
return segment_access.CommitUnmountSegment(segment_id, client_id,
metrics_dec_capacity);
}
ErrorCode MasterService::ExistKey(const std::string& key) {
@ -172,7 +224,7 @@ ErrorCode MasterService::ExistKey(const std::string& key) {
ErrorCode MasterService::GetAllKeys(std::vector<std::string> & all_keys) {
all_keys.clear();
for(int i = 0; i < kNumShards; i++) {
for(size_t i = 0; i < kNumShards; i++) {
for(const auto& item : metadata_shards_[i].metadata) {
all_keys.push_back(item.first);
}
@ -180,35 +232,16 @@ ErrorCode MasterService::GetAllKeys(std::vector<std::string> & all_keys) {
return ErrorCode::OK;
}
ErrorCode MasterService::GetAllSegments(std::vector<std::string> & all_segments) {
all_segments.clear();
std::shared_lock<std::shared_mutex> alloc_lock(
buffer_allocator_manager_->GetMutex());
const auto& allocators = buffer_allocator_manager_->GetAllocators();
for(auto & allocator : allocators) {
all_segments.push_back(allocator.first);
}
alloc_lock.unlock();
return ErrorCode::OK;
ErrorCode MasterService::GetAllSegments(
std::vector<std::string>& all_segments) {
ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
return segment_access.GetAllSegments(all_segments);
}
ErrorCode MasterService::QuerySegments(const std::string & segment,
size_t & used,
size_t & capacity) {
std::shared_lock<std::shared_mutex> alloc_lock(
buffer_allocator_manager_->GetMutex());
const auto& allocators = buffer_allocator_manager_->GetAllocators();
auto it = allocators.find(segment);
if (it != allocators.end()) {
auto& allocator = it -> second;
capacity = allocator -> capacity();
used = allocator -> size();
} else {
VLOG(1) << "### DEBUG ### MasterService::QuerySegments(" << segment << ") not found!";
return ErrorCode::AVAILABLE_SEGMENT_EMPTY;
}
alloc_lock.unlock();
return ErrorCode::OK;
ErrorCode MasterService::QuerySegments(const std::string& segment, size_t& used,
size_t& capacity) {
ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess();
return segment_access.QuerySegments(segment, used, capacity);
}
ErrorCode MasterService::GetReplicaList(
@ -309,42 +342,41 @@ ErrorCode MasterService::PutStart(
// Allocate replicas
std::vector<Replica> replicas;
replicas.reserve(config.replica_num);
for (size_t i = 0; i < config.replica_num; ++i) {
std::vector<std::unique_ptr<AllocatedBuffer>> handles;
handles.reserve(slice_lengths.size());
{
ScopedAllocatorAccess allocator_access = segment_manager_.getAllocatorAccess();
auto& allocators = allocator_access.getAllocators();
auto& allocators_by_name = allocator_access.getAllocatorsByName();
for (size_t i = 0; i < config.replica_num; ++i) {
std::vector<std::unique_ptr<AllocatedBuffer>> handles;
handles.reserve(slice_lengths.size());
// Allocate space for each slice
for (size_t j = 0; j < slice_lengths.size(); ++j) {
auto chunk_size = slice_lengths[j];
// Allocate space for each slice
for (size_t j = 0; j < slice_lengths.size(); ++j) {
auto chunk_size = slice_lengths[j];
// Use allocation strategy to select an allocator
std::shared_lock<std::shared_mutex> alloc_lock(
buffer_allocator_manager_->GetMutex());
const auto& allocators = buffer_allocator_manager_->GetAllocators();
// Use the unified allocation strategy with replica config
auto handle =
allocation_strategy_->Allocate(allocators, allocators_by_name, chunk_size, config);
// Use the unified allocation strategy with replica config
auto handle =
allocation_strategy_->Allocate(allocators, chunk_size, config);
alloc_lock.unlock();
if (!handle) {
LOG(ERROR) << "key=" << key << ", replica_id=" << i
<< ", slice_index=" << j
<< ", error=allocation_failed";
replica_list.clear();
// If the allocation failed, we need to evict some objects
// to free up space for future allocations.
need_eviction_ = true;
return ErrorCode::NO_AVAILABLE_HANDLE;
}
if (!handle) {
LOG(ERROR) << "key=" << key << ", replica_id=" << i
<< ", slice_index=" << j
<< ", error=allocation_failed";
replica_list.clear();
// If the allocation failed, we need to evict some objects
// to free up space for future allocations.
need_eviction_ = true;
return ErrorCode::NO_AVAILABLE_HANDLE;
VLOG(1) << "key=" << key << ", replica_id=" << i
<< ", slice_index=" << j << ", handle=" << *handle
<< ", action=slice_allocated";
handles.emplace_back(std::move(handle));
}
VLOG(1) << "key=" << key << ", replica_id=" << i
<< ", slice_index=" << j << ", handle=" << *handle
<< ", action=slice_allocated";
handles.emplace_back(std::move(handle));
replicas.emplace_back(std::move(handles), ReplicaStatus::PROCESSING);
}
replicas.emplace_back(std::move(handles), ReplicaStatus::PROCESSING);
}
metadata.replicas = std::move(replicas);
@ -552,6 +584,32 @@ size_t MasterService::GetKeyCount() const {
return total;
}
ErrorCode MasterService::Ping(const UUID& client_id,
ViewVersionId& view_version,
ClientStatus& client_status) {
if (!enable_ha_) {
LOG(ERROR) << "Ping is only available in HA mode";
return ErrorCode::UNAVAILABLE_IN_CURRENT_MODE;
}
std::shared_lock<std::shared_mutex> lock(client_mutex_);
auto it = ok_client_.find(client_id);
if (it != ok_client_.end()) {
client_status = ClientStatus::OK;
} else {
client_status = ClientStatus::NEED_REMOUNT;
}
view_version = view_version_;
PodUUID pod_client_id = {client_id.first, client_id.second};
if (!client_ping_queue_.push(pod_client_id)) {
// Queue is full
LOG(ERROR) << "client_id=" << client_id
<< ", error=client_ping_queue_full";
return ErrorCode::INTERNAL_ERROR;
}
return ErrorCode::OK;
}
void MasterService::GCThreadFunc() {
VLOG(1) << "action=gc_thread_started";
@ -695,4 +753,93 @@ void MasterService::BatchEvict(double eviction_ratio) {
<< ", total_freed_size=" << total_freed_size;
}
void MasterService::ClientMonitorFunc() {
std::unordered_map<UUID, std::chrono::steady_clock::time_point,
boost::hash<UUID>>
client_ttl;
while (client_monitor_running_) {
auto now = std::chrono::steady_clock::now();
// Update the client ttl
PodUUID pod_client_id;
while (client_ping_queue_.pop(pod_client_id)) {
UUID client_id = {pod_client_id.first, pod_client_id.second};
client_ttl[client_id] = now + std::chrono::seconds(client_live_ttl_sec_);
}
// Find out expired clients
std::vector<UUID> expired_clients;
for (auto it = client_ttl.begin(); it != client_ttl.end();) {
if (it->second < now) {
LOG(INFO) << "client_id=" << it->first << ", action=client_expired";
expired_clients.push_back(it->first);
it = client_ttl.erase(it);
} else {
++it;
}
}
// Update the client status to NEED_REMOUNT
if (!expired_clients.empty()) {
// Record which segments are unmounted, will be used in the commit
// phase.
std::vector<UUID> unmount_segments;
std::vector<size_t> dec_capacities;
std::vector<UUID> client_ids;
std::vector<std::string> segment_names;
{
// Lock client_mutex and segment_mutex
std::unique_lock<std::shared_mutex> lock(client_mutex_);
for (auto& client_id : expired_clients) {
auto it = ok_client_.find(client_id);
if (it != ok_client_.end()) {
ok_client_.erase(it);
MasterMetricManager::instance().dec_active_clients();
}
}
ScopedSegmentAccess segment_access =
segment_manager_.getSegmentAccess();
for (auto& client_id : expired_clients) {
std::vector<Segment> segments;
segment_access.GetClientSegments(client_id, segments);
for (auto& seg : segments) {
size_t metrics_dec_capacity = 0;
if (segment_access.PrepareUnmountSegment(
seg.id, metrics_dec_capacity) ==
ErrorCode::OK) {
unmount_segments.push_back(seg.id);
dec_capacities.push_back(metrics_dec_capacity);
client_ids.push_back(client_id);
segment_names.push_back(seg.name);
} else {
LOG(ERROR) << "client_id=" << client_id
<< ", segment_name=" << seg.name
<< ", error=prepare_unmount_expired_segment_failed";
}
}
}
} // Release the mutex before long-running ClearInvalidHandles and
// avoid deadlocks
if (!unmount_segments.empty()) {
ClearInvalidHandles();
ScopedSegmentAccess segment_access =
segment_manager_.getSegmentAccess();
for (size_t i = 0; i < unmount_segments.size(); i++) {
segment_access.CommitUnmountSegment(
unmount_segments[i], client_ids[i], dec_capacities[i]);
LOG(INFO) << "client_id=" << client_ids[i]
<< ", segment_name=" << segment_names[i]
<< ", action=unmount_expired_segment";
}
}
}
std::this_thread::sleep_for(
std::chrono::milliseconds(kClientMonitorSleepMs));
}
}
} // namespace mooncake

229
mooncake-store/src/segment.cpp Normal file
View File

@ -0,0 +1,229 @@
#include "segment.h"
#include "master_metric_manager.h"
namespace mooncake {
ErrorCode ScopedSegmentAccess::MountSegment(const Segment& segment,
const UUID& client_id) {
const uintptr_t buffer = segment.base;
const size_t size = segment.size;
// Check if parameters are valid before allocating memory.
if (buffer == 0 || size == 0 || buffer % facebook::cachelib::Slab::kSize ||
size % facebook::cachelib::Slab::kSize) {
LOG(ERROR) << "buffer=" << buffer << " or size=" << size
<< " is not aligned to " << facebook::cachelib::Slab::kSize;
return ErrorCode::INVALID_PARAMS;
}
// Check if segment already exists
auto exist_segment_it =
segment_manager_->mounted_segments_.find(segment.id);
if (exist_segment_it != segment_manager_->mounted_segments_.end()) {
auto& exist_segment = exist_segment_it->second;
if (exist_segment.status == SegmentStatus::OK) {
LOG(WARNING) << "segment_name=" << segment.name
<< ", warn=segment_already_exists";
return ErrorCode::SEGMENT_ALREADY_EXISTS;
} else {
LOG(ERROR) << "segment_name=" << segment.name
<< ", error=segment_already_exists_but_not_ok"
<< ", status=" << exist_segment.status;
return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS;
}
}
std::shared_ptr<BufferAllocator> allocator;
try {
// SlabAllocator may throw an exception if the size or base is invalid
// for the slab allocator.
allocator =
std::make_shared<BufferAllocator>(segment.name, buffer, size);
if (!allocator) {
LOG(ERROR) << "segment_name=" << segment.name
<< ", error=failed_to_create_allocator";
return ErrorCode::INVALID_PARAMS;
}
} catch (...) {
LOG(ERROR) << "segment_name=" << segment.name
<< ", error=unknown_exception_during_allocator_creation";
return ErrorCode::INVALID_PARAMS;
}
segment_manager_->allocators_.push_back(allocator);
segment_manager_->allocators_by_name_[segment.name].push_back(allocator);
segment_manager_->client_segments_[client_id].push_back(segment.id);
segment_manager_->mounted_segments_[segment.id] = {
segment, SegmentStatus::OK, std::move(allocator)};
MasterMetricManager::instance().inc_total_capacity(size);
return ErrorCode::OK;
}
ErrorCode ScopedSegmentAccess::ReMountSegment(
const std::vector<Segment>& segments, const UUID& client_id) {
for (const auto& segment : segments) {
ErrorCode err = MountSegment(segment, client_id);
if (err == ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS ||
err == ErrorCode::INTERNAL_ERROR) {
LOG(ERROR) << "segment_name=" << segment.name
<< ", error=fail_to_remount_segment";
return err;
} else if (err == ErrorCode::INVALID_PARAMS) {
// Ignore INVALID_PARAMS. This error cannot be solved by a new
// remount request.
LOG(WARNING) << "segment_name=" << segment.name
<< ", warn=invalid_params";
} else if (err == ErrorCode::SEGMENT_ALREADY_EXISTS) {
// Segment already exists, no need to remount.
LOG(WARNING) << "segment_name=" << segment.name
<< ", warn=segment_already_exists";
} else if (err != ErrorCode::OK) {
// Ignore other errors. The error may not be solvable by a new
// remount request.
LOG(ERROR) << "segment_name=" << segment.name
<< ", error=unexpected_error (" << err << ")";
}
}
return ErrorCode::OK;
}
ErrorCode ScopedSegmentAccess::PrepareUnmountSegment(
const UUID& segment_id, size_t& metrics_dec_capacity) {
auto it = segment_manager_->mounted_segments_.find(segment_id);
if (it == segment_manager_->mounted_segments_.end()) {
LOG(WARNING) << "segment_id=" << segment_id
<< ", warn=segment_not_found";
return ErrorCode::SEGMENT_NOT_FOUND;
}
if (it->second.status == SegmentStatus::UNMOUNTING) {
LOG(ERROR) << "segment_id=" << segment_id
<< ", error=segment_is_unmounting";
return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS;
}
auto& mounted_segment = it->second;
auto& segment = mounted_segment.segment;
metrics_dec_capacity = segment.size;
// Remove the allocator from the segment manager
std::shared_ptr<BufferAllocator> allocator = mounted_segment.buf_allocator;
// 1. Remove from allocators
auto alloc_it = std::find(segment_manager_->allocators_.begin(),
segment_manager_->allocators_.end(), allocator);
if (alloc_it != segment_manager_->allocators_.end()) {
segment_manager_->allocators_.erase(alloc_it);
} else {
LOG(ERROR) << "segment_name=" << segment.name
<< ", error=allocator_not_found_in_allocators";
}
// 2. Remove from allocators_by_name
bool found_in_allocators_by_name = false;
auto name_it = segment_manager_->allocators_by_name_.find(segment.name);
if (name_it != segment_manager_->allocators_by_name_.end()) {
auto& allocators = name_it->second;
auto alloc_it =
std::find(allocators.begin(), allocators.end(), allocator);
if (alloc_it != allocators.end()) {
allocators.erase(alloc_it);
found_in_allocators_by_name = true;
}
if (allocators.empty()) {
segment_manager_->allocators_by_name_.erase(name_it);
}
}
if (!found_in_allocators_by_name) {
LOG(ERROR) << "segment_name=" << segment.name
<< ", error=allocator_not_found_in_allocators_by_name";
}
// 3. Remove from mounted_segment
mounted_segment.buf_allocator.reset();
// Set the segment status to UNMOUNTING
mounted_segment.status = SegmentStatus::UNMOUNTING;
return ErrorCode::OK;
}
ErrorCode ScopedSegmentAccess::CommitUnmountSegment(
const UUID& segment_id, const UUID& client_id,
const size_t& metrics_dec_capacity) {
// Remove from client_segments_
bool found_in_client_segments = false;
auto client_it = segment_manager_->client_segments_.find(client_id);
if (client_it != segment_manager_->client_segments_.end()) {
auto& segments = client_it->second;
auto segment_it =
std::find(segments.begin(), segments.end(), segment_id);
if (segment_it != segments.end()) {
segments.erase(segment_it);
found_in_client_segments = true;
}
if (segments.empty()) {
segment_manager_->client_segments_.erase(client_it);
}
}
if (!found_in_client_segments) {
LOG(ERROR) << "segment_id=" << segment_id
<< ", error=segment_not_found_in_client_segments";
}
// Remove from mounted_segments_
segment_manager_->mounted_segments_.erase(segment_id);
// Decrease the total capacity
MasterMetricManager::instance().dec_total_capacity(metrics_dec_capacity);
return ErrorCode::OK;
}
ErrorCode ScopedSegmentAccess::GetClientSegments(
const UUID& client_id, std::vector<Segment>& segments) const {
auto it = segment_manager_->client_segments_.find(client_id);
if (it == segment_manager_->client_segments_.end()) {
return ErrorCode::SEGMENT_NOT_FOUND;
}
segments.clear();
for (auto& segment_id : it->second) {
auto segment_it = segment_manager_->mounted_segments_.find(segment_id);
if (segment_it != segment_manager_->mounted_segments_.end()) {
segments.emplace_back(segment_it->second.segment);
}
}
return ErrorCode::OK;
}
ErrorCode ScopedSegmentAccess::GetAllSegments(
std::vector<std::string>& all_segments) {
all_segments.clear();
for (auto& segment : segment_manager_->mounted_segments_) {
if (segment.second.status == SegmentStatus::OK) {
all_segments.push_back(segment.second.segment.name);
}
}
return ErrorCode::OK;
}
ErrorCode ScopedSegmentAccess::QuerySegments(const std::string& segment,
size_t& used, size_t& capacity) {
const auto& allocators =
segment_manager_->allocators_by_name_.find(segment);
if (allocators != segment_manager_->allocators_by_name_.end()) {
// Allocators Only contains the segments with OK status, so just return
// the first one.
capacity = allocators->second[0]->capacity();
used = allocators->second[0]->size();
} else {
VLOG(1) << "### DEBUG ### MasterService::QuerySegments(" << segment
<< ") not found!";
return ErrorCode::SEGMENT_NOT_FOUND;
}
return ErrorCode::OK;
}
} // namespace mooncake

18
mooncake-store/src/types.cpp
View File

@ -1,5 +1,9 @@
#include "types.h"
#include <boost/uuid/uuid.hpp>
#include <boost/uuid/uuid_generators.hpp>
#include <boost/uuid/uuid_io.hpp>
namespace mooncake {
const std::string& toString(ErrorCode errorCode) noexcept {
@ -8,7 +12,8 @@ const std::string& toString(ErrorCode errorCode) noexcept {
{ErrorCode::INTERNAL_ERROR, "INTERNAL_ERROR"},
{ErrorCode::BUFFER_OVERFLOW, "BUFFER_OVERFLOW"},
{ErrorCode::SHARD_INDEX_OUT_OF_RANGE, "SHARD_INDEX_OUT_OF_RANGE"},
{ErrorCode::AVAILABLE_SEGMENT_EMPTY, "AVAILABLE_SEGMENT_EMPTY"},
{ErrorCode::SEGMENT_NOT_FOUND, "SEGMENT_NOT_FOUND"},
{ErrorCode::SEGMENT_ALREADY_EXISTS, "SEGMENT_ALREADY_EXISTS"},
{ErrorCode::NO_AVAILABLE_HANDLE, "NO_AVAILABLE_HANDLE"},
{ErrorCode::INVALID_VERSION, "INVALID_VERSION"},
{ErrorCode::INVALID_KEY, "INVALID_KEY"},
@ -27,6 +32,8 @@ const std::string& toString(ErrorCode errorCode) noexcept {
{ErrorCode::ETCD_KEY_NOT_EXIST, "ETCD_KEY_NOT_EXIST"},
{ErrorCode::ETCD_TRANSACTION_FAIL, "ETCD_TRANSACTION_FAIL"},
{ErrorCode::ETCD_CTX_CANCELLED, "ETCD_CTX_CANCELLED"},
{ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS, "UNAVAILABLE_IN_CURRENT_STATUS"},
{ErrorCode::UNAVAILABLE_IN_CURRENT_MODE, "UNAVAILABLE_IN_CURRENT_MODE"},
};
auto it = errorCodeMap.find(errorCode);
@ -42,4 +49,13 @@ ErrorCode fromInt(int32_t errorCode) noexcept {
return static_cast<ErrorCode>(errorCode);
}
UUID generate_uuid() {
UUID pair_uuid;
boost::uuids::random_generator gen;
boost::uuids::uuid uuid = gen();
std::memcpy(&pair_uuid.first, uuid.data, sizeof(uint64_t));
std::memcpy(&pair_uuid.second, uuid.data + sizeof(uint64_t), sizeof(uint64_t));
return pair_uuid;
}
} // namespace mooncake

15
mooncake-store/tests/CMakeLists.txt
View File

@ -54,4 +54,17 @@ target_link_libraries(transfer_task_test PUBLIC
gtest_main
pthread
)
add_test(NAME transfer_task_test COMMAND transfer_task_test)
add_test(NAME transfer_task_test COMMAND transfer_task_test)
add_executable(segment_test segment_test.cpp)
target_link_libraries(segment_test PUBLIC
mooncake_store
cachelib_memory_allocator
glog
gtest
gtest_main
pthread
)
add_test(NAME segment_test COMMAND segment_test)
add_subdirectory(e2e)

115
mooncake-store/tests/allocation_strategy_test.cpp
View File

@ -31,21 +31,23 @@ class AllocationStrategyTest : public ::testing::Test {
// Test basic functionality with empty allocators map
TEST_F(AllocationStrategyTest, EmptyAllocatorsMap) {
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
empty_allocators;
std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>
empty_allocators_by_name;
std::vector<std::shared_ptr<BufferAllocator>> empty_allocators;
ReplicateConfig config{1, "local"};
auto result = strategy_->Allocate(empty_allocators, 100, config);
auto result = strategy_->Allocate(empty_allocators, empty_allocators_by_name, 100, config);
EXPECT_EQ(result, nullptr);
}
// Test preferred segment behavior with empty allocators
TEST_F(AllocationStrategyTest, PreferredSegmentWithEmptyAllocators) {
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
allocators;
std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>
empty_allocators_by_name;
std::vector<std::shared_ptr<BufferAllocator>> empty_allocators;
ReplicateConfig config{1, "preferred_segment"};
auto result = strategy_->Allocate(allocators, 100, config);
auto result = strategy_->Allocate(empty_allocators, empty_allocators_by_name, 100, config);
EXPECT_EQ(result, nullptr); // Should return nullptr for empty allocators
}
@ -54,15 +56,19 @@ TEST_F(AllocationStrategyTest, PreferredSegmentAllocation) {
auto allocator1 = CreateTestAllocator("segment1", 0);
auto allocator2 = CreateTestAllocator("preferred", 0x10000000ULL);
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
allocators;
allocators["segment1"] = allocator1;
allocators["preferred"] = allocator2;
std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>
allocators_by_name;
std::vector<std::shared_ptr<BufferAllocator>> allocators;
allocators_by_name["segment1"].push_back(allocator1);
allocators_by_name["preferred"].push_back(allocator2);
allocators.push_back(allocator1);
allocators.push_back(allocator2);
ReplicateConfig config{1, "preferred"};
size_t alloc_size = 1024;
auto result = strategy_->Allocate(allocators, alloc_size, config);
auto result = strategy_->Allocate(allocators, allocators_by_name, alloc_size, config);
ASSERT_NE(result, nullptr);
EXPECT_EQ(result->get_descriptor().segment_name_, "preferred");
EXPECT_EQ(result->get_descriptor().size_, alloc_size);
@ -73,15 +79,19 @@ TEST_F(AllocationStrategyTest, PreferredSegmentNotFound) {
auto allocator1 = CreateTestAllocator("segment1", 0);
auto allocator2 = CreateTestAllocator("segment2", 0x10000000ULL);
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
allocators;
allocators["segment1"] = allocator1;
allocators["segment2"] = allocator2;
std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>
allocators_by_name;
std::vector<std::shared_ptr<BufferAllocator>> allocators;
allocators_by_name["segment1"].push_back(allocator1);
allocators_by_name["segment2"].push_back(allocator2);
allocators.push_back(allocator1);
allocators.push_back(allocator2);
ReplicateConfig config{1, "nonexistent"};
size_t alloc_size = 1024;
auto result = strategy_->Allocate(allocators, alloc_size, config);
auto result = strategy_->Allocate(allocators, allocators_by_name, alloc_size, config);
ASSERT_NE(result, nullptr);
// Should allocate from one of the available segments
std::string segment_name = result->get_descriptor().segment_name_;
@ -95,11 +105,16 @@ TEST_F(AllocationStrategyTest, MultipleAllocatorsRandomSelection) {
auto allocator2 = CreateTestAllocator("segment2", 0x10000000ULL);
auto allocator3 = CreateTestAllocator("segment3", 0x20000000ULL);
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
allocators;
allocators["segment1"] = allocator1;
allocators["segment2"] = allocator2;
allocators["segment3"] = allocator3;
std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>
allocators_by_name;
std::vector<std::shared_ptr<BufferAllocator>> allocators;
allocators_by_name["segment1"].push_back(allocator1);
allocators_by_name["segment2"].push_back(allocator2);
allocators_by_name["segment3"].push_back(allocator3);
allocators.push_back(allocator1);
allocators.push_back(allocator2);
allocators.push_back(allocator3);
ReplicateConfig config{1, ""}; // No preferred segment
size_t alloc_size = 1024;
@ -107,7 +122,7 @@ TEST_F(AllocationStrategyTest, MultipleAllocatorsRandomSelection) {
// Perform multiple allocations to test randomness
std::vector<std::string> allocated_segments;
for (int i = 0; i < 10; ++i) {
auto result = strategy_->Allocate(allocators, alloc_size, config);
auto result = strategy_->Allocate(allocators, allocators_by_name, alloc_size, config);
ASSERT_NE(result, nullptr);
allocated_segments.push_back(result->get_descriptor().segment_name_);
EXPECT_EQ(result->get_descriptor().size_, alloc_size);
@ -125,10 +140,14 @@ TEST_F(AllocationStrategyTest, PreferredSegmentInsufficientSpace) {
auto allocator1 = CreateTestAllocator("segment1", 0);
auto allocator2 = CreateTestAllocator("preferred", 0x10000000ULL);
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
allocators;
allocators["segment1"] = allocator1;
allocators["preferred"] = allocator2;
std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>
allocators_by_name;
std::vector<std::shared_ptr<BufferAllocator>> allocators;
allocators_by_name["segment1"].push_back(allocator1);
allocators_by_name["preferred"].push_back(allocator2);
allocators.push_back(allocator1);
allocators.push_back(allocator2);
// First, fill up the preferred allocator
ReplicateConfig config{1, "preferred"};
@ -136,14 +155,14 @@ TEST_F(AllocationStrategyTest, PreferredSegmentInsufficientSpace) {
// Allocate most of the space in preferred segment
size_t large_alloc = 15 * 1024 * 1024; // 15MB out of 16MB
auto large_buffer = strategy_->Allocate(allocators, large_alloc, config);
auto large_buffer = strategy_->Allocate(allocators, allocators_by_name, large_alloc, config);
ASSERT_NE(large_buffer, nullptr);
EXPECT_EQ(large_buffer->get_descriptor().segment_name_, "preferred");
buffers.push_back(std::move(large_buffer));
// Now try to allocate more than remaining space in preferred segment
size_t small_alloc = 2 * 1024 * 1024; // 2MB (more than remaining ~1MB)
auto result = strategy_->Allocate(allocators, small_alloc, config);
auto result = strategy_->Allocate(allocators, allocators_by_name, small_alloc, config);
ASSERT_NE(result, nullptr);
// Should fall back to segment1 since preferred doesn't have enough space
EXPECT_EQ(result->get_descriptor().segment_name_, "segment1");
@ -155,18 +174,22 @@ TEST_F(AllocationStrategyTest, AllAllocatorsFull) {
auto allocator1 = CreateTestAllocator("segment1", 0);
auto allocator2 = CreateTestAllocator("segment2", 0x10000000ULL);
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
allocators;
allocators["segment1"] = allocator1;
allocators["segment2"] = allocator2;
std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>
allocators_by_name;
std::vector<std::shared_ptr<BufferAllocator>> allocators;
allocators_by_name["segment1"].push_back(allocator1);
allocators_by_name["segment2"].push_back(allocator2);
allocators.push_back(allocator1);
allocators.push_back(allocator2);
ReplicateConfig config{1, ""};
std::vector<std::unique_ptr<AllocatedBuffer>> buffers;
// Fill up both allocators
size_t large_alloc = 15 * 1024 * 1024; // 15MB each
auto buffer1 = strategy_->Allocate(allocators, large_alloc, config);
auto buffer2 = strategy_->Allocate(allocators, large_alloc, config);
auto buffer1 = strategy_->Allocate(allocators, allocators_by_name, large_alloc, config);
auto buffer2 = strategy_->Allocate(allocators, allocators_by_name, large_alloc, config);
ASSERT_NE(buffer1, nullptr);
ASSERT_NE(buffer2, nullptr);
buffers.push_back(std::move(buffer1));
@ -174,20 +197,23 @@ TEST_F(AllocationStrategyTest, AllAllocatorsFull) {
// Try to allocate more than remaining space
size_t impossible_alloc = 5 * 1024 * 1024; // 5MB (more than remaining)
auto result = strategy_->Allocate(allocators, impossible_alloc, config);
auto result = strategy_->Allocate(allocators, allocators_by_name, impossible_alloc, config);
EXPECT_EQ(result, nullptr); // Should fail
}
// Test allocation with zero size
TEST_F(AllocationStrategyTest, ZeroSizeAllocation) {
auto allocator = CreateTestAllocator("segment1");
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
allocators;
allocators["segment1"] = allocator;
std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>
allocators_by_name;
std::vector<std::shared_ptr<BufferAllocator>> allocators;
allocators_by_name["segment1"].push_back(allocator);
allocators.push_back(allocator);
ReplicateConfig config{1, ""};
auto result = strategy_->Allocate(allocators, 0, config);
auto result = strategy_->Allocate(allocators, allocators_by_name, 0, config);
// Zero-size allocation behavior depends on BufferAllocator implementation
// This test documents the current behavior
if (result != nullptr) {
@ -198,14 +224,17 @@ TEST_F(AllocationStrategyTest, ZeroSizeAllocation) {
// Test allocation with very large size
TEST_F(AllocationStrategyTest, VeryLargeSizeAllocation) {
auto allocator = CreateTestAllocator("segment1");
std::unordered_map<std::string, std::shared_ptr<BufferAllocator>>
allocators;
allocators["segment1"] = allocator;
std::unordered_map<std::string, std::vector<std::shared_ptr<BufferAllocator>>>
allocators_by_name;
std::vector<std::shared_ptr<BufferAllocator>> allocators;
allocators_by_name["segment1"].push_back(allocator);
allocators.push_back(allocator);
ReplicateConfig config{1, ""};
size_t huge_size = 100 * 1024 * 1024; // 100MB (larger than 16MB capacity)
auto result = strategy_->Allocate(allocators, huge_size, config);
auto result = strategy_->Allocate(allocators, allocators_by_name, huge_size, config);
EXPECT_EQ(result, nullptr); // Should fail due to insufficient capacity
}

30
mooncake-store/tests/client_integration_test.cpp
View File

@ -72,11 +72,11 @@ class ClientIntegrationTest : public ::testing::Test {
}
static void InitializeSegment() {
const size_t ram_buffer_size = 512 * 1024 * 1024; // 512 MB
segment_ptr_ = allocate_buffer_allocator_memory(ram_buffer_size);
ram_buffer_size_ = 512 * 1024 * 1024; // 512 MB
segment_ptr_ = allocate_buffer_allocator_memory(ram_buffer_size_);
LOG_ASSERT(segment_ptr_);
ErrorCode rc = segment_provider_client_->MountSegment(
"localhost:17812", segment_ptr_, ram_buffer_size);
ErrorCode rc = segment_provider_client_->MountSegment(segment_ptr_,
ram_buffer_size_);
if (rc != ErrorCode::OK) {
LOG(ERROR) << "Failed to mount segment: " << toString(rc);
}
@ -103,13 +103,12 @@ class ClientIntegrationTest : public ::testing::Test {
}
// Mount segment for test_client_ as well
const size_t test_client_ram_buffer_size = 512 * 1024 * 1024; // 512 MB
test_client_ram_buffer_size_ = 512 * 1024 * 1024; // 512 MB
test_client_segment_ptr_ =
allocate_buffer_allocator_memory(test_client_ram_buffer_size);
allocate_buffer_allocator_memory(test_client_ram_buffer_size_);
LOG_ASSERT(test_client_segment_ptr_);
ErrorCode rc = test_client_->MountSegment("localhost:17813",
test_client_segment_ptr_,
test_client_ram_buffer_size);
ErrorCode rc = test_client_->MountSegment(test_client_segment_ptr_,
test_client_ram_buffer_size_);
if (rc != ErrorCode::OK) {
LOG(ERROR) << "Failed to mount segment for test_client_: "
<< toString(rc);
@ -120,8 +119,8 @@ class ClientIntegrationTest : public ::testing::Test {
static void CleanupClients() {
// Unmount test client segment first
if (test_client_ && test_client_segment_ptr_) {
if (test_client_->UnmountSegment("localhost:17813",
test_client_segment_ptr_) !=
if (test_client_->UnmountSegment(test_client_segment_ptr_,
test_client_ram_buffer_size_) !=
ErrorCode::OK) {
LOG(ERROR) << "Failed to unmount test client segment";
}
@ -136,8 +135,9 @@ class ClientIntegrationTest : public ::testing::Test {
}
static void CleanupSegment() {
if (segment_provider_client_->UnmountSegment(
"localhost:17812", segment_ptr_) != ErrorCode::OK) {
if (segment_provider_client_->UnmountSegment(segment_ptr_,
ram_buffer_size_) !=
ErrorCode::OK) {
LOG(ERROR) << "Failed to unmount segment";
}
}
@ -149,7 +149,9 @@ class ClientIntegrationTest : public ::testing::Test {
// themselves.
static std::unique_ptr<SimpleAllocator> client_buffer_allocator_;
static void* segment_ptr_;
static size_t ram_buffer_size_;
static void* test_client_segment_ptr_;
static size_t test_client_ram_buffer_size_;
};
// Static members initialization
@ -160,6 +162,8 @@ void* ClientIntegrationTest::segment_ptr_ = nullptr;
void* ClientIntegrationTest::test_client_segment_ptr_ = nullptr;
std::unique_ptr<SimpleAllocator>
ClientIntegrationTest::client_buffer_allocator_ = nullptr;
size_t ClientIntegrationTest::ram_buffer_size_ = 0;
size_t ClientIntegrationTest::test_client_ram_buffer_size_ = 0;
// Test basic Put/Get operations through the client
TEST_F(ClientIntegrationTest, BasicPutGetOperations) {

7
mooncake-store/tests/e2e/CMakeLists.txt Normal file
View File

@ -0,0 +1,7 @@
add_executable(clientctl clientctl.cpp)
target_link_libraries(clientctl PUBLIC
mooncake_store
cachelib_memory_allocator
glog
pthread
)

12
mooncake-store/tests/e2e/cases/case1.txt Normal file
View File

@ -0,0 +1,12 @@
# Test case for None HA mode
create c1 9888
mount c1 s1 1073741824
put c1 key1 val1
put c1 key2 val2
create c2 9889
mount c2 s2 1073741824
get c2 key2
remove c1
put c2 key1 val1b
get c2 key1
remove c2

18
mooncake-store/tests/e2e/cases/case2.txt Normal file
View File

@ -0,0 +1,18 @@
# Test case for HA mode
# master cmd: ./mooncake-store/src/mooncake_master --enable-ha=true --local-hostname=0.0.0.0:50051 --etcd-endpoints=0.0.0.0:2379 --client-ttl=3
# clientctl cmd: ./mooncake-store/tests/clientctl --master_server_entry=etcd://0.0.0.0:2379
create c1 9888
mount c1 s1 1073741824
put c1 key1 val1
create c2 9889
mount c2 s2 1073741824
put c1 key2 val2
get c2 key1
get c1 key2
# For quick test, the master should set client ttl to 3 sec
sleep 10
# The client should not expired after the sleep
get c1 key1
get c2 key2
terminate
# The client should expired and segments are auto-unmounted

295
mooncake-store/tests/e2e/clientctl.cpp Normal file
View File

@ -0,0 +1,295 @@
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <iostream>
#include <memory>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>
#include "client.h"
#include "types.h"
#include "utils.h"
// Command line flags
DEFINE_string(metadata_connstring, "http://127.0.0.1:8080/metadata",
"Metadata connection string for transfer engine");
DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp");
DEFINE_string(device_name, "ibp6s0",
"Device name to use, valid if protocol=rdma");
DEFINE_string(master_server_entry, "localhost:50051", "Master server address");
namespace mooncake {
namespace testing {
struct SegmentInfo {
void* base;
size_t size;
};
struct ClientInfo {
std::shared_ptr<Client> client;
std::unordered_map<std::string, SegmentInfo> segments;
std::string hostname;
~ClientInfo() {
for (auto& [name, segment] : segments) {
free(segment.base);
}
}
};
class ClientCtl {
public:
void Run() {
std::string line;
while (std::getline(std::cin, line)) {
std::istringstream iss(line);
std::string cmd;
iss >> cmd;
if (cmd == "create") {
HandleCreate(iss);
} else if (cmd == "put") {
HandlePut(iss);
} else if (cmd == "get") {
HandleGet(iss);
} else if (cmd == "mount") {
HandleMount(iss);
} else if (cmd == "remove") {
HandleRemove(iss);
} else if (cmd == "sleep") {
HandleSleep(iss);
} else if (cmd[0] == '#') {
// Ignore comment lines
continue;
} else if (cmd == "terminate") {
std::exit(0);
} else {
std::cout << "Unknown command: " << cmd << std::endl;
}
}
}
private:
void HandleCreate(std::istringstream& iss) {
std::string name;
std::string port;
iss >> name >> port;
if (name.empty() || port.empty()) {
std::cout << "Invalid create command format. Expected: create "
"[name] [port]"
<< std::endl;
return;
}
void** args =
(FLAGS_protocol == "rdma") ? rdma_args(FLAGS_device_name) : nullptr;
std::string hostname = "localhost:" + port;
auto client_opt =
Client::Create(hostname, // Local hostname
FLAGS_metadata_connstring, FLAGS_protocol, args,
FLAGS_master_server_entry);
if (!client_opt.has_value()) {
std::cout << "Failed to create client: " << name << std::endl;
return;
}
clients_[name] = ClientInfo{client_opt.value(), {}, hostname};
std::cout << "Successfully created client: " << name << std::endl;
}
void HandlePut(std::istringstream& iss) {
std::string name, key, value;
iss >> name >> key >> value;
auto it = clients_.find(name);
if (it == clients_.end()) {
std::cout << "Client not found: " << name << std::endl;
return;
}
// Allocate buffer for the value
void* buffer = malloc(value.size());
if (!buffer) {
std::cout << "Failed to allocate memory for value" << std::endl;
return;
}
// Copy value to buffer
memcpy(buffer, value.data(), value.size());
// Create slices
std::vector<Slice> slices;
slices.emplace_back(Slice{buffer, value.size()});
// Configure replication
ReplicateConfig config;
config.replica_num = 1;
// Perform put operation
ErrorCode error_code = it->second.client->Put(key, slices, config);
// Free the buffer
free(buffer);
if (error_code != ErrorCode::OK) {
std::cout << "Failed to put value: " << toString(error_code)
<< std::endl;
return;
}
std::cout << "Successfully put value for key: " << key << std::endl;
}
void HandleGet(std::istringstream& iss) {
std::string name, key;
iss >> name >> key;
auto it = clients_.find(name);
if (it == clients_.end()) {
std::cout << "Client not found: " << name << std::endl;
return;
}
Client::ObjectInfo object_info;
if (it->second.client->Query(key, object_info) != ErrorCode::OK) {
std::cout << "Key not found: " << key << std::endl;
return;
}
// Create slices
std::vector<AllocatedBuffer::Descriptor>& descriptors =
object_info.replica_list[0].buffer_descriptors;
std::vector<Slice> slices(descriptors.size());
for (size_t i = 0; i < descriptors.size(); i++) {
void* buffer = malloc(descriptors[i].size_);
slices[i] = Slice{buffer, descriptors[i].size_};
}
auto free_slices = [&]() {
for (auto& slice : slices) {
free(slice.ptr);
}
};
// Perform get operation
ErrorCode error_code = it->second.client->Get(key, object_info, slices);
if (error_code != ErrorCode::OK) {
free_slices();
std::cout << "Failed to get value: " << toString(error_code)
<< std::endl;
return;
}
// Print the value
std::string value;
for (const auto& slice : slices) {
value.append(static_cast<const char*>(slice.ptr), slice.size);
}
std::cout << "Get value: " << value << std::endl;
// Free the buffer
free_slices();
}
void HandleMount(std::istringstream& iss) {
std::string client_name;
std::string segment_name;
size_t size;
iss >> client_name >> segment_name >> size;
if (segment_name.empty() || client_name.empty() || size == 0) {
std::cout << "Invalid mount command format. Expected: mount "
"[client_name] [segment_name] [size]"
<< std::endl;
return;
}
auto it = clients_.find(client_name);
if (it == clients_.end()) {
std::cout << "Client not found: " << client_name << std::endl;
return;
}
if (it->second.segments.find(segment_name) !=
it->second.segments.end()) {
std::cout << "Segment " << segment_name << " already mounted"
<< std::endl;
return;
}
void* buffer;
buffer = allocate_buffer_allocator_memory(size);
if (!buffer) {
std::cout << "Failed to allocate memory for segment" << std::endl;
return;
}
ErrorCode error_code = it->second.client->MountSegment(buffer, size);
if (error_code != ErrorCode::OK) {
std::cout << "Failed to mount segment: " << toString(error_code)
<< std::endl;
free(buffer);
return;
}
SegmentInfo segment_info{buffer, size};
it->second.segments[segment_name] = segment_info;
std::cout << "Successfully mounted segment on client " << client_name
<< std::endl;
}
void HandleRemove(std::istringstream& iss) {
std::string name;
iss >> name;
auto it = clients_.find(name);
if (it == clients_.end()) {
std::cout << "Client not found: " << name << std::endl;
return;
}
clients_.erase(it);
std::cout << "Successfully removed client: " << name << std::endl;
}
void HandleSleep(std::istringstream& iss) {
int seconds;
iss >> seconds;
if (seconds <= 0) {
std::cout << "Invalid sleep command format. Expected: sleep [seconds]"
<< std::endl;
return;
}
std::this_thread::sleep_for(std::chrono::seconds(seconds));
std::cout << "Slept for " << seconds << " seconds"
<< std::endl;
}
std::unordered_map<std::string, ClientInfo> clients_;
};
} // namespace testing
} // namespace mooncake
int main(int argc, char** argv) {
// Initialize Google's flags library
gflags::ParseCommandLineFlags(&argc, &argv, true);
// Initialize Google logging
google::InitGoogleLogging(argv[0]);
FLAGS_logtostderr = 1;
mooncake::testing::ClientCtl ctl;
ctl.Run();
return 0;
}

12
mooncake-store/tests/master_metrics_test.cpp
View File

@ -73,6 +73,14 @@ TEST_F(MasterMetricsTest, BasicRequestTest) {
constexpr size_t kBufferAddress = 0x300000000;
constexpr size_t kSegmentSize = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
UUID segment_id = generate_uuid();
Segment segment;
segment.id = segment_id;
segment.name = segment_name;
segment.base = kBufferAddress;
segment.size = kSegmentSize;
UUID client_id = generate_uuid();
std::string key = "test_key";
uint64_t value_length = 1024;
std::vector<uint64_t> slice_lengths = {value_length};
@ -81,7 +89,7 @@ TEST_F(MasterMetricsTest, BasicRequestTest) {
// Test MountSegment request
ASSERT_EQ(ErrorCode::OK,
service_.MountSegment(kBufferAddress, kSegmentSize, segment_name).error_code);
service_.MountSegment(segment, client_id).error_code);
ASSERT_EQ(metrics.get_allocated_size(), 0);
ASSERT_EQ(metrics.get_total_capacity(), kSegmentSize);
ASSERT_DOUBLE_EQ(metrics.get_global_used_ratio(), 0.0);
@ -147,7 +155,7 @@ TEST_F(MasterMetricsTest, BasicRequestTest) {
ASSERT_EQ(ErrorCode::OK,
service_.PutStart(key, value_length, slice_lengths, config).error_code);
ASSERT_EQ(ErrorCode::OK, service_.PutEnd(key).error_code);
ASSERT_EQ(ErrorCode::OK, service_.UnmountSegment(segment_name).error_code);
ASSERT_EQ(ErrorCode::OK, service_.UnmountSegment(segment_id, client_id).error_code);
ASSERT_EQ(metrics.get_unmount_segment_requests(), 1);
ASSERT_EQ(metrics.get_unmount_segment_failures(), 0);
ASSERT_EQ(metrics.get_key_count(), 0);

206
mooncake-store/tests/master_service_test.cpp
View File

@ -67,43 +67,57 @@ TEST_F(MasterServiceTest, MountUnmountSegment) {
constexpr size_t kSegmentSize = 1024 * 1024 * 16;
// Define the name of the test segment.
std::string segment_name = "test_segment";
Segment segment(generate_uuid(), segment_name, kBufferAddress, kSegmentSize);
UUID client_id = generate_uuid();
// Test invalid parameters.
// Invalid buffer address (0).
segment.base = 0;
segment.size = kSegmentSize;
EXPECT_EQ(ErrorCode::INVALID_PARAMS,
service_->MountSegment(0, kSegmentSize, segment_name));
service_->MountSegment(segment, client_id));
// Invalid segment size (0).
segment.base = kBufferAddress;
segment.size = 0;
EXPECT_EQ(ErrorCode::INVALID_PARAMS,
service_->MountSegment(kBufferAddress, 0, segment_name));
service_->MountSegment(segment, client_id));
// Base is not aligned
EXPECT_EQ(
ErrorCode::INVALID_PARAMS,
service_->MountSegment(kBufferAddress + 1, kSegmentSize, segment_name));
segment.base = kBufferAddress + 1;
segment.size = kSegmentSize;
EXPECT_EQ(ErrorCode::INVALID_PARAMS,
service_->MountSegment(segment, client_id));
// Size is not aligned
EXPECT_EQ(
ErrorCode::INVALID_PARAMS,
service_->MountSegment(kBufferAddress, kSegmentSize + 1, segment_name));
segment.base = kBufferAddress;
segment.size = kSegmentSize + 1;
EXPECT_EQ(ErrorCode::INVALID_PARAMS,
service_->MountSegment(segment, client_id));
// Test normal mount operation.
EXPECT_EQ(ErrorCode::OK, service_->MountSegment(
kBufferAddress, kSegmentSize, segment_name));
segment.base = kBufferAddress;
segment.size = kSegmentSize;
EXPECT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
// Test mounting the same segment again (should fail).
EXPECT_EQ(
ErrorCode::INVALID_PARAMS,
service_->MountSegment(kBufferAddress, kSegmentSize, segment_name));
// Test mounting the same segment again (idempotent request should succeed).
EXPECT_EQ(ErrorCode::OK,
service_->MountSegment(segment, client_id));
// Test unmounting the segment.
EXPECT_EQ(ErrorCode::OK, service_->UnmountSegment(segment_name));
EXPECT_EQ(ErrorCode::OK, service_->UnmountSegment(segment.id, client_id));
// Test unmounting a non-existent segment (should fail).
EXPECT_EQ(ErrorCode::INVALID_PARAMS,
service_->UnmountSegment("non_existent"));
// Test unmounting the same segment again (idempotent request should succeed).
EXPECT_EQ(ErrorCode::OK, service_->UnmountSegment(segment.id, client_id));
// Test unmounting a non-existent segment (idempotent request should succeed).
UUID non_existent_id = generate_uuid();
EXPECT_EQ(ErrorCode::OK,
service_->UnmountSegment(non_existent_id, client_id));
// Test remounting after unmount.
EXPECT_EQ(ErrorCode::OK, service_->MountSegment(
kBufferAddress, kSegmentSize, segment_name));
EXPECT_EQ(ErrorCode::OK, service_->UnmountSegment(segment_name));
EXPECT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
EXPECT_EQ(ErrorCode::OK, service_->UnmountSegment(segment.id, client_id));
}
TEST_F(MasterServiceTest, RandomMountUnmountSegment) {
@ -113,6 +127,8 @@ TEST_F(MasterServiceTest, RandomMountUnmountSegment) {
constexpr size_t kBufferAddress = 0x300000000;
// Define the name of the test segment.
std::string segment_name = "test_random_segment";
UUID segment_id = generate_uuid();
UUID client_id = generate_uuid();
size_t times = 10;
std::random_device rd;
std::mt19937 gen(rd());
@ -121,11 +137,12 @@ TEST_F(MasterServiceTest, RandomMountUnmountSegment) {
int random_number = dis(gen);
// Define the size of the segment (16MB).
size_t kSegmentSize = 1024 * 1024 * 16 * random_number;
Segment segment(segment_id, segment_name, kBufferAddress, kSegmentSize);
// Test remounting after unmount.
EXPECT_EQ(
ErrorCode::OK,
service_->MountSegment(kBufferAddress, kSegmentSize, segment_name));
EXPECT_EQ(ErrorCode::OK, service_->UnmountSegment(segment_name));
EXPECT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
EXPECT_EQ(ErrorCode::OK, service_->UnmountSegment(segment.id, client_id));
}
}
@ -139,15 +156,18 @@ TEST_F(MasterServiceTest, ConcurrentMountUnmount) {
// Launch multiple threads to mount/unmount segments concurrently
for (size_t i = 0; i < num_threads; i++) {
threads.emplace_back([&service_, i, &success_count]() {
std::string segment_name = "segment_" + std::to_string(i);
size_t buffer = 0x300000000 + i * 0x10000000;
constexpr size_t size = 16 * 1024 * 1024;
Segment segment;
segment.name = "segment_" + std::to_string(i);
segment.id = generate_uuid();
segment.base = 0x300000000 + i * 0x10000000;
segment.size = 16 * 1024 * 1024;
UUID client_id = generate_uuid();
for (size_t j = 0; j < iterations; j++) {
if (service_->MountSegment(buffer, size, segment_name) ==
if (service_->MountSegment(segment, client_id) ==
ErrorCode::OK) {
EXPECT_EQ(ErrorCode::OK,
service_->UnmountSegment(segment_name));
service_->UnmountSegment(segment.id, client_id));
success_count++;
}
}
@ -168,8 +188,11 @@ TEST_F(MasterServiceTest, PutStartInvalidParams) {
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
std::string key = "test_key";
ReplicateConfig config;
@ -197,8 +220,10 @@ TEST_F(MasterServiceTest, PutStartEndFlow) {
constexpr size_t size = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
// Test PutStart
std::string key = "test_key";
@ -233,8 +258,10 @@ TEST_F(MasterServiceTest, RandomPutStartEndFlow) {
constexpr size_t size = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
// Test PutStart
std::string key = "test_key";
@ -277,8 +304,11 @@ TEST_F(MasterServiceTest, GetReplicaList) {
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
std::string key = "test_key";
std::vector<uint64_t> slice_lengths = {1024};
@ -300,8 +330,11 @@ TEST_F(MasterServiceTest, RemoveObject) {
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
std::string key = "test_key";
std::vector<uint64_t> slice_lengths = {1024};
@ -330,8 +363,11 @@ TEST_F(MasterServiceTest, RandomRemoveObject) {
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
int times = 10;
std::random_device rd;
std::mt19937 gen(rd());
@ -363,8 +399,11 @@ TEST_F(MasterServiceTest, RemoveAll) {
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
int times = 10;
while (times--) {
std::string key = "test_key" + std::to_string(times);
@ -395,8 +434,11 @@ TEST_F(MasterServiceTest, MultiSliceMultiReplicaFlow) {
constexpr size_t segment_size =
1024 * 1024 * 64; // 64MB to accommodate multiple replicas
std::string segment_name = "test_segment_multi";
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, segment_size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, segment_size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
// Test parameters
std::string key = "multi_slice_object";
@ -483,8 +525,9 @@ TEST_F(MasterServiceTest, ConcurrentGarbageCollectionTest) {
constexpr size_t size =
1024 * 1024 * 256; // Larger segment for concurrent use
std::string segment_name = "concurrent_gc_segment";
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
constexpr size_t num_threads = 4;
constexpr size_t objects_per_thread = 25;
@ -555,6 +598,7 @@ TEST_F(MasterServiceTest, ConcurrentGarbageCollectionTest) {
// All objects should have been garbage collected
EXPECT_EQ(0, found_count);
}
TEST_F(MasterServiceTest, CleanupStaleHandlesTest) {
std::unique_ptr<MasterService> service_(new MasterService());
@ -563,9 +607,11 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) {
constexpr size_t size = 1024 * 1024 * 16; // 16MB
std::string segment_name = "test_segment";
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
// Mount the segment
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
// Create an object that will be stored in the segment
std::string key = "segment_object";
@ -585,7 +631,7 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) {
ASSERT_EQ(1, retrieved_replicas.size());
// Unmount the segment
ASSERT_EQ(ErrorCode::OK, service_->UnmountSegment(segment_name));
ASSERT_EQ(ErrorCode::OK, service_->UnmountSegment(segment.id, client_id));
// Try to get the object - it should be automatically removed since the
// replica is invalid
@ -595,8 +641,7 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) {
EXPECT_TRUE(retrieved_replicas.empty());
// Mount the segment again
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
// Create another object
std::string key2 = "another_segment_object";
@ -611,19 +656,20 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) {
service_->GetReplicaList(key2, retrieved_replicas));
// Unmount the segment
ASSERT_EQ(ErrorCode::OK, service_->UnmountSegment(segment_name));
ASSERT_EQ(ErrorCode::OK, service_->UnmountSegment(segment.id, client_id));
// Try to remove the object that should already be cleaned up
EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, service_->Remove(key2));
}
TEST_F(MasterServiceTest, ConcurrentWriteAndRemoveAll) {
std::unique_ptr<MasterService> service_(new MasterService());
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 256; // 256MB for concurrent testing
std::string segment_name = "concurrent_segment";
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
constexpr int num_threads = 4;
constexpr int objects_per_thread = 100;
@ -681,7 +727,6 @@ TEST_F(MasterServiceTest, ConcurrentWriteAndRemoveAll) {
ASSERT_EQ(total_removed, num_threads * objects_per_thread);
}
TEST_F(MasterServiceTest, ConcurrentReadAndRemoveAll) {
// set a large kv_lease_ttl so the granted lease will not quickly expire
const uint64_t kv_lease_ttl = 200;
@ -689,7 +734,9 @@ TEST_F(MasterServiceTest, ConcurrentReadAndRemoveAll) {
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 256; // 256MB for concurrent testing
std::string segment_name = "concurrent_segment";
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
// Pre-populate with test data
constexpr int num_objects = 1000;
@ -761,7 +808,9 @@ TEST_F(MasterServiceTest, ConcurrentRemoveAllOperations) {
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 16 * 100; // 256MB for concurrent testing
std::string segment_name = "concurrent_segment";
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(buffer, size, segment_name));
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment, client_id));
// Pre-populate with test data
constexpr int num_objects = 1000000;
@ -811,21 +860,22 @@ TEST_F(MasterServiceTest, UnmountSegmentImmediateCleanup) {
constexpr size_t buffer1 = 0x300000000;
constexpr size_t buffer2 = 0x400000000;
constexpr size_t size = 1024 * 1024 * 16;
std::string segment1 = "segment1";
std::string segment2 = "segment2";
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(buffer1, size, segment1));
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(buffer2, size, segment2));
Segment segment1(generate_uuid(), "segment1", buffer1, size);
Segment segment2(generate_uuid(), "segment2", buffer2, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment1, client_id));
ASSERT_EQ(ErrorCode::OK, service_->MountSegment(segment2, client_id));
// Create two objects in the two segments
std::string key1 = GenerateKeyForSegment(service_, segment1);
std::string key2 = GenerateKeyForSegment(service_, segment2);
std::string key1 = GenerateKeyForSegment(service_, segment1.name);
std::string key2 = GenerateKeyForSegment(service_, segment2.name);
std::vector<uint64_t> slice_lengths = {1024};
ReplicateConfig config;
config.replica_num = 1;
// Unmount segment1
ASSERT_EQ(ErrorCode::OK, service_->UnmountSegment(segment1));
ASSERT_EQ(ErrorCode::OK, service_->UnmountSegment(segment1.id, client_id));
// Umount will remove all objects in the segment, include the key1
ASSERT_EQ(1, service_->GetKeyCount());
// Verify objects in segment1 is gone
@ -843,7 +893,7 @@ TEST_F(MasterServiceTest, UnmountSegmentImmediateCleanup) {
ASSERT_EQ(ErrorCode::OK, service_->PutEnd(key1));
ASSERT_EQ(ErrorCode::OK,
service_->GetReplicaList(key1, retrieved));
ASSERT_EQ(replica_list[0].buffer_descriptors[0].segment_name_, segment2);
ASSERT_EQ(replica_list[0].buffer_descriptors[0].segment_name_, segment2.name);
}
TEST_F(MasterServiceTest, UnmountSegmentPerformance) {
@ -851,10 +901,12 @@ TEST_F(MasterServiceTest, UnmountSegmentPerformance) {
constexpr size_t kBufferAddress = 0x300000000;
constexpr size_t kSegmentSize = 1024 * 1024 * 256; // 256MB
std::string segment_name = "perf_test_segment";
Segment segment(generate_uuid(), segment_name, kBufferAddress, kSegmentSize);
UUID client_id = generate_uuid();
// Mount a segment for testing
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(kBufferAddress, kSegmentSize, segment_name));
service_->MountSegment(segment, client_id));
// Create 10000 keys for testing
constexpr int kNumKeys = 1000;
@ -873,7 +925,7 @@ TEST_F(MasterServiceTest, UnmountSegmentPerformance) {
// Execute unmount operation and record operation time
auto unmount_start = std::chrono::steady_clock::now();
EXPECT_EQ(ErrorCode::OK, service_->UnmountSegment(segment_name));
EXPECT_EQ(ErrorCode::OK, service_->UnmountSegment(segment.id, client_id));
auto unmount_end = std::chrono::steady_clock::now();
auto unmount_duration = std::chrono::duration_cast<std::chrono::milliseconds>(
@ -907,8 +959,10 @@ TEST_F(MasterServiceTest, RemoveLeasedObject) {
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
service_->MountSegment(segment, client_id));
std::string key = "test_key";
std::vector<uint64_t> slice_lengths = {1024};
@ -968,8 +1022,10 @@ TEST_F(MasterServiceTest, RemoveAllLeasedObject) {
constexpr size_t buffer = 0x300000000;
constexpr size_t size = 1024 * 1024 * 16;
std::string segment_name = "test_segment";
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
service_->MountSegment(segment, client_id));
for (int i = 0; i < 10; ++i) {
std::string key = "test_key" + std::to_string(i);
std::vector<uint64_t> slice_lengths = {1024};
@ -1009,8 +1065,10 @@ TEST_F(MasterServiceTest, EvictObject) {
constexpr size_t size = 1024 * 1024 * 16 * 15;
constexpr size_t object_size = 1024 * 15;
std::string segment_name = "test_segment";
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
service_->MountSegment(segment, client_id));
// Verify if we can put objects more than the segment can hold
int success_puts = 0;
@ -1042,8 +1100,10 @@ TEST_F(MasterServiceTest, TryEvictLeasedObject) {
constexpr size_t size = 1024 * 1024 * 16;
constexpr size_t object_size = 1024 * 1024;
std::string segment_name = "test_segment";
Segment segment(generate_uuid(), segment_name, buffer, size);
UUID client_id = generate_uuid();
ASSERT_EQ(ErrorCode::OK,
service_->MountSegment(buffer, size, segment_name));
service_->MountSegment(segment, client_id));
// Verify leased object will not be evicted.
int success_puts = 0;

443
mooncake-store/tests/segment_test.cpp Normal file
View File

@ -0,0 +1,443 @@
#include "segment.h"
#include <glog/logging.h>
#include <gtest/gtest.h>
#include <boost/functional/hash.hpp>
namespace mooncake {
// Test fixture for Segment tests
class SegmentTest : public ::testing::Test {
protected:
void SetUp() override {
// Initialize glog for logging
google::InitGoogleLogging("EvictionStrategyTest");
FLAGS_logtostderr = 1; // Output logs to stderr
}
void TearDown() override {
// Cleanup glog
google::ShutdownGoogleLogging();
}
void ValidateMountedSegments(const SegmentManager& segment_manager,
const std::vector<Segment>& segments,
const std::vector<UUID>& client_ids) {
// validate client_segments_ and mounted_segments_
size_t total_num = 0;
for (const auto& it : segment_manager.client_segments_) {
total_num += it.second.size();
}
ASSERT_EQ(total_num, segments.size());
ASSERT_EQ(segment_manager.mounted_segments_.size(), segments.size());
for (size_t i = 0; i < client_ids.size(); i++) {
auto client_it =
segment_manager.client_segments_.find(client_ids[i]);
ASSERT_NE(client_it, segment_manager.client_segments_.end());
auto segment_it =
std::find(client_it->second.begin(), client_it->second.end(),
segments[i].id);
ASSERT_NE(segment_it, client_it->second.end());
ASSERT_EQ(*segment_it, segments[i].id);
ASSERT_NE(segment_manager.mounted_segments_.find(segments[i].id),
segment_manager.mounted_segments_.end());
MountedSegment seg =
segment_manager.mounted_segments_.at(segments[i].id);
ASSERT_EQ(seg.segment.id, segments[i].id);
ASSERT_EQ(seg.segment.name, segments[i].name);
ASSERT_EQ(seg.segment.size, segments[i].size);
ASSERT_EQ(seg.segment.base, segments[i].base);
ASSERT_EQ(seg.status, SegmentStatus::OK);
ASSERT_EQ(seg.buf_allocator->getSegmentName(), segments[i].name);
ASSERT_EQ(seg.buf_allocator->capacity(), segments[i].size);
}
// validate allocators and allocators_by_name
total_num = 0;
for (const auto& it : segment_manager.allocators_by_name_) {
total_num += it.second.size();
}
ASSERT_EQ(total_num, segments.size());
ASSERT_EQ(segment_manager.allocators_.size(), segments.size());
for (const auto& segment : segments) {
MountedSegment mounted_segment =
segment_manager.mounted_segments_.at(segment.id);
auto allocator = mounted_segment.buf_allocator;
// validate allocators_
ASSERT_NE(std::find(segment_manager.allocators_.begin(),
segment_manager.allocators_.end(),
mounted_segment.buf_allocator),
segment_manager.allocators_.end());
// validate allocators_by_name
auto map_it =
segment_manager.allocators_by_name_.find(segment.name);
ASSERT_NE(map_it, segment_manager.allocators_by_name_.end());
auto name_allocator_it = map_it->second.begin();
for (; name_allocator_it != map_it->second.end();
name_allocator_it++) {
if (*name_allocator_it == allocator) {
break;
}
}
ASSERT_NE(name_allocator_it, map_it->second.end());
}
}
void ValidateMountedSegment(const SegmentManager& segment_manager,
const Segment segment, const UUID& client_id) {
std::vector<Segment> segments;
segments.push_back(segment);
std::vector<UUID> client_ids;
client_ids.push_back(client_id);
ValidateMountedSegments(segment_manager, segments, client_ids);
}
};
// Mount Segment Operations Tests:
TEST_F(SegmentTest, MountSegmentSuccess) {
SegmentManager segment_manager;
// Create a valid segment and client ID
Segment segment;
segment.id = generate_uuid();
segment.name = "test_segment";
segment.size = 1024 * 1024 * 16;
segment.base = 0x100000000;
UUID client_id = generate_uuid();
// Get segment access and attempt to mount
auto segment_access = segment_manager.getSegmentAccess();
ASSERT_EQ(segment_access.MountSegment(segment, client_id), ErrorCode::OK);
// Verify segment is properly mounted
ValidateMountedSegment(segment_manager, segment, client_id);
}
// MountSegmentDuplicate Tests:
// 1. MountSegment with the same segment id. The second mount operation return
// SEGMENT_ALREADY_EXISTS.
// 2. MountSegment with different segment id and the same segment name should be
// considered as different segments. Validate the status of SegmentManager use
// ValidateMountedSegments function.
TEST_F(SegmentTest, MountSegmentDuplicate) {
SegmentManager segment_manager;
// Create a valid segment and client ID
Segment segment;
segment.id = generate_uuid();
segment.name = "test_segment";
segment.size = 1024 * 1024 * 16;
segment.base = 0x100000000;
UUID client_id = generate_uuid();
// Get segment access and mount first time
auto segment_access = segment_manager.getSegmentAccess();
ASSERT_EQ(segment_access.MountSegment(segment, client_id), ErrorCode::OK);
// Verify first mount
ValidateMountedSegment(segment_manager, segment, client_id);
// Test duplicate mount - mount the same segment again
ASSERT_EQ(segment_access.MountSegment(segment, client_id),
ErrorCode::SEGMENT_ALREADY_EXISTS);
// Verify state remains the same after duplicate mount
ValidateMountedSegment(segment_manager, segment, client_id);
// Create a new segment with same name but different ID
Segment segment2;
segment2.id = generate_uuid(); // Different ID
segment2.name = segment.name; // Same name
segment2.size = segment.size * 2;
segment2.base = segment.base + segment.size;
// Mount the second segment
ASSERT_EQ(segment_access.MountSegment(segment2, client_id), ErrorCode::OK);
// Verify both segments are mounted correctly
std::vector<Segment> segments = {segment, segment2};
std::vector<UUID> client_ids = {client_id, client_id};
ValidateMountedSegments(segment_manager, segments, client_ids);
}
// UnmountSegmentSuccess:
// 1. Mount a segment and then unmount it. Unmount operation return success.
// 2. Use ValidateMountedSegments function to validate the status of
// SegmentManager.
TEST_F(SegmentTest, UnmountSegmentSuccess) {
SegmentManager segment_manager;
// Create and mount a segment
Segment segment;
segment.id = generate_uuid();
segment.name = "test_segment";
segment.size = 1024 * 1024 * 16;
segment.base = 0x100000000;
UUID client_id = generate_uuid();
// Get segment access and mount
auto segment_access = segment_manager.getSegmentAccess();
ASSERT_EQ(segment_access.MountSegment(segment, client_id), ErrorCode::OK);
// Verify segment is mounted correctly
ValidateMountedSegment(segment_manager, segment, client_id);
// Prepare unmount
size_t metrics_dec_capacity = 0;
ASSERT_EQ(
segment_access.PrepareUnmountSegment(segment.id, metrics_dec_capacity),
ErrorCode::OK);
ASSERT_EQ(metrics_dec_capacity, segment.size);
// Commit unmount
ASSERT_EQ(segment_access.CommitUnmountSegment(segment.id, client_id,
metrics_dec_capacity),
ErrorCode::OK);
// Verify segment is unmounted correctly
std::vector<Segment> empty_segment_vec;
std::vector<UUID> empty_client_ids_vec;
ValidateMountedSegments(segment_manager, empty_segment_vec,
empty_client_ids_vec);
}
// UnmountSegmentDuplicate:
// 1. Mount a segment and then unmount it twice. The second unmount operation
// returns SEGMENT_NOT_FOUND.
// 2. Only use ValidateMountedSegments function to validate the status of
// SegmentManager. Do not use other interfaces for validation.
TEST_F(SegmentTest, UnmountSegmentDuplicate) {
SegmentManager segment_manager;
// Create and mount a segment
Segment segment;
segment.id = generate_uuid();
segment.name = "test_segment";
segment.size = 1024 * 1024 * 16;
segment.base = 0x100000000;
UUID client_id = generate_uuid();
// Get segment access and mount
auto segment_access = segment_manager.getSegmentAccess();
ASSERT_EQ(segment_access.MountSegment(segment, client_id), ErrorCode::OK);
// Verify initial mounted state
ValidateMountedSegment(segment_manager, segment, client_id);
// First unmount
size_t metrics_dec_capacity = 0;
ASSERT_EQ(
segment_access.PrepareUnmountSegment(segment.id, metrics_dec_capacity),
ErrorCode::OK);
ASSERT_EQ(segment_access.CommitUnmountSegment(segment.id, client_id,
metrics_dec_capacity),
ErrorCode::OK);
// Verify segment is unmounted after first unmount
std::vector<Segment> empty_segment_vec;
std::vector<UUID> empty_client_ids_vec;
ValidateMountedSegments(segment_manager, empty_segment_vec,
empty_client_ids_vec);
// Second unmount attempt
metrics_dec_capacity = 0;
ASSERT_EQ(
segment_access.PrepareUnmountSegment(segment.id, metrics_dec_capacity),
ErrorCode::SEGMENT_NOT_FOUND);
// Verify segment remains unmounted after second unmount
ValidateMountedSegments(segment_manager, empty_segment_vec,
empty_client_ids_vec);
}
// ReMountSegmentSuccess:
// 1. Mount a segment A;
// 2. Remount two segments: A and B where A is already mounted and B is a new
// segment. The remount operation return success.
// 3. Only use ValidateMountedSegments function to validate the status of
// SegmentManager. Do not use other interfaces for validation.
TEST_F(SegmentTest, ReMountSegmentSuccess) {
SegmentManager segment_manager;
// Create and mount segment A
Segment segment_a;
segment_a.id = generate_uuid();
segment_a.name = "test_segment_a";
segment_a.size = 1024 * 1024 * 16;
segment_a.base = 0x100000000;
UUID client_id = generate_uuid();
// Get segment access and mount segment A
auto segment_access = segment_manager.getSegmentAccess();
ASSERT_EQ(segment_access.MountSegment(segment_a, client_id), ErrorCode::OK);
// Verify segment A is mounted correctly
ValidateMountedSegment(segment_manager, segment_a, client_id);
// Create segment B
Segment segment_b;
segment_b.id = generate_uuid();
segment_b.name = "test_segment_b";
segment_b.size = 1024 * 1024 * 32;
segment_b.base = 0x200000000;
// Remount both segments A and B
std::vector<Segment> segments_to_remount = {segment_a, segment_b};
ASSERT_EQ(segment_access.ReMountSegment(segments_to_remount, client_id),
ErrorCode::OK);
// Verify both segments are mounted correctly
std::vector<UUID> client_ids = {client_id, client_id};
ValidateMountedSegments(segment_manager, segments_to_remount, client_ids);
}
// ReMountUnmountingSegment:
// 1. Mount a segment A;
// 2. PrepareUnmount segment A;
// 3. Remount segment A. The remount operation return
// UNAVAILABLE_IN_CURRENT_STATUS.
// 4. CommitUnmount segment A;
// 5. Only use ValidateMountedSegments function to validate the status of
// SegmentManager. Do not use other interfaces for validation.
TEST_F(SegmentTest, ReMountUnmountingSegment) {
SegmentManager segment_manager;
// Create and mount segment A
Segment segment_a;
segment_a.id = generate_uuid();
segment_a.name = "test_segment_a";
segment_a.size = 1024 * 1024 * 16;
segment_a.base = 0x100000000;
UUID client_id = generate_uuid();
// Get segment access and mount segment A
auto segment_access = segment_manager.getSegmentAccess();
ASSERT_EQ(segment_access.MountSegment(segment_a, client_id), ErrorCode::OK);
// Verify segment A is mounted correctly
ValidateMountedSegment(segment_manager, segment_a, client_id);
// Prepare unmount segment A
size_t metrics_dec_capacity = 0;
ASSERT_EQ(segment_access.PrepareUnmountSegment(segment_a.id,
metrics_dec_capacity),
ErrorCode::OK);
// Attempt to remount segment A while it's in UNMOUNTING state
std::vector<Segment> segments_to_remount = {segment_a};
ASSERT_EQ(segment_access.ReMountSegment(segments_to_remount, client_id),
ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS);
// Complete the unmount process
ASSERT_EQ(segment_access.CommitUnmountSegment(segment_a.id, client_id,
metrics_dec_capacity),
ErrorCode::OK);
// Verify segment is completely unmounted
std::vector<Segment> empty_segment_vec;
std::vector<UUID> empty_client_ids_vec;
ValidateMountedSegments(segment_manager, empty_segment_vec,
empty_client_ids_vec);
}
// QuerySegments:
// 1. Create and mount 10 different segments with different names and different
// client ids;
// 2. Test GetClientSegments, verify the return value is correct.
// 3. Test GetAllSegments, verify the return value is correct.
// 4. Test QuerySegments, verify the return value is correct.
TEST_F(SegmentTest, QuerySegments) {
SegmentManager segment_manager;
auto segment_access = segment_manager.getSegmentAccess();
// Create 10 different segments with different names and client IDs
std::vector<Segment> segments;
std::vector<UUID> client_ids;
std::unordered_map<UUID, UUID, boost::hash<UUID>>
expected_client_segments;
for (int i = 0; i < 10; i++) {
// Create segment
Segment segment;
segment.id = generate_uuid();
segment.name = "test_segment_" + std::to_string(i);
segment.size = 1024 * 1024 * 16;
segment.base =
0x100000000 + (i * 0x100000000); // Different base addresses
// Create client ID
UUID client_id = generate_uuid();
// Mount segment
ASSERT_EQ(segment_access.MountSegment(segment, client_id),
ErrorCode::OK);
// Store for verification
segments.push_back(segment);
client_ids.push_back(client_id);
expected_client_segments[client_id] = segment.id;
}
// Verify all segments are mounted correctly
ValidateMountedSegments(segment_manager, segments, client_ids);
// Test GetClientSegments for each client
for (size_t i = 0; i < client_ids.size(); i++) {
std::vector<Segment> client_segments;
ASSERT_EQ(
segment_access.GetClientSegments(client_ids[i], client_segments),
ErrorCode::OK);
// Verify correct number of segments
ASSERT_EQ(client_segments.size(), 1);
// Verify all expected segments are present
ASSERT_EQ(client_segments[0].id,
expected_client_segments[client_ids[i]]);
}
// Test GetAllSegments
std::vector<std::string> all_segments;
ASSERT_EQ(segment_access.GetAllSegments(all_segments), ErrorCode::OK);
// Verify correct number of segments
ASSERT_EQ(all_segments.size(), segments.size());
// Verify all segment names are present
for (const auto& segment : segments) {
ASSERT_NE(
std::find(all_segments.begin(), all_segments.end(), segment.name),
all_segments.end());
}
// Test QuerySegments for each segment
for (const auto& segment : segments) {
size_t used = 0, capacity = 0;
ASSERT_EQ(segment_access.QuerySegments(segment.name, used, capacity),
ErrorCode::OK);
// Verify capacity matches segment size
ASSERT_EQ(capacity, segment.size);
// Verify used space is 0 for newly mounted segments
ASSERT_EQ(used, 0);
}
// Test QuerySegments for non-existent segment
size_t used = 0, capacity = 0;
ASSERT_EQ(
segment_access.QuerySegments("non_existent_segment", used, capacity),
ErrorCode::SEGMENT_NOT_FOUND);
ASSERT_EQ(used, 0);
ASSERT_EQ(capacity, 0);
}
} // namespace mooncake

12
mooncake-store/tests/stress_workload_test.cpp
View File

@ -112,18 +112,18 @@ class ClientIntegrationTest : public ::testing::Test {
}
static void InitializeSegment() {
const size_t ram_buffer_size = 3200ull * 1024 * 1024;
segment_ptr_ = allocate_buffer_allocator_memory(ram_buffer_size);
ram_buffer_size_ = 3200ull * 1024 * 1024;
segment_ptr_ = allocate_buffer_allocator_memory(ram_buffer_size_);
ASSERT_TRUE(segment_ptr_);
ErrorCode rc = client_->MountSegment("localhost:12345", segment_ptr_,
ram_buffer_size);
ErrorCode rc = client_->MountSegment(segment_ptr_,
ram_buffer_size_);
if (rc != ErrorCode::OK) {
LOG(ERROR) << "Failed to mount segment: " << toString(rc);
}
}
static void CleanupSegment() {
if (client_->UnmountSegment("localhost:12345", segment_ptr_) !=
if (client_->UnmountSegment(segment_ptr_, ram_buffer_size_) !=
ErrorCode::OK) {
LOG(ERROR) << "Failed to unmount segment";
}
@ -162,6 +162,7 @@ class ClientIntegrationTest : public ::testing::Test {
static std::shared_ptr<Client> client_;
static std::unique_ptr<SimpleAllocator> client_buffer_allocator_;
static void* segment_ptr_;
static size_t ram_buffer_size_;
};
// Static members initialization
@ -169,6 +170,7 @@ std::shared_ptr<Client> ClientIntegrationTest::client_ = nullptr;
void* ClientIntegrationTest::segment_ptr_ = nullptr;
std::unique_ptr<SimpleAllocator>
ClientIntegrationTest::client_buffer_allocator_ = nullptr;
size_t ClientIntegrationTest::ram_buffer_size_ = 0;
// Test basic Put/Get operations through the client
TEST_F(ClientIntegrationTest, StressPutOperations) {