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

[Store] One Replica Has One Slice (#1032)

This commit is contained in:
ykwd 2025-11-10 16:35:08 +08:00 committed by GitHub
parent 196af0d4a9
commit d9fb46342f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
20 changed files with 468 additions and 941 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,13 +1,12 @@
#pragma once
#include <algorithm>
#include <atomic>
#include <memory>
#include <optional>
#include <random>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <iterator>
#include <time.h>
#include <ylt/util/tl/expected.hpp>
#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<std::vector<Replica>, ErrorCode> containing
@ -60,8 +58,7 @@ class AllocationStrategy {
const std::unordered_map<
std::string, std::vector<std::shared_ptr<BufferAllocatorBase>>>&
allocators_by_name,
const std::vector<size_t>& 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<std::shared_ptr<BufferAllocatorBase>>>&
allocators_by_name,
const std::vector<size_t>& 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<std::vector<std::unique_ptr<AllocatedBuffer>>>
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<Replica> result;
result.emplace_back(std::move(buffer),
ReplicaStatus::PROCESSING);
return result;
}
return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
}
std::vector<Replica> 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<ErrorCode> validateInput(
const std::vector<size_t>& 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<std::unique_ptr<AllocatedBuffer>> allocateSlice(
const std::vector<std::shared_ptr<BufferAllocatorBase>>& allocators,
const std::unordered_map<
std::string, std::vector<std::shared_ptr<BufferAllocatorBase>>>&
allocators_by_name,
size_t slice_size, size_t replica_num, const ReplicateConfig& config,
std::unordered_set<std::string>& used_segments) {
std::vector<std::unique_ptr<AllocatedBuffer>> 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<std::unique_ptr<AllocatedBuffer>> allocateSlice(
const std::vector<std::shared_ptr<BufferAllocatorBase>>& allocators,
const std::unordered_map<
std::string, std::vector<std::shared_ptr<BufferAllocatorBase>>>&
allocators_by_name,
size_t slice_size, size_t replica_num, const ReplicateConfig& config) {
std::unordered_set<std::string> 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<AllocatedBuffer> allocateSingleBuffer(
const std::vector<std::shared_ptr<BufferAllocatorBase>>& allocators,
const std::unordered_map<
std::string, std::vector<std::shared_ptr<BufferAllocatorBase>>>&
allocators_by_name,
size_t size, const ReplicateConfig& config,
const std::unordered_set<std::string>& 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<AllocatedBuffer> tryRandomAllocate(
const std::vector<std::shared_ptr<BufferAllocatorBase>>& allocators,
size_t size, const std::unordered_set<std::string>& excluded_segments) {
std::vector<size_t> 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<size_t> 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

4
mooncake-store/include/client_buffer.hpp
View File

@ -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<Slice>& slices,
const Replica::Descriptor& replica,
BufferHandle& buffer_handle);
void* buffer_ptr);
} // namespace mooncake

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

@ -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<uint64_t>& slice_lengths,
const ReplicateConfig& config)
const uint64_t slice_length, const ReplicateConfig& config)
-> tl::expected<std::vector<Replica::Descriptor>, ErrorCode>;
/**
@ -583,4 +583,4 @@ class MasterService {
GUARDED_BY(discarded_replicas_mutex_);
};
} // namespace mooncake
} // namespace mooncake

59
mooncake-store/include/replica.h
View File

@ -87,7 +87,7 @@ struct ReplicateConfig {
};
struct MemoryReplicaData {
std::vector<std::unique_ptr<AllocatedBuffer>> buffers;
std::unique_ptr<AllocatedBuffer> buffer;
};
struct DiskReplicaData {
@ -96,8 +96,8 @@ struct DiskReplicaData {
};
struct MemoryDescriptor {
std::vector<AllocatedBuffer::Descriptor> 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<std::unique_ptr<AllocatedBuffer>> buffers,
ReplicaStatus status)
: data_(MemoryReplicaData{std::move(buffers)}), status_(status) {}
Replica(std::unique_ptr<AllocatedBuffer> 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<MemoryReplicaData>(data_);
return std::any_of(
mem_data.buffers.begin(), mem_data.buffers.end(),
[](const std::unique_ptr<AllocatedBuffer>& 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<MemoryReplicaData>(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<std::optional<std::string>> get_segment_names()
@ -292,12 +286,13 @@ inline Replica::Descriptor Replica::get_descriptor() const {
if (is_memory_replica()) {
const auto& mem_data = std::get<MemoryReplicaData>(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<std::optional<std::string>> Replica::get_segment_names()
const {
if (is_memory_replica()) {
const auto& mem_data = std::get<MemoryReplicaData>(data_);
std::vector<std::optional<std::string>> 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<std::optional<std::string>> 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<MemoryReplicaData>(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()) {

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

@ -42,8 +42,7 @@ class WrappedMasterService {
tl::expected<std::vector<Replica::Descriptor>, ErrorCode> PutStart(
const UUID& client_id, const std::string& key,
const std::vector<uint64_t>& slice_lengths,
const ReplicateConfig& config);
const uint64_t slice_length, const ReplicateConfig& config);
tl::expected<void, ErrorCode> PutEnd(const UUID& client_id,
const std::string& key,
@ -55,7 +54,7 @@ class WrappedMasterService {
std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
BatchPutStart(const UUID& client_id, const std::vector<std::string>& keys,
const std::vector<std::vector<uint64_t>>& slice_lengths,
const std::vector<uint64_t>& slice_lengths,
const ReplicateConfig& config);
std::vector<tl::expected<void, ErrorCode>> BatchPutEnd(
@ -95,4 +94,4 @@ class WrappedMasterService {
void RegisterRpcService(coro_rpc::coro_rpc_server& server,
mooncake::WrappedMasterService& wrapped_master_service);
} // namespace mooncake
} // namespace mooncake

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

@ -400,36 +400,35 @@ class TransferSubmitter {
/**
* @brief Select the optimal transfer strategy
*/
TransferStrategy selectStrategy(
const std::vector<AllocatedBuffer::Descriptor>& handles,
const std::vector<Slice>& slices) const;
TransferStrategy selectStrategy(const AllocatedBuffer::Descriptor& handle,
const std::vector<Slice>& slices) const;
/**
* @brief Check if all handles refer to local segments
*/
bool isLocalTransfer(
const std::vector<AllocatedBuffer::Descriptor>& handles) const;
bool isLocalTransfer(const AllocatedBuffer::Descriptor& handle) const;
/**
* @brief Validate transfer parameters
*/
bool validateTransferParams(
const std::vector<AllocatedBuffer::Descriptor>& handles,
const std::vector<Slice>& slices, bool is_multi_buffers = false) const;
bool validateTransferParams(const AllocatedBuffer::Descriptor& handle,
const std::vector<Slice>& slices) const;
/**
* @brief Submit memcpy operation asynchronously
*/
std::optional<TransferFuture> submitMemcpyOperation(
const std::vector<AllocatedBuffer::Descriptor>& handles,
std::vector<Slice>& slices, TransferRequest::OpCode op_code);
const AllocatedBuffer::Descriptor& handle,
const std::vector<Slice>& slices,
const TransferRequest::OpCode op_code);
/**
* @brief Submit transfer engine operation asynchronously
*/
std::optional<TransferFuture> submitTransferEngineOperation(
const std::vector<AllocatedBuffer::Descriptor>& handles,
std::vector<Slice>& slices, TransferRequest::OpCode op_code);
const AllocatedBuffer::Descriptor& handle,
const std::vector<Slice>& slices,
const TransferRequest::OpCode op_code);
std::optional<TransferFuture> submitFileReadOperation(
const Replica::Descriptor& replica, std::vector<Slice>& slices,

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

@ -589,11 +589,11 @@ std::vector<tl::expected<void, ErrorCode>> 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<tl::expected<void, ErrorCode>> 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<tl::expected<void, ErrorCode>> 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<tl::expected<void, ErrorCode>> 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<EmptyOperationState>();
@ -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;

21
mooncake-store/src/client_buffer.cpp
View File

@ -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<Slice>& 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<char*>(buffer_handle.ptr()) + offset;
void* chunk_ptr = static_cast<char*>(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<char*>(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;
}

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

@ -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<uint64_t> 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<Replica::Descriptor>>(
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<uint64_t> 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<Replica::Descriptor>>(
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;
}

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

@ -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<uint64_t>& slice_lengths,
const uint64_t slice_length,
const ReplicateConfig& config)
-> tl::expected<std::vector<Replica::Descriptor>, 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

39
mooncake-store/src/pybind_client.cpp
View File

@ -614,7 +614,7 @@ std::shared_ptr<BufferHandle> PyClient::get_buffer(const std::string &key) {
// Create slices for the allocated buffer
std::vector<Slice> 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<std::shared_ptr<BufferHandle>> PyClient::batch_get_buffer_internal(
auto buffer_handle =
std::make_unique<BufferHandle>(std::move(*alloc_result));
std::vector<Slice> 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<int64_t, ErrorCode> PyClient::get_into_internal(
// Step 2: Split user buffer according to object info and create
// slices
std::vector<mooncake::Slice> 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<char *>(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<char *>(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<tl::expected<int64_t, ErrorCode>> PyClient::batch_get_into_internal(
// Create slices for this key's buffer
std::vector<Slice> 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<char *>(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<char *>(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(

43
mooncake-store/src/rpc_service.cpp
View File

@ -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<std::vector<Replica::Descriptor>, ErrorCode>
WrappedMasterService::PutStart(const UUID& client_id, const std::string& key,
const std::vector<uint64_t>& 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<void, ErrorCode> WrappedMasterService::PutRevoke(
}
std::vector<tl::expected<std::vector<Replica::Descriptor>, ErrorCode>>
WrappedMasterService::BatchPutStart(
const UUID& client_id, const std::vector<std::string>& keys,
const std::vector<std::vector<uint64_t>>& slice_lengths,
const ReplicateConfig& config) {
WrappedMasterService::BatchPutStart(const UUID& client_id,
const std::vector<std::string>& keys,
const std::vector<uint64_t>& 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<uint64_t> 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

134
mooncake-store/src/transfer_task.cpp
View File

@ -392,23 +392,21 @@ std::optional<TransferFuture> TransferSubmitter::submit(
std::optional<TransferFuture> future;
if (replica.is_memory_replica()) {
std::vector<AllocatedBuffer::Descriptor> 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<TransferFuture> 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<uint64_t>(ERR_INVALID_ARGUMENT)) {
@ -470,16 +467,17 @@ std::optional<TransferFuture> TransferSubmitter::submit_batch(
}
std::optional<TransferFuture> TransferSubmitter::submitMemcpyOperation(
const std::vector<AllocatedBuffer::Descriptor>& handles,
std::vector<Slice>& slices, TransferRequest::OpCode op_code) {
const AllocatedBuffer::Descriptor& handle, const std::vector<Slice>& slices,
const TransferRequest::OpCode op_code) {
auto state = std::make_shared<MemcpyOperationState>();
// Create memcpy operations
std::vector<MemcpyOperation> operations;
operations.reserve(handles.size());
operations.reserve(slices.size());
uint64_t base_address = static_cast<uint64_t>(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<TransferFuture> TransferSubmitter::submitMemcpyOperation(
// READ: from handle (remote buffer) to slice (local
// buffer)
dest = slice.ptr;
src = reinterpret_cast<const void*>(handle.buffer_address_);
src = reinterpret_cast<const void*>(base_address + offset);
} else {
// WRITE: from slice (local buffer) to handle (remote
// buffer)
dest = reinterpret_cast<void*>(handle.buffer_address_);
dest = reinterpret_cast<void*>(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<TransferFuture> TransferSubmitter::submitTransfer(
}
std::optional<TransferFuture> TransferSubmitter::submitTransferEngineOperation(
const std::vector<AllocatedBuffer::Descriptor>& handles,
std::vector<Slice>& slices, TransferRequest::OpCode op_code) {
const AllocatedBuffer::Descriptor& handle, const std::vector<Slice>& 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<uint64_t>(ERR_INVALID_ARGUMENT)) {
LOG(ERROR) << "Failed to open segment for endpoint='"
<< handle.transport_endpoint_ << "'";
return std::nullopt;
}
// Create transfer requests
std::vector<TransferRequest> requests;
requests.reserve(handles.size());
requests.reserve(slices.size());
uint64_t base_address = static_cast<uint64_t>(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<uint64_t>(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<char*>(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<TransferFuture> TransferSubmitter::submitFileReadOperation(
}
TransferStrategy TransferSubmitter::selectStrategy(
const std::vector<AllocatedBuffer::Descriptor>& handles,
const AllocatedBuffer::Descriptor& handle,
const std::vector<Slice>& 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<AllocatedBuffer::Descriptor>& 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<AllocatedBuffer::Descriptor>& handles,
const std::vector<Slice>& slices, bool is_multi_buffers) const {
if (handles.empty()) {
LOG(ERROR) << "handles is empty";
return false;
const AllocatedBuffer::Descriptor& handle,
const std::vector<Slice>& 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;
}

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

@ -3,8 +3,10 @@
#include <gtest/gtest.h>
#include <memory>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#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<RandomAllocationStrategy>();
}
// Helper function to create test allocators
std::shared_ptr<BufferAllocatorBase> 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<CachelibBufferAllocator>(
segment_name, base, size, segment_name);
case BufferAllocatorType::OFFSET:
return std::make_shared<OffsetBufferAllocator>(
segment_name, base, size, segment_name);
default:
throw std::invalid_argument("Invalid allocator type");
}
}
std::unique_ptr<RandomAllocationStrategy> strategy_;
};
// Test basic functionality with empty allocators map (non-parameterized)
TEST_F(AllocationStrategyTest, EmptyAllocatorsMap) {
std::unordered_map<std::string,
@ -107,9 +82,9 @@ TEST_F(AllocationStrategyTest, EmptyAllocatorsMap) {
std::vector<std::shared_ptr<BufferAllocatorBase>> empty_allocators;
ReplicateConfig config{1, false, "local"};
std::vector<size_t> 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<std::shared_ptr<BufferAllocatorBase>> empty_allocators;
ReplicateConfig config{1, false, "preferred_segment"};
std::vector<size_t> 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<size_t> 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<size_t> 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<size_t> 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<size_t> 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<size_t> 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<std::vector<Replica>> 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<size_t> 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<size_t> 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<std::vector<Replica>> 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<size_t> 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<size_t> 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<size_t> 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<OffsetBufferAllocator>(
"segment1", 0x100000000ULL, 64 * MB, "segment1");
std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocatorBase>>>
allocators_by_name;
std::vector<std::shared_ptr<BufferAllocatorBase>> allocators;
allocators_by_name["segment1"].push_back(allocator);
allocators.push_back(allocator);
ReplicateConfig config{1, false, ""};
std::vector<size_t> 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<size_t> 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<size_t> 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<std::shared_ptr<BufferAllocatorBase>> allocators = {allocator1,
allocator2};
std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocatorBase>>>
allocators_by_name;
allocators_by_name["segment1"] = {allocator1};
allocators_by_name["segment2"] = {allocator2};
ReplicateConfig config{1, false, "nonexistent"};
std::unordered_set<std::string> 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<std::shared_ptr<BufferAllocatorBase>> allocators = {allocator1};
std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocatorBase>>>
allocators_by_name;
allocators_by_name["segment1"] = {allocator1};
ReplicateConfig config{1, false, ""}; // Empty preferred segment
std::unordered_set<std::string> 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<std::shared_ptr<BufferAllocatorBase>> allocators = {allocator1,
allocator2};
std::unordered_set<std::string> 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<std::shared_ptr<BufferAllocatorBase>> allocators = {allocator1,
allocator2};
std::unordered_set<std::string> 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<std::shared_ptr<BufferAllocatorBase>> allocators = {allocator};
std::unordered_set<std::string> 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<std::shared_ptr<BufferAllocatorBase>> allocators = {allocator1};
std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocatorBase>>>
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<std::shared_ptr<BufferAllocatorBase>> allocators = {
allocator1, allocator2, allocator3};
std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocatorBase>>>
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<std::string> 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<std::shared_ptr<BufferAllocatorBase>> allocators = {allocator1};
std::unordered_map<std::string,
std::vector<std::shared_ptr<BufferAllocatorBase>>>
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<std::unique_ptr<AllocatedBuffer>> 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<std::shared_ptr<BufferAllocatorBase>> allocators = {
allocator1, allocator2};
std::unordered_set<std::string> 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

71
mooncake-store/tests/client_buffer_test.cpp
View File

@ -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<Slice> 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<char*>(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<Slice> 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<Slice> 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

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

@ -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);

6
mooncake-store/tests/e2e/client_wrapper.cpp
View File

@ -106,9 +106,9 @@ ErrorCode ClientTestWrapper::Get(const std::string& key, std::string& value) {
}
// Create slices
const std::vector<AllocatedBuffer::Descriptor>& 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 =

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

@ -119,7 +119,6 @@ TEST_F(MasterMetricsTest, BasicRequestTest) {
std::string key = "test_key";
uint64_t value_length = 1024;
std::vector<uint64_t> 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<std::string> keys = {"test_key1", "test_key2", "test_key3"};
std::vector<std::vector<uint64_t>> slice_lengths = {{1024}, {2048}, {512}};
std::vector<uint64_t> 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);

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

@ -45,12 +45,12 @@ TEST_F(MasterServiceSSDTest, PutEndBothReplica) {
ASSERT_TRUE(mount_result.has_value());
std::string key = "disk_key";
std::vector<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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);

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

@ -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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> slice_lengths = {1024};
uint64_t slice_length = 1024;
ReplicateConfig config;
config.replica_num = 1;
std::vector<Replica::Descriptor> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<std::string> leased_keys;
for (int i = 0; i < 16 + 10; ++i) {
std::string key = "test_key" + std::to_string(i);
std::vector<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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<std::string> success_keys;
for (int i = 0; i < 16 + 50; ++i) {
std::string key = "test_key" + std::to_string(i);
std::vector<uint64_t> 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<MasterService> 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<uint64_t> 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<std::string> 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<std::string> 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<uint64_t> 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<uint64_t> 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<uint64_t> 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);