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

[Store][Feature] Add CXL storage for mooncake_store (#1365) 

* [Store] feat: add cxl storage for mooncake store

* Update extern/pybind11 to match main

* fix: use fake cxl device to bypass ci-test error

* Fix code formatting in segment.cpp

---------

Co-authored-by: Teng Ma <sima.mt@alibaba-inc.com>
This commit is contained in:
qiuweit7 2026-01-28 19:51:20 +08:00 committed by GitHub
parent c8c7b0a025
commit 67d0afc4a9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
23 changed files with 847 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/ci.yml vendored
View File

@ -86,7 +86,7 @@ jobs:
sudo bash -x dependencies.sh -y
mkdir build
cd build
cmake .. -DUSE_HTTP=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Debug
cmake .. -DUSE_HTTP=ON -DUSE_CXL=ON -DUSE_ETCD=ON -DSTORE_USE_ETCD=ON -DENABLE_ASAN=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Debug
shell: bash
- name: Build project

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

@ -359,4 +359,63 @@ class RandomAllocationStrategy : public AllocationStrategy {
static constexpr size_t kMaxRetryLimit = 100;
};
class CxlAllocationStrategy : public AllocationStrategy {
public:
CxlAllocationStrategy() = default;
tl::expected<std::vector<Replica>, ErrorCode> Allocate(
const AllocatorManager& allocator_manager, const size_t slice_length,
const size_t replica_num = 1,
const std::vector<std::string>& preferred_segments =
std::vector<std::string>(),
const std::set<std::string>& excluded_segments =
std::set<std::string>()) {
if (slice_length == 0 || replica_num == 0) {
return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
}
if (preferred_segments.empty()) {
LOG(ERROR) << "Preferred_segments is empty.";
return tl::make_unexpected(ErrorCode::INVALID_PARAMS);
}
const std::string& cxl_segment_name = preferred_segments[0];
VLOG(1) << "Do cxl allocate, overwritten segment=" << cxl_segment_name;
const auto cxl_allocators =
allocator_manager.getAllocators(cxl_segment_name);
if (cxl_allocators == nullptr || cxl_allocators->size() == 0) {
return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
}
std::shared_ptr<BufferAllocatorBase> cxl_allocator =
(*cxl_allocators)[0];
if (!cxl_allocator) {
LOG(ERROR) << "No CXL allocator in preferred_segment";
return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
}
std::vector<Replica> replicas;
replicas.reserve(replica_num);
auto buffer = cxl_allocator->allocate(slice_length);
if (!buffer) {
return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
}
buffer->change_to_cxl(cxl_segment_name);
replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING);
VLOG(1) << "Successfully allocated " << replicas.size()
<< " CXL replica.";
return replicas;
}
tl::expected<Replica, ErrorCode> AllocateFrom(
const AllocatorManager& allocator_manager, const size_t slice_length,
const std::string& segment_name) {
return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE);
}
};
} // namespace mooncake

9
mooncake-store/include/allocator.h
View File

@ -66,14 +66,21 @@ class AllocatedBuffer {
struct Descriptor {
uint64_t size_;
uintptr_t buffer_address_;
std::string protocol_;
std::string transport_endpoint_;
YLT_REFL(Descriptor, size_, buffer_address_, transport_endpoint_);
YLT_REFL(Descriptor, size_, buffer_address_, protocol_,
transport_endpoint_);
};
void change_to_cxl(std::string client_segment_name);
void* get_vaddr_from_cxl();
private:
std::weak_ptr<BufferAllocatorBase> allocator_;
std::string segment_name_;
void* buffer_ptr_{nullptr};
std::size_t size_{0};
std::string protocol{"tcp"};
// RAII handle for buffer allocated by offset allocator
std::optional<offset_allocator::OffsetAllocationHandle> offset_handle_{
std::nullopt};

12
mooncake-store/include/client_service.h
View File

@ -225,7 +225,8 @@ class Client {
* @param size Size of the buffer in bytes
* @return ErrorCode indicating success/failure
*/
tl::expected<void, ErrorCode> MountSegment(const void* buffer, size_t size);
tl::expected<void, ErrorCode> MountSegment(
const void* buffer, size_t size, const std::string& protocol = "tcp");
/**
* @brief Unregisters a memory segment from master
@ -305,6 +306,12 @@ class Client {
*/
tl::expected<QueryTaskResponse, ErrorCode> QueryTask(const UUID& task_id);
/**
* @brief Get global segment base address for cxl protocol
* @return Global segment base address
*/
void* GetBaseAddr();
/**
* @brief Mounts a local disk segment into the master.
* @param enable_offloading If true, enables offloading (write-to-file).
@ -404,7 +411,7 @@ class Client {
* @brief Private constructor to enforce creation through Create() method
*/
Client(const std::string& local_hostname,
const std::string& metadata_connstring,
const std::string& metadata_connstring, const std::string& protocol,
const std::map<std::string, std::string>& labels = {});
/**
@ -486,6 +493,7 @@ class Client {
// Configuration
const std::string local_hostname_;
const std::string metadata_connstring_;
const std::string protocol_;
// Client persistent thread pool for async operations
ThreadPool write_thread_pool_;

76
mooncake-store/include/master_config.h
View File

@ -52,6 +52,9 @@ struct MasterConfig {
uint32_t max_total_processing_tasks;
uint64_t pending_task_timeout_sec;
uint64_t processing_task_timeout_sec;
std::string cxl_path;
size_t cxl_size;
bool enable_cxl = false;
};
class MasterServiceSupervisorConfig {
@ -94,6 +97,9 @@ class MasterServiceSupervisorConfig {
uint64_t processing_task_timeout_sec =
DEFAULT_PROCESSING_TASK_TIMEOUT_SEC; // 0 = no timeout(infinite)
std::string cxl_path = DEFAULT_CXL_PATH;
size_t cxl_size = DEFAULT_CXL_SIZE;
bool enable_cxl = false;
MasterServiceSupervisorConfig() = default;
// From MasterConfig
@ -141,6 +147,9 @@ class MasterServiceSupervisorConfig {
pending_task_timeout_sec = config.pending_task_timeout_sec;
processing_task_timeout_sec = config.processing_task_timeout_sec;
cxl_path = config.cxl_path;
cxl_size = config.cxl_size;
enable_cxl = config.enable_cxl;
validate();
}
@ -219,6 +228,9 @@ class WrappedMasterServiceConfig {
uint64_t processing_task_timeout_sec =
DEFAULT_PROCESSING_TASK_TIMEOUT_SEC; // 0 = no timeout(infinite)
std::string cxl_path = DEFAULT_CXL_PATH;
size_t cxl_size = DEFAULT_CXL_SIZE;
bool enable_cxl = false;
WrappedMasterServiceConfig() = default;
// From MasterConfig
@ -260,6 +272,9 @@ class WrappedMasterServiceConfig {
max_total_processing_tasks = config.max_total_processing_tasks;
pending_task_timeout_sec = config.pending_task_timeout_sec;
processing_task_timeout_sec = config.processing_task_timeout_sec;
cxl_path = config.cxl_path;
cxl_size = config.cxl_size;
enable_cxl = config.enable_cxl;
}
// From MasterServiceSupervisorConfig, enable_ha is set to true
@ -295,6 +310,10 @@ class WrappedMasterServiceConfig {
max_total_processing_tasks = config.max_total_processing_tasks;
pending_task_timeout_sec = config.pending_task_timeout_sec;
processing_task_timeout_sec = config.processing_task_timeout_sec;
cxl_path = config.cxl_path;
cxl_size = config.cxl_size;
enable_cxl = config.enable_cxl;
}
};
@ -329,6 +348,10 @@ class MasterServiceConfigBuilder {
uint64_t pending_task_timeout_sec_ = DEFAULT_PENDING_TASK_TIMEOUT_SEC;
uint64_t processing_task_timeout_sec_ = DEFAULT_PROCESSING_TASK_TIMEOUT_SEC;
std::string cxl_path_ = DEFAULT_CXL_PATH;
size_t cxl_size_ = DEFAULT_CXL_SIZE;
bool enable_cxl_ = false;
public:
MasterServiceConfigBuilder() = default;
@ -441,6 +464,21 @@ class MasterServiceConfigBuilder {
return *this;
}
MasterServiceConfigBuilder& set_cxl_path(const std::string& path) {
cxl_path_ = path;
return *this;
}
MasterServiceConfigBuilder& set_cxl_size(size_t size) {
cxl_size_ = size;
return *this;
}
MasterServiceConfigBuilder& set_enable_cxl(bool enable) {
enable_cxl_ = enable;
return *this;
}
MasterServiceConfig build() const;
};
@ -483,10 +521,15 @@ class MasterServiceConfig {
.processing_task_timeout_sec = DEFAULT_PROCESSING_TASK_TIMEOUT_SEC,
};
std::string cxl_path = DEFAULT_CXL_PATH;
size_t cxl_size = DEFAULT_CXL_SIZE;
bool enable_cxl = false;
MasterServiceConfig() = default;
// From WrappedMasterServiceConfig
MasterServiceConfig(const WrappedMasterServiceConfig& config) {
auto cxl_allocator_type = BufferAllocatorType::CACHELIB;
default_kv_lease_ttl = config.default_kv_lease_ttl;
default_kv_soft_pin_ttl = config.default_kv_soft_pin_ttl;
allow_evict_soft_pinned_objects =
@ -500,7 +543,8 @@ class MasterServiceConfig {
cluster_id = config.cluster_id;
root_fs_dir = config.root_fs_dir;
global_file_segment_size = config.global_file_segment_size;
memory_allocator = config.memory_allocator;
memory_allocator =
config.enable_cxl ? cxl_allocator_type : config.memory_allocator;
enable_disk_eviction = config.enable_disk_eviction;
quota_bytes = config.quota_bytes;
put_start_discard_timeout_sec = config.put_start_discard_timeout_sec;
@ -515,6 +559,9 @@ class MasterServiceConfig {
config.pending_task_timeout_sec;
task_manager_config.processing_task_timeout_sec =
config.processing_task_timeout_sec;
cxl_path = config.cxl_path;
cxl_size = config.cxl_size;
enable_cxl = config.enable_cxl;
}
// Static factory method to create a builder
@ -551,6 +598,9 @@ inline MasterServiceConfig MasterServiceConfigBuilder::build() const {
pending_task_timeout_sec_;
config.task_manager_config.processing_task_timeout_sec =
processing_task_timeout_sec_;
config.cxl_path = cxl_path_;
config.cxl_size = cxl_size_;
config.enable_cxl = enable_cxl_;
return config;
}
@ -565,6 +615,9 @@ struct InProcMasterConfig {
std::optional<int> http_metrics_port;
std::optional<int> http_metadata_port;
std::optional<uint64_t> default_kv_lease_ttl;
std::optional<bool> enable_cxl;
std::optional<std::string> cxl_path;
std::optional<size_t> cxl_size;
};
// Builder class for InProcMasterConfig
@ -574,6 +627,9 @@ class InProcMasterConfigBuilder {
std::optional<int> http_metrics_port_ = std::nullopt;
std::optional<int> http_metadata_port_ = std::nullopt;
std::optional<uint64_t> default_kv_lease_ttl_ = std::nullopt;
std::optional<bool> enable_cxl_ = std::nullopt;
std::optional<std::string> cxl_path_ = std::nullopt;
std::optional<size_t> cxl_size_ = std::nullopt;
public:
InProcMasterConfigBuilder() = default;
@ -598,6 +654,21 @@ class InProcMasterConfigBuilder {
return *this;
}
InProcMasterConfigBuilder& set_enable_cxl(bool enable) {
enable_cxl_ = enable;
return *this;
}
InProcMasterConfigBuilder& set_cxl_path(const std::string& path) {
cxl_path_ = path;
return *this;
}
InProcMasterConfigBuilder& set_cxl_size(size_t size) {
cxl_size_ = size;
return *this;
}
InProcMasterConfig build() const;
};
@ -608,6 +679,9 @@ inline InProcMasterConfig InProcMasterConfigBuilder::build() const {
config.http_metrics_port = http_metrics_port_;
config.http_metadata_port = http_metadata_port_;
config.default_kv_lease_ttl = default_kv_lease_ttl_;
config.enable_cxl = enable_cxl_;
config.cxl_path = cxl_path_;
config.cxl_size = cxl_size_;
return config;
}

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

@ -391,7 +391,6 @@ class MasterService {
// Resolve the key to a sanitized format for storage
std::string SanitizeKey(const std::string& key) const;
std::string ResolvePath(const std::string& key) const;
// BatchEvict evicts objects in a near-LRU way, i.e., prioritizes to evict
// object with smaller lease timeout. It has two passes. The first pass only
// evicts objects without soft pin. The second pass prioritizes objects
@ -926,6 +925,10 @@ class MasterService {
// Discarded replicas management
const std::chrono::seconds put_start_discard_timeout_sec_;
const std::chrono::seconds put_start_release_timeout_sec_;
const std::string cxl_path_;
const size_t cxl_size_;
bool enable_cxl_;
class DiscardedReplicas {
public:
DiscardedReplicas() = delete;

12
mooncake-store/include/segment.h
View File

@ -194,8 +194,9 @@ class SegmentManager {
* @param memory_allocator Type of buffer allocator to use for new segments
*/
explicit SegmentManager(
BufferAllocatorType memory_allocator = BufferAllocatorType::CACHELIB)
: memory_allocator_(memory_allocator) {}
BufferAllocatorType memory_allocator = BufferAllocatorType::CACHELIB,
bool enable_cxl = false)
: memory_allocator_(memory_allocator), enable_cxl_(enable_cxl) {}
/**
* @brief Get RAII-style access to segment management operations
@ -218,11 +219,18 @@ class SegmentManager {
client_by_name_, client_local_disk_segment_, segment_mutex_);
}
void initializeCxlAllocator(const std::string& cxl_path,
const size_t cxl_size);
private:
mutable std::shared_mutex segment_mutex_;
std::shared_ptr<AllocationStrategy> allocation_strategy_;
const BufferAllocatorType
memory_allocator_; // Type of buffer allocator to use
// This singleton allocator is managed by the master
// Used for unified allocation and recycling of CXL shared memory.
const bool enable_cxl_;
std::shared_ptr<BufferAllocatorBase> cxl_global_allocator_;
// allocator_manager_ only contains allocators whose segment status is OK.
AllocatorManager allocator_manager_;
std::unordered_map<UUID, MountedSegment, boost::hash<UUID>>

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

@ -16,7 +16,6 @@
#ifdef STORE_USE_ETCD
#include "libetcd_wrapper.h"
#endif
namespace mooncake {
// Constants
@ -33,6 +32,9 @@ static constexpr double DEFAULT_EVICTION_HIGH_WATERMARK_RATIO = 0.95;
static constexpr int64_t ETCD_MASTER_VIEW_LEASE_TTL = 5; // in seconds
static constexpr int64_t DEFAULT_CLIENT_LIVE_TTL_SEC = 10; // in seconds
constexpr const char* DEFAULT_CLUSTER_ID = "mooncake_cluster";
static const std::string DEFAULT_CXL_PATH = "/dev/dax0.0";
static const size_t DEFAULT_CXL_BASE = 0x100000000ULL;
static const size_t DEFAULT_CXL_SIZE = 8ULL * 1024 * 1024 * 1024;
constexpr const char* DEFAULT_ROOT_FS_DIR = "";
// default do not limit DFS usage, and use
// int64_t to make it compaitable to file metrics monitor
@ -240,9 +242,10 @@ struct Segment {
size_t size{0};
// TE p2p endpoint (ip:port) for transport-only addressing
std::string te_endpoint{};
std::string protocol;
Segment() = default;
};
YLT_REFL(Segment, id, name, base, size, te_endpoint);
YLT_REFL(Segment, id, name, base, size, te_endpoint, protocol);
/**
* @brief Client status from the master's perspective

24
mooncake-store/src/allocator.cpp
View File

@ -37,8 +37,25 @@ AllocatedBuffer::Descriptor AllocatedBuffer::get_descriptor() const {
} else {
LOG(ERROR) << "allocator=expired_or_null in get_descriptor";
}
if (this->protocol == "cxl") {
endpoint = this->segment_name_;
}
return {static_cast<uint64_t>(size()),
reinterpret_cast<uintptr_t>(buffer_ptr_), endpoint};
reinterpret_cast<uintptr_t>(buffer_ptr_), this->protocol, endpoint};
}
void AllocatedBuffer::change_to_cxl(std::string client_segment_name) {
uint64_t offset_raw = reinterpret_cast<uintptr_t>(buffer_ptr_);
buffer_ptr_ = reinterpret_cast<void*>(offset_raw - DEFAULT_CXL_BASE);
protocol = "cxl";
segment_name_ = client_segment_name;
}
void* AllocatedBuffer::get_vaddr_from_cxl() {
uint64_t offset_raw = reinterpret_cast<uintptr_t>(buffer_ptr_);
return reinterpret_cast<void*>(offset_raw + DEFAULT_CXL_BASE);
}
// Define operator<< using public accessors or get_descriptor if appropriate
@ -126,8 +143,11 @@ std::unique_ptr<AllocatedBuffer> CachelibBufferAllocator::allocate(
void CachelibBufferAllocator::deallocate(AllocatedBuffer* handle) {
try {
void* buffer = handle->get_descriptor().protocol_ == "cxl"
? handle->get_vaddr_from_cxl()
: handle->buffer_ptr_;
// Deallocate memory using CacheLib.
memory_allocator_->free(handle->buffer_ptr_);
memory_allocator_->free(buffer);
size_t freed_size =
handle->size_; // Store size before handle might become invalid
cur_size_.fetch_sub(freed_size);

48
mooncake-store/src/client_service.cpp
View File

@ -13,6 +13,7 @@
#include "transfer_engine.h"
#include "transfer_task.h"
#include "transport/transport.h"
#include "config.h"
#include "types.h"
@ -36,6 +37,7 @@ namespace mooncake {
Client::Client(const std::string& local_hostname,
const std::string& metadata_connstring,
const std::string& protocol,
const std::map<std::string, std::string>& labels)
: client_id_(generate_uuid()),
metrics_(ClientMetric::Create(merge_labels(labels))),
@ -43,6 +45,7 @@ Client::Client(const std::string& local_hostname,
metrics_ ? &metrics_->master_client_metric : nullptr),
local_hostname_(local_hostname),
metadata_connstring_(metadata_connstring),
protocol_(protocol),
write_thread_pool_(2) {
LOG(INFO) << "client_id=" << client_id_;
@ -369,6 +372,23 @@ ErrorCode Client::InitTransferEngine(
LOG(ERROR) << "Failed to install Ascend transport";
return ErrorCode::INTERNAL_ERROR;
}
} else if (protocol == "cxl") {
if (device_names.has_value()) {
LOG(WARNING) << "CXL protocol does not use device "
"names, ignoring";
}
try {
transport = transfer_engine_->installTransport("cxl", nullptr);
} catch (std::exception& e) {
LOG(ERROR) << "cxl_transport_install_failed error_message=\""
<< e.what() << "\"";
return ErrorCode::INTERNAL_ERROR;
}
if (!transport) {
LOG(ERROR) << "Failed to install CXL transport";
return ErrorCode::INTERNAL_ERROR;
}
} else {
LOG(ERROR) << "unsupported_protocol protocol=" << protocol;
return ErrorCode::INVALID_PARAMS;
@ -394,7 +414,7 @@ std::optional<std::shared_ptr<Client>> Client::Create(
const std::shared_ptr<TransferEngine>& transfer_engine,
std::map<std::string, std::string> labels) {
auto client = std::shared_ptr<Client>(
new Client(local_hostname, metadata_connstring, labels));
new Client(local_hostname, metadata_connstring, protocol, labels));
ErrorCode err = client->ConnectToMaster(master_server_entry);
if (err != ErrorCode::OK) {
@ -849,8 +869,13 @@ tl::expected<void, ErrorCode> Client::Put(const ObjectKey& key,
slice_lengths.emplace_back(slices[i].size);
}
ReplicateConfig client_cfg = config;
if (protocol_ == "cxl") {
client_cfg.preferred_segment = local_hostname_;
}
// Start put operation
auto start_result = master_client_.PutStart(key, slice_lengths, config);
auto start_result = master_client_.PutStart(key, slice_lengths, client_cfg);
if (!start_result) {
ErrorCode err = start_result.error();
if (err == ErrorCode::OBJECT_ALREADY_EXISTS) {
@ -1404,18 +1429,22 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchPut(
const std::vector<ObjectKey>& keys,
std::vector<std::vector<Slice>>& batched_slices,
const ReplicateConfig& config) {
ReplicateConfig client_cfg = config;
if (protocol_ == "cxl") {
client_cfg.preferred_segment = local_hostname_;
}
std::vector<PutOperation> ops = CreatePutOperations(keys, batched_slices);
if (config.prefer_alloc_in_same_node) {
if (config.replica_num != 1) {
if (client_cfg.prefer_alloc_in_same_node) {
if (client_cfg.replica_num != 1) {
LOG(ERROR) << "prefer_alloc_in_same_node is not supported with "
"replica_num != 1";
return std::vector<tl::expected<void, ErrorCode>>(
keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS));
}
StartBatchPut(ops, config);
StartBatchPut(ops, client_cfg);
return BatchPutWhenPreferSameNode(ops);
}
StartBatchPut(ops, config);
StartBatchPut(ops, client_cfg);
auto t0 = std::chrono::steady_clock::now();
SubmitTransfers(ops);
@ -1460,8 +1489,8 @@ tl::expected<long, ErrorCode> Client::RemoveAll() {
return master_client_.RemoveAll();
}
tl::expected<void, ErrorCode> Client::MountSegment(const void* buffer,
size_t size) {
tl::expected<void, ErrorCode> Client::MountSegment(
const void* buffer, size_t size, const std::string& protocol) {
auto check_result = CheckRegisterMemoryParams(buffer, size);
if (!check_result) {
return tl::unexpected(check_result.error());
@ -1498,6 +1527,7 @@ tl::expected<void, ErrorCode> Client::MountSegment(const void* buffer,
segment.name = local_hostname_;
segment.base = reinterpret_cast<uintptr_t>(buffer);
segment.size = size;
segment.protocol = protocol;
// For P2P handshake mode, publish the actual transport endpoint that was
// negotiated by the transfer engine. Otherwise, keep the logical hostname
// so metadata backends (HTTP/etcd/redis) can resolve the segment by name.
@ -1612,6 +1642,8 @@ std::vector<tl::expected<bool, ErrorCode>> Client::BatchIsExist(
return response;
}
void* Client::GetBaseAddr() { return transfer_engine_->getBaseAddr(); }
tl::expected<void, ErrorCode> Client::MountLocalDiskSegment(
bool enable_offloading) {
auto response =

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

@ -113,9 +113,19 @@ DEFINE_uint64(pending_task_timeout_sec, 300,
DEFINE_uint64(processing_task_timeout_sec, 300,
"Timeout in seconds for processing tasks (0 = no timeout)");
DEFINE_string(cxl_path, mooncake::DEFAULT_CXL_PATH,
"DAX device path for CXL memory");
DEFINE_uint64(cxl_size, mooncake::DEFAULT_CXL_SIZE, "CXL memory size in bytes");
DEFINE_bool(enable_cxl, false, "Whether to enable CXL memory support");
void InitMasterConf(const mooncake::DefaultConfig& default_config,
mooncake::MasterConfig& master_config) {
// Initialize the master service configuration from the default config
default_config.GetBool("enable_cxl", &master_config.enable_cxl,
FLAGS_enable_cxl);
default_config.GetString("cxl_path", &master_config.cxl_path,
FLAGS_cxl_path);
default_config.GetUInt64("cxl_size", &master_config.cxl_size,
FLAGS_cxl_size);
default_config.GetBool("enable_metric_reporting",
&master_config.enable_metric_reporting,
FLAGS_enable_metric_reporting);
@ -251,6 +261,21 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config,
}
google::CommandLineFlagInfo info;
if ((google::GetCommandLineFlagInfo("enable_cxl", &info) &&
!info.is_default) ||
!conf_set) {
master_config.enable_cxl = FLAGS_enable_cxl;
}
if ((google::GetCommandLineFlagInfo("cxl_path", &info) &&
!info.is_default) ||
!conf_set) {
master_config.cxl_path = FLAGS_cxl_path;
}
if ((google::GetCommandLineFlagInfo("cxl_size", &info) &&
!info.is_default) ||
!conf_set) {
master_config.cxl_size = FLAGS_cxl_size;
}
if ((google::GetCommandLineFlagInfo("rpc_address", &info) &&
!info.is_default) ||
!conf_set) {
@ -532,7 +557,10 @@ int main(int argc, char* argv[]) {
<< ", pending_task_timeout_sec="
<< master_config.pending_task_timeout_sec
<< ", processing_task_timeout_sec="
<< master_config.processing_task_timeout_sec;
<< master_config.processing_task_timeout_sec
<< ", enable_cxl=" << master_config.enable_cxl
<< ", cxl_path=" << master_config.cxl_path
<< ", cxl_size=" << master_config.cxl_size;
// Start HTTP metadata server if enabled
std::unique_ptr<mooncake::HttpMetadataServer> http_metadata_server;

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

@ -29,12 +29,15 @@ MasterService::MasterService(const MasterServiceConfig& config)
global_file_segment_size_(config.global_file_segment_size),
enable_disk_eviction_(config.enable_disk_eviction),
quota_bytes_(config.quota_bytes),
segment_manager_(config.memory_allocator),
segment_manager_(config.memory_allocator, config.enable_cxl),
memory_allocator_type_(config.memory_allocator),
allocation_strategy_(std::make_shared<RandomAllocationStrategy>()),
put_start_discard_timeout_sec_(config.put_start_discard_timeout_sec),
put_start_release_timeout_sec_(config.put_start_release_timeout_sec),
task_manager_(config.task_manager_config) {
task_manager_(config.task_manager_config),
cxl_path_(config.cxl_path),
cxl_size_(config.cxl_size),
enable_cxl_(config.enable_cxl) {
if (eviction_ratio_ < 0.0 || eviction_ratio_ > 1.0) {
LOG(ERROR) << "Eviction ratio must be between 0.0 and 1.0, "
<< "current value: " << eviction_ratio_;
@ -79,6 +82,13 @@ MasterService::MasterService(const MasterServiceConfig& config)
MasterMetricManager::instance().inc_total_file_capacity(
global_file_segment_size_);
}
if (enable_cxl_) {
allocation_strategy_ = std::make_shared<CxlAllocationStrategy>();
segment_manager_.initializeCxlAllocator(cxl_path_, cxl_size_);
VLOG(1) << "action=start_cxl_global_allocator";
} else {
allocation_strategy_ = std::make_shared<RandomAllocationStrategy>();
}
}
MasterService::~MasterService() {

94
mooncake-store/src/real_client.cpp
View File

@ -267,47 +267,75 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
// If global_segment_size is 0, skip mount segment;
// If global_segment_size is larger than max_mr_size, split to multiple
// mapped_shms.
auto max_mr_size = globalConfig().max_mr_size; // Max segment size
uint64_t total_glbseg_size = global_segment_size; // For logging
uint64_t current_glbseg_size = 0; // For logging
while (global_segment_size > 0) {
size_t segment_size = std::min(global_segment_size, max_mr_size);
global_segment_size -= segment_size;
current_glbseg_size += segment_size;
LOG(INFO) << "Mounting segment: " << segment_size << " bytes, "
<< current_glbseg_size << " of " << total_glbseg_size;
size_t mapped_size = segment_size;
void *ptr = nullptr;
if (should_use_hugepage) {
mapped_size = align_up(segment_size, get_hugepage_size_from_env());
ptr = allocate_buffer_mmap_memory(mapped_size,
get_hugepage_size_from_env());
if (protocol == "cxl") {
size_t cxl_dev_size = 0;
const char *env = std::getenv("MC_CXL_DEV_SIZE");
if (env) {
char *end = nullptr;
unsigned long long val = strtoull(env, &end, 10);
if (end != env && *end == '\0')
cxl_dev_size = static_cast<size_t>(val);
} else {
ptr =
allocate_buffer_allocator_memory(segment_size, this->protocol);
}
if (!ptr) {
LOG(ERROR) << "Failed to allocate segment memory";
LOG(FATAL) << "MC_CXL_DEV_SIZE not set";
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
if (this->protocol == "ascend") {
ascend_segment_ptrs_.emplace_back(ptr);
} else if (should_use_hugepage) {
hugepage_segment_ptrs_.emplace_back(
ptr, HugepageSegmentDeleter{mapped_size});
} else {
segment_ptrs_.emplace_back(ptr);
}
auto mount_result = client_->MountSegment(ptr, mapped_size);
void *ptr = client_->GetBaseAddr();
LOG(INFO) << "Mounting CXL segment: " << cxl_dev_size << " bytes, "
<< ptr;
auto mount_result = client_->MountSegment(ptr, cxl_dev_size, protocol);
if (!mount_result.has_value()) {
LOG(ERROR) << "Failed to mount segment: "
<< toString(mount_result.error());
return tl::unexpected(mount_result.error());
}
}
if (total_glbseg_size == 0) {
LOG(INFO) << "Global segment size is 0, skip mounting segment";
} else {
auto max_mr_size = globalConfig().max_mr_size; // Max segment size
uint64_t total_glbseg_size = global_segment_size; // For logging
uint64_t current_glbseg_size = 0; // For logging
while (global_segment_size > 0) {
size_t segment_size = std::min(global_segment_size, max_mr_size);
global_segment_size -= segment_size;
current_glbseg_size += segment_size;
LOG(INFO) << "Mounting segment: " << segment_size << " bytes, "
<< current_glbseg_size << " of " << total_glbseg_size;
size_t mapped_size = segment_size;
void *ptr = nullptr;
if (should_use_hugepage) {
mapped_size =
align_up(segment_size, get_hugepage_size_from_env());
ptr = allocate_buffer_mmap_memory(mapped_size,
get_hugepage_size_from_env());
} else {
ptr = allocate_buffer_allocator_memory(segment_size,
this->protocol);
}
if (!ptr) {
LOG(ERROR) << "Failed to allocate segment memory";
return tl::unexpected(ErrorCode::INVALID_PARAMS);
}
if (this->protocol == "ascend") {
ascend_segment_ptrs_.emplace_back(ptr);
} else if (should_use_hugepage) {
hugepage_segment_ptrs_.emplace_back(
ptr, HugepageSegmentDeleter{mapped_size});
} else {
segment_ptrs_.emplace_back(ptr);
}
auto mount_result =
client_->MountSegment(ptr, mapped_size, protocol);
if (!mount_result.has_value()) {
LOG(ERROR) << "Failed to mount segment: "
<< toString(mount_result.error());
return tl::unexpected(mount_result.error());
}
}
if (total_glbseg_size == 0) {
LOG(INFO) << "Global segment size is 0, skip mounting segment";
}
}
// Start IPC server to accept FD from dummy clients

50
mooncake-store/src/segment.cpp
View File

@ -9,6 +9,33 @@ ErrorCode ScopedSegmentAccess::MountSegment(const Segment& segment,
const uintptr_t buffer = segment.base;
const size_t size = segment.size;
// Check if cxl storage is enable
if (segment_manager_->enable_cxl_ && segment.protocol == "cxl") {
LOG(INFO) << "Start Mounting CXL Segment.";
if (segment_manager_->memory_allocator_ ==
BufferAllocatorType::CACHELIB) {
auto allocator = segment_manager_->cxl_global_allocator_;
if (segment_manager_->cxl_global_allocator_ == nullptr) {
LOG(ERROR) << "Cxl global allocator has not been initialized.";
return ErrorCode::INTERNAL_ERROR;
}
segment_manager_->allocator_manager_.addAllocator(segment.name,
allocator);
segment_manager_->client_segments_[client_id].push_back(segment.id);
segment_manager_->mounted_segments_[segment.id] = {
segment, SegmentStatus::OK, allocator};
segment_manager_->client_by_name_[segment.name] = client_id;
MasterMetricManager::instance().inc_total_mem_capacity(segment.name,
size);
LOG(INFO) << "[CXL Segment Mounted Successfully] Segment name: "
<< segment.name
<< ", Mount size: " << (size / 1024 / 1024 / 1024)
<< " GB";
return ErrorCode::OK;
}
return ErrorCode::INTERNAL_ERROR;
}
// Check if parameters are valid before allocating memory.
if (buffer == 0 || size == 0) {
LOG(ERROR) << "buffer=" << buffer << " or size=" << size
@ -163,7 +190,6 @@ ErrorCode ScopedSegmentAccess::PrepareUnmountSegment(
// Set the segment status to UNMOUNTING
mounted_segment.status = SegmentStatus::UNMOUNTING;
return ErrorCode::OK;
}
@ -192,18 +218,22 @@ ErrorCode ScopedSegmentAccess::CommitUnmountSegment(
// segment_id -> segment_name
std::string segment_name;
bool is_cxl = false;
auto&& segment = segment_manager_->mounted_segments_.find(segment_id);
if (segment != segment_manager_->mounted_segments_.end()) {
segment_name = segment->second.segment.name;
// Also remove from segment_name_client_id_map_
segment_manager_->client_by_name_.erase(segment_name);
is_cxl = (segment->second.segment.protocol == "cxl");
}
// Remove from mounted_segments_
segment_manager_->mounted_segments_.erase(segment_id);
// Decrease the total capacity
MasterMetricManager::instance().dec_total_mem_capacity(
segment_name, metrics_dec_capacity);
if (!is_cxl) {
MasterMetricManager::instance().dec_total_mem_capacity(
segment_name, metrics_dec_capacity);
}
return ErrorCode::OK;
}
@ -277,4 +307,16 @@ bool ScopedSegmentAccess::ExistsSegmentName(
return it != segment_manager_->client_by_name_.end();
}
} // namespace mooncake
void SegmentManager::initializeCxlAllocator(const std::string& cxl_path,
const size_t cxl_size) {
LOG(INFO) << "Init CXL global allocator.";
LOG(INFO) << "[CXL] create allocator with " << "path=" << cxl_path
<< " base=0x" << std::hex << DEFAULT_CXL_BASE << std::dec
<< " size=" << cxl_size << " (" << std::fixed
<< std::setprecision(2) << cxl_size / (1024.0 * 1024 * 1024)
<< " GB)";
cxl_global_allocator_ = std::make_shared<CachelibBufferAllocator>(
cxl_path, DEFAULT_CXL_BASE, cxl_size, cxl_path);
}
} // namespace mooncake

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

@ -20,6 +20,7 @@ add_store_test(eviction_strategy_test eviction_strategy_test.cpp)
add_store_test(master_service_test master_service_test.cpp)
add_store_test(master_service_ssd_test master_service_ssd_test.cpp)
add_store_test(client_integration_test client_integration_test.cpp)
add_store_test(cxl_client_integration_test cxl_client_integration_test.cpp)
add_store_test(master_metrics_test master_metrics_test.cpp)
add_store_test(posix_file_test posix_file_test.cpp)
add_store_test(thread_pool_test thread_pool_test.cpp)

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

@ -142,7 +142,7 @@ class ClientIntegrationTest : public ::testing::Test {
segment_ptr_ = allocate_buffer_allocator_memory(ram_buffer_size_);
LOG_ASSERT(segment_ptr_);
auto mount_result = segment_provider_client_->MountSegment(
segment_ptr_, ram_buffer_size_);
segment_ptr_, ram_buffer_size_, FLAGS_protocol);
if (!mount_result.has_value()) {
LOG(ERROR) << "Failed to mount segment: "
<< toString(mount_result.error());
@ -213,7 +213,8 @@ class ClientIntegrationTest : public ::testing::Test {
allocate_buffer_allocator_memory(test_client_ram_buffer_size_);
LOG_ASSERT(test_client_segment_ptr_);
auto test_client_mount_result = test_client_->MountSegment(
test_client_segment_ptr_, test_client_ram_buffer_size_);
test_client_segment_ptr_, test_client_ram_buffer_size_,
FLAGS_protocol);
if (!test_client_mount_result.has_value()) {
LOG(ERROR) << "Failed to mount segment for test_client_: "
<< toString(test_client_mount_result.error());

420
mooncake-store/tests/cxl_client_integration_test.cpp Normal file
View File

@ -0,0 +1,420 @@
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <gtest/gtest.h>
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include <regex>
#include <unordered_set>
#include <unordered_map>
#include <thread>
#include <chrono>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include "allocator.h"
#include "client_service.h"
#include "types.h"
#include "utils.h"
#include "test_server_helpers.h"
#include "default_config.h"
DEFINE_string(protocol, "cxl", "Transfer protocol: rdma|tcp|cxl");
DEFINE_string(device_name, "", "Device name to use, valid if protocol=rdma");
DEFINE_uint64(default_kv_lease_ttl, mooncake::DEFAULT_DEFAULT_KV_LEASE_TTL,
"Default lease time for kv objects, must be set to the "
"same as the master's default_kv_lease_ttl");
DEFINE_string(cxl_device_name, "tmp_dax_sim", "Device name for cxl");
DEFINE_uint64(cxl_device_size, 1073741824, "Device Size for cxl");
DEFINE_bool(auto_disc, false, "Auto discover tcp devices");
DEFINE_string(transfer_engine_metadata_url, "127.0.0.1:2379",
"Metadata connection string for transfer engine");
namespace mooncake {
namespace testing {
// Helper functions for client_id parsing
std::string FormatClientId(const UUID& client_id) {
return std::to_string(client_id.first) + "-" +
std::to_string(client_id.second);
}
UUID ParseClientId(const std::string& client_id_str) {
UUID client_id{0, 0};
size_t dash_pos = client_id_str.find('-');
if (dash_pos != std::string::npos) {
try {
client_id.first = std::stoull(client_id_str.substr(0, dash_pos));
client_id.second = std::stoull(client_id_str.substr(dash_pos + 1));
} catch (const std::exception& e) {
LOG(ERROR) << "Failed to parse client_id: " << e.what();
}
} else {
LOG(ERROR) << "Invalid client_id format. Expected format: first-second";
}
return client_id;
}
class ClientIdCaptureSink : public google::LogSink {
public:
std::string captured_client_id;
void send(google::LogSeverity severity, const char* full_filename,
const char* base_filename, int line, const struct ::tm* tm_time,
const char* message, size_t message_len) override {
(void)severity;
(void)full_filename;
(void)base_filename;
(void)line;
(void)tm_time;
std::string msg(message, message_len);
size_t pos = msg.find("client_id=");
if (pos != std::string::npos) {
std::string client_id_str = msg.substr(pos + 10);
client_id_str.erase(0, client_id_str.find_first_not_of(" \t\n\r"));
client_id_str.erase(client_id_str.find_last_not_of(" \t\n\r") + 1);
std::regex uuid_pattern(R"((\d+)-(\d+))");
std::smatch match;
if (std::regex_search(client_id_str, match, uuid_pattern)) {
captured_client_id = match[0].str();
}
}
}
};
class ClientIntegrationTestCxl : public ::testing::Test {
protected:
static std::shared_ptr<Client> CreateClient(const std::string& host_name) {
auto client_opt = Client::Create(
host_name, // Local hostname
FLAGS_transfer_engine_metadata_url, // Metadata connection string
FLAGS_protocol, // Transfer protocol
std::nullopt, // RDMA device names (auto-discovery)
master_address_);
EXPECT_TRUE(client_opt.has_value())
<< "Failed to create client with host_name: " << host_name;
if (!client_opt.has_value()) {
return nullptr;
}
return client_opt.value();
}
static void SetUpTestSuite() {
// Initialize glog
google::InitGoogleLogging("ClientIntegrationTestCxl");
FLAGS_logtostderr = 1;
tmp_fd = open(FLAGS_cxl_device_name.c_str(), O_RDWR | O_CREAT, 0666);
ASSERT_GE(tmp_fd, 0);
ASSERT_EQ(ftruncate(tmp_fd, FLAGS_cxl_device_size), 0);
// Override flags from environment variables if present
setenv("MC_CXL_DEV_PATH", FLAGS_cxl_device_name.c_str(), 1);
setenv("MC_CXL_DEV_SIZE", std::to_string(FLAGS_cxl_device_size).c_str(),
1);
setenv("MC_MS_AUTO_DISC", std::to_string(FLAGS_auto_disc).c_str(), 0);
if (getenv("DEFAULT_KV_LEASE_TTL")) {
default_kv_lease_ttl_ = std::stoul(getenv("DEFAULT_KV_LEASE_TTL"));
} else {
default_kv_lease_ttl_ = FLAGS_default_kv_lease_ttl;
}
LOG(INFO) << "Default KV lease TTL: " << default_kv_lease_ttl_;
// Start in-proc master
InProcMasterConfig config = InProcMasterConfigBuilder()
.set_enable_cxl(true)
.set_cxl_path(FLAGS_cxl_device_name)
.set_cxl_size(FLAGS_cxl_device_size)
.build();
ASSERT_TRUE(master_.Start(config)) << "Failed to start InProcMaster!";
master_address_ = master_.master_address();
LOG(INFO) << "Started in-proc master at " << master_address_;
InitializeClients();
InitializeSegment();
}
static void TearDownTestSuite() {
CleanupSegment();
CleanupClients();
master_.Stop();
google::ShutdownGoogleLogging();
if (tmp_fd >= 0) {
close(tmp_fd);
unlink(FLAGS_cxl_device_name.c_str());
}
}
static void InitializeSegment() {
// init local buffer allocator
client_buffer_allocator_ =
std::make_unique<SimpleAllocator>(128 * 1024 * 1024);
// Mount segment for test_client_ as well
ASSERT_TRUE(FLAGS_protocol == "cxl");
size_t cxl_dev_size = 0;
const char* env_cxl_size = getenv("MC_CXL_DEV_SIZE");
if (env_cxl_size) {
char* end = nullptr;
unsigned long long val = std::strtoull(env_cxl_size, &end, 10);
if (end != env_cxl_size && *end == '\0')
cxl_dev_size = static_cast<size_t>(val);
} else {
LOG(FATAL) << "MC_CXL_DEV_SIZE environment variable not set "
"(required for CXL protocol)";
return;
}
test_client_ram_buffer_size_ = cxl_dev_size;
test_client_segment_ptr_ = test_client_->GetBaseAddr();
LOG_ASSERT(test_client_segment_ptr_);
LOG(INFO) << "test_client_segment_ptr_: " << test_client_segment_ptr_;
auto test_client_mount_result = test_client_->MountSegment(
test_client_segment_ptr_, test_client_ram_buffer_size_,
FLAGS_protocol);
if (!test_client_mount_result.has_value()) {
LOG(ERROR) << "Failed to mount CXL segment for test_client_: "
<< toString(test_client_mount_result.error());
}
LOG(INFO) << "Test client CXL segment mounted successfully";
}
static void InitializeClients() {
// This client is used for testing purposes.
// Capture test_client_ client_id from logs
ClientIdCaptureSink* test_client_sink = new ClientIdCaptureSink();
google::AddLogSink(test_client_sink);
test_client_ = CreateClient("localhost:17813");
ASSERT_TRUE(test_client_ != nullptr);
// Wait for logs to flush
std::this_thread::sleep_for(std::chrono::milliseconds(200));
google::RemoveLogSink(test_client_sink);
if (!test_client_sink->captured_client_id.empty()) {
UUID extracted_id =
ParseClientId(test_client_sink->captured_client_id);
if (extracted_id.first != 0 || extracted_id.second != 0) {
test_client_id_ = extracted_id;
LOG(INFO) << "Captured test_client_id: "
<< FormatClientId(test_client_id_);
}
}
delete test_client_sink;
}
static void CleanupClients() {
if (test_client_) {
test_client_.reset();
}
}
static void CleanupSegment() {
// Unmount test client segment first
if (test_client_ && test_client_segment_ptr_) {
if (!test_client_
->UnmountSegment(test_client_segment_ptr_,
test_client_ram_buffer_size_)
.has_value()) {
LOG(ERROR) << "Failed to unmount test client CXL segment";
}
}
}
static std::shared_ptr<Client> test_client_;
// Here we use a simple allocator for the client buffer. In a real
// application, user should manage the memory allocation and deallocation
// themselves.
static std::unique_ptr<SimpleAllocator> client_buffer_allocator_;
static void* segment_ptr_;
static size_t ram_buffer_size_;
static void* test_client_segment_ptr_;
static size_t test_client_ram_buffer_size_;
static uint64_t default_kv_lease_ttl_;
static InProcMaster master_;
static std::string master_address_;
static std::string metadata_url_;
static UUID test_client_id_;
static inline bool is_cxl = false;
static int tmp_fd;
};
// Static members initialization
std::shared_ptr<Client> ClientIntegrationTestCxl::test_client_ = nullptr;
void* ClientIntegrationTestCxl::segment_ptr_ = nullptr;
void* ClientIntegrationTestCxl::test_client_segment_ptr_ = nullptr;
std::unique_ptr<SimpleAllocator>
ClientIntegrationTestCxl::client_buffer_allocator_ = nullptr;
size_t ClientIntegrationTestCxl::ram_buffer_size_ = 0;
size_t ClientIntegrationTestCxl::test_client_ram_buffer_size_ = 0;
uint64_t ClientIntegrationTestCxl::default_kv_lease_ttl_ = 0;
InProcMaster ClientIntegrationTestCxl::master_;
std::string ClientIntegrationTestCxl::master_address_;
std::string ClientIntegrationTestCxl::metadata_url_;
UUID ClientIntegrationTestCxl::test_client_id_{0, 0};
int ClientIntegrationTestCxl::tmp_fd = -1;
// Test basic Put/Get operations through the client
TEST_F(ClientIntegrationTestCxl, BasicPutGetOperations) {
const std::string test_data = "Hello, World!";
const std::string key = "test_key";
void* buffer = client_buffer_allocator_->allocate(test_data.size());
// write
memcpy(buffer, test_data.data(), test_data.size());
std::vector<Slice> slices;
slices.emplace_back(Slice{buffer, test_data.size()});
// Test Put operation
ReplicateConfig config;
config.replica_num = 1;
auto put_result = test_client_->Put(key, slices, config);
ASSERT_TRUE(put_result.has_value())
<< "Put operation failed: " << toString(put_result.error());
client_buffer_allocator_->deallocate(buffer, test_data.size());
buffer = client_buffer_allocator_->allocate(1 * 1024 * 1024);
slices.clear();
slices.emplace_back(Slice{buffer, test_data.size()});
// Verify data through Get operation
auto get_result = test_client_->Get(key, slices);
ASSERT_TRUE(get_result.has_value())
<< "Get operation failed: " << toString(get_result.error());
ASSERT_EQ(slices.size(), 1);
ASSERT_EQ(slices[0].size, test_data.size());
ASSERT_EQ(slices[0].ptr, buffer);
ASSERT_EQ(memcmp(slices[0].ptr, test_data.data(), test_data.size()), 0);
client_buffer_allocator_->deallocate(buffer, test_data.size());
// Put again with the same key, should succeed
buffer = client_buffer_allocator_->allocate(test_data.size());
memcpy(buffer, test_data.data(), test_data.size());
slices.clear();
slices.emplace_back(Slice{buffer, test_data.size()});
auto put_result2 = test_client_->Put(key, slices, config);
ASSERT_TRUE(put_result2.has_value())
<< "Second Put operation failed: " << toString(put_result2.error());
std::this_thread::sleep_for(
std::chrono::milliseconds(default_kv_lease_ttl_));
auto remove_result = test_client_->Remove(key);
ASSERT_TRUE(remove_result.has_value())
<< "Remove operation failed: " << toString(remove_result.error());
client_buffer_allocator_->deallocate(buffer, test_data.size());
}
// Test batch Put/Get operations through the client
TEST_F(ClientIntegrationTestCxl, BatchPutGetOperations) {
int batch_sz = 10;
std::vector<std::string> keys;
std::vector<std::string> test_data_list;
std::vector<std::vector<Slice>> batched_slices;
for (int i = 0; i < batch_sz; i++) {
keys.push_back("test_key_batch_put_" + std::to_string(i));
test_data_list.push_back("test_data_" + std::to_string(i));
}
void* buffer = nullptr;
void* target_buffer = nullptr;
batched_slices.reserve(batch_sz);
for (int i = 0; i < batch_sz; i++) {
std::vector<Slice> slices;
buffer = client_buffer_allocator_->allocate(test_data_list[i].size());
memcpy(buffer, test_data_list[i].data(), test_data_list[i].size());
slices.emplace_back(Slice{buffer, test_data_list[i].size()});
batched_slices.push_back(std::move(slices));
}
// Test Batch Put operation
ReplicateConfig config;
config.replica_num = 1;
auto start = std::chrono::high_resolution_clock::now();
auto batch_put_results =
test_client_->BatchPut(keys, batched_slices, config);
// Check that all operations succeeded
for (const auto& result : batch_put_results) {
ASSERT_TRUE(result.has_value()) << "BatchPut operation failed";
}
auto end = std::chrono::high_resolution_clock::now();
LOG(INFO) << "Time taken for BatchPut: "
<< std::chrono::duration_cast<std::chrono::microseconds>(end -
start)
.count()
<< "us";
start = std::chrono::high_resolution_clock::now();
for (int i = 0; i < batch_sz; i++) {
std::vector<Slice> slices;
target_buffer =
client_buffer_allocator_->allocate(test_data_list[i].size());
slices.emplace_back(Slice{target_buffer, test_data_list[i].size()});
auto get_result = test_client_->Get(keys[i], slices);
ASSERT_TRUE(get_result.has_value())
<< "Get operation failed: " << toString(get_result.error());
client_buffer_allocator_->deallocate(target_buffer,
test_data_list[i].size());
}
end = std::chrono::high_resolution_clock::now();
LOG(INFO) << "Time taken for single Get: "
<< std::chrono::duration_cast<std::chrono::microseconds>(end -
start)
.count()
<< "us";
start = std::chrono::high_resolution_clock::now();
std::unordered_map<std::string, std::vector<Slice>> target_batched_slices;
for (int i = 0; i < batch_sz; i++) {
std::vector<Slice> target_slices;
target_buffer =
client_buffer_allocator_->allocate(test_data_list[i].size());
target_slices.emplace_back(
Slice{target_buffer, test_data_list[i].size()});
target_batched_slices.emplace(keys[i], target_slices);
}
auto batch_get_results =
test_client_->BatchGet(keys, target_batched_slices);
for (const auto& result : batch_get_results) {
ASSERT_TRUE(result.has_value()) << "BatchGet operation failed";
}
end = std::chrono::high_resolution_clock::now();
LOG(INFO) << "Time taken for BatchGet: "
<< std::chrono::duration_cast<std::chrono::microseconds>(end -
start)
.count()
<< "us";
for (int i = 0; i < batch_sz; i++) {
ASSERT_EQ(target_batched_slices[keys[i]][0].size,
test_data_list[i].size());
ASSERT_EQ(memcmp(target_batched_slices[keys[i]][0].ptr,
test_data_list[i].data(), test_data_list[i].size()),
0);
client_buffer_allocator_->deallocate(
target_batched_slices[keys[i]][0].ptr, test_data_list[i].size());
}
}
} // namespace testing
} // namespace mooncake
int main(int argc, char** argv) {
// Initialize Google Test
::testing::InitGoogleTest(&argc, argv);
// Initialize Google's flags library
gflags::ParseCommandLineFlags(&argc, &argv, false);
mooncake::init_ylt_log_level();
// Run all tests
return RUN_ALL_TESTS();
}

22
mooncake-store/tests/test_server_helpers.h
View File

@ -85,6 +85,28 @@ class InProcMaster {
wms_cfg.root_fs_dir = DEFAULT_ROOT_FS_DIR;
wms_cfg.memory_allocator = BufferAllocatorType::OFFSET;
wms_cfg.enable_cxl = config.enable_cxl.has_value()
? config.enable_cxl.value()
: false;
if (config.cxl_path.has_value()) {
wms_cfg.cxl_path = config.cxl_path.value();
} else if (const char* cxl_path_env =
std::getenv("MC_CXL_DEV_PATH")) {
wms_cfg.cxl_path = cxl_path_env;
}
if (config.cxl_size.has_value()) {
wms_cfg.cxl_size = config.cxl_size.value();
} else if (const char* cxl_size_env =
std::getenv("MC_CXL_DEV_SIZE")) {
char* endptr = nullptr;
unsigned long long val =
std::strtoull(cxl_size_env, &endptr, 10);
if (endptr != cxl_size_env && *endptr == '\0') {
wms_cfg.cxl_size = static_cast<size_t>(val);
}
}
wrapped_ = std::make_unique<WrappedMasterService>(wms_cfg);
RegisterRpcService(*server_, *wrapped_);

2
mooncake-transfer-engine/include/multi_transport.h
View File

@ -51,6 +51,8 @@ class MultiTransport {
std::vector<Transport *> listTransports();
void *getBaseAddr();
private:
Status selectTransport(const TransferRequest &entry, Transport *&transport);

2
mooncake-transfer-engine/include/transfer_engine.h
View File

@ -111,6 +111,8 @@ class TransferEngine {
void setAutoDiscover(bool auto_discover);
void* getBaseAddr();
void setWhitelistFilters(std::vector<std::string>&& filters);
int numContexts() const;

2
mooncake-transfer-engine/include/transfer_engine_impl.h
View File

@ -279,6 +279,8 @@ class TransferEngineImpl {
void setAutoDiscover(bool auto_discover) { auto_discover_ = auto_discover; }
void* getBaseAddr() { return multi_transports_->getBaseAddr(); }
void setWhitelistFilters(std::vector<std::string>&& filters) {
filter_ = std::move(filters);
}

11
mooncake-transfer-engine/src/multi_transport.cpp
View File

@ -368,4 +368,15 @@ std::vector<Transport *> MultiTransport::listTransports() {
return transport_list;
}
void *MultiTransport::getBaseAddr() {
#ifdef USE_CXL
Transport *transport = getTransport("cxl");
if (transport) {
auto *cxl_transport = dynamic_cast<CxlTransport *>(transport);
return cxl_transport ? cxl_transport->getCxlBaseAddr() : 0;
}
#endif
return 0;
}
} // namespace mooncake

2
mooncake-transfer-engine/src/transfer_engine.cpp
View File

@ -161,6 +161,8 @@ void TransferEngine::setAutoDiscover(bool auto_discover) {
impl_->setAutoDiscover(auto_discover);
}
void* TransferEngine::getBaseAddr() { return impl_->getBaseAddr(); }
void TransferEngine::setWhitelistFilters(std::vector<std::string>&& filters) {
impl_->setWhitelistFilters(std::move(filters));
}