diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index d45a9a71..ee0bdb16 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -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; diff --git a/mooncake-store/include/allocation_strategy.h b/mooncake-store/include/allocation_strategy.h index 17448125..83b57c7a 100644 --- a/mooncake-store/include/allocation_strategy.h +++ b/mooncake-store/include/allocation_strategy.h @@ -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 Allocate( - const std::unordered_map>& - allocators, + const std::vector>& allocators, + const std::unordered_map>>& + 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 Allocate( - const std::unordered_map>& - allocators, + const std::vector>& allocators, + const std::unordered_map>>& + 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 TryPreferredAllocation( - const std::unordered_map>& + std::unique_ptr TryPreferredAllocate( + const std::unordered_map>>& 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 RandomAllocateFromEligible( - const std::unordered_map>& - 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> CollectEligibleAllocators( - const std::unordered_map>& - allocators, - size_t objectSize) { - std::vector> 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 TryAllocateWithRetry( - std::vector>& eligible, + std::unique_ptr TryRandomAllocate( + const std::vector>& 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 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 dist(0, eligible.size() - 1); - const size_t random_index = dist(rng_); + std::uniform_int_distribution 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& 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>& allocators, - size_t index) { - if (index + 1 != allocators.size()) { - std::swap(allocators[index], allocators.back()); - } - allocators.pop_back(); - } }; } // namespace mooncake diff --git a/mooncake-store/include/client.h b/mooncake-store/include/client.h index 2f2abbf2..f8635693 100644 --- a/mooncake-store/include/client.h +++ b/mooncake-store/include/client.h @@ -5,6 +5,7 @@ #include #include #include +#include #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 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 mounted_segments_; + std::unordered_map> mounted_segments_; // Configuration const std::string local_hostname_; @@ -250,7 +244,10 @@ class Client { MasterViewHelper master_view_helper_; std::thread ping_thread_; std::atomic ping_running_{false}; - void PingThreadFunc(int current_version); + void PingThreadFunc(); + + // Client identification + UUID client_id_; }; } // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/ha_helper.h b/mooncake-store/include/ha_helper.h index 0ced521d..e831d9d4 100644 --- a/mooncake-store/include/ha_helper.h +++ b/mooncake-store/include/ha_helper.h @@ -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_; diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 77b97c3a..b44d65b0 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -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& 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_; diff --git a/mooncake-store/include/master_metric_manager.h b/mooncake-store/include/master_metric_manager.h index 14229ad7..01e08de1 100644 --- a/mooncake-store/include/master_metric_manager.h +++ b/mooncake-store/include/master_metric_manager.h @@ -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 diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index f6a1bf29..8f2b6551 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -9,6 +10,7 @@ #include #include #include +#include #include #include @@ -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>& - 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> - 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& 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 replicas; @@ -284,11 +287,10 @@ class MasterService { } }; - // Buffer allocator management - std::shared_ptr buffer_allocator_manager_; + // Segment management + SegmentManager segment_manager_; std::shared_ptr 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> ok_client_; // client with ok status + void ClientMonitorFunc(); + std::thread client_monitor_thread_; + std::atomic 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 client_ping_queue_{kClientPingQueueSize}; + const int64_t client_live_ttl_sec_; + + // if high availability features enabled + const bool enable_ha_; }; } // namespace mooncake diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index bef7981b..45b28755 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -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& 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; diff --git a/mooncake-store/include/segment.h b/mooncake-store/include/segment.h new file mode 100644 index 00000000..3b889079 --- /dev/null +++ b/mooncake-store/include/segment.h @@ -0,0 +1,182 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#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 + 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 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& 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& segments) const; + + /** + * @brief Get the names of all the segments + */ + ErrorCode GetAllSegments(std::vector& 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 lock_; +}; + +/** + * @brief RAII-style access to allocators for thread-safe allocator usage + */ +class ScopedAllocatorAccess { + public: + explicit ScopedAllocatorAccess( + std::unordered_map>>& + allocators_by_name, + std::vector>& allocators, + std::shared_mutex& mutex) + : allocators_by_name_(allocators_by_name), + allocators_(allocators), + lock_(mutex) {} + + const std::unordered_map>>& + getAllocatorsByName() { + return allocators_by_name_; + } + + const std::vector>& getAllocators() { + return allocators_; + } + + private: + const std::unordered_map>>& + allocators_by_name_; // segment name -> allocators + const std::vector>& allocators_; + std::shared_lock 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 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>> + allocators_by_name_; // segment name -> allocators + std::vector> allocators_; // allocators + std::unordered_map> + mounted_segments_; // segment_id -> mounted segment + std::unordered_map, boost::hash> + client_segments_; // client_id -> segment_ids + + friend class ScopedSegmentAccess; + friend class SegmentTest; // for unit tests +}; + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h index 09636f85..ef8c14bf 100644 --- a/mooncake-store/include/types.h +++ b/mooncake-store/include/types.h @@ -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; + +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 + 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 diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 1a75c8ff..73c34b60 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -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 diff --git a/mooncake-store/src/client.cpp b/mooncake-store/src/client.cpp index 1508b3b7..7c7970c4 100644 --- a/mooncake-store/src/client.cpp +++ b/mooncake-store/src/client.cpp @@ -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 segments_to_unmount = - mounted_segments_; + std::vector segments_to_unmount; + { + std::lock_guard 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(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 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> 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(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 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 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(mtseg.base); + uintptr_t r1 = reinterpret_cast(mtseg.size) + l1; + uintptr_t l2 = reinterpret_cast(buffer); + uintptr_t r2 = reinterpret_cast(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(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 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(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 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 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(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 lock(mounted_segments_mutex_); + std::vector 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 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(); + } + + // 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(); } } diff --git a/mooncake-store/src/ha_helper.cpp b/mooncake-store/src/ha_helper.cpp index 61e58100..ecc28d22 100644 --- a/mooncake-store/src/ha_helper.cpp +++ b/mooncake-store/src/ha_helper.cpp @@ -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. diff --git a/mooncake-store/src/master.cpp b/mooncake-store/src/master.cpp index cf7b3b45..3a219427 100644 --- a/mooncake-store/src/master.cpp +++ b/mooncake-store/src/master.cpp @@ -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(); diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index 6cd3b83c..c0c0d115 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -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 result = syncAwait([&]() -> coro::Lazy> { Lazy> handler = co_await client_ .send_request<&WrappedMasterService::MountSegment>( - reinterpret_cast(buffer), - static_cast(size), segment_name); + segment, client_id); async_rpc_result 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& segments, const UUID& client_id) { + ScopedVLogTimer timer(1, "MasterClient::ReMountSegment"); + timer.LogRequest("segments_num=", segments.size(), ", client_id=", client_id); + + std::optional result = + syncAwait([&]() -> coro::Lazy> { + Lazy> handler = + co_await client_ + .send_request<&WrappedMasterService::ReMountSegment>( + segments, client_id); + async_rpc_result 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 result = coro::syncAwait( [&]() -> coro::Lazy> { 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 result = coro::syncAwait([&]() -> coro::Lazy> { 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; } diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp index 9632a4b4..f3285c24 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -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: " diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 102bdda2..0de161a4 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -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(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 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 allocator; - try { - // SlabAllocator may throw an exception if the size or base is invalid - // for the slab allocator. - allocator = std::make_shared(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 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()), - allocation_strategy_(std::make_shared()), + double eviction_high_watermark_ratio, + ViewVersionId view_version, + int64_t client_live_ttl_sec, bool enable_ha) + : allocation_strategy_(std::make_shared()), 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& 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 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 & 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 & all_keys) { return ErrorCode::OK; } -ErrorCode MasterService::GetAllSegments(std::vector & all_segments) { - all_segments.clear(); - std::shared_lock 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& 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 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 replicas; replicas.reserve(config.replica_num); - for (size_t i = 0; i < config.replica_num; ++i) { - std::vector> 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> 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 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 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> + 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 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 unmount_segments; + std::vector dec_capacities; + std::vector client_ids; + std::vector segment_names; + { + // Lock client_mutex and segment_mutex + std::unique_lock 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 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 diff --git a/mooncake-store/src/segment.cpp b/mooncake-store/src/segment.cpp new file mode 100644 index 00000000..95fb068b --- /dev/null +++ b/mooncake-store/src/segment.cpp @@ -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 allocator; + try { + // SlabAllocator may throw an exception if the size or base is invalid + // for the slab allocator. + allocator = + std::make_shared(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& 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 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& 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& 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 \ No newline at end of file diff --git a/mooncake-store/src/types.cpp b/mooncake-store/src/types.cpp index c8b616f3..8f884641 100644 --- a/mooncake-store/src/types.cpp +++ b/mooncake-store/src/types.cpp @@ -1,5 +1,9 @@ #include "types.h" +#include +#include +#include + 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); } +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 diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index c0652666..ae85cf53 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -54,4 +54,17 @@ target_link_libraries(transfer_task_test PUBLIC gtest_main pthread ) -add_test(NAME transfer_task_test COMMAND transfer_task_test) \ No newline at end of file +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) \ No newline at end of file diff --git a/mooncake-store/tests/allocation_strategy_test.cpp b/mooncake-store/tests/allocation_strategy_test.cpp index 79c0f732..49d67c7c 100644 --- a/mooncake-store/tests/allocation_strategy_test.cpp +++ b/mooncake-store/tests/allocation_strategy_test.cpp @@ -31,21 +31,23 @@ class AllocationStrategyTest : public ::testing::Test { // Test basic functionality with empty allocators map TEST_F(AllocationStrategyTest, EmptyAllocatorsMap) { - std::unordered_map> - empty_allocators; + std::unordered_map>> + empty_allocators_by_name; + std::vector> 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> - allocators; + std::unordered_map>> + empty_allocators_by_name; + std::vector> 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> - allocators; - allocators["segment1"] = allocator1; - allocators["preferred"] = allocator2; + std::unordered_map>> + allocators_by_name; + std::vector> 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> - allocators; - allocators["segment1"] = allocator1; - allocators["segment2"] = allocator2; + std::unordered_map>> + allocators_by_name; + std::vector> 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> - allocators; - allocators["segment1"] = allocator1; - allocators["segment2"] = allocator2; - allocators["segment3"] = allocator3; + std::unordered_map>> + allocators_by_name; + std::vector> 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 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> - allocators; - allocators["segment1"] = allocator1; - allocators["preferred"] = allocator2; + std::unordered_map>> + allocators_by_name; + std::vector> 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> - allocators; - allocators["segment1"] = allocator1; - allocators["segment2"] = allocator2; + std::unordered_map>> + allocators_by_name; + std::vector> 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> 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> - allocators; - allocators["segment1"] = allocator; + std::unordered_map>> + allocators_by_name; + std::vector> 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> - allocators; - allocators["segment1"] = allocator; + std::unordered_map>> + allocators_by_name; + std::vector> 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 } diff --git a/mooncake-store/tests/client_integration_test.cpp b/mooncake-store/tests/client_integration_test.cpp index 473f5e4f..15a54638 100644 --- a/mooncake-store/tests/client_integration_test.cpp +++ b/mooncake-store/tests/client_integration_test.cpp @@ -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 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 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) { diff --git a/mooncake-store/tests/e2e/CMakeLists.txt b/mooncake-store/tests/e2e/CMakeLists.txt new file mode 100644 index 00000000..1843f862 --- /dev/null +++ b/mooncake-store/tests/e2e/CMakeLists.txt @@ -0,0 +1,7 @@ +add_executable(clientctl clientctl.cpp) +target_link_libraries(clientctl PUBLIC + mooncake_store + cachelib_memory_allocator + glog + pthread +) \ No newline at end of file diff --git a/mooncake-store/tests/e2e/cases/case1.txt b/mooncake-store/tests/e2e/cases/case1.txt new file mode 100644 index 00000000..e0eafaed --- /dev/null +++ b/mooncake-store/tests/e2e/cases/case1.txt @@ -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 \ No newline at end of file diff --git a/mooncake-store/tests/e2e/cases/case2.txt b/mooncake-store/tests/e2e/cases/case2.txt new file mode 100644 index 00000000..d2ad616a --- /dev/null +++ b/mooncake-store/tests/e2e/cases/case2.txt @@ -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 \ No newline at end of file diff --git a/mooncake-store/tests/e2e/clientctl.cpp b/mooncake-store/tests/e2e/clientctl.cpp new file mode 100644 index 00000000..7daa4651 --- /dev/null +++ b/mooncake-store/tests/e2e/clientctl.cpp @@ -0,0 +1,295 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +#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; + std::unordered_map 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 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& descriptors = + object_info.replica_list[0].buffer_descriptors; + std::vector 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(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 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; +} diff --git a/mooncake-store/tests/master_metrics_test.cpp b/mooncake-store/tests/master_metrics_test.cpp index d4b01bd6..ebe58a39 100644 --- a/mooncake-store/tests/master_metrics_test.cpp +++ b/mooncake-store/tests/master_metrics_test.cpp @@ -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 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); diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp index 0703173f..0e3b1272 100644 --- a/mooncake-store/tests/master_service_test.cpp +++ b/mooncake-store/tests/master_service_test.cpp @@ -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 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 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 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 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 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( @@ -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 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 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; diff --git a/mooncake-store/tests/segment_test.cpp b/mooncake-store/tests/segment_test.cpp new file mode 100644 index 00000000..988f36ba --- /dev/null +++ b/mooncake-store/tests/segment_test.cpp @@ -0,0 +1,443 @@ +#include "segment.h" + +#include +#include + +#include + +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& segments, + const std::vector& 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 segments; + segments.push_back(segment); + std::vector 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 segments = {segment, segment2}; + std::vector 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 empty_segment_vec; + std::vector 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 empty_segment_vec; + std::vector 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 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 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 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 empty_segment_vec; + std::vector 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 segments; + std::vector client_ids; + std::unordered_map> + 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 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 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 \ No newline at end of file diff --git a/mooncake-store/tests/stress_workload_test.cpp b/mooncake-store/tests/stress_workload_test.cpp index dcc9d100..93d2799e 100644 --- a/mooncake-store/tests/stress_workload_test.cpp +++ b/mooncake-store/tests/stress_workload_test.cpp @@ -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_; static std::unique_ptr client_buffer_allocator_; static void* segment_ptr_; + static size_t ram_buffer_size_; }; // Static members initialization @@ -169,6 +170,7 @@ std::shared_ptr ClientIntegrationTest::client_ = nullptr; void* ClientIntegrationTest::segment_ptr_ = nullptr; std::unique_ptr ClientIntegrationTest::client_buffer_allocator_ = nullptr; +size_t ClientIntegrationTest::ram_buffer_size_ = 0; // Test basic Put/Get operations through the client TEST_F(ClientIntegrationTest, StressPutOperations) {