diff --git a/mooncake-store/include/allocation_strategy.h b/mooncake-store/include/allocation_strategy.h index 3e7fa65d..4c2b1949 100644 --- a/mooncake-store/include/allocation_strategy.h +++ b/mooncake-store/include/allocation_strategy.h @@ -1,13 +1,12 @@ #pragma once #include -#include #include -#include #include #include #include -#include +#include +#include #include #include "allocator.h" // Contains BufferAllocator declaration @@ -18,7 +17,7 @@ namespace mooncake { /** * @brief Abstract interface for allocation strategy, responsible for - * allocating multiple slices across multiple replicas using available + * allocating a slice (with one or more replicas) using available * BufferAllocators. * * The allocation strategy follows best-effort semantics: if the requested @@ -31,9 +30,8 @@ class AllocationStrategy { virtual ~AllocationStrategy() = default; /** - * @brief Allocates multiple slices across the requested number of replicas - * using best-effort semantics. Each replica will contain all - * requested slices. + * @brief Allocates a slice across the requested number of replicas + * using best-effort semantics. * * The allocation follows best-effort semantics: if the full requested * replica count cannot be satisfied, the method will allocate as many @@ -44,7 +42,7 @@ class AllocationStrategy { * @param allocators_by_name Container of mounted allocators, key is * segment_name, value is the corresponding * allocators - * @param slice_sizes Sizes of slices to be allocated in each replica + * @param slice_length Length of the slice to be allocated * @param config Replica configuration containing number of replicas and * placement constraints * @return tl::expected, ErrorCode> containing @@ -60,8 +58,7 @@ class AllocationStrategy { const std::unordered_map< std::string, std::vector>>& allocators_by_name, - const std::vector& slice_sizes, - const ReplicateConfig& config) = 0; + const size_t slice_length, const ReplicateConfig& config) = 0; }; /** @@ -87,189 +84,110 @@ class RandomAllocationStrategy : public AllocationStrategy { const std::unordered_map< std::string, std::vector>>& allocators_by_name, - const std::vector& slice_sizes, const ReplicateConfig& config) { - if (auto validation_error = - validateInput(slice_sizes, config.replica_num)) { - return tl::make_unexpected(*validation_error); + const size_t slice_length, const ReplicateConfig& config) { + // Validate input parameters + if (slice_length == 0 || config.replica_num == 0) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } - std::vector>> - replica_buffers(config.replica_num); - for (auto& replica_buffer : replica_buffers) { - replica_buffer.reserve(slice_sizes.size()); - } - - // Track the actual number of replicas we can allocate - size_t actual_replica_count = config.replica_num; - - // Allocate each slice across replicas - for (size_t slice_idx = 0; slice_idx < slice_sizes.size(); - ++slice_idx) { - auto slice_replicas = allocateSlice(allocators, allocators_by_name, - slice_sizes[slice_idx], - actual_replica_count, config); - - if (slice_replicas.empty()) { - return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); - } - - if (slice_replicas.size() < actual_replica_count) { - actual_replica_count = slice_replicas.size(); - // NOTE: replica allocation is best effort - VLOG(1) << "Failed to allocate all replicas for slice " - << slice_idx << ", reducing replica count to " - << actual_replica_count; - - // Resize replica_buffers to match the new count - replica_buffers.resize(actual_replica_count); - } - - for (size_t replica_idx = 0; replica_idx < actual_replica_count; - ++replica_idx) { - replica_buffers[replica_idx].push_back( - std::move(slice_replicas[replica_idx])); + // Fast path: single allocator case + if (allocators.size() == 1) { + if (auto buffer = allocators[0]->allocate(slice_length)) { + std::vector result; + result.emplace_back(std::move(buffer), + ReplicaStatus::PROCESSING); + return result; } + return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } std::vector replicas; - replicas.reserve(actual_replica_count); - for (size_t replica_idx = 0; replica_idx < actual_replica_count; - ++replica_idx) { - replicas.emplace_back(std::move(replica_buffers[replica_idx]), - ReplicaStatus::PROCESSING); - } + replicas.reserve(config.replica_num); - return replicas; - } - - std::optional validateInput( - const std::vector& slice_sizes, size_t replica_num) const { - if (replica_num == 0 || slice_sizes.empty() || - std::count(slice_sizes.begin(), slice_sizes.end(), 0) > 0) { - return ErrorCode::INVALID_PARAMS; - } - - return std::nullopt; - } - - /** - * @brief Allocates replicas for a single slice across different segments - */ - std::vector> allocateSlice( - const std::vector>& allocators, - const std::unordered_map< - std::string, std::vector>>& - allocators_by_name, - size_t slice_size, size_t replica_num, const ReplicateConfig& config, - std::unordered_set& used_segments) { - std::vector> buffers; - buffers.reserve(replica_num); - - for (size_t i = 0; i < replica_num; ++i) { - auto buffer = - allocateSingleBuffer(allocators, allocators_by_name, slice_size, - config, used_segments); - - if (!buffer) { - break; - } - - used_segments.insert(buffer->getSegmentName()); - buffers.push_back(std::move(buffer)); - } - - return buffers; - } - - std::vector> allocateSlice( - const std::vector>& allocators, - const std::unordered_map< - std::string, std::vector>>& - allocators_by_name, - size_t slice_size, size_t replica_num, const ReplicateConfig& config) { - std::unordered_set empty_segments; - return allocateSlice(allocators, allocators_by_name, slice_size, - replica_num, config, empty_segments); - } - - /** - * @brief Allocates a single buffer respecting preferences and exclusions - */ - std::unique_ptr allocateSingleBuffer( - const std::vector>& allocators, - const std::unordered_map< - std::string, std::vector>>& - allocators_by_name, - size_t size, const ReplicateConfig& config, - const std::unordered_set& excluded_segments) { - // Try preferred segment first - if (!config.preferred_segment.empty() && - !excluded_segments.contains(config.preferred_segment)) { + // Try preferred segment first if specified + if (!config.preferred_segment.empty()) { auto preferred_it = allocators_by_name.find(config.preferred_segment); if (preferred_it != allocators_by_name.end()) { for (auto& allocator : preferred_it->second) { - if (auto buffer = allocator->allocate(size)) { - return buffer; + if (auto buffer = allocator->allocate(slice_length)) { + replicas.emplace_back(std::move(buffer), + ReplicaStatus::PROCESSING); + break; } } } } - return tryRandomAllocate(allocators, size, excluded_segments); - } + if (replicas.size() == config.replica_num) { + return replicas; + } - /** - * @brief Attempts allocation with random selection from allocators that can - * fit the size - */ - std::unique_ptr tryRandomAllocate( - const std::vector>& allocators, - size_t size, const std::unordered_set& excluded_segments) { - std::vector eligible_indices; - eligible_indices.reserve(allocators.size()); - for (size_t i = 0; i < allocators.size(); ++i) { - if (!excluded_segments.contains(allocators[i]->getSegmentName()) && - allocators[i]->getLargestFreeRegion() >= size) { - eligible_indices.push_back(i); + // If replica_num is not satisfied, allocate the remaining replicas + // randomly Randomly select a starting point from allocators_by_name + if (allocators_by_name.empty()) { + if (replicas.empty()) { + return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); + } + return replicas; + } + + static thread_local std::mt19937 generator(clock()); + std::uniform_int_distribution distribution( + 0, allocators_by_name.size() - 1); + size_t start_idx = distribution(generator); + + // Get iterator to the starting point + auto start_it = allocators_by_name.begin(); + std::advance(start_it, start_idx); + + auto it = start_it; + size_t max_retry = std::min(kMaxRetryLimit, allocators_by_name.size()); + size_t retry_count = 0; + + // Try to allocate remaining replicas, starting from random position + // TODO: Change the segment data structure to avoid traversing the + // entire map every time + while (replicas.size() < config.replica_num && + retry_count < max_retry) { + // Skip preferred segment if it was already allocated + if (it->first != config.preferred_segment) { + // Try each allocator in this segment + bool allocated = false; + for (auto& allocator : it->second) { + if (auto buffer = allocator->allocate(slice_length)) { + replicas.emplace_back(std::move(buffer), + ReplicaStatus::PROCESSING); + // Allocate at most one replica per segment + allocated = true; + break; + } + } + if (!allocated) { + ++retry_count; + } + } + // Move to next segment (circular) + ++it; + if (it == allocators_by_name.end()) { + it = allocators_by_name.begin(); + } + + // If we have cycled through all segments, break + if (it == start_it) { + break; } } - if (eligible_indices.empty()) { - return nullptr; + // Return allocated replicas (may be fewer than requested) + if (replicas.empty()) { + return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } - - // Thread-local random number generator for thread safety - thread_local std::mt19937 rng(std::random_device{}()); - std::shuffle(eligible_indices.begin(), eligible_indices.end(), rng); - - const size_t max_tries = - std::min(kMaxRetryLimit, eligible_indices.size()); - for (size_t i = 0; i < max_tries; ++i) { - auto& allocator = allocators[eligible_indices[i]]; - if (auto buffer = allocator->allocate(size)) { - return buffer; - } - retry_counter_.fetch_add(1); // Track allocation attempts - } - - return nullptr; + return replicas; } - /** - * @brief Get the number of allocation retry attempts - */ - uint64_t getRetryCount() const { return retry_counter_.load(); } - - /** - * @brief Reset the retry counter - */ - void resetRetryCount() { retry_counter_.store(0); } - private: - static constexpr size_t kMaxRetryLimit = 10; - // Observer for allocation retries - std::atomic_uint64_t retry_counter_{0}; + static constexpr size_t kMaxRetryLimit = 100; }; -} // namespace mooncake +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/client_buffer.hpp b/mooncake-store/include/client_buffer.hpp index 0a1ca1e4..79c39360 100644 --- a/mooncake-store/include/client_buffer.hpp +++ b/mooncake-store/include/client_buffer.hpp @@ -103,11 +103,11 @@ uint64_t calculate_total_size(const Replica::Descriptor& replica); * @brief Allocate slices from a buffer handle based on replica descriptor * @param slices Output vector to store the allocated slices * @param replica The replica descriptor defining the slice structure - * @param buffer_handle The buffer handle to allocate slices from + * @param buffer_ptr The buffer pointer to allocate slices from * @return 0 on success, non-zero on error */ int allocateSlices(std::vector& slices, const Replica::Descriptor& replica, - BufferHandle& buffer_handle); + void* buffer_ptr); } // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index ba1848c2..348407c9 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -133,14 +133,14 @@ class MasterService { /** * @brief Start a put operation for an object - * @param[out] replica_list Vector to store replica information for slices + * @param[out] replica_list Vector to store replica information for the + * slice * @return ErrorCode::OK on success, ErrorCode::OBJECT_NOT_FOUND if exists, * ErrorCode::NO_AVAILABLE_HANDLE if allocation fails, * ErrorCode::INVALID_PARAMS if slice size is invalid */ auto PutStart(const UUID& client_id, const std::string& key, - const std::vector& slice_lengths, - const ReplicateConfig& config) + const uint64_t slice_length, const ReplicateConfig& config) -> tl::expected, ErrorCode>; /** @@ -583,4 +583,4 @@ class MasterService { GUARDED_BY(discarded_replicas_mutex_); }; -} // namespace mooncake +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/replica.h b/mooncake-store/include/replica.h index 0386eaf2..740d0f6d 100644 --- a/mooncake-store/include/replica.h +++ b/mooncake-store/include/replica.h @@ -87,7 +87,7 @@ struct ReplicateConfig { }; struct MemoryReplicaData { - std::vector> buffers; + std::unique_ptr buffer; }; struct DiskReplicaData { @@ -96,8 +96,8 @@ struct DiskReplicaData { }; struct MemoryDescriptor { - std::vector buffer_descriptors; - YLT_REFL(MemoryDescriptor, buffer_descriptors); + AllocatedBuffer::Descriptor buffer_descriptor; + YLT_REFL(MemoryDescriptor, buffer_descriptor); }; struct DiskDescriptor { @@ -111,9 +111,8 @@ class Replica { struct Descriptor; // memory replica constructor - Replica(std::vector> buffers, - ReplicaStatus status) - : data_(MemoryReplicaData{std::move(buffers)}), status_(status) {} + Replica(std::unique_ptr buffer, ReplicaStatus status) + : data_(MemoryReplicaData{std::move(buffer)}), status_(status) {} // disk replica constructor Replica(std::string file_path, uint64_t object_size, ReplicaStatus status) @@ -183,24 +182,19 @@ class Replica { [[nodiscard]] bool has_invalid_mem_handle() const { if (is_memory_replica()) { const auto& mem_data = std::get(data_); - return std::any_of( - mem_data.buffers.begin(), mem_data.buffers.end(), - [](const std::unique_ptr& buf_ptr) { - return !buf_ptr->isAllocatorValid(); - }); + return !mem_data.buffer->isAllocatorValid(); } return false; // DiskReplicaData does not have handles } [[nodiscard]] size_t get_memory_buffer_size() const { - size_t size = 0; if (is_memory_replica()) { const auto& mem_data = std::get(data_); - for (auto& buffer : mem_data.buffers) { - size += buffer->size(); - } + return mem_data.buffer->size(); + } else { + LOG(ERROR) << "Invalid replica type: " << type(); + return 0; } - return size; } [[nodiscard]] std::vector> get_segment_names() @@ -292,12 +286,13 @@ inline Replica::Descriptor Replica::get_descriptor() const { if (is_memory_replica()) { const auto& mem_data = std::get(data_); MemoryDescriptor mem_desc; - mem_desc.buffer_descriptors.reserve(mem_data.buffers.size()); - for (const auto& buf_ptr : mem_data.buffers) { - if (buf_ptr) { - mem_desc.buffer_descriptors.push_back( - buf_ptr->get_descriptor()); - } + if (mem_data.buffer) { + mem_desc.buffer_descriptor = mem_data.buffer->get_descriptor(); + } else { + mem_desc.buffer_descriptor.size_ = 0; + mem_desc.buffer_descriptor.buffer_address_ = 0; + mem_desc.buffer_descriptor.transport_endpoint_ = ""; + LOG(ERROR) << "Trying to get invalid memory replica descriptor"; } desc.descriptor_variant = std::move(mem_desc); } else if (is_disk_replica()) { @@ -315,15 +310,11 @@ inline std::vector> Replica::get_segment_names() const { if (is_memory_replica()) { const auto& mem_data = std::get(data_); - std::vector> segment_names( - mem_data.buffers.size()); - for (size_t i = 0; i < mem_data.buffers.size(); ++i) { - if (mem_data.buffers[i] && - mem_data.buffers[i]->isAllocatorValid()) { - segment_names[i] = mem_data.buffers[i]->getSegmentName(); - } else { - segment_names[i] = std::nullopt; - } + std::vector> segment_names; + if (mem_data.buffer && mem_data.buffer->isAllocatorValid()) { + segment_names.push_back(mem_data.buffer->getSegmentName()); + } else { + segment_names.push_back(std::nullopt); } return segment_names; } @@ -336,10 +327,8 @@ inline std::ostream& operator<<(std::ostream& os, const Replica& replica) { if (replica.is_memory_replica()) { const auto& mem_data = std::get(replica.data_); os << "type: MEMORY, buffers: ["; - for (const auto& buf_ptr : mem_data.buffers) { - if (buf_ptr) { - os << *buf_ptr; - } + if (mem_data.buffer) { + os << *mem_data.buffer; } os << "]"; } else if (replica.is_disk_replica()) { diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index f76267be..96943fbe 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -42,8 +42,7 @@ class WrappedMasterService { tl::expected, ErrorCode> PutStart( const UUID& client_id, const std::string& key, - const std::vector& slice_lengths, - const ReplicateConfig& config); + const uint64_t slice_length, const ReplicateConfig& config); tl::expected PutEnd(const UUID& client_id, const std::string& key, @@ -55,7 +54,7 @@ class WrappedMasterService { std::vector, ErrorCode>> BatchPutStart(const UUID& client_id, const std::vector& keys, - const std::vector>& slice_lengths, + const std::vector& slice_lengths, const ReplicateConfig& config); std::vector> BatchPutEnd( @@ -95,4 +94,4 @@ class WrappedMasterService { void RegisterRpcService(coro_rpc::coro_rpc_server& server, mooncake::WrappedMasterService& wrapped_master_service); -} // namespace mooncake +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/include/transfer_task.h b/mooncake-store/include/transfer_task.h index a97b47bc..6b8fc82c 100644 --- a/mooncake-store/include/transfer_task.h +++ b/mooncake-store/include/transfer_task.h @@ -400,36 +400,35 @@ class TransferSubmitter { /** * @brief Select the optimal transfer strategy */ - TransferStrategy selectStrategy( - const std::vector& handles, - const std::vector& slices) const; + TransferStrategy selectStrategy(const AllocatedBuffer::Descriptor& handle, + const std::vector& slices) const; /** * @brief Check if all handles refer to local segments */ - bool isLocalTransfer( - const std::vector& handles) const; + bool isLocalTransfer(const AllocatedBuffer::Descriptor& handle) const; /** * @brief Validate transfer parameters */ - bool validateTransferParams( - const std::vector& handles, - const std::vector& slices, bool is_multi_buffers = false) const; + bool validateTransferParams(const AllocatedBuffer::Descriptor& handle, + const std::vector& slices) const; /** * @brief Submit memcpy operation asynchronously */ std::optional submitMemcpyOperation( - const std::vector& handles, - std::vector& slices, TransferRequest::OpCode op_code); + const AllocatedBuffer::Descriptor& handle, + const std::vector& slices, + const TransferRequest::OpCode op_code); /** * @brief Submit transfer engine operation asynchronously */ std::optional submitTransferEngineOperation( - const std::vector& handles, - std::vector& slices, TransferRequest::OpCode op_code); + const AllocatedBuffer::Descriptor& handle, + const std::vector& slices, + const TransferRequest::OpCode op_code); std::optional submitFileReadOperation( const Replica::Descriptor& replica, std::vector& slices, diff --git a/mooncake-store/src/client.cpp b/mooncake-store/src/client.cpp index ee50c097..263d8cbf 100644 --- a/mooncake-store/src/client.cpp +++ b/mooncake-store/src/client.cpp @@ -589,11 +589,11 @@ std::vector> Client::BatchGetWhenPreferSameNode( continue; } auto& memory_descriptor = replica.get_memory_descriptor(); - if (memory_descriptor.buffer_descriptors.empty()) { + if (memory_descriptor.buffer_descriptor.size_ == 0) { results[i] = tl::unexpected(ErrorCode::INVALID_REPLICA); continue; } - auto& buffer_descriptor = memory_descriptor.buffer_descriptors[0]; + auto& buffer_descriptor = memory_descriptor.buffer_descriptor; auto seg = buffer_descriptor.transport_endpoint_; auto& op = seg_to_op_map[seg]; op.replicas.emplace_back(replica); @@ -1240,12 +1240,11 @@ std::vector> Client::BatchPutWhenPreferSameNode( continue; } auto& memory_descriptor = replica.get_memory_descriptor(); - if (memory_descriptor.buffer_descriptors.empty()) { - op.SetError(ErrorCode::INVALID_PARAMS, - "buffer descriptors is empty."); + if (memory_descriptor.buffer_descriptor.size_ == 0) { + op.SetError(ErrorCode::INVALID_PARAMS, "buffer size is 0."); continue; } - auto& buffer_descriptor = memory_descriptor.buffer_descriptors[0]; + auto& buffer_descriptor = memory_descriptor.buffer_descriptor; auto seg = buffer_descriptor.transport_endpoint_; if (seg_to_ops.find(seg) == seg_to_ops.end()) { seg_to_ops.emplace(seg, PutOperation(op.key, op.slices)); @@ -1286,7 +1285,7 @@ std::vector> Client::BatchPutWhenPreferSameNode( WaitForTransfers(merged_ops); for (auto& op : merged_ops) { auto& memory_descriptor = op.replicas[0].get_memory_descriptor(); - auto& buffer_descriptor = memory_descriptor.buffer_descriptors[0]; + auto& buffer_descriptor = memory_descriptor.buffer_descriptor; auto seg = buffer_descriptor.transport_endpoint_; seg_to_ops.at(seg).state = op.state; } @@ -1295,7 +1294,7 @@ std::vector> Client::BatchPutWhenPreferSameNode( continue; } auto& memory_descriptor = op.replicas[0].get_memory_descriptor(); - auto& buffer_descriptor = memory_descriptor.buffer_descriptors[0]; + auto& buffer_descriptor = memory_descriptor.buffer_descriptor; auto seg = buffer_descriptor.transport_endpoint_; op.state = seg_to_ops.at(seg).state; auto state = std::make_shared(); @@ -1611,9 +1610,7 @@ ErrorCode Client::TransferRead(const Replica::Descriptor& replica_descriptor, size_t total_size = 0; if (replica_descriptor.is_memory_replica()) { auto& mem_desc = replica_descriptor.get_memory_descriptor(); - for (const auto& handle : mem_desc.buffer_descriptors) { - total_size += handle.size_; - } + total_size = mem_desc.buffer_descriptor.size_; } else { auto& disk_desc = replica_descriptor.get_disk_descriptor(); total_size = disk_desc.object_size; diff --git a/mooncake-store/src/client_buffer.cpp b/mooncake-store/src/client_buffer.cpp index 5acfc293..23fad52d 100644 --- a/mooncake-store/src/client_buffer.cpp +++ b/mooncake-store/src/client_buffer.cpp @@ -72,36 +72,29 @@ uint64_t calculate_total_size(const Replica::Descriptor& replica) { auto& disk_descriptor = replica.get_disk_descriptor(); total_length = disk_descriptor.object_size; } else { - for (auto& handle : - replica.get_memory_descriptor().buffer_descriptors) { - total_length += handle.size_; - } + total_length = replica.get_memory_descriptor().buffer_descriptor.size_; } return total_length; } int allocateSlices(std::vector& slices, - const Replica::Descriptor& replica, - BufferHandle& buffer_handle) { - uint64_t offset = 0; + const Replica::Descriptor& replica, void* buffer_ptr) { if (replica.is_memory_replica() == false) { // For disk-based replica, split into slices based on file size + uint64_t offset = 0; uint64_t total_length = replica.get_disk_descriptor().object_size; while (offset < total_length) { auto chunk_size = std::min(total_length - offset, kMaxSliceSize); - void* chunk_ptr = static_cast(buffer_handle.ptr()) + offset; + void* chunk_ptr = static_cast(buffer_ptr) + offset; slices.emplace_back(Slice{chunk_ptr, chunk_size}); offset += chunk_size; } } else { // For memory-based replica, split into slices based on buffer // descriptors - for (auto& handle : - replica.get_memory_descriptor().buffer_descriptors) { - void* chunk_ptr = static_cast(buffer_handle.ptr()) + offset; - slices.emplace_back(Slice{chunk_ptr, handle.size_}); - offset += handle.size_; - } + auto& handle = replica.get_memory_descriptor().buffer_descriptor; + void* chunk_ptr = buffer_ptr; + slices.emplace_back(Slice{chunk_ptr, handle.size_}); } return 0; } diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index 094ca94b..6731153b 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -297,16 +297,14 @@ MasterClient::PutStart(const std::string& key, ScopedVLogTimer timer(1, "MasterClient::PutStart"); timer.LogRequest("key=", key, ", slice_count=", slice_lengths.size()); - // Convert size_t to uint64_t for RPC - std::vector rpc_slice_lengths; - rpc_slice_lengths.reserve(slice_lengths.size()); - for (const auto& length : slice_lengths) { - rpc_slice_lengths.push_back(length); + uint64_t total_slice_length = 0; + for (const auto& slice_length : slice_lengths) { + total_slice_length += slice_length; } auto result = invoke_rpc<&WrappedMasterService::PutStart, std::vector>( - client_id_, key, rpc_slice_lengths, config); + client_id_, key, total_slice_length, config); timer.LogResponseExpected(result); return result; } @@ -319,9 +317,19 @@ MasterClient::BatchPutStart( ScopedVLogTimer timer(1, "MasterClient::BatchPutStart"); timer.LogRequest("keys_count=", keys.size()); + std::vector total_slice_lengths; + total_slice_lengths.reserve(slice_lengths.size()); + for (const auto& slice_lengths : slice_lengths) { + uint64_t total_slice_length = 0; + for (const auto& slice_length : slice_lengths) { + total_slice_length += slice_length; + } + total_slice_lengths.emplace_back(total_slice_length); + } + auto result = invoke_batch_rpc<&WrappedMasterService::BatchPutStart, std::vector>( - keys.size(), client_id_, keys, slice_lengths, config); + keys.size(), client_id_, keys, total_slice_lengths, config); timer.LogResponse("result=", result.size(), " operations"); return result; } diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 31ce4ae8..437561e4 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -353,32 +353,29 @@ auto MasterService::GetReplicaList(std::string_view key) } auto MasterService::PutStart(const UUID& client_id, const std::string& key, - const std::vector& slice_lengths, + const uint64_t slice_length, const ReplicateConfig& config) -> tl::expected, ErrorCode> { - if (config.replica_num == 0 || key.empty() || slice_lengths.empty()) { + if (config.replica_num == 0 || key.empty() || slice_length == 0) { LOG(ERROR) << "key=" << key << ", replica_num=" << config.replica_num - << ", slice_count=" << slice_lengths.size() + << ", slice_length=" << slice_length << ", key_size=" << key.size() << ", error=invalid_params"; return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } // Validate slice lengths uint64_t total_length = 0; - for (size_t i = 0; i < slice_lengths.size(); ++i) { - if ((memory_allocator_type_ == BufferAllocatorType::CACHELIB) && - (slice_lengths[i] > kMaxSliceSize)) { - LOG(ERROR) << "key=" << key << ", slice_index=" << i - << ", slice_size=" << slice_lengths[i] - << ", max_size=" << kMaxSliceSize - << ", error=invalid_slice_size"; - return tl::make_unexpected(ErrorCode::INVALID_PARAMS); - } - total_length += slice_lengths[i]; + if ((memory_allocator_type_ == BufferAllocatorType::CACHELIB) && + (slice_length > kMaxSliceSize)) { + LOG(ERROR) << "key=" << key << ", slice_length=" << slice_length + << ", max_size=" << kMaxSliceSize + << ", error=invalid_slice_size"; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } + total_length += slice_length; VLOG(1) << "key=" << key << ", value_length=" << total_length - << ", slice_count=" << slice_lengths.size() << ", config=" << config + << ", slice_length=" << slice_length << ", config=" << config << ", action=put_start_begin"; // Lock the shard and check if object already exists @@ -419,7 +416,7 @@ auto MasterService::PutStart(const UUID& client_id, const std::string& key, auto& allocators_by_name = allocator_access.getAllocatorsByName(); auto allocation_result = allocation_strategy_->Allocate( - allocators, allocators_by_name, slice_lengths, config); + allocators, allocators_by_name, slice_length, config); if (!allocation_result.has_value()) { VLOG(1) << "Failed to allocate all replicas for key=" << key @@ -1195,4 +1192,4 @@ std::string MasterService::ResolvePath(const std::string& key) const { return full_path.lexically_normal().string(); } -} // namespace mooncake +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/src/pybind_client.cpp b/mooncake-store/src/pybind_client.cpp index 996d4aa7..06a00dc1 100644 --- a/mooncake-store/src/pybind_client.cpp +++ b/mooncake-store/src/pybind_client.cpp @@ -614,7 +614,7 @@ std::shared_ptr PyClient::get_buffer(const std::string &key) { // Create slices for the allocated buffer std::vector slices; - allocateSlices(slices, replica, buffer_handle); + allocateSlices(slices, replica, buffer_handle.ptr()); // Get the object data auto get_result = client_->Get(key, query_result.value(), slices); @@ -691,7 +691,7 @@ std::vector> PyClient::batch_get_buffer_internal( auto buffer_handle = std::make_unique(std::move(*alloc_result)); std::vector slices; - allocateSlices(slices, replica, *buffer_handle); + allocateSlices(slices, replica, buffer_handle->ptr()); valid_ops.emplace_back( KeyOp{.original_index = i, @@ -819,23 +819,7 @@ tl::expected PyClient::get_into_internal( // Step 2: Split user buffer according to object info and create // slices std::vector slices; - uint64_t offset = 0; - - if (replica.is_memory_replica() == false) { - while (offset < total_size) { - auto chunk_size = std::min(total_size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - } else { - for (auto &handle : - replica.get_memory_descriptor().buffer_descriptors) { - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, handle.size_}); - offset += handle.size_; - } - } + allocateSlices(slices, replica, buffer); // Step 3: Read data directly into user buffer auto get_result = client_->Get(key, query_result.value(), slices); @@ -1064,22 +1048,7 @@ std::vector> PyClient::batch_get_into_internal( // Create slices for this key's buffer std::vector key_slices; - uint64_t offset = 0; - if (replica.is_memory_replica() == false) { - while (offset < total_size) { - auto chunk_size = std::min(total_size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffers[i]) + offset; - key_slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - } else { - for (auto &handle : - replica.get_memory_descriptor().buffer_descriptors) { - void *chunk_ptr = static_cast(buffers[i]) + offset; - key_slices.emplace_back(Slice{chunk_ptr, handle.size_}); - offset += handle.size_; - } - } + allocateSlices(key_slices, replica, buffers[i]); // Store operation info for batch processing valid_operations.push_back( diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index 442269c2..f9b5caf3 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -86,13 +86,11 @@ void WrappedMasterService::init_http_server() { if (replicas[i].is_memory_replica()) { auto& memory_descriptors = replicas[i].get_memory_descriptor(); - for (const auto& handle : - memory_descriptors.buffer_descriptors) { - std::string tmp = ""; - struct_json::to_json(handle, tmp); - ss += tmp; - ss += "\n"; - } + std::string tmp = ""; + struct_json::to_json( + memory_descriptors.buffer_descriptor, tmp); + ss += tmp; + ss += "\n"; } } resp.set_status_and_content(status_type::ok, std::move(ss)); @@ -292,17 +290,17 @@ WrappedMasterService::BatchGetReplicaList( tl::expected, ErrorCode> WrappedMasterService::PutStart(const UUID& client_id, const std::string& key, - const std::vector& slice_lengths, + const uint64_t slice_length, const ReplicateConfig& config) { return execute_rpc( "PutStart", [&] { - return master_service_.PutStart(client_id, key, slice_lengths, + return master_service_.PutStart(client_id, key, slice_length, config); }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, - ", slice_lengths=", slice_lengths.size()); + ", slice_length=", slice_length); }, [&] { MasterMetricManager::instance().inc_put_start_requests(); }, [] { MasterMetricManager::instance().inc_put_start_failures(); }); @@ -335,10 +333,10 @@ tl::expected WrappedMasterService::PutRevoke( } std::vector, ErrorCode>> -WrappedMasterService::BatchPutStart( - const UUID& client_id, const std::vector& keys, - const std::vector>& slice_lengths, - const ReplicateConfig& config) { +WrappedMasterService::BatchPutStart(const UUID& client_id, + const std::vector& keys, + const std::vector& slice_lengths, + const ReplicateConfig& config) { ScopedVLogTimer timer(1, "BatchPutStart"); const size_t total_keys = keys.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); @@ -351,24 +349,17 @@ WrappedMasterService::BatchPutStart( if (config.prefer_alloc_in_same_node) { ReplicateConfig new_config = config; for (size_t i = 0; i < keys.size(); ++i) { - auto& slice_lens = slice_lengths[i]; - std::vector alloc_slice_lens; - size_t all_slice_len = 0; - for (auto& slice_len : slice_lens) { - all_slice_len += slice_len; - } - alloc_slice_lens.emplace_back(all_slice_len); auto result = master_service_.PutStart( - client_id, keys[i], alloc_slice_lens, new_config); + client_id, keys[i], slice_lengths[i], new_config); results.emplace_back(result); if ((i == 0) && result.has_value()) { std::string preferred_segment; for (const auto& replica : result.value()) { if (replica.is_memory_replica()) { auto handles = - replica.get_memory_descriptor().buffer_descriptors; - if (!handles.empty()) { - preferred_segment = handles[0].transport_endpoint_; + replica.get_memory_descriptor().buffer_descriptor; + if (!handles.transport_endpoint_.empty()) { + preferred_segment = handles.transport_endpoint_; } } } @@ -640,4 +631,4 @@ void RegisterRpcService( &wrapped_master_service); } -} // namespace mooncake +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/src/transfer_task.cpp b/mooncake-store/src/transfer_task.cpp index 158487d5..0299bbbe 100644 --- a/mooncake-store/src/transfer_task.cpp +++ b/mooncake-store/src/transfer_task.cpp @@ -392,23 +392,21 @@ std::optional TransferSubmitter::submit( std::optional future; if (replica.is_memory_replica()) { - std::vector handles; auto& mem_desc = replica.get_memory_descriptor(); - handles = mem_desc.buffer_descriptors; + auto& handle = mem_desc.buffer_descriptor; - if (!validateTransferParams(handles, slices)) { + if (!validateTransferParams(handle, slices)) { return std::nullopt; } - TransferStrategy strategy = selectStrategy(handles, slices); + TransferStrategy strategy = selectStrategy(handle, slices); switch (strategy) { case TransferStrategy::LOCAL_MEMCPY: - future = submitMemcpyOperation(handles, slices, op_code); + future = submitMemcpyOperation(handle, slices, op_code); break; case TransferStrategy::TRANSFER_ENGINE: - future = - submitTransferEngineOperation(handles, slices, op_code); + future = submitTransferEngineOperation(handle, slices, op_code); break; default: LOG(ERROR) << "Unknown transfer strategy: " << strategy; @@ -436,11 +434,10 @@ std::optional TransferSubmitter::submit_batch( auto& replica = replicas[i]; auto& slices = all_slices[i]; auto& mem_desc = replica.get_memory_descriptor(); - if (!validateTransferParams(mem_desc.buffer_descriptors, slices, - true)) { + if (!validateTransferParams(mem_desc.buffer_descriptor, slices)) { return std::nullopt; } - auto handle = mem_desc.buffer_descriptors[0]; + auto& handle = mem_desc.buffer_descriptor; uint64_t offset = 0; SegmentHandle seg = engine_.openSegment(handle.transport_endpoint_); if (seg == static_cast(ERR_INVALID_ARGUMENT)) { @@ -470,16 +467,17 @@ std::optional TransferSubmitter::submit_batch( } std::optional TransferSubmitter::submitMemcpyOperation( - const std::vector& handles, - std::vector& slices, TransferRequest::OpCode op_code) { + const AllocatedBuffer::Descriptor& handle, const std::vector& slices, + const TransferRequest::OpCode op_code) { auto state = std::make_shared(); // Create memcpy operations std::vector operations; - operations.reserve(handles.size()); + operations.reserve(slices.size()); + uint64_t base_address = static_cast(handle.buffer_address_); + uint64_t offset = 0; - for (size_t i = 0; i < handles.size(); ++i) { - const auto& handle = handles[i]; + for (size_t i = 0; i < slices.size(); ++i) { const auto& slice = slices[i]; if (slice.ptr == nullptr) continue; @@ -491,23 +489,24 @@ std::optional TransferSubmitter::submitMemcpyOperation( // READ: from handle (remote buffer) to slice (local // buffer) dest = slice.ptr; - src = reinterpret_cast(handle.buffer_address_); + src = reinterpret_cast(base_address + offset); } else { // WRITE: from slice (local buffer) to handle (remote // buffer) - dest = reinterpret_cast(handle.buffer_address_); + dest = reinterpret_cast(base_address + offset); src = slice.ptr; } + offset += slice.size; - operations.emplace_back(dest, src, handle.size_); + operations.emplace_back(dest, src, slice.size); } // Submit memcpy operations to worker pool for async execution MemcpyTask task(std::move(operations), state); memcpy_pool_->submitTask(std::move(task)); - VLOG(1) << "Memcpy transfer submitted to worker pool with " - << handles.size() << " operations"; + VLOG(1) << "Memcpy transfer submitted to worker pool with " << slices.size() + << " operations"; return TransferFuture(state); } @@ -548,39 +547,39 @@ std::optional TransferSubmitter::submitTransfer( } std::optional TransferSubmitter::submitTransferEngineOperation( - const std::vector& handles, - std::vector& slices, TransferRequest::OpCode op_code) { + const AllocatedBuffer::Descriptor& handle, const std::vector& slices, + const TransferRequest::OpCode op_code) { + if (handle.transport_endpoint_.empty()) { + LOG(ERROR) << "Transport endpoint is empty for handle with address " + << handle.buffer_address_; + return std::nullopt; + } + SegmentHandle seg = engine_.openSegment(handle.transport_endpoint_); + + if (seg == static_cast(ERR_INVALID_ARGUMENT)) { + LOG(ERROR) << "Failed to open segment for endpoint='" + << handle.transport_endpoint_ << "'"; + return std::nullopt; + } + // Create transfer requests std::vector requests; - requests.reserve(handles.size()); + requests.reserve(slices.size()); + uint64_t base_address = static_cast(handle.buffer_address_); + uint64_t offset = 0; - for (size_t i = 0; i < handles.size(); ++i) { - const auto& handle = handles[i]; + for (size_t i = 0; i < slices.size(); ++i) { const auto& slice = slices[i]; - if (slice.ptr == nullptr) continue; - if (handle.transport_endpoint_.empty()) { - LOG(ERROR) << "Transport endpoint is empty for handle with address " - << handle.buffer_address_; - return std::nullopt; - } - - SegmentHandle seg = engine_.openSegment(handle.transport_endpoint_); - - if (seg == static_cast(ERR_INVALID_ARGUMENT)) { - LOG(ERROR) << "Failed to open segment for endpoint='" - << handle.transport_endpoint_ << "'"; - return std::nullopt; - } - TransferRequest request; request.opcode = op_code; request.source = static_cast(slice.ptr); request.target_id = seg; - request.target_offset = handle.buffer_address_; - request.length = handle.size_; + request.target_offset = base_address + offset; + request.length = slice.size; + offset += slice.size; requests.emplace_back(request); } return submitTransfer(requests); @@ -604,7 +603,7 @@ std::optional TransferSubmitter::submitFileReadOperation( } TransferStrategy TransferSubmitter::selectStrategy( - const std::vector& handles, + const AllocatedBuffer::Descriptor& handle, const std::vector& slices) const { // Check if memcpy operations are enabled via environment variable if (!memcpy_enabled_) { @@ -614,7 +613,7 @@ TransferStrategy TransferSubmitter::selectStrategy( } // Check conditions for local memcpy optimization - if (isLocalTransfer(handles)) { + if (isLocalTransfer(handle)) { return TransferStrategy::LOCAL_MEMCPY; } @@ -622,15 +621,12 @@ TransferStrategy TransferSubmitter::selectStrategy( } bool TransferSubmitter::isLocalTransfer( - const std::vector& handles) const { + const AllocatedBuffer::Descriptor& handle) const { std::string local_ep = engine_.getLocalIpAndPort(); if (!local_ep.empty()) { - return std::all_of(handles.begin(), handles.end(), - [&local_ep](const auto& h) { - return !h.transport_endpoint_.empty() && - h.transport_endpoint_ == local_ep; - }); + return !handle.transport_endpoint_.empty() && + handle.transport_endpoint_ == local_ep; } // Without a local endpoint we cannot prove locality; disable memcpy. @@ -638,39 +634,17 @@ bool TransferSubmitter::isLocalTransfer( } bool TransferSubmitter::validateTransferParams( - const std::vector& handles, - const std::vector& slices, bool is_multi_buffers) const { - if (handles.empty()) { - LOG(ERROR) << "handles is empty"; - return false; + const AllocatedBuffer::Descriptor& handle, + const std::vector& slices) const { + uint64_t all_slice_len = 0; + for (auto slice : slices) { + all_slice_len += slice.size; } - - if (handles.size() > slices.size()) { - LOG(ERROR) << "invalid_partition_count handles_size=" << handles.size() - << " slices_size=" << slices.size(); + if (handle.size_ != all_slice_len) { + LOG(ERROR) << "handles len:" << handle.size_ + << ", all_slice_len:" << all_slice_len; return false; } - if (is_multi_buffers) { - uint64_t all_slice_len = 0; - for (auto slice : slices) { - all_slice_len += slice.size; - } - if (handles[0].size_ != all_slice_len) { - LOG(ERROR) << "handles len:" << handles[0].size_ - << ", all_slice_len:" << all_slice_len; - return false; - } - } else { - for (size_t i = 0; i < handles.size(); ++i) { - if (handles[i].size_ != slices[i].size) { - LOG(ERROR) << "Size of replica partition " << i << " (" - << handles[i].size_ - << ") does not match provided buffer (" - << slices[i].size << ")"; - return false; - } - } - } return true; } diff --git a/mooncake-store/tests/allocation_strategy_test.cpp b/mooncake-store/tests/allocation_strategy_test.cpp index 679ca806..778dcb8f 100644 --- a/mooncake-store/tests/allocation_strategy_test.cpp +++ b/mooncake-store/tests/allocation_strategy_test.cpp @@ -3,8 +3,10 @@ #include #include +#include #include #include +#include #include #include "allocator.h" @@ -72,33 +74,6 @@ INSTANTIATE_TEST_SUITE_P( } }); -// Unit test class for testing individual functions -class AllocationStrategyUnitTest : public ::testing::Test { - protected: - void SetUp() override { - strategy_ = std::make_unique(); - } - - // Helper function to create test allocators - std::shared_ptr CreateTestAllocator( - const std::string& segment_name, size_t base_offset, - BufferAllocatorType type, size_t size = 64 * MB) { - const size_t base = 0x100000000ULL + base_offset; // 4GB + offset - switch (type) { - case BufferAllocatorType::CACHELIB: - return std::make_shared( - segment_name, base, size, segment_name); - case BufferAllocatorType::OFFSET: - return std::make_shared( - segment_name, base, size, segment_name); - default: - throw std::invalid_argument("Invalid allocator type"); - } - } - - std::unique_ptr strategy_; -}; - // Test basic functionality with empty allocators map (non-parameterized) TEST_F(AllocationStrategyTest, EmptyAllocatorsMap) { std::unordered_map> empty_allocators; ReplicateConfig config{1, false, "local"}; - std::vector slice_sizes = {100}; + size_t slice_length = 100; auto result = strategy_->Allocate( - empty_allocators, empty_allocators_by_name, slice_sizes, config); + empty_allocators, empty_allocators_by_name, slice_length, config); EXPECT_FALSE(result.has_value()); EXPECT_EQ(result.error(), ErrorCode::NO_AVAILABLE_HANDLE); } @@ -122,9 +97,9 @@ TEST_F(AllocationStrategyTest, PreferredSegmentWithEmptyAllocators) { std::vector> empty_allocators; ReplicateConfig config{1, false, "preferred_segment"}; - std::vector slice_sizes = {100}; + size_t slice_length = 100; auto result = strategy_->Allocate( - empty_allocators, empty_allocators_by_name, slice_sizes, config); + empty_allocators, empty_allocators_by_name, slice_length, config); EXPECT_FALSE(result.has_value()); EXPECT_EQ(result.error(), ErrorCode::NO_AVAILABLE_HANDLE); } @@ -145,10 +120,10 @@ TEST_P(AllocationStrategyParameterizedTest, PreferredSegmentAllocation) { allocators.push_back(allocator2); ReplicateConfig config{1, false, "preferred"}; - std::vector slice_sizes = {1024}; + size_t slice_length = 1024; auto result = strategy_->Allocate(allocators, allocators_by_name, - slice_sizes, config); + slice_length, config); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 1); ASSERT_FALSE(result.value().empty()); @@ -157,9 +132,8 @@ TEST_P(AllocationStrategyParameterizedTest, PreferredSegmentAllocation) { auto descriptor = replica.get_descriptor(); ASSERT_TRUE(descriptor.is_memory_replica()); const auto& mem_desc = descriptor.get_memory_descriptor(); - ASSERT_EQ(mem_desc.buffer_descriptors.size(), 1); - EXPECT_EQ(mem_desc.buffer_descriptors[0].transport_endpoint_, "preferred"); - EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 1024); + EXPECT_EQ(mem_desc.buffer_descriptor.transport_endpoint_, "preferred"); + EXPECT_EQ(mem_desc.buffer_descriptor.size_, 1024); } // Test fallback to random allocation when preferred segment doesn't exist @@ -178,10 +152,10 @@ TEST_P(AllocationStrategyParameterizedTest, PreferredSegmentNotFound) { allocators.push_back(allocator2); ReplicateConfig config{1, false, "nonexistent"}; - std::vector slice_sizes = {1024}; + size_t slice_length = 1024; auto result = strategy_->Allocate(allocators, allocators_by_name, - slice_sizes, config); + slice_length, config); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 1); @@ -189,14 +163,13 @@ TEST_P(AllocationStrategyParameterizedTest, PreferredSegmentNotFound) { auto descriptor = replica.get_descriptor(); ASSERT_TRUE(descriptor.is_memory_replica()); const auto& mem_desc = descriptor.get_memory_descriptor(); - ASSERT_EQ(mem_desc.buffer_descriptors.size(), 1); - std::string segment_ep = mem_desc.buffer_descriptors[0].transport_endpoint_; + std::string segment_ep = mem_desc.buffer_descriptor.transport_endpoint_; EXPECT_TRUE(segment_ep == "segment1" || segment_ep == "segment2"); - EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 1024); + EXPECT_EQ(mem_desc.buffer_descriptor.size_, 1024); } -// Test multiple slices allocation -TEST_P(AllocationStrategyParameterizedTest, MultipleSlicesAllocation) { +// Test single slice allocation +TEST_P(AllocationStrategyParameterizedTest, SingleSliceAllocation) { auto allocator1 = CreateTestAllocator("segment1", 0); auto allocator2 = CreateTestAllocator("segment2", 0x10000000ULL); @@ -211,10 +184,10 @@ TEST_P(AllocationStrategyParameterizedTest, MultipleSlicesAllocation) { allocators.push_back(allocator2); ReplicateConfig config{1, false, ""}; - std::vector slice_sizes = {1024, 2048, 512}; + size_t slice_length = 1024; auto result = strategy_->Allocate(allocators, allocators_by_name, - slice_sizes, config); + slice_length, config); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 1); @@ -222,10 +195,7 @@ TEST_P(AllocationStrategyParameterizedTest, MultipleSlicesAllocation) { auto descriptor = replica.get_descriptor(); ASSERT_TRUE(descriptor.is_memory_replica()); const auto& mem_desc = descriptor.get_memory_descriptor(); - ASSERT_EQ(mem_desc.buffer_descriptors.size(), 3); - EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 1024); - EXPECT_EQ(mem_desc.buffer_descriptors[1].size_, 2048); - EXPECT_EQ(mem_desc.buffer_descriptors[2].size_, 512); + EXPECT_EQ(mem_desc.buffer_descriptor.size_, 1024); } // Test multiple replicas allocation @@ -247,21 +217,19 @@ TEST_P(AllocationStrategyParameterizedTest, MultipleReplicasAllocation) { allocators.push_back(allocator3); ReplicateConfig config{3, false, ""}; // Request 3 replicas - std::vector slice_sizes = {1024, 2048}; + size_t slice_length = 1024; auto result = strategy_->Allocate(allocators, allocators_by_name, - slice_sizes, config); + slice_length, config); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 3); - // Check each replica has all slices + // Check each replica has the correct slice size for (const auto& replica : result.value()) { auto descriptor = replica.get_descriptor(); ASSERT_TRUE(descriptor.is_memory_replica()); const auto& mem_desc = descriptor.get_memory_descriptor(); - ASSERT_EQ(mem_desc.buffer_descriptors.size(), 2); - EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 1024); - EXPECT_EQ(mem_desc.buffer_descriptors[1].size_, 2048); + EXPECT_EQ(mem_desc.buffer_descriptor.size_, 1024); } // Check that replicas are on different segments @@ -293,31 +261,33 @@ TEST_P(AllocationStrategyParameterizedTest, PreferredSegmentInsufficientSpace) { // First, fill up the preferred allocator ReplicateConfig config{1, false, "preferred"}; - std::vector large_slices = {10 * 1024 * 1024, 10 * 1024 * 1024, - 10 * 1024 * 1024, 10 * 1024 * 1024, - 10 * 1024 * 1024, 10 * 1024 * 1024, - 3 * 1024 * 1024}; // 63MB out of 64MB - - auto large_result = strategy_->Allocate(allocators, allocators_by_name, - large_slices, config); - ASSERT_TRUE(large_result.has_value()); - auto large_desc = large_result.value()[0].get_descriptor(); - ASSERT_TRUE(large_desc.is_memory_replica()); - EXPECT_EQ(large_desc.get_memory_descriptor() - .buffer_descriptors[0] - .transport_endpoint_, - "preferred"); + // Store the results of the allocations to avoid deallocation of the buffers + // before the test is done + std::vector> results; + // Allocate multiple times to fill up the preferred allocator + for (int i = 0; i < 4; ++i) { + size_t large_slice = 15 * 1024 * 1024; // 10MB + auto large_result = strategy_->Allocate(allocators, allocators_by_name, + large_slice, config); + ASSERT_TRUE(large_result.has_value()); + auto last_desc = large_result.value()[0].get_descriptor(); + ASSERT_TRUE(last_desc.is_memory_replica()); + EXPECT_EQ(last_desc.get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + "preferred"); + results.emplace_back(std::move(large_result.value())); + } // Now try to allocate more than remaining space in preferred segment - std::vector small_slice = {2 * 1024 * 1024}; + size_t small_slice = 5 * 1024 * 1024; // 2MB auto result = strategy_->Allocate(allocators, allocators_by_name, small_slice, config); ASSERT_TRUE(result.has_value()); auto small_desc = result.value()[0].get_descriptor(); ASSERT_TRUE(small_desc.is_memory_replica()); const auto& mem_desc = small_desc.get_memory_descriptor(); - EXPECT_EQ(mem_desc.buffer_descriptors[0].transport_endpoint_, "segment1"); - EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 2 * 1024 * 1024); + EXPECT_EQ(mem_desc.buffer_descriptor.transport_endpoint_, "segment1"); + EXPECT_EQ(mem_desc.buffer_descriptor.size_, small_slice); } // Test allocation when all allocators are full @@ -338,19 +308,20 @@ TEST_P(AllocationStrategyParameterizedTest, AllAllocatorsFull) { ReplicateConfig config{1, false, ""}; // Fill up both allocators - std::vector large_slices = {15 * 1024 * 1024, 15 * 1024 * 1024, - 15 * 1024 * 1024, - 15 * 1024 * 1024}; // 60MB - auto result1 = strategy_->Allocate(allocators, allocators_by_name, - large_slices, config); - ASSERT_TRUE(result1.has_value()); - auto result2 = strategy_->Allocate(allocators, allocators_by_name, - large_slices, config); - ASSERT_TRUE(result2.has_value()); + size_t large_slice = 15 * 1024 * 1024; // 15MB + // Store the results of the allocations to avoid deallocation of the buffers + // before the test is done + std::vector> results; + // Allocate 8 times to use 120MB total + for (int i = 0; i < 8; ++i) { + auto result = strategy_->Allocate(allocators, allocators_by_name, + large_slice, config); + ASSERT_TRUE(result.has_value()); + results.emplace_back(std::move(result.value())); + } // Try to allocate more than remaining space - std::vector impossible_slice = {5 * 1024 * - 1024}; // 5MB (more than remaining) + size_t impossible_slice = 5 * 1024 * 1024; // 5MB (more than remaining) auto result = strategy_->Allocate(allocators, allocators_by_name, impossible_slice, config); EXPECT_FALSE(result.has_value()); @@ -369,7 +340,7 @@ TEST_P(AllocationStrategyParameterizedTest, ZeroSizeAllocation) { allocators.push_back(allocator); ReplicateConfig config{1, false, ""}; - std::vector zero_slice = {0}; + size_t zero_slice = 0; auto result = strategy_->Allocate(allocators, allocators_by_name, zero_slice, config); @@ -389,8 +360,7 @@ TEST_P(AllocationStrategyParameterizedTest, VeryLargeSizeAllocation) { allocators.push_back(allocator); ReplicateConfig config{1, false, ""}; - std::vector huge_slice = { - 100 * 1024 * 1024}; // 100MB (larger than 64MB capacity) + size_t huge_slice = 100 * 1024 * 1024; // 100MB (larger than 64MB capacity) auto result = strategy_->Allocate(allocators, allocators_by_name, huge_slice, config); @@ -398,26 +368,7 @@ TEST_P(AllocationStrategyParameterizedTest, VeryLargeSizeAllocation) { EXPECT_EQ(result.error(), ErrorCode::NO_AVAILABLE_HANDLE); } -// Test empty slice sizes -TEST_F(AllocationStrategyTest, EmptySliceSizes) { - auto allocator = std::make_shared( - "segment1", 0x100000000ULL, 64 * MB, "segment1"); - std::unordered_map>> - allocators_by_name; - std::vector> allocators; - - allocators_by_name["segment1"].push_back(allocator); - allocators.push_back(allocator); - - ReplicateConfig config{1, false, ""}; - std::vector empty_slices; - - auto result = strategy_->Allocate(allocators, allocators_by_name, - empty_slices, config); - EXPECT_FALSE(result.has_value()); - EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); -} +// Test zero slice length (already covered by ZeroSizeAllocation test) // Test invalid replication count TEST_F(AllocationStrategyTest, InvalidReplicationCount) { @@ -432,10 +383,10 @@ TEST_F(AllocationStrategyTest, InvalidReplicationCount) { allocators.push_back(allocator); ReplicateConfig config{0, false, ""}; // Invalid: 0 replicas - std::vector slice_sizes = {1024}; + size_t slice_length = 1024; auto result = strategy_->Allocate(allocators, allocators_by_name, - slice_sizes, config); + slice_length, config); EXPECT_FALSE(result.has_value()); EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); } @@ -460,10 +411,10 @@ TEST_F(AllocationStrategyTest, InsufficientAllocatorsForReplicas) { ReplicateConfig config{ 5, false, ""}; // Request 5 replicas, but only 2 segments available - std::vector slice_sizes = {1024}; + size_t slice_length = 1024; auto result = strategy_->Allocate(allocators, allocators_by_name, - slice_sizes, config); + slice_length, config); // With best-effort semantics, should succeed with available replicas EXPECT_TRUE(result.has_value()); // Should get 2 replicas (limited by number of segments) @@ -474,8 +425,7 @@ TEST_F(AllocationStrategyTest, InsufficientAllocatorsForReplicas) { auto descriptor = replica.get_descriptor(); ASSERT_TRUE(descriptor.is_memory_replica()); const auto& mem_desc = descriptor.get_memory_descriptor(); - ASSERT_EQ(mem_desc.buffer_descriptors.size(), 1u); - EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 1024u); + EXPECT_EQ(mem_desc.buffer_descriptor.size_, 1024u); } // Verify replicas are on different segments @@ -483,216 +433,15 @@ TEST_F(AllocationStrategyTest, InsufficientAllocatorsForReplicas) { for (const auto& replica : result.value()) { auto descriptor = replica.get_descriptor(); const auto& mem_desc = descriptor.get_memory_descriptor(); - segment_names.insert( - mem_desc.buffer_descriptors[0].transport_endpoint_); + segment_names.insert(mem_desc.buffer_descriptor.transport_endpoint_); } EXPECT_EQ(2u, segment_names.size()); } -TEST_F(AllocationStrategyUnitTest, - AllocateSingleBuffer_PreferredSegmentNotFound) { - auto allocator1 = - CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET); - auto allocator2 = CreateTestAllocator("segment2", 0x10000000ULL, - BufferAllocatorType::OFFSET); - - std::vector> allocators = {allocator1, - allocator2}; - std::unordered_map>> - allocators_by_name; - allocators_by_name["segment1"] = {allocator1}; - allocators_by_name["segment2"] = {allocator2}; - - ReplicateConfig config{1, false, "nonexistent"}; - std::unordered_set excluded_segments; - - auto buffer = strategy_->allocateSingleBuffer( - allocators, allocators_by_name, 1024, config, excluded_segments); - - ASSERT_TRUE(buffer != nullptr); - std::string segment_name = buffer->getSegmentName(); - EXPECT_TRUE(segment_name == "segment1" || segment_name == "segment2"); -} - -TEST_F(AllocationStrategyUnitTest, AllocateSingleBuffer_EmptyPreferredSegment) { - auto allocator1 = - CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET); - - std::vector> allocators = {allocator1}; - std::unordered_map>> - allocators_by_name; - allocators_by_name["segment1"] = {allocator1}; - - ReplicateConfig config{1, false, ""}; // Empty preferred segment - std::unordered_set excluded_segments; - - auto buffer = strategy_->allocateSingleBuffer( - allocators, allocators_by_name, 1024, config, excluded_segments); - - ASSERT_TRUE(buffer != nullptr); - EXPECT_EQ(buffer->getSegmentName(), "segment1"); -} - -// Test tryRandomAllocate function -TEST_F(AllocationStrategyUnitTest, TryRandomAllocate_Success) { - auto allocator1 = - CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET); - auto allocator2 = CreateTestAllocator("segment2", 0x10000000ULL, - BufferAllocatorType::OFFSET); - - std::vector> allocators = {allocator1, - allocator2}; - std::unordered_set excluded_segments; - - auto buffer = - strategy_->tryRandomAllocate(allocators, 1024, excluded_segments); - ASSERT_TRUE(buffer != nullptr); - EXPECT_EQ(buffer->size(), 1024); -} - -TEST_F(AllocationStrategyUnitTest, TryRandomAllocate_AllSegmentsExcluded) { - auto allocator1 = - CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET); - auto allocator2 = CreateTestAllocator("segment2", 0x10000000ULL, - BufferAllocatorType::OFFSET); - - std::vector> allocators = {allocator1, - allocator2}; - std::unordered_set excluded_segments = {"segment1", - "segment2"}; - - auto buffer = - strategy_->tryRandomAllocate(allocators, 1024, excluded_segments); - EXPECT_TRUE(buffer == nullptr); -} - -TEST_F(AllocationStrategyUnitTest, TryRandomAllocate_InsufficientSpace) { - auto allocator = CreateTestAllocator( - "segment1", 0, BufferAllocatorType::OFFSET, 1024); // Only 1KB - std::vector> allocators = {allocator}; - std::unordered_set excluded_segments; - - auto buffer = strategy_->tryRandomAllocate( - allocators, 2048, excluded_segments); // Request 2KB - EXPECT_TRUE(buffer == nullptr); -} - -// Test allocateSlice function -TEST_F(AllocationStrategyUnitTest, AllocateSlice_SingleReplica) { - auto allocator1 = - CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET); - - std::vector> allocators = {allocator1}; - std::unordered_map>> - allocators_by_name; - allocators_by_name["segment1"] = {allocator1}; - - ReplicateConfig config{1, false, ""}; - auto buffers = strategy_->allocateSlice(allocators, allocators_by_name, - 1024, 1, config); - - ASSERT_EQ(buffers.size(), 1); - EXPECT_EQ(buffers[0]->size(), 1024); - EXPECT_EQ(buffers[0]->getSegmentName(), "segment1"); -} - -TEST_F(AllocationStrategyUnitTest, AllocateSlice_MultipleReplicas) { - auto allocator1 = - CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET); - auto allocator2 = CreateTestAllocator("segment2", 0x10000000ULL, - BufferAllocatorType::OFFSET); - auto allocator3 = CreateTestAllocator("segment3", 0x20000000ULL, - BufferAllocatorType::OFFSET); - - std::vector> allocators = { - allocator1, allocator2, allocator3}; - std::unordered_map>> - allocators_by_name; - allocators_by_name["segment1"] = {allocator1}; - allocators_by_name["segment2"] = {allocator2}; - allocators_by_name["segment3"] = {allocator3}; - - ReplicateConfig config{3, false, ""}; - auto buffers = strategy_->allocateSlice(allocators, allocators_by_name, - 1024, 3, config); - - ASSERT_EQ(buffers.size(), 3); - - // Verify all buffers have correct size - for (const auto& buffer : buffers) { - EXPECT_EQ(buffer->size(), 1024); - } - - // Verify replicas are on different segments - std::unordered_set used_segments; - for (const auto& buffer : buffers) { - used_segments.insert(buffer->getSegmentName()); - } - EXPECT_EQ(used_segments.size(), 3); -} - -TEST_F(AllocationStrategyUnitTest, AllocateSlice_InsufficientAllocators) { - auto allocator1 = - CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET); - - std::vector> allocators = {allocator1}; - std::unordered_map>> - allocators_by_name; - allocators_by_name["segment1"] = {allocator1}; - - ReplicateConfig config{3, false, - ""}; // Request 3 replicas but only 1 allocator - auto buffers = strategy_->allocateSlice(allocators, allocators_by_name, - 1024, 3, config); - - // Should allocate as many as possible (best-effort) - ASSERT_EQ(buffers.size(), 1); - EXPECT_EQ(buffers[0]->getSegmentName(), "segment1"); -} - -// Test getLargestFreeRegion() filtering logic with fragmented allocators -TEST_F(AllocationStrategyUnitTest, - TryRandomAllocate_LargestFreeRegionFiltering) { - // Run the test 10 times to account for randomness - for (int run = 0; run < 10; ++run) { - // Create two OffsetBufferAllocators with 10MB each - auto allocator1 = CreateTestAllocator( - "segment1", 0, BufferAllocatorType::OFFSET, 10 * MB); - auto allocator2 = CreateTestAllocator( - "segment2", 0x10000000ULL, BufferAllocatorType::OFFSET, 10 * MB); - - // Fragment allocator1 heavily - leave only small free regions - std::vector> fragments1; - for (int i = 0; i < 9; ++i) { - fragments1.push_back(allocator1->allocate(1 * MB)); - } - // allocator1: 9MB allocated, only 1MB free - - // Leave allocator2 with enough contiguous space - auto fragment2 = allocator2->allocate(5 * MB); - // allocator2: 5MB allocated, 5MB contiguous free - - std::vector> allocators = { - allocator1, allocator2}; - std::unordered_set excluded_segments; - - // Reset retry counter before test - strategy_->resetRetryCount(); - - auto buffer = - strategy_->tryRandomAllocate(allocators, 4 * MB, excluded_segments); - - ASSERT_TRUE(buffer != nullptr) << "Failed on run " << run; - EXPECT_EQ(buffer->size(), 4 * MB) << "Failed on run " << run; - EXPECT_EQ(buffer->getSegmentName(), "segment2") - << "Failed on run " << run; - EXPECT_EQ(strategy_->getRetryCount(), 0) << "Failed on run " << run; - } -} +// Note: The following unit tests for internal helper methods have been removed +// because those methods (allocateSingleBuffer, tryRandomAllocate, +// allocateSlice, resetRetryCount, getRetryCount) are no longer part of the +// public API. The functionality is now encapsulated within the Allocate() +// method. } // namespace mooncake diff --git a/mooncake-store/tests/client_buffer_test.cpp b/mooncake-store/tests/client_buffer_test.cpp index 6cdd25f2..818c15d5 100644 --- a/mooncake-store/tests/client_buffer_test.cpp +++ b/mooncake-store/tests/client_buffer_test.cpp @@ -269,25 +269,15 @@ TEST_F(ClientBufferTest, CalculateTotalSizeMemoryReplica) { Replica::Descriptor replica; MemoryDescriptor mem_desc; - // Add some buffer descriptors with proper initialization - AllocatedBuffer::Descriptor buf1; - buf1.size_ = 1024; - buf1.buffer_address_ = 0x1000; + // Set buffer descriptor with proper initialization + mem_desc.buffer_descriptor.size_ = 4096; + mem_desc.buffer_descriptor.buffer_address_ = 0x1000; - AllocatedBuffer::Descriptor buf2; - buf2.size_ = 2048; - buf2.buffer_address_ = 0x2000; - - AllocatedBuffer::Descriptor buf3; - buf3.size_ = 512; - buf3.buffer_address_ = 0x3000; - - mem_desc.buffer_descriptors = {buf1, buf2, buf3}; replica.descriptor_variant = mem_desc; replica.status = ReplicaStatus::COMPLETE; uint64_t total_size = calculate_total_size(replica); - EXPECT_EQ(total_size, 1024 + 2048 + 512); + EXPECT_EQ(total_size, 4096); } // Test calculate_total_size function with disk replica @@ -304,12 +294,13 @@ TEST_F(ClientBufferTest, CalculateTotalSizeDiskReplica) { EXPECT_EQ(total_size, 4096); } -// Test calculate_total_size function with empty memory replica -TEST_F(ClientBufferTest, CalculateTotalSizeEmptyMemoryReplica) { - // Create an empty memory replica descriptor +// Test calculate_total_size function with zero-size memory replica +TEST_F(ClientBufferTest, CalculateTotalSizeZeroSizeMemoryReplica) { + // Create a memory replica descriptor with zero size Replica::Descriptor replica; MemoryDescriptor mem_desc; - // Empty buffer_descriptors vector + mem_desc.buffer_descriptor.size_ = 0; + mem_desc.buffer_descriptor.buffer_address_ = 0x1000; replica.descriptor_variant = mem_desc; replica.status = ReplicaStatus::COMPLETE; @@ -334,34 +325,23 @@ TEST_F(ClientBufferTest, AllocateSlicesMemoryReplica) { // Create a memory replica descriptor Replica::Descriptor replica; MemoryDescriptor mem_desc; + mem_desc.buffer_descriptor.size_ = 4096; + mem_desc.buffer_descriptor.buffer_address_ = 0x1000; - AllocatedBuffer::Descriptor buf1; - buf1.size_ = 1024; - AllocatedBuffer::Descriptor buf2; - buf2.size_ = 2048; - AllocatedBuffer::Descriptor buf3; - buf3.size_ = 1024; - - mem_desc.buffer_descriptors = {buf1, buf2, buf3}; replica.descriptor_variant = mem_desc; replica.status = ReplicaStatus::COMPLETE; std::vector slices; - int result = allocateSlices(slices, replica, handle); + int result = allocateSlices(slices, replica, handle.ptr()); EXPECT_EQ(result, 0); - EXPECT_EQ(slices.size(), 3); + EXPECT_EQ(slices.size(), 1); - // Verify slice sizes match buffer descriptors - EXPECT_EQ(slices[0].size, 1024); - EXPECT_EQ(slices[1].size, 2048); - EXPECT_EQ(slices[2].size, 1024); + // Verify slice size matches buffer descriptor + EXPECT_EQ(slices[0].size, 4096); - // Verify slices are contiguous - char* base_ptr = static_cast(handle.ptr()); - EXPECT_EQ(slices[0].ptr, base_ptr); - EXPECT_EQ(slices[1].ptr, base_ptr + 1024); - EXPECT_EQ(slices[2].ptr, base_ptr + 1024 + 2048); + // Verify slice pointer matches buffer pointer + EXPECT_EQ(slices[0].ptr, handle.ptr()); } // Test allocateSlices function with disk replica @@ -386,7 +366,7 @@ TEST_F(ClientBufferTest, AllocateSlicesDiskReplica) { replica.status = ReplicaStatus::COMPLETE; std::vector slices; - int result = allocateSlices(slices, replica, handle); + int result = allocateSlices(slices, replica, handle.ptr()); EXPECT_EQ(result, 0); EXPECT_GE(slices.size(), 1); @@ -403,8 +383,8 @@ TEST_F(ClientBufferTest, AllocateSlicesDiskReplica) { EXPECT_EQ(total_slice_size, 8192); } -// Test allocateSlices function with empty memory replica -TEST_F(ClientBufferTest, AllocateSlicesEmptyMemoryReplica) { +// Test allocateSlices function with zero-size memory replica +TEST_F(ClientBufferTest, AllocateSlicesZeroSizeMemoryReplica) { const size_t buffer_size = 1024 * 1024; // 1MB const size_t alloc_size = 1024; // 1KB @@ -416,19 +396,22 @@ TEST_F(ClientBufferTest, AllocateSlicesEmptyMemoryReplica) { BufferHandle handle = std::move(handle_opt.value()); - // Create an empty memory replica descriptor + // Create a memory replica descriptor with zero size Replica::Descriptor replica; MemoryDescriptor mem_desc; - // Empty buffer_descriptors vector + mem_desc.buffer_descriptor.size_ = 0; + mem_desc.buffer_descriptor.buffer_address_ = 0x1000; replica.descriptor_variant = mem_desc; replica.status = ReplicaStatus::COMPLETE; std::vector slices; - int result = allocateSlices(slices, replica, handle); + int result = allocateSlices(slices, replica, handle.ptr()); EXPECT_EQ(result, 0); - EXPECT_EQ(slices.size(), 0); + EXPECT_EQ(slices.size(), 1); + EXPECT_EQ(slices[0].size, 0); + EXPECT_EQ(slices[0].ptr, handle.ptr()); } } // namespace mooncake diff --git a/mooncake-store/tests/client_integration_test.cpp b/mooncake-store/tests/client_integration_test.cpp index e42d50a3..8df3176e 100644 --- a/mooncake-store/tests/client_integration_test.cpp +++ b/mooncake-store/tests/client_integration_test.cpp @@ -314,12 +314,9 @@ TEST_F(ClientIntegrationTest, LocalPreferredAllocationTest) { << "Query operation failed: " << toString(query_result.error()); auto replica_list = query_result.value().replicas; ASSERT_EQ(replica_list.size(), 1); - ASSERT_EQ(replica_list[0].get_memory_descriptor().buffer_descriptors.size(), - 1); ASSERT_EQ(replica_list[0] .get_memory_descriptor() - .buffer_descriptors[0] - .transport_endpoint_, + .buffer_descriptor.transport_endpoint_, segment_provider_client_->GetTransportEndpoint()); auto get_result = test_client_->Get(key, query_result.value(), slices); diff --git a/mooncake-store/tests/e2e/client_wrapper.cpp b/mooncake-store/tests/e2e/client_wrapper.cpp index ce3451bc..74e66637 100644 --- a/mooncake-store/tests/e2e/client_wrapper.cpp +++ b/mooncake-store/tests/e2e/client_wrapper.cpp @@ -106,9 +106,9 @@ ErrorCode ClientTestWrapper::Get(const std::string& key, std::string& value) { } // Create slices - const std::vector& descriptors = - replica_list[0].get_memory_descriptor().buffer_descriptors; - SliceGuard slice_guard(descriptors, allocator_); + const AllocatedBuffer::Descriptor& descriptor = + replica_list[0].get_memory_descriptor().buffer_descriptor; + SliceGuard slice_guard(descriptor.size_, allocator_); // Perform get operation auto get_result = diff --git a/mooncake-store/tests/master_metrics_test.cpp b/mooncake-store/tests/master_metrics_test.cpp index 62952145..f17ad801 100644 --- a/mooncake-store/tests/master_metrics_test.cpp +++ b/mooncake-store/tests/master_metrics_test.cpp @@ -119,7 +119,6 @@ TEST_F(MasterMetricsTest, BasicRequestTest) { std::string key = "test_key"; uint64_t value_length = 1024; - std::vector slice_lengths = {value_length}; ReplicateConfig config; config.replica_num = 1; @@ -138,7 +137,7 @@ TEST_F(MasterMetricsTest, BasicRequestTest) { // Test PutStart and PutRevoke request auto put_start_result1 = - service_.PutStart(client_id, key, slice_lengths, config); + service_.PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result1.has_value()); ASSERT_EQ(metrics.get_key_count(), 1); ASSERT_EQ(metrics.get_allocated_mem_size(), value_length); @@ -157,7 +156,7 @@ TEST_F(MasterMetricsTest, BasicRequestTest) { // Test PutStart and PutEnd request auto put_start_result2 = - service_.PutStart(client_id, key, slice_lengths, config); + service_.PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result2.has_value()); ASSERT_EQ(metrics.get_key_count(), 1); ASSERT_EQ(metrics.get_allocated_mem_size(), value_length); @@ -199,7 +198,7 @@ TEST_F(MasterMetricsTest, BasicRequestTest) { // Test RemoveAll request auto put_start_result3 = - service_.PutStart(client_id, key, slice_lengths, config); + service_.PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result3.has_value()); auto put_end_result2 = service_.PutEnd(client_id, key, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result2.has_value()); @@ -213,7 +212,7 @@ TEST_F(MasterMetricsTest, BasicRequestTest) { // Test UnmountSegment request auto put_start_result4 = - service_.PutStart(client_id, key, slice_lengths, config); + service_.PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result4.has_value()); auto put_end_result3 = service_.PutEnd(client_id, key, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result3.has_value()); @@ -253,7 +252,7 @@ TEST_F(MasterMetricsTest, BatchRequestTest) { UUID client_id = generate_uuid(); std::vector keys = {"test_key1", "test_key2", "test_key3"}; - std::vector> slice_lengths = {{1024}, {2048}, {512}}; + std::vector value_lengths = {1024, 2048, 512}; ReplicateConfig config; config.replica_num = 1; @@ -272,7 +271,7 @@ TEST_F(MasterMetricsTest, BatchRequestTest) { // Test BatchPutStart request auto batch_put_start_result = - service_.BatchPutStart(client_id, keys, slice_lengths, config); + service_.BatchPutStart(client_id, keys, value_lengths, config); ASSERT_EQ(batch_put_start_result.size(), 3); ASSERT_EQ(metrics.get_batch_put_start_requests(), 1); ASSERT_EQ(metrics.get_batch_put_start_partial_successes(), 0); @@ -327,7 +326,7 @@ TEST_F(MasterMetricsTest, BatchRequestTest) { // Test partial success keys.push_back("test_key4"); - slice_lengths.push_back({512}); + value_lengths.push_back(512); auto batch_get_replica_result3 = service_.BatchGetReplicaList(keys); ASSERT_EQ(batch_get_replica_result3.size(), 4); ASSERT_EQ(metrics.get_batch_get_replica_list_requests(), 3); @@ -337,7 +336,7 @@ TEST_F(MasterMetricsTest, BatchRequestTest) { ASSERT_EQ(metrics.get_batch_get_replica_list_failed_items(), 4); auto batch_put_start_result2 = - service_.BatchPutStart(client_id, keys, slice_lengths, config); + service_.BatchPutStart(client_id, keys, value_lengths, config); ASSERT_EQ(batch_put_start_result2.size(), 4); ASSERT_EQ(metrics.get_batch_put_start_requests(), 2); ASSERT_EQ(metrics.get_batch_put_start_partial_successes(), 1); diff --git a/mooncake-store/tests/master_service_ssd_test.cpp b/mooncake-store/tests/master_service_ssd_test.cpp index 16dc24e9..422e1aa1 100644 --- a/mooncake-store/tests/master_service_ssd_test.cpp +++ b/mooncake-store/tests/master_service_ssd_test.cpp @@ -45,12 +45,12 @@ TEST_F(MasterServiceSSDTest, PutEndBothReplica) { ASSERT_TRUE(mount_result.has_value()); std::string key = "disk_key"; - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto replicas = put_start_result.value(); ASSERT_EQ(2, replicas.size()); @@ -100,12 +100,12 @@ TEST_F(MasterServiceSSDTest, PutRevokeDiskReplica) { ASSERT_TRUE(mount_result.has_value()); std::string key = "revoke_key"; - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; ASSERT_TRUE( - service_->PutStart(client_id, key, slice_lengths, config).has_value()); + service_->PutStart(client_id, key, slice_length, config).has_value()); EXPECT_TRUE( service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value()); @@ -141,12 +141,12 @@ TEST_F(MasterServiceSSDTest, PutRevokeMemoryReplica) { ASSERT_TRUE(mount_result.has_value()); std::string key = "revoke_key"; - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; ASSERT_TRUE( - service_->PutStart(client_id, key, slice_lengths, config).has_value()); + service_->PutStart(client_id, key, slice_length, config).has_value()); EXPECT_TRUE( service_->PutRevoke(client_id, key, ReplicaType::MEMORY).has_value()); @@ -180,12 +180,12 @@ TEST_F(MasterServiceSSDTest, PutRevokeBothReplica) { ASSERT_TRUE(mount_result.has_value()); std::string key = "revoke_key"; - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; ASSERT_TRUE( - service_->PutStart(client_id, key, slice_lengths, config).has_value()); + service_->PutStart(client_id, key, slice_length, config).has_value()); EXPECT_TRUE( service_->PutRevoke(client_id, key, ReplicaType::DISK).has_value()); @@ -218,12 +218,12 @@ TEST_F(MasterServiceSSDTest, RemoveKey) { ASSERT_TRUE(mount_result.has_value()); std::string key = "remove_key"; - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; ASSERT_TRUE( - service_->PutStart(client_id, key, slice_lengths, config).has_value()); + service_->PutStart(client_id, key, slice_length, config).has_value()); EXPECT_TRUE( service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value()); EXPECT_TRUE( @@ -260,11 +260,11 @@ TEST_F(MasterServiceSSDTest, EvictObject) { int success_puts = 0; for (int i = 0; i < 1024 * 16 + 50; ++i) { std::string key = "test_key" + std::to_string(i); - std::vector slice_lengths = {object_size}; + uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); if (put_start_result.has_value()) { auto put_end_mem_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -325,7 +325,7 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) { std::string key = "test_key"; uint64_t value_length = 16 * 1024 * 1024; // 16MB - std::vector slice_lengths = {value_length}; + uint64_t slice_length = value_length; ReplicateConfig config; auto test_discard_replica = [&](ReplicaType discard_type) { @@ -335,7 +335,7 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) { // Put key, should success. auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); EXPECT_TRUE(put_start_result.has_value()); auto replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -362,7 +362,7 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) { // Put key again, should fail because the object has had an completed // replica. put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS); diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp index 1706fdc6..91ffb8a6 100644 --- a/mooncake-store/tests/master_service_test.cpp +++ b/mooncake-store/tests/master_service_test.cpp @@ -96,8 +96,7 @@ std::string GenerateKeyForSegment(const UUID& client_id, } if (replica_list[0] .get_memory_descriptor() - .buffer_descriptors[0] - .transport_endpoint_ == segment_name) { + .buffer_descriptor.transport_endpoint_ == segment_name) { return key; } // Clean up failed attempt @@ -314,14 +313,13 @@ TEST_F(MasterServiceTest, PutStartInvalidParams) { // Test invalid replica_num config.replica_num = 0; - auto put_result1 = service_->PutStart(client_id, key, {1024}, config); + auto put_result1 = service_->PutStart(client_id, key, 1024, config); EXPECT_FALSE(put_result1.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result1.error()); - // Test empty slice_lengths + // Test zero slice_length config.replica_num = 1; - std::vector empty_slices; - auto put_result2 = service_->PutStart(client_id, key, empty_slices, config); + auto put_result2 = service_->PutStart(client_id, key, 0, config); EXPECT_FALSE(put_result2.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result2.error()); } @@ -336,12 +334,11 @@ TEST_F(MasterServiceTest, PutStartEndFlow) { // Test PutStart std::string key = "test_key"; uint64_t value_length = 1024; - std::vector slice_lengths = {value_length}; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, value_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_FALSE(replica_list.empty()); @@ -395,7 +392,6 @@ TEST_F(MasterServiceTest, RandomPutStartEndFlow) { // Test PutStart std::string key = "test_key"; uint64_t value_length = 1024; - std::vector slice_lengths = {value_length}; ReplicateConfig config; std::random_device rd; std::mt19937 gen(rd()); @@ -403,7 +399,7 @@ TEST_F(MasterServiceTest, RandomPutStartEndFlow) { int random_number = dis(gen); config.replica_num = random_number; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, value_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_FALSE(replica_list.empty()); @@ -445,11 +441,11 @@ TEST_F(MasterServiceTest, GetReplicaListByRegex) { int times = 10; while (times--) { std::string key = "test_key" + std::to_string(times); - std::vector slice_lengths = {1024}; + uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -470,11 +466,11 @@ TEST_F(MasterServiceTest, GetReplicaListByRegex) { // Helper function to put an object, making the test cleaner void put_object(MasterService& service, const UUID& client_id, const std::string& key) { - std::vector slice_lengths = {1024}; + uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service.PutStart(client_id, key, slice_lengths, config); + service.PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result.has_value()) << "Failed to PutStart for key: " << key; auto put_end_result = service.PutEnd(client_id, key, ReplicaType::MEMORY); @@ -615,11 +611,11 @@ TEST_F(MasterServiceTest, GetReplicaList) { [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_); std::string key = "test_key"; - std::vector slice_lengths = {1024}; + uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); @@ -637,11 +633,11 @@ TEST_F(MasterServiceTest, RemoveObject) { const UUID client_id = generate_uuid(); std::string key = "test_key"; - std::vector slice_lengths = {1024}; + uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); @@ -671,11 +667,11 @@ TEST_F(MasterServiceTest, RandomRemoveObject) { std::uniform_int_distribution<> dis(1, 1000); while (times--) { std::string key = "test_key" + std::to_string(dis(gen)); - std::vector slice_lengths = {1024}; + uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -703,11 +699,11 @@ TEST_F(MasterServiceTest, RemoveByRegex) { int times = 10; while (times--) { std::string key = "test_key" + std::to_string(times); - std::vector slice_lengths = {1024}; + uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -916,11 +912,11 @@ TEST_F(MasterServiceTest, RemoveAll) { int times = 10; while (times--) { std::string key = "test_key" + std::to_string(times); - std::vector slice_lengths = {1024}; + uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -940,7 +936,7 @@ TEST_F(MasterServiceTest, RemoveAll) { } } -TEST_F(MasterServiceTest, MultiSliceMultiReplicaFlow) { +TEST_F(MasterServiceTest, SingleSliceMultiReplicaFlow) { const uint64_t kv_lease_ttl = 50; auto service_config = MasterServiceConfig::builder() .set_default_kv_lease_ttl(kv_lease_ttl) @@ -960,22 +956,7 @@ TEST_F(MasterServiceTest, MultiSliceMultiReplicaFlow) { // Test parameters std::string key = "multi_slice_object"; constexpr size_t num_replicas = 3; - constexpr size_t total_size = 1024 * 1024 * 5; // 5MB total size - - // Create multiple slices of different sizes - std::vector slice_lengths = { - 1024 * 1024 * 2, // 2MB - 1024 * 1024 * 1, // 1MB - 1024 * 1024 * 1, // 1MB - 1024 * 1024 * 1 // 1MB - }; - - // Verify total size matches sum of slices - uint64_t sum_slices = 0; - for (const auto& size : slice_lengths) { - sum_slices += size; - } - ASSERT_EQ(total_size, sum_slices); + constexpr size_t slice_length = 1024 * 1024 * 5; // 5MB // Configure replication ReplicateConfig config; @@ -984,7 +965,7 @@ TEST_F(MasterServiceTest, MultiSliceMultiReplicaFlow) { // Test PutStart with multiple slices and replicas auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); @@ -994,19 +975,9 @@ TEST_F(MasterServiceTest, MultiSliceMultiReplicaFlow) { // Verify replica status EXPECT_EQ(ReplicaStatus::PROCESSING, replica.status); - // Verify number of handles matches number of slices - ASSERT_EQ(slice_lengths.size(), - replica.get_memory_descriptor().buffer_descriptors.size()); - - // Verify each handle's properties - for (size_t i = 0; - i < replica.get_memory_descriptor().buffer_descriptors.size(); - i++) { - const auto& handle = - replica.get_memory_descriptor().buffer_descriptors[i]; - - EXPECT_EQ(slice_lengths[i], handle.size_); - } + // Verify slice length matches buffer descriptor + EXPECT_EQ(slice_length, + replica.get_memory_descriptor().buffer_descriptor.size_); } // Test GetReplicaList during processing (should fail) @@ -1027,8 +998,8 @@ TEST_F(MasterServiceTest, MultiSliceMultiReplicaFlow) { // Verify final state of all replicas for (const auto& replica : retrieved_replicas) { EXPECT_EQ(ReplicaStatus::COMPLETE, replica.status); - ASSERT_EQ(slice_lengths.size(), - replica.get_memory_descriptor().buffer_descriptors.size()); + ASSERT_EQ(slice_length, + replica.get_memory_descriptor().buffer_descriptor.size_); } } @@ -1047,13 +1018,13 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) { // Create an object that will be stored in the segment std::string key = "segment_object"; - std::vector slice_lengths = {1024 * 1024}; // One 1MB slice + uint64_t slice_length = 1024 * 1024; // One 1MB slice ReplicateConfig config; config.replica_num = 1; // One replica // Create the object auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); @@ -1081,7 +1052,7 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) { // Create another object std::string key2 = "another_segment_object"; auto put_start_result2 = - service_->PutStart(client_id, key2, slice_lengths, config); + service_->PutStart(client_id, key2, slice_length, config); ASSERT_TRUE(put_start_result2.has_value()); auto put_end_result2 = service_->PutEnd(client_id, key2, ReplicaType::MEMORY); @@ -1123,13 +1094,13 @@ TEST_F(MasterServiceTest, ConcurrentWriteAndRemoveAll) { for (int j = 0; j < objects_per_thread; ++j) { std::string key = "key_" + std::to_string(i) + "_" + std::to_string(j); - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; std::vector replica_list; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); if (put_start_result.has_value()) { auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -1193,12 +1164,12 @@ TEST_F(MasterServiceTest, ConcurrentReadAndRemoveAll) { constexpr int num_objects = 1000; for (int i = 0; i < num_objects; ++i) { std::string key = "pre_key_" + std::to_string(i); - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -1274,12 +1245,12 @@ TEST_F(MasterServiceTest, ConcurrentRemoveAllOperations) { constexpr int num_objects = 1000; for (int i = 0; i < num_objects; ++i) { std::string key = "pre_key_" + std::to_string(i); - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -1336,7 +1307,7 @@ TEST_F(MasterServiceTest, UnmountSegmentImmediateCleanup) { GenerateKeyForSegment(client_id, service_, segment1.name); std::string key2 = GenerateKeyForSegment(client_id, service_, segment2.name); - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; @@ -1356,7 +1327,7 @@ TEST_F(MasterServiceTest, UnmountSegmentImmediateCleanup) { // Verify put key1 will put into segment2 rather than segment1 auto put_start_result = - service_->PutStart(client_id, key1, slice_lengths, config); + service_->PutStart(client_id, key1, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); auto put_end_result = @@ -1367,8 +1338,7 @@ TEST_F(MasterServiceTest, UnmountSegmentImmediateCleanup) { auto retrieved = get_result3.value(); ASSERT_EQ(replica_list[0] .get_memory_descriptor() - .buffer_descriptors[0] - .transport_endpoint_, + .buffer_descriptor.transport_endpoint_, segment2.name); } @@ -1391,12 +1361,12 @@ TEST_F(MasterServiceTest, ReadableAfterPartialUnmountWithReplication) { // Put a key with 2 replicas std::string key = "replicated_key"; - std::vector slice_lengths = {object_size}; + uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 2; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); ASSERT_EQ(2u, put_start_result->size()); ASSERT_TRUE( @@ -1411,8 +1381,8 @@ TEST_F(MasterServiceTest, ReadableAfterPartialUnmountWithReplication) { for (const auto& rep : replicas) { ASSERT_EQ(ReplicaStatus::COMPLETE, rep.status); const auto& mem = rep.get_memory_descriptor(); - ASSERT_EQ(1u, mem.buffer_descriptors.size()); - seg_names.insert(mem.buffer_descriptors[0].transport_endpoint_); + ASSERT_EQ(slice_length, mem.buffer_descriptor.size_); + seg_names.insert(mem.buffer_descriptor.transport_endpoint_); } ASSERT_EQ(2u, seg_names.size()) << "Replicas should be on different segments"; @@ -1496,13 +1466,13 @@ TEST_F(MasterServiceTest, RemoveLeasedObject) { const UUID client_id = generate_uuid(); std::string key = "test_key"; - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; // Verify lease is granted on ExistsKey auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); @@ -1517,7 +1487,7 @@ TEST_F(MasterServiceTest, RemoveLeasedObject) { // Verify lease is extended on successive ExistsKey auto put_start_result2 = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result2.has_value()); auto put_end_result2 = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -1536,7 +1506,7 @@ TEST_F(MasterServiceTest, RemoveLeasedObject) { // Verify lease is granted on GetReplicaList auto put_start_result3 = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result3.has_value()); auto put_end_result3 = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -1552,7 +1522,7 @@ TEST_F(MasterServiceTest, RemoveLeasedObject) { // Verify lease is extended on successive GetReplicaList auto put_start_result4 = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result4.has_value()); auto put_end_result4 = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -1585,11 +1555,11 @@ TEST_F(MasterServiceTest, RemoveAllLeasedObject) { const UUID client_id = generate_uuid(); for (int i = 0; i < 10; ++i) { std::string key = "test_key" + std::to_string(i); - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -1637,11 +1607,11 @@ TEST_F(MasterServiceTest, EvictObject) { int success_puts = 0; for (int i = 0; i < 1024 * 16 + 50; ++i) { std::string key = "test_key" + std::to_string(i); - std::vector slice_lengths = {object_size}; + uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); if (put_start_result.has_value()) { auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -1677,11 +1647,11 @@ TEST_F(MasterServiceTest, TryEvictLeasedObject) { std::vector leased_keys; for (int i = 0; i < 16 + 10; ++i) { std::string key = "test_key" + std::to_string(i); - std::vector slice_lengths = {object_size}; + uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 1; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); if (put_start_result.has_value()) { auto put_end_result = service_->PutEnd(client_id, key, ReplicaType::MEMORY); @@ -1728,21 +1698,21 @@ TEST_F(MasterServiceTest, RemoveSoftPinObject) { PrepareSimpleSegment(*service_, "test_segment", buffer, size); std::string key = "test_key"; - std::vector slice_lengths = {1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; config.with_soft_pin = true; // Verify soft pin does not block remove ASSERT_TRUE( - service_->PutStart(client_id, key, slice_lengths, config).has_value()); + service_->PutStart(client_id, key, slice_length, config).has_value()); ASSERT_TRUE( service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value()); EXPECT_TRUE(service_->Remove(key).has_value()); // Verify soft pin does not block RemoveAll ASSERT_TRUE( - service_->PutStart(client_id, key, slice_lengths, config).has_value()); + service_->PutStart(client_id, key, slice_length, config).has_value()); ASSERT_TRUE( service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value()); EXPECT_EQ(1, service_->RemoveAll()); @@ -1776,13 +1746,13 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotEvictedBeforeOtherObjects) { // Put pin_key first for (int i = 0; i < 2; i++) { std::string pin_key = "pin_key" + std::to_string(i); - std::vector slice_lengths = {value_size}; + uint64_t slice_length = value_size; ReplicateConfig soft_pin_config; soft_pin_config.replica_num = 1; soft_pin_config.with_soft_pin = true; ASSERT_TRUE(service_ - ->PutStart(client_id, pin_key, slice_lengths, + ->PutStart(client_id, pin_key, slice_length, soft_pin_config) .has_value()); ASSERT_TRUE( @@ -1794,10 +1764,10 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotEvictedBeforeOtherObjects) { int failed_puts = 0; for (int i = 0; i < 20; i++) { std::string key = "key" + std::to_string(i); - std::vector slice_lengths = {value_size}; + uint64_t slice_length = value_size; ReplicateConfig config; config.replica_num = 1; - if (service_->PutStart(client_id, key, slice_lengths, config) + if (service_->PutStart(client_id, key, slice_length, config) .has_value()) { ASSERT_TRUE( service_->PutEnd(client_id, key, ReplicaType::MEMORY) @@ -1848,11 +1818,11 @@ TEST_F(MasterServiceTest, SoftPinObjectsCanBeEvicted) { int success_puts = 0; for (int i = 0; i < 16 + 50; ++i) { std::string key = "test_key" + std::to_string(i); - std::vector slice_lengths = {value_size}; + uint64_t slice_length = value_size; ReplicateConfig config; config.replica_num = 1; config.with_soft_pin = true; - if (service_->PutStart(client_id, key, slice_lengths, config) + if (service_->PutStart(client_id, key, slice_length, config) .has_value()) { ASSERT_TRUE(service_->PutEnd(client_id, key, ReplicaType::MEMORY) .has_value()); @@ -1899,12 +1869,12 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) { // Put pin_key first for (int i = 0; i < 2; i++) { std::string pin_key = "pin_key" + std::to_string(i); - std::vector slice_lengths = {value_size}; + uint64_t slice_length = value_size; ReplicateConfig soft_pin_config; soft_pin_config.replica_num = 1; soft_pin_config.with_soft_pin = true; - ASSERT_TRUE(service_->PutStart(client_id, pin_key, slice_lengths, + ASSERT_TRUE(service_->PutStart(client_id, pin_key, slice_length, soft_pin_config)); ASSERT_TRUE( service_->PutEnd(client_id, pin_key, ReplicaType::MEMORY) @@ -1924,10 +1894,10 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) { int failed_puts = 0; for (int i = 0; i < 16; i++) { std::string key = "key" + std::to_string(i); - std::vector slice_lengths = {value_size}; + uint64_t slice_length = value_size; ReplicateConfig config; config.replica_num = 1; - if (service_->PutStart(client_id, key, slice_lengths, config) + if (service_->PutStart(client_id, key, slice_length, config) .has_value()) { ASSERT_TRUE( service_->PutEnd(client_id, key, ReplicaType::MEMORY) @@ -1981,11 +1951,11 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotAllowEvict) { std::vector success_keys; for (int i = 0; i < 16 + 50; ++i) { std::string key = "test_key" + std::to_string(i); - std::vector slice_lengths = {value_size}; + uint64_t slice_length = value_size; ReplicateConfig config; config.replica_num = 1; config.with_soft_pin = true; - if (service_->PutStart(client_id, key, slice_lengths, config) + if (service_->PutStart(client_id, key, slice_length, config) .has_value()) { ASSERT_TRUE(service_->PutEnd(client_id, key, ReplicaType::MEMORY) .has_value()); @@ -2004,7 +1974,7 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotAllowEvict) { service_->RemoveAll(); } -TEST_F(MasterServiceTest, PerSliceReplicaSegmentsAreUnique) { +TEST_F(MasterServiceTest, ReplicaSegmentsAreUnique) { std::unique_ptr service_(new MasterService()); const UUID client_id = generate_uuid(); @@ -2019,29 +1989,26 @@ TEST_F(MasterServiceTest, PerSliceReplicaSegmentsAreUnique) { // Object with 16 slices of ~1MB and replication factor 10 const std::string key = "replica_uniqueness_test_key"; - std::vector slice_lengths(16, 1024 * 1024 - 16); + uint64_t slice_length = 1024 * 1024 - 16; ReplicateConfig config; config.replica_num = 10; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto replica_list_local = put_start_result.value(); ASSERT_EQ(config.replica_num, replica_list_local.size()); - // For each slice index, segment names across replicas must be unique - for (size_t slice_idx = 0; slice_idx < slice_lengths.size(); ++slice_idx) { - std::unordered_set segment_names; - for (const auto& replica : replica_list_local) { - ASSERT_TRUE(replica.is_memory_replica()); - const auto& mem = replica.get_memory_descriptor(); - ASSERT_EQ(slice_lengths.size(), mem.buffer_descriptors.size()); - segment_names.insert( - mem.buffer_descriptors[slice_idx].transport_endpoint_); - } - EXPECT_EQ(segment_names.size(), config.replica_num) - << "Duplicate segment found for slice index " << slice_idx; + // Segment names across replicas must be unique + std::unordered_set segment_names; + for (const auto& replica : replica_list_local) { + ASSERT_TRUE(replica.is_memory_replica()); + const auto& mem = replica.get_memory_descriptor(); + ASSERT_EQ(slice_length, mem.buffer_descriptor.size_); + segment_names.insert(mem.buffer_descriptor.transport_endpoint_); } + EXPECT_EQ(segment_names.size(), config.replica_num) + << "Duplicate segment found"; ASSERT_TRUE( service_->PutEnd(client_id, key, ReplicaType::MEMORY).has_value()); @@ -2060,12 +2027,12 @@ TEST_F(MasterServiceTest, ReplicationFactorTwoWithSingleSegment) { // Request replication factor 2 with a single 1KB slice // With best-effort semantics, should succeed with 1 replica const std::string key = "replication_factor_two_single_segment"; - std::vector slice_lengths{1024}; + uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 2; auto put_start_result = - service_->PutStart(client_id, key, slice_lengths, config); + service_->PutStart(client_id, key, slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto replicas = put_start_result.value(); @@ -2075,10 +2042,8 @@ TEST_F(MasterServiceTest, ReplicationFactorTwoWithSingleSegment) { // Verify the replica is properly allocated on the single segment auto mem_desc = replicas[0].get_memory_descriptor(); - EXPECT_EQ(1u, mem_desc.buffer_descriptors.size()); - EXPECT_EQ("single_segment", - mem_desc.buffer_descriptors[0].transport_endpoint_); - EXPECT_EQ(1024u, mem_desc.buffer_descriptors[0].size_); + EXPECT_EQ("single_segment", mem_desc.buffer_descriptor.transport_endpoint_); + EXPECT_EQ(1024u, mem_desc.buffer_descriptor.size_); } TEST_F(MasterServiceTest, BatchExistKeyTest) { @@ -2098,9 +2063,9 @@ TEST_F(MasterServiceTest, BatchExistKeyTest) { test_keys.push_back("test_key" + std::to_string(i)); ReplicateConfig config; config.replica_num = 1; - std::vector slice_lengths = {value_size}; + uint64_t slice_length = value_size; auto put_start_result = - service_->PutStart(client_id, test_keys[i], slice_lengths, config); + service_->PutStart(client_id, test_keys[i], slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto put_end_result = service_->PutEnd(client_id, test_keys[i], ReplicaType::MEMORY); @@ -2150,13 +2115,13 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { auto client_id = generate_uuid(); std::string key_1 = "test_key_1", key_2 = "test_key_2"; uint64_t value_length = 6 * 1024 * 1024; // 6MB - std::vector slice_lengths = {value_length}; + uint64_t slice_length = value_length; ReplicateConfig config; config.replica_num = kReplicaCnt; // Put key_1, should success. auto put_start_result = - service_->PutStart(client_id, key_1, slice_lengths, config); + service_->PutStart(client_id, key_1, slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -2166,7 +2131,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { // Put key_1 again, should fail because the key exists. put_start_result = - service_->PutStart(client_id, key_1, slice_lengths, config); + service_->PutStart(client_id, key_1, slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS); @@ -2182,7 +2147,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { // Put key_1 again, should success because the old one has expired and will // be discarded by this put. put_start_result = - service_->PutStart(client_id, key_1, slice_lengths, config); + service_->PutStart(client_id, key_1, slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -2202,7 +2167,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { // Put key_2, should fail because the key_1 occupied 12MB (6MB processing, // 6MB discarded but not yet released) on each segment. put_start_result = - service_->PutStart(client_id, key_2, slice_lengths, config); + service_->PutStart(client_id, key_2, slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE); @@ -2223,7 +2188,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { // Put key_2 again, should success because the discarded replica has been // released. put_start_result = - service_->PutStart(client_id, key_2, slice_lengths, config); + service_->PutStart(client_id, key_2, slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -2246,7 +2211,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { // Put key_2 again, should fail because eviction has not been triggered. And // this PutStart should trigger the eviction. put_start_result = - service_->PutStart(client_id, key_2, slice_lengths, config); + service_->PutStart(client_id, key_2, slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE); @@ -2256,7 +2221,7 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { // Put key_2 again, should success because the previous one has been // discarded and released. put_start_result = - service_->PutStart(client_id, key_2, slice_lengths, config); + service_->PutStart(client_id, key_2, slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt);