forked from mooncake-track/Mooncake
[Store] Add replication guarantees (#744)
* add replication guarantees Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * refactor implementation Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * fix merge aftermath Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * less copies Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * try fix wheel tests Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * fix wheel tests Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * remove old comment Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * fix wheel tests and add replication fault tolerance wheel test Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * split wheel tests Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * simplify code Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * replica allocation is best effort operation Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * fix tests for new best-effort behaviour Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * fix Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * update docs Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> * update another docs Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com> --------- Signed-off-by: Vladislav Nosivskoy <vladnosiv@gmail.com>
This commit is contained in:
Vladislav Nosivskoy
GitHub
parent
f76c92295b
commit
aa9d4471a1
|
|
@ -4,7 +4,7 @@ Mooncake aims to enhance the inference efficiency of large language models (LLMs
|
|||
|
||||
Mooncake:
|
||||
- provides object-level data storage services
|
||||
- supports data replication in the cache layer, with a lightweight design due to not guaranteeing high availability
|
||||
- supports data replication in the cache layer with slice-level placement guarantees and best-effort allocation, with a lightweight design due to not guaranteeing high availability
|
||||
- ensures the atomicity of object write operations, meaning a `Get` operation will always read one consistent version, but not necessarily the latest one
|
||||
- supports striping and parallel I/O transfer for larger objects to utilize the aggregated bandwidth of multiple network cards
|
||||
- supports multiple modes for flushing slow object storage
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ Mooncake Store provides low-level object storage and management capabilities, in
|
|||
|
||||
Key features of Mooncake Store include:
|
||||
- **Object-level storage operations**: Mooncake Store provides simple and easy-to-use object-level APIs, including `Put`, `Get`, and `Remove` operations.
|
||||
- **Multi-replica support**: Mooncake Store supports storing multiple data replicas for the same object, effectively alleviating hotspots in access pressure.
|
||||
- **Multi-replica support**: Mooncake Store supports storing multiple data replicas for the same object, effectively alleviating hotspots in access pressure. Each slice within an object is guaranteed to be placed in different segments, while different objects' slices may share segments. Replication operates on a best-effort basis.
|
||||
- **Strong consistency**: Mooncake Store Guarantees that `Get` operations always read accurate and complete data, and after a successful write, all subsequent Gets will return the most recent value.
|
||||
- **Zero-copy, bandwidth-saturating transfers**: Powered by the Transfer Engine, Mooncake Store eliminates redundant memory copies and exploits multi-NIC GPUDirect RDMA pooling to drive data across the network at full line rate while keeping CPU overhead negligible.
|
||||
- **High bandwidth utilization**: Mooncake Store supports striping and parallel I/O transfer of large objects, fully utilizing multi-NIC aggregated bandwidth for high-speed data reads and writes.
|
||||
|
|
@ -89,7 +89,14 @@ tl::expected<void, ErrorCode> Put(const ObjectKey& key,
|
|||
|
||||
|
||||
|
||||
Used to store the value corresponding to `key`. The required number of replicas can be set via the `config` parameter.(When persistence is enabled, Put not only writes to the memory pool but also asynchronously initiates a data persistence operation to the SSD.) The data structure details of `ReplicateConfig` are as follows:
|
||||
Used to store the value corresponding to `key`. The required number of replicas can be set via the `config` parameter.(When persistence is enabled, Put not only writes to the memory pool but also asynchronously initiates a data persistence operation to the SSD.)
|
||||
|
||||
**Replication Guarantees and Best Effort Behavior:**
|
||||
- Each slice of an object is guaranteed to be replicated to different segments, ensuring distribution across separate storage nodes
|
||||
- Different slices from different objects may be placed in the same segment
|
||||
- Replication operates on a best-effort basis: if insufficient space is available for all requested replicas, the object will still be written with as many replicas as possible
|
||||
|
||||
The data structure details of `ReplicateConfig` are as follows:
|
||||
|
||||
```C++
|
||||
struct ReplicateConfig {
|
||||
|
|
@ -246,7 +253,7 @@ message PutStartResponse {
|
|||
|
||||
- **Request**: `PutStartRequest` containing the key, data length, and replica configuration config.
|
||||
- **Response**: `PutStartResponse` containing the status code status_code and the allocated replica information replica_list.
|
||||
- **Description**: Before writing an object, the Client must call PutStart to request storage space from the Master Service. The Master Service allocates space based on the config and returns the allocation results (`replica_list`) to the Client. The Client then writes data to the storage nodes where the allocated replicas are located. The need for both start and end steps ensures that other Clients do not read partially written values, preventing dirty reads.
|
||||
- **Description**: Before writing an object, the Client must call PutStart to request storage space from the Master Service. The Master Service allocates space based on the config and returns the allocation results (`replica_list`) to the Client. The allocation strategy ensures that each slice of the object is placed in different segments, while operating on a best-effort basis - if insufficient space is available for all requested replicas, as many replicas as possible will be allocated. The Client then writes data to the storage nodes where the allocated replicas are located. The need for both start and end steps ensures that other Clients do not read partially written values, preventing dirty reads.
|
||||
|
||||
4. PutEnd
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
Mooncake 旨在通过在高速互联的 DRAM/SSD 资源上构建一个多级缓存池，来提升大型语言模型（LLM）等场景下的推理效率，尤其是在慢速对象存储环境中。与传统缓存相比，Mooncake 的独特之处在于，它能够利用 (GPUDirect) RDMA 技术，以零拷贝的方式将数据从发起端的 DRAM/VRAM 直接传输到接收端的 DRAM/SSD，同时最大限度地利用单机多网卡资源。
|
||||
|
||||
- 提供对象级别的数据存储服务
|
||||
- 支持在缓存层多副本保存数据，由于不保证绝对的高可用，因此系统设计更为轻量化
|
||||
- 支持在缓存层多副本保存数据，提供slice级别的分布保证和尽力而为的分配策略，由于不保证绝对的高可用，因此系统设计更为轻量化
|
||||
- 保证对象写操作的原子性，即 Get 一定会读到某次 Put 生成的完整的数据，但不一定是最新的
|
||||
- 支持对较大的对象进行条带化和并行 I/O 传输，以利用多网卡的聚合带宽
|
||||
- 支持 Eager/Lazy/None 三种下刷慢速对象存储的模式，分别对应持久化要求从高到低的对象
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ Mooncake Store 提供了底层的对象存储和管理能力，包括可配置
|
|||
Mooncake Store 的主要特性包括：
|
||||
|
||||
* **对象级存储操作**：提供简单易用的对象级 API，包括 Put、Get 和 Remove 等操作，方便用户进行数据管理。
|
||||
* **多副本支持**：支持为同一对象保存多个数据副本，有效缓解热点访问压力。
|
||||
* **多副本支持**：支持为同一对象保存多个数据副本，有效缓解热点访问压力。保证同一对象的每个slice被放置在不同的segment中，不同对象的slice可以共享segment。采用尽力而为的副本分配策略。
|
||||
* **强一致性**：保证 `Get` 操作读取到完整且正确的数据，并且数据写入成功后，后续的 `Get` 操作一定能读取到最新写入的值。
|
||||
* **零拷贝、高带宽利用**：由 Transfer Engine 提供支持, 数据链路零拷贝，对大型对象进行条带化和并行 I/O 传输，充分利用多网卡聚合带宽，实现高速数据读写。
|
||||
* **动态资源伸缩**：支持动态添加和删除节点，灵活应对系统负载变化，实现资源的弹性管理。
|
||||
|
|
@ -95,6 +95,12 @@ tl::expected<void, ErrorCode> Put(const ObjectKey& key,
|
|||
|
||||
|
||||
用于存储 `key` 对应的值。可通过 `config` 参数设置所需的副本数量。（当启用了持久化功能时，`Put`除了对memory pool的写入之外，还会异步发起一次向SSD的数据持久化操作）
|
||||
|
||||
**副本保证和尽力而为行为：**
|
||||
- 保证对象的每个slice被复制到不同的segment，确保分布在不同的存储节点上
|
||||
- 不同对象的slice可能被放置在同一个segment中
|
||||
- 副本采用尽力而为的方式运行：如果没有足够的空间来分配所有请求的副本，对象仍将被写入，副本数量为实际能够分配的数量
|
||||
|
||||
其中`ReplicateConfig` 的数据结构细节如下：
|
||||
|
||||
```C++
|
||||
|
|
@ -253,7 +259,7 @@ message PutStartResponse {
|
|||
* 请求: PutStartRequest，包含 key、数据长度和副本配置config。
|
||||
* 响应: PutStartResponse，包含状态码 status_code 和分配好的副本信息 replica_list。
|
||||
|
||||
说明: Client 在写入对象前，需要先调用 PutStart 向 `Master Service` 申请存储空间。`Master Service` 会根据 config 分配空间，并将分配结果（replica_list）返回给 Client。Client 随后将数据写入到分配副本所在的存储节点。 之所以需要 start 和 end 两步，是为确保其他Client不会读到正在写的值，进而造成脏读。
|
||||
说明: Client 在写入对象前，需要先调用 PutStart 向 `Master Service` 申请存储空间。`Master Service` 会根据 config 分配空间，并将分配结果（replica_list）返回给 Client。分配策略确保对象的每个slice被放置在不同的segment中，同时采用尽力而为的方式运行——如果没有足够的空间来分配所有请求的副本，将分配尽可能多的副本。Client 随后将数据写入到分配副本所在的存储节点。 之所以需要 start 和 end 两步，是为确保其他Client不会读到正在写的值，进而造成脏读。
|
||||
|
||||
4. PutEnd
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ Mooncake aims to enhance the inference efficiency of large language models (LLMs
|
|||
|
||||
Mooncake:
|
||||
- provides object-level data storage services
|
||||
- supports data replication in the cache layer, with a lightweight design due to not guaranteeing high availability
|
||||
- supports data replication in the cache layer with slice-level placement guarantees and best-effort allocation, with a lightweight design due to not guaranteeing high availability
|
||||
- ensures the atomicity of object write operations, meaning a `Get` operation will always read one consistent version, but not necessarily the latest one
|
||||
- supports striping and parallel I/O transfer for larger objects to utilize the aggregated bandwidth of multiple network cards
|
||||
- supports multiple modes for flushing slow object storage
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ Mooncake Store provides low-level object storage and management capabilities, in
|
|||
|
||||
Key features of Mooncake Store include:
|
||||
- **Object-level storage operations**: Mooncake Store provides simple and easy-to-use object-level APIs, including `Put`, `Get`, and `Remove` operations.
|
||||
- **Multi-replica support**: Mooncake Store supports storing multiple data replicas for the same object, effectively alleviating hotspots in access pressure.
|
||||
- **Multi-replica support**: Mooncake Store supports storing multiple data replicas for the same object, effectively alleviating hotspots in access pressure. Each slice within an object is guaranteed to be placed in different segments, while different objects' slices may share segments. Replication operates on a best-effort basis.
|
||||
- **Strong consistency**: Mooncake Store Guarantees that `Get` operations always read accurate and complete data, and after a successful write, all subsequent Gets will return the most recent value.
|
||||
- **Zero-copy, bandwidth-saturating transfers**: Powered by the Transfer Engine, Mooncake Store eliminates redundant memory copies and exploits multi-NIC GPUDirect RDMA pooling to drive data across the network at full line rate while keeping CPU overhead negligible.
|
||||
- **High bandwidth utilization**: Mooncake Store supports striping and parallel I/O transfer of large objects, fully utilizing multi-NIC aggregated bandwidth for high-speed data reads and writes.
|
||||
|
|
@ -89,7 +89,14 @@ tl::expected<void, ErrorCode> Put(const ObjectKey& key,
|
|||
|
||||
|
||||
|
||||
Used to store the value corresponding to `key`. The required number of replicas can be set via the `config` parameter.(When persistence is enabled, Put not only writes to the memory pool but also asynchronously initiates a data persistence operation to the SSD.) The data structure details of `ReplicateConfig` are as follows:
|
||||
Used to store the value corresponding to `key`. The required number of replicas can be set via the `config` parameter.(When persistence is enabled, Put not only writes to the memory pool but also asynchronously initiates a data persistence operation to the SSD.)
|
||||
|
||||
**Replication Guarantees and Best Effort Behavior:**
|
||||
- Each slice of an object is guaranteed to be replicated to different segments, ensuring distribution across separate storage nodes
|
||||
- Different slices from different objects may be placed in the same segment
|
||||
- Replication operates on a best-effort basis: if insufficient space is available for all requested replicas, the object will still be written with as many replicas as possible
|
||||
|
||||
The data structure details of `ReplicateConfig` are as follows:
|
||||
|
||||
```C++
|
||||
struct ReplicateConfig {
|
||||
|
|
@ -246,7 +253,7 @@ message PutStartResponse {
|
|||
|
||||
- **Request**: `PutStartRequest` containing the key, data length, and replica configuration config.
|
||||
- **Response**: `PutStartResponse` containing the status code status_code and the allocated replica information replica_list.
|
||||
- **Description**: Before writing an object, the Client must call PutStart to request storage space from the Master Service. The Master Service allocates space based on the config and returns the allocation results (`replica_list`) to the Client. The Client then writes data to the storage nodes where the allocated replicas are located. The need for both start and end steps ensures that other Clients do not read partially written values, preventing dirty reads.
|
||||
- **Description**: Before writing an object, the Client must call PutStart to request storage space from the Master Service. The Master Service allocates space based on the config and returns the allocation results (`replica_list`) to the Client. The allocation strategy ensures that each slice of the object is placed in different segments, while operating on a best-effort basis - if insufficient space is available for all requested replicas, as many replicas as possible will be allocated. The Client then writes data to the storage nodes where the allocated replicas are located. The need for both start and end steps ensures that other Clients do not read partially written values, preventing dirty reads.
|
||||
|
||||
4. PutEnd
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,13 @@
|
|||
#pragma once
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <random>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
#include <ylt/util/tl/expected.hpp>
|
||||
|
||||
#include "allocator.h" // Contains BufferAllocator declaration
|
||||
#include "types.h"
|
||||
|
|
@ -11,127 +15,241 @@
|
|||
namespace mooncake {
|
||||
|
||||
/**
|
||||
* @brief Abstract interface for allocation strategy, responsible for choosing
|
||||
* among multiple BufferAllocators.
|
||||
* @brief Abstract interface for allocation strategy, responsible for
|
||||
* allocating multiple slices across multiple replicas using available
|
||||
* BufferAllocators.
|
||||
*
|
||||
* The allocation strategy follows best-effort semantics: if the requested
|
||||
* number of replicas cannot be fully satisfied due to resource constraints,
|
||||
* it will allocate as many replicas as possible rather than failing entirely.
|
||||
* Only returns an error if no replicas can be allocated at all.
|
||||
*/
|
||||
class AllocationStrategy {
|
||||
public:
|
||||
virtual ~AllocationStrategy() = default;
|
||||
|
||||
/**
|
||||
* @brief Given all mounted BufferAllocators and required object size,
|
||||
* the strategy can freely choose a suitable BufferAllocator.
|
||||
* @brief Allocates multiple slices across the requested number of replicas
|
||||
* using best-effort semantics. Each replica will contain all
|
||||
* requested slices.
|
||||
*
|
||||
* The allocation follows best-effort semantics: if the full requested
|
||||
* replica count cannot be satisfied, the method will allocate as many
|
||||
* replicas as possible across different segments. For each slice, replicas
|
||||
* are guaranteed to be placed on different segments to ensure redundancy.
|
||||
*
|
||||
* @param allocators Container of mounted allocators
|
||||
* @param allocators_by_name Container of mounted allocators, key is
|
||||
* segment_name, value is the corresponding allocator
|
||||
* @param objectSize Size of object to be allocated
|
||||
* @param config Replica configuration
|
||||
* @return Selected allocator; returns nullptr if allocation is not possible
|
||||
* or no suitable allocator is found
|
||||
* segment_name, value is the corresponding
|
||||
* allocators
|
||||
* @param slice_sizes Sizes of slices to be allocated in each replica
|
||||
* @param config Replica configuration containing number of replicas and
|
||||
* placement constraints
|
||||
* @return tl::expected<std::vector<Replica>, ErrorCode> containing
|
||||
* allocated replicas.
|
||||
* - On success: vector of allocated replicas (may be fewer than
|
||||
* requested due to resource constraints, but at least 1)
|
||||
* - On failure: ErrorCode::NO_AVAILABLE_HANDLE if no replicas can
|
||||
* be allocated, ErrorCode::INVALID_PARAMS for invalid
|
||||
* configuration
|
||||
*/
|
||||
virtual std::unique_ptr<AllocatedBuffer> Allocate(
|
||||
virtual tl::expected<std::vector<Replica>, ErrorCode> Allocate(
|
||||
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 objectSize, const ReplicateConfig& config) = 0;
|
||||
const std::vector<size_t>& slice_sizes,
|
||||
const ReplicateConfig& config) = 0;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Random allocation strategy with local preference support.
|
||||
* @brief Random batch allocation strategy with local preference and
|
||||
* replication guarantees support using best-effort semantics.
|
||||
*
|
||||
* This strategy first attempts to allocate from a preferred segment if
|
||||
* specified, then falls back to random allocation among all available
|
||||
* allocators.
|
||||
* This strategy ensures that for each slice, its replicas are placed in
|
||||
* different segments. Different slices may use the same segments.
|
||||
*
|
||||
* Best-effort behavior:
|
||||
* - Attempts to allocate the requested number of replicas
|
||||
* - If insufficient segments are available, allocates as many replicas as
|
||||
* possible (limited by the number of available segments)
|
||||
* - Only fails if no replicas can be allocated at all
|
||||
* - Preferred segment allocation is attempted first if specified
|
||||
*/
|
||||
class RandomAllocationStrategy : public AllocationStrategy {
|
||||
public:
|
||||
RandomAllocationStrategy() : rng_(std::random_device{}()) {}
|
||||
|
||||
std::unique_ptr<AllocatedBuffer> Allocate(
|
||||
tl::expected<std::vector<Replica>, ErrorCode> Allocate(
|
||||
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 objectSize, const ReplicateConfig& config) override {
|
||||
// Fast path: single allocator case
|
||||
if (allocators.size() == 1) {
|
||||
return allocators[0]->allocate(objectSize);
|
||||
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);
|
||||
}
|
||||
|
||||
// Try preferred segment first if specified
|
||||
if (auto preferred_buffer =
|
||||
TryPreferredAllocate(allocators_by_name, objectSize, config)) {
|
||||
return preferred_buffer;
|
||||
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());
|
||||
}
|
||||
|
||||
// Fall back to random allocation among all eligible allocators
|
||||
return TryRandomAllocate(allocators, objectSize);
|
||||
// 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
|
||||
LOG(WARNING)
|
||||
<< "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]));
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
return replicas;
|
||||
}
|
||||
|
||||
private:
|
||||
static constexpr size_t kMaxRetryLimit = 10;
|
||||
std::mt19937 rng_;
|
||||
|
||||
std::mt19937 rng_; // Mersenne Twister random number generator
|
||||
|
||||
/**
|
||||
* @brief Attempts allocation from preferred segment if available and
|
||||
* eligible
|
||||
*/
|
||||
std::unique_ptr<AllocatedBuffer> TryPreferredAllocate(
|
||||
const std::unordered_map<
|
||||
std::string, std::vector<std::shared_ptr<BufferAllocatorBase>>>&
|
||||
allocators,
|
||||
size_t objectSize, const ReplicateConfig& config) {
|
||||
if (config.preferred_segment.empty()) {
|
||||
return nullptr;
|
||||
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;
|
||||
}
|
||||
|
||||
auto preferred_it = allocators.find(config.preferred_segment);
|
||||
if (preferred_it == allocators.end()) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
auto& preferred_allocators = preferred_it->second;
|
||||
for (auto& allocator : preferred_allocators) {
|
||||
auto buffer = allocator->allocate(objectSize);
|
||||
if (buffer != nullptr) {
|
||||
return buffer;
|
||||
}
|
||||
}
|
||||
|
||||
return nullptr;
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Attempts allocation with random selection and retry logic
|
||||
* @brief Allocates replicas for a single slice across different segments
|
||||
*/
|
||||
std::unique_ptr<AllocatedBuffer> TryRandomAllocate(
|
||||
std::vector<std::unique_ptr<AllocatedBuffer>> allocateSlice(
|
||||
const std::vector<std::shared_ptr<BufferAllocatorBase>>& allocators,
|
||||
size_t objectSize) {
|
||||
const size_t max_tries = std::min(kMaxRetryLimit, allocators.size());
|
||||
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);
|
||||
|
||||
std::vector<size_t> allocator_indices(allocators.size());
|
||||
std::iota(allocator_indices.begin(), allocator_indices.end(), 0);
|
||||
for (size_t i = 0; i < replica_num; ++i) {
|
||||
auto buffer =
|
||||
allocateSingleBuffer(allocators, allocators_by_name, slice_size,
|
||||
config, used_segments);
|
||||
|
||||
for (size_t try_count = 0; try_count < max_tries; ++try_count) {
|
||||
// Randomly select an allocator
|
||||
std::uniform_int_distribution<size_t> dist(
|
||||
0, allocator_indices.size() - 1);
|
||||
const size_t random_index = allocator_indices[dist(rng_)];
|
||||
if (!buffer) {
|
||||
break;
|
||||
}
|
||||
|
||||
auto& allocator = allocators[random_index];
|
||||
if (auto buffer = allocator->allocate(objectSize)) {
|
||||
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)) {
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return tryRandomAllocate(allocators, size, excluded_segments);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Attempts allocation with random selection
|
||||
*/
|
||||
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())) {
|
||||
eligible_indices.push_back(i);
|
||||
}
|
||||
}
|
||||
|
||||
if (eligible_indices.empty()) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
// Remove failed allocator and continue with remaining ones
|
||||
if (random_index + 1 != allocator_indices.size()) {
|
||||
std::swap(allocator_indices[random_index],
|
||||
allocator_indices[allocator_indices.size() - 1]);
|
||||
}
|
||||
allocator_indices.pop_back();
|
||||
}
|
||||
|
||||
return nullptr;
|
||||
}
|
||||
};
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
#include <cstdint>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <variant>
|
||||
|
|
@ -289,6 +290,10 @@ class AllocatedBuffer {
|
|||
// Serialize the buffer into a descriptor for transfer
|
||||
[[nodiscard]] Descriptor get_descriptor() const;
|
||||
|
||||
[[nodiscard]] std::string getSegmentName() const noexcept {
|
||||
return segment_name_;
|
||||
}
|
||||
|
||||
// Friend declaration for operator<<
|
||||
friend std::ostream& operator<<(std::ostream& os,
|
||||
const AllocatedBuffer& buffer);
|
||||
|
|
@ -393,6 +398,9 @@ class Replica {
|
|||
return false; // DiskReplicaData does not have handles
|
||||
}
|
||||
|
||||
[[nodiscard]] std::vector<std::optional<std::string>> get_segment_names()
|
||||
const;
|
||||
|
||||
void mark_complete() {
|
||||
if (status_ == ReplicaStatus::PROCESSING) {
|
||||
status_ = ReplicaStatus::COMPLETE;
|
||||
|
|
@ -504,6 +512,25 @@ inline Replica::Descriptor Replica::get_descriptor() const {
|
|||
return desc;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
return segment_names;
|
||||
}
|
||||
return std::vector<std::optional<std::string>>();
|
||||
}
|
||||
|
||||
inline std::ostream& operator<<(std::ostream& os, const Replica& replica) {
|
||||
os << "Replica: { status: " << replica.status_ << ", ";
|
||||
|
||||
|
|
|
|||
|
|
@ -392,40 +392,26 @@ auto MasterService::PutStart(const std::string& key,
|
|||
|
||||
// Allocate replicas
|
||||
std::vector<Replica> replicas;
|
||||
replicas.reserve(config.replica_num + use_disk_replica_);
|
||||
{
|
||||
ScopedAllocatorAccess allocator_access =
|
||||
segment_manager_.getAllocatorAccess();
|
||||
auto& allocators = allocator_access.getAllocators();
|
||||
auto& allocators_by_name = allocator_access.getAllocatorsByName();
|
||||
for (size_t i = 0; i < config.replica_num; ++i) {
|
||||
std::vector<std::unique_ptr<AllocatedBuffer>> handles;
|
||||
handles.reserve(slice_lengths.size());
|
||||
|
||||
// Allocate space for each slice
|
||||
for (size_t j = 0; j < slice_lengths.size(); ++j) {
|
||||
auto chunk_size = slice_lengths[j];
|
||||
auto allocation_result = allocation_strategy_->Allocate(
|
||||
allocators, allocators_by_name, slice_lengths, config);
|
||||
|
||||
// Use the unified allocation strategy with replica config
|
||||
auto handle = allocation_strategy_->Allocate(
|
||||
allocators, allocators_by_name, chunk_size, config);
|
||||
|
||||
if (!handle) {
|
||||
// If the allocation failed, we need to evict some objects
|
||||
// to free up space for future allocations.
|
||||
need_eviction_ = true;
|
||||
return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
|
||||
}
|
||||
|
||||
VLOG(1) << "key=" << key << ", replica_id=" << i
|
||||
<< ", slice_index=" << j << ", handle=" << *handle
|
||||
<< ", action=slice_allocated";
|
||||
handles.emplace_back(std::move(handle));
|
||||
if (!allocation_result.has_value()) {
|
||||
LOG(ERROR) << "Failed to allocate all replicas for key=" << key
|
||||
<< ", error: " << allocation_result.error();
|
||||
if (allocation_result.error() == ErrorCode::INVALID_PARAMS) {
|
||||
return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
|
||||
}
|
||||
|
||||
replicas.emplace_back(std::move(handles),
|
||||
ReplicaStatus::PROCESSING);
|
||||
need_eviction_ = true;
|
||||
return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
|
||||
}
|
||||
|
||||
replicas = std::move(allocation_result.value());
|
||||
}
|
||||
|
||||
// If disk replica is enabled, allocate a disk replica
|
||||
|
|
|
|||
|
|
@ -23,7 +23,8 @@ class AllocationStrategyTest : public ::testing::Test {
|
|||
const std::string& segment_name, size_t base_offset,
|
||||
BufferAllocatorType allocator_type) {
|
||||
const size_t base = 0x100000000ULL + base_offset; // 4GB + offset
|
||||
const size_t size = 1024 * 1024 * 16; // 16MB (multiple of 4MB)
|
||||
const size_t size =
|
||||
1024 * 1024 * 64; // 64MB (for multiple slabs in cachelib)
|
||||
switch (allocator_type) {
|
||||
case BufferAllocatorType::CACHELIB:
|
||||
return std::make_shared<CachelibBufferAllocator>(segment_name,
|
||||
|
|
@ -50,9 +51,11 @@ TEST_F(AllocationStrategyTest, EmptyAllocatorsMap) {
|
|||
std::vector<std::shared_ptr<BufferAllocatorBase>> empty_allocators;
|
||||
ReplicateConfig config{1, false, "local"};
|
||||
|
||||
auto result = strategy_->Allocate(empty_allocators,
|
||||
empty_allocators_by_name, 100, config);
|
||||
EXPECT_EQ(result, nullptr);
|
||||
std::vector<size_t> slice_sizes = {100};
|
||||
auto result = strategy_->Allocate(
|
||||
empty_allocators, empty_allocators_by_name, slice_sizes, config);
|
||||
EXPECT_FALSE(result.has_value());
|
||||
EXPECT_EQ(result.error(), ErrorCode::NO_AVAILABLE_HANDLE);
|
||||
}
|
||||
|
||||
// Test preferred segment behavior with empty allocators
|
||||
|
|
@ -63,9 +66,11 @@ TEST_F(AllocationStrategyTest, PreferredSegmentWithEmptyAllocators) {
|
|||
std::vector<std::shared_ptr<BufferAllocatorBase>> empty_allocators;
|
||||
ReplicateConfig config{1, false, "preferred_segment"};
|
||||
|
||||
auto result = strategy_->Allocate(empty_allocators,
|
||||
empty_allocators_by_name, 100, config);
|
||||
EXPECT_EQ(result, nullptr); // Should return nullptr for empty allocators
|
||||
std::vector<size_t> slice_sizes = {100};
|
||||
auto result = strategy_->Allocate(
|
||||
empty_allocators, empty_allocators_by_name, slice_sizes, config);
|
||||
EXPECT_FALSE(result.has_value());
|
||||
EXPECT_EQ(result.error(), ErrorCode::NO_AVAILABLE_HANDLE);
|
||||
}
|
||||
|
||||
// Test preferred segment allocation when available
|
||||
|
|
@ -86,13 +91,21 @@ TEST_F(AllocationStrategyTest, PreferredSegmentAllocation) {
|
|||
allocators.push_back(allocator2);
|
||||
|
||||
ReplicateConfig config{1, false, "preferred"};
|
||||
size_t alloc_size = 1024;
|
||||
std::vector<size_t> slice_sizes = {1024};
|
||||
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
alloc_size, config);
|
||||
ASSERT_NE(result, nullptr);
|
||||
EXPECT_EQ(result->get_descriptor().segment_name_, "preferred");
|
||||
EXPECT_EQ(result->get_descriptor().size_, alloc_size);
|
||||
slice_sizes, config);
|
||||
ASSERT_TRUE(result.has_value());
|
||||
EXPECT_EQ(result.value().size(), 1);
|
||||
ASSERT_FALSE(result.value().empty());
|
||||
|
||||
const auto& replica = result.value()[0];
|
||||
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].segment_name_, "preferred");
|
||||
EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 1024);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -114,20 +127,62 @@ TEST_F(AllocationStrategyTest, PreferredSegmentNotFound) {
|
|||
allocators.push_back(allocator2);
|
||||
|
||||
ReplicateConfig config{1, false, "nonexistent"};
|
||||
size_t alloc_size = 1024;
|
||||
std::vector<size_t> slice_sizes = {1024};
|
||||
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
alloc_size, config);
|
||||
ASSERT_NE(result, nullptr);
|
||||
// Should allocate from one of the available segments
|
||||
std::string segment_name = result->get_descriptor().segment_name_;
|
||||
slice_sizes, config);
|
||||
ASSERT_TRUE(result.has_value());
|
||||
EXPECT_EQ(result.value().size(), 1);
|
||||
|
||||
const auto& replica = result.value()[0];
|
||||
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_name = mem_desc.buffer_descriptors[0].segment_name_;
|
||||
EXPECT_TRUE(segment_name == "segment1" || segment_name == "segment2");
|
||||
EXPECT_EQ(result->get_descriptor().size_, alloc_size);
|
||||
EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 1024);
|
||||
}
|
||||
}
|
||||
|
||||
// Test multiple allocators with random selection
|
||||
TEST_F(AllocationStrategyTest, MultipleAllocatorsRandomSelection) {
|
||||
// Test multiple slices allocation
|
||||
TEST_F(AllocationStrategyTest, MultipleSlicesAllocation) {
|
||||
for (const auto& allocator_type : allocator_types_) {
|
||||
auto allocator1 = CreateTestAllocator("segment1", 0, allocator_type);
|
||||
auto allocator2 =
|
||||
CreateTestAllocator("segment2", 0x10000000ULL, allocator_type);
|
||||
|
||||
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(allocator1);
|
||||
allocators_by_name["segment2"].push_back(allocator2);
|
||||
allocators.push_back(allocator1);
|
||||
allocators.push_back(allocator2);
|
||||
|
||||
ReplicateConfig config{1, false, ""};
|
||||
std::vector<size_t> slice_sizes = {1024, 2048, 512};
|
||||
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
slice_sizes, config);
|
||||
ASSERT_TRUE(result.has_value());
|
||||
EXPECT_EQ(result.value().size(), 1);
|
||||
|
||||
const auto& replica = result.value()[0];
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
// Test multiple replicas allocation
|
||||
TEST_F(AllocationStrategyTest, MultipleReplicasAllocation) {
|
||||
for (const auto& allocator_type : allocator_types_) {
|
||||
auto allocator1 = CreateTestAllocator("segment1", 0, allocator_type);
|
||||
auto allocator2 =
|
||||
|
|
@ -147,24 +202,33 @@ TEST_F(AllocationStrategyTest, MultipleAllocatorsRandomSelection) {
|
|||
allocators.push_back(allocator2);
|
||||
allocators.push_back(allocator3);
|
||||
|
||||
ReplicateConfig config{1, false, ""}; // No preferred segment
|
||||
size_t alloc_size = 1024;
|
||||
ReplicateConfig config{3, false, ""}; // Request 3 replicas
|
||||
std::vector<size_t> slice_sizes = {1024, 2048};
|
||||
|
||||
// Perform multiple allocations to test randomness
|
||||
std::vector<std::string> allocated_segments;
|
||||
for (int i = 0; i < 10; ++i) {
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
alloc_size, config);
|
||||
ASSERT_NE(result, nullptr);
|
||||
allocated_segments.push_back(
|
||||
result->get_descriptor().segment_name_);
|
||||
EXPECT_EQ(result->get_descriptor().size_, alloc_size);
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
slice_sizes, config);
|
||||
ASSERT_TRUE(result.has_value());
|
||||
EXPECT_EQ(result.value().size(), 3);
|
||||
|
||||
// Check each replica has all slices
|
||||
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);
|
||||
}
|
||||
|
||||
// Verify that allocations happened on available segments
|
||||
for (const auto& segment : allocated_segments) {
|
||||
EXPECT_TRUE(segment == "segment1" || segment == "segment2" ||
|
||||
segment == "segment3");
|
||||
// Check that replicas are on different segments
|
||||
std::set<std::string> used_segments;
|
||||
for (const auto& replica : result.value()) {
|
||||
auto segment_names = replica.get_segment_names();
|
||||
for (const auto& name_ptr : segment_names) {
|
||||
if (name_ptr) {
|
||||
used_segments.insert(*name_ptr);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -188,25 +252,31 @@ TEST_F(AllocationStrategyTest, PreferredSegmentInsufficientSpace) {
|
|||
|
||||
// First, fill up the preferred allocator
|
||||
ReplicateConfig config{1, false, "preferred"};
|
||||
std::vector<std::unique_ptr<AllocatedBuffer>> buffers;
|
||||
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
|
||||
|
||||
// Allocate most of the space in preferred segment
|
||||
size_t large_alloc = 15 * 1024 * 1024; // 15MB out of 16MB
|
||||
auto large_buffer = strategy_->Allocate(allocators, allocators_by_name,
|
||||
large_alloc, config);
|
||||
ASSERT_NE(large_buffer, nullptr);
|
||||
EXPECT_EQ(large_buffer->get_descriptor().segment_name_, "preferred");
|
||||
buffers.push_back(std::move(large_buffer));
|
||||
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]
|
||||
.segment_name_,
|
||||
"preferred");
|
||||
|
||||
// Now try to allocate more than remaining space in preferred segment
|
||||
size_t small_alloc = 2 * 1024 * 1024; // 2MB (more than remaining ~1MB)
|
||||
std::vector<size_t> small_slice = {2 * 1024 * 1024};
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
small_alloc, config);
|
||||
ASSERT_NE(result, nullptr);
|
||||
// Should fall back to segment1 since preferred doesn't have enough
|
||||
// space
|
||||
EXPECT_EQ(result->get_descriptor().segment_name_, "segment1");
|
||||
EXPECT_EQ(result->get_descriptor().size_, small_alloc);
|
||||
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].segment_name_, "segment1");
|
||||
EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 2 * 1024 * 1024);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -228,24 +298,25 @@ TEST_F(AllocationStrategyTest, AllAllocatorsFull) {
|
|||
allocators.push_back(allocator2);
|
||||
|
||||
ReplicateConfig config{1, false, ""};
|
||||
std::vector<std::unique_ptr<AllocatedBuffer>> buffers;
|
||||
|
||||
// Fill up both allocators
|
||||
size_t large_alloc = 15 * 1024 * 1024; // 15MB each
|
||||
auto buffer1 = strategy_->Allocate(allocators, allocators_by_name,
|
||||
large_alloc, config);
|
||||
auto buffer2 = strategy_->Allocate(allocators, allocators_by_name,
|
||||
large_alloc, config);
|
||||
ASSERT_NE(buffer1, nullptr);
|
||||
ASSERT_NE(buffer2, nullptr);
|
||||
buffers.push_back(std::move(buffer1));
|
||||
buffers.push_back(std::move(buffer2));
|
||||
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());
|
||||
|
||||
// Try to allocate more than remaining space
|
||||
size_t impossible_alloc = 5 * 1024 * 1024; // 5MB (more than remaining)
|
||||
std::vector<size_t> impossible_slice = {
|
||||
5 * 1024 * 1024}; // 5MB (more than remaining)
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
impossible_alloc, config);
|
||||
EXPECT_EQ(result, nullptr); // Should fail
|
||||
impossible_slice, config);
|
||||
EXPECT_FALSE(result.has_value());
|
||||
EXPECT_EQ(result.error(), ErrorCode::NO_AVAILABLE_HANDLE);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -262,14 +333,12 @@ TEST_F(AllocationStrategyTest, ZeroSizeAllocation) {
|
|||
allocators.push_back(allocator);
|
||||
|
||||
ReplicateConfig config{1, false, ""};
|
||||
std::vector<size_t> zero_slice = {0};
|
||||
|
||||
auto result =
|
||||
strategy_->Allocate(allocators, allocators_by_name, 0, config);
|
||||
// Zero-size allocation behavior depends on BufferAllocator
|
||||
// implementation This test documents the current behavior
|
||||
if (result != nullptr) {
|
||||
EXPECT_EQ(result->get_descriptor().segment_name_, "segment1");
|
||||
}
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
zero_slice, config);
|
||||
EXPECT_FALSE(result.has_value());
|
||||
EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -286,13 +355,104 @@ TEST_F(AllocationStrategyTest, VeryLargeSizeAllocation) {
|
|||
allocators.push_back(allocator);
|
||||
|
||||
ReplicateConfig config{1, false, ""};
|
||||
size_t huge_size =
|
||||
100 * 1024 * 1024; // 100MB (larger than 16MB capacity)
|
||||
std::vector<size_t> huge_slice = {
|
||||
100 * 1024 * 1024}; // 100MB (larger than 64MB capacity)
|
||||
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
huge_size, config);
|
||||
EXPECT_EQ(result, nullptr); // Should fail due to insufficient capacity
|
||||
huge_slice, config);
|
||||
EXPECT_FALSE(result.has_value());
|
||||
EXPECT_EQ(result.error(), ErrorCode::NO_AVAILABLE_HANDLE);
|
||||
}
|
||||
}
|
||||
|
||||
// Test empty slice sizes
|
||||
TEST_F(AllocationStrategyTest, EmptySliceSizes) {
|
||||
auto allocator =
|
||||
CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET);
|
||||
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 invalid replication count
|
||||
TEST_F(AllocationStrategyTest, InvalidReplicationCount) {
|
||||
auto allocator =
|
||||
CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET);
|
||||
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{0, false, ""}; // Invalid: 0 replicas
|
||||
std::vector<size_t> slice_sizes = {1024};
|
||||
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
slice_sizes, config);
|
||||
EXPECT_FALSE(result.has_value());
|
||||
EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS);
|
||||
}
|
||||
|
||||
// Test best-effort behavior when insufficient allocators for requested replica
|
||||
// count
|
||||
TEST_F(AllocationStrategyTest, InsufficientAllocatorsForReplicas) {
|
||||
auto allocator1 =
|
||||
CreateTestAllocator("segment1", 0, BufferAllocatorType::OFFSET);
|
||||
auto allocator2 = CreateTestAllocator("segment2", 0x10000000ULL,
|
||||
BufferAllocatorType::OFFSET);
|
||||
|
||||
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(allocator1);
|
||||
allocators_by_name["segment2"].push_back(allocator2);
|
||||
allocators.push_back(allocator1);
|
||||
allocators.push_back(allocator2);
|
||||
|
||||
ReplicateConfig config{
|
||||
5, false, ""}; // Request 5 replicas, but only 2 segments available
|
||||
std::vector<size_t> slice_sizes = {1024};
|
||||
|
||||
auto result = strategy_->Allocate(allocators, allocators_by_name,
|
||||
slice_sizes, config);
|
||||
// With best-effort semantics, should succeed with available replicas
|
||||
EXPECT_TRUE(result.has_value());
|
||||
// Should get 2 replicas (limited by number of segments)
|
||||
EXPECT_EQ(2u, result.value().size());
|
||||
|
||||
// Verify each replica has the expected slice structure
|
||||
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(), 1u);
|
||||
EXPECT_EQ(mem_desc.buffer_descriptors[0].size_, 1024u);
|
||||
}
|
||||
|
||||
// Verify replicas are on different segments
|
||||
std::unordered_set<std::string> segment_names;
|
||||
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].segment_name_);
|
||||
}
|
||||
EXPECT_EQ(2u, segment_names.size());
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
@ -8,6 +8,7 @@
|
|||
#include <random>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
#include <unordered_set>
|
||||
|
||||
#include "types.h"
|
||||
|
||||
|
|
@ -355,15 +356,17 @@ TEST_F(MasterServiceTest, PutStartEndFlow) {
|
|||
|
||||
TEST_F(MasterServiceTest, RandomPutStartEndFlow) {
|
||||
std::unique_ptr<MasterService> service_(new MasterService());
|
||||
constexpr size_t buffer = 0x300000000;
|
||||
constexpr size_t size = 1024 * 1024 * 16;
|
||||
std::string segment_name = "test_segment";
|
||||
|
||||
Segment segment(generate_uuid(), segment_name, buffer, size);
|
||||
// Mount 5 segments, each 16MB
|
||||
constexpr size_t kBaseAddr = 0x300000000;
|
||||
constexpr size_t kSegmentSize = 1024 * 1024 * 16; // 16MB
|
||||
UUID client_id = generate_uuid();
|
||||
|
||||
auto mount_result = service_->MountSegment(segment, client_id);
|
||||
ASSERT_TRUE(mount_result.has_value());
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
Segment segment(generate_uuid(), "segment_" + std::to_string(i),
|
||||
kBaseAddr + static_cast<size_t>(i) * kSegmentSize,
|
||||
kSegmentSize);
|
||||
ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
|
||||
}
|
||||
|
||||
// Test PutStart
|
||||
std::string key = "test_key";
|
||||
|
|
@ -954,20 +957,22 @@ TEST_F(MasterServiceTest, RemoveAll) {
|
|||
}
|
||||
|
||||
TEST_F(MasterServiceTest, MultiSliceMultiReplicaFlow) {
|
||||
auto service_config = MasterServiceConfig::builder().build();
|
||||
const uint64_t kv_lease_ttl = 50;
|
||||
auto service_config = MasterServiceConfig::builder()
|
||||
.set_default_kv_lease_ttl(kv_lease_ttl)
|
||||
.build();
|
||||
std::unique_ptr<MasterService> service_(new MasterService(service_config));
|
||||
|
||||
// Mount a segment with sufficient size for multiple replicas
|
||||
constexpr size_t buffer = 0x300000000;
|
||||
constexpr size_t segment_size =
|
||||
1024 * 1024 * 64; // 64MB to accommodate multiple replicas
|
||||
std::string segment_name = "test_segment_multi";
|
||||
|
||||
Segment segment(generate_uuid(), segment_name, buffer, segment_size);
|
||||
// Mount 3 segments, each 64MB
|
||||
constexpr size_t kBaseAddr = 0x300000000;
|
||||
constexpr size_t kSegmentSize = 1024 * 1024 * 64; // 64MB
|
||||
UUID client_id = generate_uuid();
|
||||
|
||||
auto mount_result = service_->MountSegment(segment, client_id);
|
||||
ASSERT_TRUE(mount_result.has_value());
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
Segment segment(generate_uuid(), "segment_" + std::to_string(i),
|
||||
kBaseAddr + static_cast<size_t>(i) * kSegmentSize,
|
||||
kSegmentSize);
|
||||
ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
|
||||
}
|
||||
|
||||
// Test parameters
|
||||
std::string key = "multi_slice_object";
|
||||
|
|
@ -1386,13 +1391,11 @@ TEST_F(MasterServiceTest, UnmountSegmentImmediateCleanup) {
|
|||
TEST_F(MasterServiceTest, ReadableAfterPartialUnmountWithReplication) {
|
||||
std::unique_ptr<MasterService> service_(new MasterService());
|
||||
|
||||
// TODO: mount two larger segments when replication affinity fixed
|
||||
// Mount two segments sized to fit exactly one replica each
|
||||
// Mount two large segments
|
||||
constexpr size_t buffer1 = 0x300000000;
|
||||
constexpr size_t buffer2 = 0x400000000;
|
||||
constexpr size_t segment_size = 1024 * 1024 * 16;
|
||||
constexpr size_t object_size =
|
||||
segment_size / 2 + 16; // force at most 1 replica per segment
|
||||
constexpr size_t segment_size = 1024 * 1024 * 64; // 64MB
|
||||
constexpr size_t object_size = 1024 * 1024; // 1MB
|
||||
|
||||
Segment segment1(generate_uuid(), "segment1", buffer1, segment_size);
|
||||
Segment segment2(generate_uuid(), "segment2", buffer2, segment_size);
|
||||
|
|
@ -2008,6 +2011,80 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotAllowEvict) {
|
|||
service_->RemoveAll();
|
||||
}
|
||||
|
||||
TEST_F(MasterServiceTest, PerSliceReplicaSegmentsAreUnique) {
|
||||
std::unique_ptr<MasterService> service_(new MasterService());
|
||||
|
||||
// Mount 20 segments, each 16MB and slab-aligned
|
||||
constexpr size_t kBaseAddr = 0x300000000;
|
||||
constexpr size_t kSegmentSize = 1024 * 1024 * 16; // 16MB
|
||||
UUID client_id = generate_uuid();
|
||||
for (int i = 0; i < 20; ++i) {
|
||||
Segment segment(generate_uuid(), "segment_" + std::to_string(i),
|
||||
kBaseAddr + static_cast<size_t>(i) * kSegmentSize,
|
||||
kSegmentSize);
|
||||
ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
|
||||
}
|
||||
|
||||
// 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);
|
||||
ReplicateConfig config;
|
||||
config.replica_num = 10;
|
||||
|
||||
auto put_start_result = service_->PutStart(key, slice_lengths, 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].segment_name_);
|
||||
}
|
||||
EXPECT_EQ(segment_names.size(), config.replica_num)
|
||||
<< "Duplicate segment found for slice index " << slice_idx;
|
||||
}
|
||||
|
||||
ASSERT_TRUE(service_->PutEnd(key, ReplicaType::MEMORY).has_value());
|
||||
}
|
||||
|
||||
TEST_F(MasterServiceTest, ReplicationFactorTwoWithSingleSegment) {
|
||||
std::unique_ptr<MasterService> service_(new MasterService());
|
||||
|
||||
// Mount a single 16MB segment
|
||||
constexpr size_t kBaseAddr = 0x300000000;
|
||||
constexpr size_t kSegmentSize = 1024 * 1024 * 16; // 16MB
|
||||
Segment segment(generate_uuid(), "single_segment", kBaseAddr, kSegmentSize);
|
||||
UUID client_id = generate_uuid();
|
||||
ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value());
|
||||
|
||||
// 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};
|
||||
ReplicateConfig config;
|
||||
config.replica_num = 2;
|
||||
|
||||
auto put_start_result = service_->PutStart(key, slice_lengths, config);
|
||||
ASSERT_TRUE(put_start_result.has_value());
|
||||
auto replicas = put_start_result.value();
|
||||
|
||||
// Should get 1 replica instead of the requested 2 (best-effort)
|
||||
EXPECT_EQ(1u, replicas.size());
|
||||
EXPECT_TRUE(replicas[0].is_memory_replica());
|
||||
|
||||
// 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].segment_name_);
|
||||
EXPECT_EQ(1024u, mem_desc.buffer_descriptors[0].size_);
|
||||
}
|
||||
|
||||
TEST_F(MasterServiceTest, BatchExistKeyTest) {
|
||||
std::unique_ptr<MasterService> service_(new MasterService());
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ DEFAULT_DEFAULT_KV_LEASE_TTL = 5000 # 5000 milliseconds
|
|||
# Use environment variable if set, otherwise use default
|
||||
default_kv_lease_ttl = int(os.getenv("DEFAULT_KV_LEASE_TTL", DEFAULT_DEFAULT_KV_LEASE_TTL))
|
||||
|
||||
|
||||
def get_client(store, local_buffer_size_param=None):
|
||||
"""Initialize and setup the distributed store client."""
|
||||
protocol = os.getenv("PROTOCOL", "tcp")
|
||||
|
|
@ -57,7 +58,9 @@ class TestZeroLocalBufferSize(unittest.TestCase):
|
|||
result = zero_buffer_store.is_exist(key)
|
||||
self.assertEqual(result, 0, "Key should not exist after failed put")
|
||||
|
||||
class TestDistributedObjectStore(unittest.TestCase):
|
||||
class TestDistributedObjectStoreSingleStore(unittest.TestCase):
|
||||
"""Test class for single store operations (no replication)."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Initialize the store once for all tests."""
|
||||
|
|
@ -567,214 +570,6 @@ class TestDistributedObjectStore(unittest.TestCase):
|
|||
self.assertIsInstance(config_str, str)
|
||||
self.assertIn("3", config_str) # Should contain replica_num
|
||||
|
||||
def test_put_with_config_parameter(self):
|
||||
"""Test put method with config parameter."""
|
||||
from mooncake.store import ReplicateConfig
|
||||
|
||||
test_data = b"Hello, Config World!"
|
||||
key = "test_put_config_key"
|
||||
|
||||
# Test with default config (backward compatibility)
|
||||
result = self.store.put(key=key, value=test_data)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, test_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
# Test with custom config
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = 2
|
||||
|
||||
key2 = "test_put_config_key2"
|
||||
result = self.store.put(key=key2, value=test_data, config=config)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
retrieved_data = self.store.get(key2)
|
||||
self.assertEqual(retrieved_data, test_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.remove(key2), 0)
|
||||
|
||||
with self.assertRaises(TypeError):
|
||||
result = self.store.put(key_arg_name_error=key, value=test_data, config=config)
|
||||
|
||||
with self.assertRaises(TypeError):
|
||||
result = self.store.put(key=key, value_arg_name_error=test_data, config=config)
|
||||
|
||||
with self.assertRaises(TypeError):
|
||||
result = self.store.put(key=key, value=test_data, config_arg_name_error=config)
|
||||
|
||||
|
||||
def test_put_batch_with_config_parameter(self):
|
||||
"""Test put_batch method with config parameter."""
|
||||
from mooncake.store import ReplicateConfig
|
||||
|
||||
keys = ["test_batch_config_key1", "test_batch_config_key2", "test_batch_config_key3"]
|
||||
values = [b"Batch Data 1", b"Batch Data 2", b"Batch Data 3"]
|
||||
|
||||
# Test with default config (backward compatibility)
|
||||
result = self.store.put_batch(keys, values)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
for key, expected_value in zip(keys, values):
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, expected_value)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
for key in keys:
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
# Test with custom config
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = 2
|
||||
|
||||
keys2 = ["test_batch_config_key4", "test_batch_config_key5", "test_batch_config_key6"]
|
||||
result = self.store.put_batch(keys=keys2, values=values, config=config)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
for key, expected_value in zip(keys2, values):
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, expected_value)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
for key in keys2:
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
def test_put_from_with_config_parameter(self):
|
||||
"""Test put_from method with config parameter."""
|
||||
import ctypes
|
||||
from mooncake.store import ReplicateConfig
|
||||
|
||||
test_data = b"Hello, put_from config world!"
|
||||
key = "test_put_from_config_key"
|
||||
buffer_size = len(test_data)
|
||||
|
||||
# Allocate and register buffer
|
||||
buffer = (ctypes.c_ubyte * buffer_size)()
|
||||
buffer_ptr = ctypes.addressof(buffer)
|
||||
|
||||
result = self.store.register_buffer(buffer_ptr, buffer_size)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Copy test data to buffer
|
||||
ctypes.memmove(buffer, test_data, len(test_data))
|
||||
|
||||
# Test with default config (backward compatibility)
|
||||
result = self.store.put_from(key=key, buffer_ptr=buffer_ptr, size=len(test_data))
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, test_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
# Test with custom config
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = 2
|
||||
config.with_soft_pin = False
|
||||
|
||||
key2 = "test_put_from_config_key2"
|
||||
result = self.store.put_from(key=key2, buffer_ptr=buffer_ptr, size=len(test_data), config=config)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
retrieved_data = self.store.get(key2)
|
||||
self.assertEqual(retrieved_data, test_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
|
||||
self.assertEqual(self.store.remove(key2), 0)
|
||||
|
||||
def test_batch_put_from_with_config_parameter(self):
|
||||
"""Test batch_put_from method with config parameter."""
|
||||
import ctypes
|
||||
from mooncake.store import ReplicateConfig
|
||||
|
||||
# Test data
|
||||
test_data = [
|
||||
b"Batch Config Data 1",
|
||||
b"Batch Config Data 2",
|
||||
b"Batch Config Data 3"
|
||||
]
|
||||
keys = ["test_batch_put_from_config_key1", "test_batch_put_from_config_key2", "test_batch_put_from_config_key3"]
|
||||
|
||||
# Use large spacing between buffers
|
||||
buffer_spacing = 1024 * 1024 # 1MB spacing
|
||||
total_buffer_size = buffer_spacing * len(test_data)
|
||||
large_buffer = (ctypes.c_ubyte * total_buffer_size)()
|
||||
large_buffer_ptr = ctypes.addressof(large_buffer)
|
||||
|
||||
# Register buffer
|
||||
result = self.store.register_buffer(large_buffer_ptr, total_buffer_size)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Prepare individual buffers
|
||||
buffer_ptrs = []
|
||||
buffer_sizes = []
|
||||
|
||||
for i, data in enumerate(test_data):
|
||||
offset = i * buffer_spacing
|
||||
buffer_ptr = large_buffer_ptr + offset
|
||||
|
||||
# Copy test data to buffer
|
||||
ctypes.memmove(ctypes.c_void_p(buffer_ptr), data, len(data))
|
||||
|
||||
buffer_ptrs.append(buffer_ptr)
|
||||
buffer_sizes.append(len(data))
|
||||
|
||||
# Test with default config (backward compatibility)
|
||||
results = self.store.batch_put_from(keys=keys, buffer_ptrs=buffer_ptrs, sizes=buffer_sizes)
|
||||
self.assertEqual(len(results), len(keys))
|
||||
for result in results:
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
for key, expected_data in zip(keys, test_data):
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, expected_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
for key in keys:
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
# Test with custom config
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = 2
|
||||
config.with_soft_pin = False
|
||||
|
||||
keys2 = ["test_batch_put_from_config_key4", "test_batch_put_from_config_key5", "test_batch_put_from_config_key6"]
|
||||
results = self.store.batch_put_from(keys=keys2, buffer_ptrs=buffer_ptrs, sizes=buffer_sizes, config=config)
|
||||
self.assertEqual(len(results), len(keys2))
|
||||
for result in results:
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
for key, expected_data in zip(keys2, test_data):
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, expected_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.unregister_buffer(large_buffer_ptr), 0)
|
||||
for key in keys2:
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
def test_batch_get_buffer_operations(self):
|
||||
"""Test batch_get_buffer operations for multiple keys."""
|
||||
# Test data
|
||||
|
|
@ -839,5 +634,6 @@ class TestDistributedObjectStore(unittest.TestCase):
|
|||
for key in keys:
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
@ -0,0 +1,306 @@
|
|||
import unittest
|
||||
import os
|
||||
import time
|
||||
from mooncake.store import MooncakeDistributedStore, ReplicateConfig
|
||||
|
||||
# The lease time of the kv object, should be set equal to
|
||||
# the master's value.
|
||||
DEFAULT_DEFAULT_KV_LEASE_TTL = 5000 # 5000 milliseconds
|
||||
# Use environment variable if set, otherwise use default
|
||||
default_kv_lease_ttl = int(os.getenv("DEFAULT_KV_LEASE_TTL", DEFAULT_DEFAULT_KV_LEASE_TTL))
|
||||
|
||||
|
||||
def get_clients(stores, local_buffer_size_param=None):
|
||||
"""Initialize and setup the distributed store clients."""
|
||||
protocol = os.getenv("PROTOCOL", "tcp")
|
||||
device_name = os.getenv("DEVICE_NAME", "ibp6s0")
|
||||
base_hostname = os.getenv("LOCAL_HOSTNAME", "localhost")
|
||||
metadata_server = os.getenv("MC_METADATA_SERVER", "http://127.0.0.1:8080/metadata")
|
||||
segment_size = 1600 * 1024 * 1024 # 1600 MB per segment
|
||||
local_buffer_size = (
|
||||
local_buffer_size_param if local_buffer_size_param is not None
|
||||
else 512 * 1024 * 1024 # 512 MB
|
||||
)
|
||||
master_server_address = os.getenv("MASTER_SERVER", "127.0.0.1:50051")
|
||||
|
||||
base_port = 12345
|
||||
for i, store in enumerate(stores):
|
||||
hostname = f"{base_hostname}:{base_port + i}"
|
||||
retcode = store.setup(
|
||||
hostname,
|
||||
metadata_server,
|
||||
segment_size,
|
||||
local_buffer_size,
|
||||
protocol,
|
||||
device_name,
|
||||
master_server_address
|
||||
)
|
||||
|
||||
if retcode:
|
||||
raise RuntimeError(f"Failed to setup segment. Return code: {retcode}")
|
||||
|
||||
def get_client(store, local_buffer_size_param=None):
|
||||
"""Initialize and setup the distributed store client."""
|
||||
return get_clients([store], local_buffer_size_param)
|
||||
|
||||
class TestDistributedObjectStoreReplication(unittest.TestCase):
|
||||
"""Test class for replication operations (multiple stores)."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Initialize the main store and additional stores for replication."""
|
||||
cls.store = MooncakeDistributedStore()
|
||||
|
||||
# Additional stores for replication testing
|
||||
cls.additional_stores = []
|
||||
cls.max_replicate_num = 2
|
||||
for _ in range(cls.max_replicate_num - 1): # -1 because main store is already created
|
||||
cls.additional_stores.append(MooncakeDistributedStore())
|
||||
|
||||
get_clients([cls.store] + cls.additional_stores)
|
||||
|
||||
def test_put_with_config_parameter(self):
|
||||
"""Test put method with config parameter."""
|
||||
|
||||
test_data = b"Hello, Config World!"
|
||||
key = "test_put_config_key"
|
||||
|
||||
# Test with default config (backward compatibility)
|
||||
result = self.store.put(key=key, value=test_data)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, test_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
# Test with custom config
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = self.max_replicate_num
|
||||
|
||||
key2 = "test_put_config_key2"
|
||||
result = self.store.put(key=key2, value=test_data, config=config)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
retrieved_data = self.store.get(key2)
|
||||
self.assertEqual(retrieved_data, test_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.remove(key2), 0)
|
||||
|
||||
with self.assertRaises(TypeError):
|
||||
result = self.store.put(key_arg_name_error=key, value=test_data, config=config)
|
||||
|
||||
with self.assertRaises(TypeError):
|
||||
result = self.store.put(key=key, value_arg_name_error=test_data, config=config)
|
||||
|
||||
with self.assertRaises(TypeError):
|
||||
result = self.store.put(key=key, value=test_data, config_arg_name_error=config)
|
||||
|
||||
def test_put_batch_with_config_parameter(self):
|
||||
"""Test put_batch method with config parameter."""
|
||||
|
||||
keys = ["test_batch_config_key1", "test_batch_config_key2", "test_batch_config_key3"]
|
||||
values = [b"Batch Data 1", b"Batch Data 2", b"Batch Data 3"]
|
||||
|
||||
# Test with default config (backward compatibility)
|
||||
result = self.store.put_batch(keys, values)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
for key, expected_value in zip(keys, values):
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, expected_value)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
for key in keys:
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
# Test with custom config
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = self.max_replicate_num
|
||||
|
||||
keys2 = ["test_batch_config_key4", "test_batch_config_key5", "test_batch_config_key6"]
|
||||
result = self.store.put_batch(keys=keys2, values=values, config=config)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
for key, expected_value in zip(keys2, values):
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, expected_value)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
for key in keys2:
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
def test_put_from_with_config_parameter(self):
|
||||
"""Test put_from method with config parameter."""
|
||||
import ctypes
|
||||
|
||||
test_data = b"Hello, put_from config world!"
|
||||
key = "test_put_from_config_key"
|
||||
buffer_size = len(test_data)
|
||||
|
||||
# Allocate and register buffer
|
||||
buffer = (ctypes.c_ubyte * buffer_size)()
|
||||
buffer_ptr = ctypes.addressof(buffer)
|
||||
|
||||
result = self.store.register_buffer(buffer_ptr, buffer_size)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Copy test data to buffer
|
||||
ctypes.memmove(buffer, test_data, len(test_data))
|
||||
|
||||
# Test with default config (backward compatibility)
|
||||
result = self.store.put_from(key=key, buffer_ptr=buffer_ptr, size=len(test_data))
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, test_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
# Test with custom config
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = self.max_replicate_num
|
||||
config.with_soft_pin = False
|
||||
|
||||
key2 = "test_put_from_config_key2"
|
||||
result = self.store.put_from(key=key2, buffer_ptr=buffer_ptr, size=len(test_data), config=config)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
retrieved_data = self.store.get(key2)
|
||||
self.assertEqual(retrieved_data, test_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.unregister_buffer(buffer_ptr), 0)
|
||||
self.assertEqual(self.store.remove(key2), 0)
|
||||
|
||||
def test_batch_put_from_with_config_parameter(self):
|
||||
"""Test batch_put_from method with config parameter."""
|
||||
import ctypes
|
||||
|
||||
# Test data
|
||||
test_data = [
|
||||
b"Batch Config Data 1",
|
||||
b"Batch Config Data 2",
|
||||
b"Batch Config Data 3"
|
||||
]
|
||||
keys = ["test_batch_put_from_config_key1", "test_batch_put_from_config_key2", "test_batch_put_from_config_key3"]
|
||||
|
||||
# Use large spacing between buffers
|
||||
buffer_spacing = 1024 * 1024 # 1MB spacing
|
||||
total_buffer_size = buffer_spacing * len(test_data)
|
||||
large_buffer = (ctypes.c_ubyte * total_buffer_size)()
|
||||
large_buffer_ptr = ctypes.addressof(large_buffer)
|
||||
|
||||
# Register buffer
|
||||
result = self.store.register_buffer(large_buffer_ptr, total_buffer_size)
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Prepare individual buffers
|
||||
buffer_ptrs = []
|
||||
buffer_sizes = []
|
||||
|
||||
for i, data in enumerate(test_data):
|
||||
offset = i * buffer_spacing
|
||||
buffer_ptr = large_buffer_ptr + offset
|
||||
|
||||
# Copy test data to buffer
|
||||
ctypes.memmove(ctypes.c_void_p(buffer_ptr), data, len(data))
|
||||
|
||||
buffer_ptrs.append(buffer_ptr)
|
||||
buffer_sizes.append(len(data))
|
||||
|
||||
# Test with default config (backward compatibility)
|
||||
results = self.store.batch_put_from(keys=keys, buffer_ptrs=buffer_ptrs, sizes=buffer_sizes)
|
||||
self.assertEqual(len(results), len(keys))
|
||||
for result in results:
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
for key, expected_data in zip(keys, test_data):
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, expected_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
for key in keys:
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
# Test with custom config
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = self.max_replicate_num
|
||||
config.with_soft_pin = False
|
||||
|
||||
keys2 = ["test_batch_put_from_config_key4", "test_batch_put_from_config_key5", "test_batch_put_from_config_key6"]
|
||||
results = self.store.batch_put_from(keys=keys2, buffer_ptrs=buffer_ptrs, sizes=buffer_sizes, config=config)
|
||||
self.assertEqual(len(results), len(keys2))
|
||||
for result in results:
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
# Verify data
|
||||
for key, expected_data in zip(keys2, test_data):
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, expected_data)
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.unregister_buffer(large_buffer_ptr), 0)
|
||||
for key in keys2:
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
def test_replication_failure_tolerance(self):
|
||||
"""Test that replicated data remains accessible after main store failure and reinit."""
|
||||
|
||||
test_data = b"Replicated failure tolerance test data!"
|
||||
key = "test_replication_failure_key"
|
||||
|
||||
# Create config with replication
|
||||
config = ReplicateConfig()
|
||||
config.replica_num = self.max_replicate_num
|
||||
|
||||
# Put data with replication
|
||||
result = self.store.put(key=key, value=test_data, config=config)
|
||||
self.assertEqual(result, 0, "Put with replication should succeed")
|
||||
|
||||
# Verify data is initially accessible
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, test_data, "Data should be accessible after put")
|
||||
|
||||
# Teardown the main store (simulate failure)
|
||||
result = self.store.close()
|
||||
self.assertEqual(result, 0, "Store teardown should succeed")
|
||||
time.sleep(1) # Allow time for teardown to complete
|
||||
|
||||
# Verify data is still accessible from replica stores
|
||||
# Since main store is down, read from one of the additional stores
|
||||
replica_store = self.additional_stores[0] # Use first replica
|
||||
retrieved_data = replica_store.get(key)
|
||||
self.assertEqual(retrieved_data, test_data, "Data should remain accessible from replica after main store teardown")
|
||||
|
||||
# Reinitialize the main store
|
||||
get_client(self.store)
|
||||
|
||||
# Verify data is still accessible after main store reinit
|
||||
retrieved_data = self.store.get(key)
|
||||
self.assertEqual(retrieved_data, test_data, "Data should remain accessible after main store reinit")
|
||||
|
||||
# Clean up
|
||||
time.sleep(default_kv_lease_ttl / 1000)
|
||||
self.assertEqual(self.store.remove(key), 0)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
@ -28,6 +28,7 @@ mooncake_master --default_kv_lease_ttl=500 &
|
|||
MASTER_PID=$!
|
||||
sleep 1
|
||||
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_distributed_object_store.py
|
||||
MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_replicated_distributed_object_store.py
|
||||
sleep 1
|
||||
|
||||
pip install torch numpy
|
||||
|
|
|
|||
Loading…
Reference in New Issue