mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

feat(Store): Support get local ssd object (#1203) 

Co-authored-by: zhang zuoyuan <zhangzuo21@mails.tsinghua.edu.cn>
This commit is contained in:
zhuxinjie-nz 2026-01-22 21:15:50 +08:00 committed by GitHub
parent 1e32df7054
commit 24c632f84a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 503 additions and 93 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -324,21 +324,21 @@ class Client {
std::unordered_map<std::string, int64_t>& offloading_objects);
/**
* @brief Performs a batched write of multiple objects using a
* @brief Performs a batched read of multiple objects using a
* high-throughput Transfer Engine.
* @param transfer_engine_addr Address of the Transfer Engine service (e.g.,
* "ip:port").
* @param keys List of keys identifying the data objects to be transferred
* @param pointers Array of destination memory addresses on the remote node
* where data will be written (one per key)
* @param batched_slices Map from object key to its data slice
* @param batch_slices Map from object key to its data slice
* (`mooncake::Slice`), containing raw bytes to be written.
*/
tl::expected<void, ErrorCode> BatchPutOffloadObject(
tl::expected<void, ErrorCode> BatchGetOffloadObject(
const std::string& transfer_engine_addr,
const std::vector<std::string>& keys,
const std::vector<uintptr_t>& pointers,
const std::unordered_map<std::string, Slice>& batched_slices);
const std::unordered_map<std::string, Slice>& batch_slices);
/**
* @brief Notifies the master that offloading of specified objects has

31
mooncake-store/include/file_storage.h
View File

@ -8,9 +8,8 @@ namespace mooncake {
class FileStorage {
public:
FileStorage(std::shared_ptr<Client> client,
const std::string& local_rpc_addr,
const FileStorageConfig& config);
FileStorage(const FileStorageConfig& config, std::shared_ptr<Client> client,
const std::string& local_rpc_addr);
~FileStorage();
tl::expected<void, ErrorCode> Init();
@ -18,25 +17,25 @@ class FileStorage {
/**
* @brief Reads multiple key-value (KV) entries from local storage and
* forwards them to a remote node.
* @param transfer_engine_addr Address of the remote transfer engine
* (format: "ip:port")
* @param keys List of keys to read from the local KV store
* @param pointers Array of remote memory base addresses (on the
* destination node) where each corresponding value will be written
* @param sizes Expected size in bytes for each value
* @return tl::expected<void, ErrorCode> indicating operation status.
* @return tl::expected<std::vector<uint64_t>, ErrorCode> indicating
* operation status.
*/
tl::expected<void, ErrorCode> BatchGet(
const std::string& transfer_engine_addr,
tl::expected<std::vector<uint64_t>, ErrorCode> BatchGet(
const std::vector<std::string>& keys,
const std::vector<uintptr_t>& pointers,
const std::vector<int64_t>& sizes);
FileStorageConfig config_;
private:
friend class FileStorageTest;
struct AllocatedBatch {
std::vector<BufferHandle> handles;
std::unordered_map<std::string, Slice> slices;
std::chrono::steady_clock::time_point lease_timeout;
std::vector<uint64_t> pointers;
uint64_t total_size;
AllocatedBatch() = default;
AllocatedBatch(AllocatedBatch&&) = default;
@ -76,20 +75,26 @@ class FileStorage {
tl::expected<void, ErrorCode> RegisterLocalMemory();
tl::expected<AllocatedBatch, ErrorCode> AllocateBatch(
tl::expected<std::shared_ptr<AllocatedBatch>, ErrorCode> AllocateBatch(
const std::vector<std::string>& keys,
const std::vector<int64_t>& sizes);
void ClientBufferGCThreadFunc();
std::shared_ptr<Client> client_;
std::string local_rpc_addr_;
FileStorageConfig config_;
std::shared_ptr<StorageBackendInterface> storage_backend_;
std::shared_ptr<ClientBufferAllocator> client_buffer_allocator_;
mutable Mutex client_buffer_mutex_;
std::vector<std::shared_ptr<AllocatedBatch>> GUARDED_BY(
client_buffer_mutex_) client_buffer_allocated_batches_;
mutable Mutex offloading_mutex_;
bool GUARDED_BY(offloading_mutex_) enable_offloading_;
std::atomic<bool> heartbeat_running_;
std::thread heartbeat_thread_;
std::atomic<bool> client_buffer_gc_running_;
std::thread client_buffer_gc_thread_;
};
} // namespace mooncake

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

@ -816,6 +816,19 @@ class MasterService {
replication_task_it_ = shard_guard_->replication_tasks.end();
}
void Create(const UUID& client_id, uint64_t total_length,
std::vector<Replica> replicas, bool enable_soft_pin) {
if (Exists()) {
throw std::logic_error("Already exists");
}
const auto now = std::chrono::steady_clock::now();
auto result = shard_guard_->metadata.emplace(
std::piecewise_construct, std::forward_as_tuple(key_),
std::forward_as_tuple(client_id, now, total_length,
std::move(replicas), enable_soft_pin));
it_ = result.first;
}
private:
MasterService* service_;
std::string key_;

60
mooncake-store/include/pyclient.h
View File

@ -25,6 +25,65 @@ struct ShmRegisterRequest {
bool is_local_buffer;
};
class ClientRequester {
public:
ClientRequester();
/**
* @brief Retrieves multiple objects from a remote Transfer Engine (TE)
* @param client_addr Network address (e.g., "ip:port") of the remote
* Transfer Engine service.
* @param keys Map from object key to size (bytes);
*/
tl::expected<BatchGetOffloadObjectResponse, ErrorCode>
batch_get_offload_object(const std::string &client_addr,
const std::vector<std::string> &keys,
const std::vector<int64_t> sizes);
private:
/**
* @brief A batch of allocated memory buffers, tracking both handles and
* mapped addresses. This struct holds a collection of buffer resources
* obtained from a memory allocator. It includes:
* - `handles`: Opaque handles used to manage lifetime and deallocation.
* - `pointers`: Direct virtual addresses where the buffers are accessible.
*/
struct AllocatedBatch {
std::vector<BufferHandle>
handles; ///< Unique handles for each buffer (used for release)
std::vector<uintptr_t>
pointers; ///< Virtual memory addresses where buffers are mapped
// Allow move semantics
AllocatedBatch() = default;
AllocatedBatch(AllocatedBatch &&) = default;
AllocatedBatch &operator=(AllocatedBatch &&) = default;
// Prevent copying (because BufferHandle is move-only)
AllocatedBatch(const AllocatedBatch &) = delete;
AllocatedBatch &operator=(const AllocatedBatch &) = delete;
~AllocatedBatch() =
default; // Automatically releases all handles via RAII
};
mutable std::shared_mutex client_pool_mutex_;
std::shared_ptr<coro_io::client_pools<coro_rpc::coro_rpc_client>>
client_pools_;
/**
* @brief Generic RPC invocation helper for single-result operations
* @tparam ServiceMethod Pointer to WrappedMasterService member function
* @tparam ReturnType The expected return type of the RPC call
* @tparam Args Parameter types for the RPC call
* @param args Arguments to pass to the RPC call
* @return The result of the RPC call
*/
template <auto ServiceMethod, typename ReturnType, typename... Args>
[[nodiscard]] tl::expected<ReturnType, ErrorCode> invoke_rpc(
const std::string &client_addr, Args &&...args);
};
// Python-specific wrapper class for client interface
class PyClient {
public:
@ -138,6 +197,7 @@ class PyClient {
const UUID &task_id) = 0;
std::shared_ptr<mooncake::Client> client_ = nullptr;
std::shared_ptr<mooncake::ClientRequester> client_requester_ = nullptr;
std::shared_ptr<mooncake::FileStorage> file_storage_ = nullptr;
std::shared_ptr<ClientBufferAllocator> client_buffer_allocator_ = nullptr;
};

18
mooncake-store/include/real_client.h
View File

@ -359,7 +359,8 @@ class RealClient : public PyClient {
const std::string &rdma_devices = "",
const std::string &master_server_addr = "127.0.0.1:50051",
const std::shared_ptr<TransferEngine> &transfer_engine = nullptr,
const std::string &ipc_socket_path = "", bool enable_offload = false);
const std::string &ipc_socket_path = "", int local_rpc_port = 50052,
bool enable_offload = false);
tl::expected<void, ErrorCode> initAll_internal(
const std::string &protocol, const std::string &device_name,
@ -450,6 +451,20 @@ class RealClient : public PyClient {
tl::expected<PingResponse, ErrorCode> ping(const UUID &client_id);
tl::expected<BatchGetOffloadObjectResponse, ErrorCode>
batch_get_offload_object(const std::vector<std::string> &keys,
const std::vector<int64_t> &sizes);
/**
* @brief Retrieves multiple stored objects from a remote service.
* @param target_rpc_service_addr Address of the remote RPC service (e.g.,
"ip:port").
*/
tl::expected<void, ErrorCode> batch_get_into_offload_object_internal(
const std::string &target_rpc_service_addr,
std::unordered_map<std::string, Slice> &objects);
std::unique_ptr<AutoPortBinder> port_binder_ = nullptr;
struct SegmentDeleter {
@ -485,6 +500,7 @@ class RealClient : public PyClient {
std::string protocol;
std::string device_name;
std::string local_hostname;
std::string local_rpc_addr;
bool use_hugepage_ = false;
struct MappedShm {

16
mooncake-store/include/rpc_types.h
View File

@ -140,4 +140,20 @@ struct TaskCompleteRequest {
};
YLT_REFL(TaskCompleteRequest, id, status, message);
struct BatchGetOffloadObjectResponse {
std::vector<uint64_t> pointers;
std::string transfer_engine_addr;
uint64_t gc_ttl_ms;
BatchGetOffloadObjectResponse() = default;
BatchGetOffloadObjectResponse(std::vector<uint64_t>&& pointers_param,
std::string transfer_engine_addr_param,
uint64_t gc_ttl_ms_param)
: pointers(std::move(pointers_param)),
transfer_engine_addr(std::move(transfer_engine_addr_param)),
gc_ttl_ms(gc_ttl_ms_param) {}
};
YLT_REFL(BatchGetOffloadObjectResponse, pointers, transfer_engine_addr,
gc_ttl_ms);
} // namespace mooncake

4
mooncake-store/include/storage_backend.h
View File

@ -91,6 +91,10 @@ struct FileStorageConfig {
// Interval between heartbeats sent to the control plane (in seconds)
uint32_t heartbeat_interval_seconds = 10;
// Interval between client_buffer_gc (in seconds)
uint32_t client_buffer_gc_interval_seconds = 1;
uint64_t client_buffer_gc_ttl_ms = 5000;
// Validates the configuration for correctness and consistency
bool Validate() const;

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

@ -390,6 +390,12 @@ class TransferSubmitter {
std::vector<std::vector<Slice>>& all_slices,
TransferRequest::OpCode op_code);
std::optional<TransferFuture> submit_batch_get_offload_object(
const std::string& transfer_engine_addr,
const std::vector<std::string>& keys,
const std::vector<uint64_t>& pointers,
const std::unordered_map<std::string, Slice>& batched_slices);
private:
TransferEngine& engine_;
std::unique_ptr<MemcpyWorkerPool> memcpy_pool_;

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

@ -111,9 +111,11 @@ std::vector<Slice> split_into_slices(BufferHandle& handle) {
uint64_t calculate_total_size(const Replica::Descriptor& replica) {
uint64_t total_length = 0;
if (replica.is_memory_replica() == false) {
if (replica.is_disk_replica()) {
auto& disk_descriptor = replica.get_disk_descriptor();
total_length = disk_descriptor.object_size;
} else if (replica.is_local_disk_replica()) {
total_length = replica.get_local_disk_descriptor().object_size;
} else {
total_length = replica.get_memory_descriptor().buffer_descriptor.size_;
}
@ -122,7 +124,7 @@ uint64_t calculate_total_size(const Replica::Descriptor& replica) {
int allocateSlices(std::vector<Slice>& slices,
const Replica::Descriptor& replica, void* buffer_ptr) {
if (replica.is_memory_replica() == false) {
if (replica.is_disk_replica()) {
// For disk-based replica, split into slices based on file size
uint64_t offset = 0;
uint64_t total_length = replica.get_disk_descriptor().object_size;
@ -132,6 +134,9 @@ int allocateSlices(std::vector<Slice>& slices,
slices.emplace_back(Slice{chunk_ptr, chunk_size});
offset += chunk_size;
}
} else if (replica.is_local_disk_replica()) {
slices.emplace_back(
Slice{buffer_ptr, replica.get_local_disk_descriptor().object_size});
} else {
// For memory-based replica, split into slices based on buffer
// descriptors

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

@ -1609,11 +1609,23 @@ tl::expected<void, ErrorCode> Client::OffloadObjectHeartbeat(
return {};
}
tl::expected<void, ErrorCode> Client::BatchPutOffloadObject(
tl::expected<void, ErrorCode> Client::BatchGetOffloadObject(
const std::string& transfer_engine_addr,
const std::vector<std::string>& keys,
const std::vector<uintptr_t>& pointers,
const std::unordered_map<std::string, Slice>& batched_slices) {
const std::unordered_map<std::string, Slice>& batch_slices) {
auto future = transfer_submitter_->submit_batch_get_offload_object(
transfer_engine_addr, keys, pointers, batch_slices);
if (!future) {
LOG(ERROR) << "Failed to submit transfer operation";
return tl::make_unexpected(ErrorCode::TRANSFER_FAIL);
}
VLOG(1) << "Using transfer strategy: " << future->strategy();
auto result = future->get();
if (result != ErrorCode::OK) {
LOG(ERROR) << "Transfer failed, error code is " << result;
return tl::make_unexpected(result);
}
return {};
}

114
mooncake-store/src/file_storage.cpp
View File

@ -44,6 +44,13 @@ FileStorageConfig FileStorageConfig::FromEnvironment() {
config.heartbeat_interval_seconds =
GetEnvOr<uint32_t>("MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS",
config.heartbeat_interval_seconds);
config.client_buffer_gc_interval_seconds =
GetEnvOr<uint32_t>("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_INTERVAL_SECONDS",
config.heartbeat_interval_seconds);
config.client_buffer_gc_ttl_ms =
GetEnvOr<uint64_t>("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS",
config.client_buffer_gc_ttl_ms);
return config;
}
@ -128,12 +135,12 @@ bool FileStorageConfig::Validate() const {
return true;
}
FileStorage::FileStorage(std::shared_ptr<Client> client,
const std::string& local_rpc_addr,
const FileStorageConfig& config)
: client_(client),
FileStorage::FileStorage(const FileStorageConfig& config,
std::shared_ptr<Client> client,
const std::string& local_rpc_addr)
: config_(config),
client_(client),
local_rpc_addr_(local_rpc_addr),
config_(config),
client_buffer_allocator_(
ClientBufferAllocator::create(config.local_buffer_size, "")) {
if (!config.Validate()) {
@ -154,6 +161,10 @@ FileStorage::~FileStorage() {
if (heartbeat_thread_.joinable()) {
heartbeat_thread_.join();
}
client_buffer_gc_running_ = false;
if (client_buffer_gc_thread_.joinable()) {
client_buffer_gc_thread_.join();
}
}
tl::expected<void, ErrorCode> FileStorage::Init() {
@ -220,41 +231,35 @@ tl::expected<void, ErrorCode> FileStorage::Init() {
std::chrono::seconds(config_.heartbeat_interval_seconds));
}
});
client_buffer_gc_running_.store(true);
client_buffer_gc_thread_ =
std::thread(&FileStorage::ClientBufferGCThreadFunc, this);
return {};
}
tl::expected<void, ErrorCode> FileStorage::BatchGet(
const std::string& transfer_engine_addr,
const std::vector<std::string>& keys,
const std::vector<uintptr_t>& pointers, const std::vector<int64_t>& sizes) {
tl::expected<std::vector<uint64_t>, ErrorCode> FileStorage::BatchGet(
const std::vector<std::string>& keys, const std::vector<int64_t>& sizes) {
auto start_time = std::chrono::steady_clock::now();
auto allocate_res = AllocateBatch(keys, sizes);
if (!allocate_res) {
LOG(ERROR) << "Failed to allocate batch objects, target = "
<< transfer_engine_addr;
LOG(ERROR) << "Failed to allocate batch objects";
return tl::make_unexpected(allocate_res.error());
}
auto result = BatchLoad(allocate_res.value().slices);
auto allocated_batch = allocate_res.value();
auto result = BatchLoad(allocated_batch->slices);
if (!result) {
LOG(ERROR) << "Batch load object failed,err_code = " << result.error();
return result;
return tl::make_unexpected(result.error());
}
auto batch_put_result = client_->BatchPutOffloadObject(
transfer_engine_addr, keys, pointers, allocate_res.value().slices);
MutexLocker locker(&client_buffer_mutex_);
client_buffer_allocated_batches_.emplace_back(std::move(allocated_batch));
auto end_time = std::chrono::steady_clock::now();
auto elapsed_time = std::chrono::duration_cast<std::chrono::microseconds>(
end_time - start_time)
.count();
VLOG(1) << "Time taken for FileStorage::BatchGet: " << elapsed_time
<< "us,with transfer_engine_addr: " << transfer_engine_addr
<< ", key size: " << keys.size();
if (!batch_put_result) {
LOG(ERROR) << "Batch write offload object failed,err_code = "
<< batch_put_result.error();
return batch_put_result;
}
return {};
<< "us, key size: " << keys.size();
return allocate_res.value<>()->pointers;
}
tl::expected<void, ErrorCode> FileStorage::OffloadObjects(
@ -434,22 +439,73 @@ tl::expected<void, ErrorCode> FileStorage::RegisterLocalMemory() {
return {};
}
tl::expected<FileStorage::AllocatedBatch, ErrorCode> FileStorage::AllocateBatch(
const std::vector<std::string>& keys, const std::vector<int64_t>& sizes) {
AllocatedBatch result;
tl::expected<std::shared_ptr<FileStorage::AllocatedBatch>, ErrorCode>
FileStorage::AllocateBatch(const std::vector<std::string>& keys,
const std::vector<int64_t>& sizes) {
auto result = std::make_shared<AllocatedBatch>();
std::chrono::steady_clock::time_point now =
std::chrono::steady_clock::now();
auto lease_timeout =
now + std::chrono::milliseconds(config_.client_buffer_gc_ttl_ms);
u_int64_t total_size = 0;
bool gc_triggered = false;
for (size_t i = 0; i < keys.size(); ++i) {
assert(sizes[i] <= kMaxSliceSize);
auto alloc_result = client_buffer_allocator_->allocate(sizes[i]);
if (!alloc_result && !gc_triggered) {
gc_triggered = true;
{
MutexLocker locker(&client_buffer_mutex_);
auto gc_now = std::chrono::steady_clock::now();
auto it = client_buffer_allocated_batches_.begin();
while (it != client_buffer_allocated_batches_.end()) {
if (gc_now >= (*it)->lease_timeout) {
it = client_buffer_allocated_batches_.erase(it);
} else {
++it;
}
}
}
alloc_result = client_buffer_allocator_->allocate(sizes[i]);
}
if (!alloc_result) {
LOG(ERROR) << "Failed to allocate slice buffer, size = " << sizes[i]
<< ", key = " << keys[i];
return tl::make_unexpected(ErrorCode::BUFFER_OVERFLOW);
}
result.slices.emplace(
total_size += sizes[i];
result->slices.emplace(
keys[i], Slice{alloc_result->ptr(), static_cast<size_t>(sizes[i])});
result.handles.emplace_back(std::move(alloc_result.value()));
result->pointers.emplace_back(
reinterpret_cast<uintptr_t>(alloc_result->ptr()));
result->handles.emplace_back(std::move(alloc_result.value()));
result->lease_timeout = lease_timeout;
}
result->total_size = total_size;
return result;
}
void FileStorage::ClientBufferGCThreadFunc() {
LOG(INFO) << "action=client_buffer_gc_thread_started";
while (client_buffer_gc_running_) {
{
MutexLocker locker(&client_buffer_mutex_);
if (!client_buffer_allocated_batches_.empty()) {
auto now = std::chrono::steady_clock::now();
client_buffer_allocated_batches_.erase(
std::remove_if(
client_buffer_allocated_batches_.begin(),
client_buffer_allocated_batches_.end(),
[&](const std::shared_ptr<AllocatedBatch>& batch) {
return now >= batch->lease_timeout;
}),
client_buffer_allocated_batches_.end());
}
}
std::this_thread::sleep_for(
std::chrono::seconds(config_.client_buffer_gc_interval_seconds));
}
LOG(INFO) << "action=client_buffer_gc_thread_stopped";
}
} // namespace mooncake

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

@ -740,8 +740,10 @@ auto MasterService::AddReplica(const UUID& client_id, const std::string& key,
-> tl::expected<void, ErrorCode> {
MetadataAccessorRW accessor(this, key);
if (!accessor.Exists()) {
LOG(ERROR) << "key=" << key << ", error=object_not_found";
return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND);
accessor.Create(
client_id,
replica.get_descriptor().get_local_disk_descriptor().object_size,
std::vector<Replica>{}, false);
}
auto& metadata = accessor.Get();
if (replica.type() != ReplicaType::LOCAL_DISK) {

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

@ -166,7 +166,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
const std::string &protocol, const std::string &rdma_devices,
const std::string &master_server_addr,
const std::shared_ptr<TransferEngine> &transfer_engine,
const std::string &ipc_socket_path, bool enable_offload) {
const std::string &ipc_socket_path, int local_rpc_port,
bool enable_offload) {
this->protocol = protocol;
this->ipc_socket_path_ = ipc_socket_path;
const bool should_use_hugepage =
@ -184,6 +185,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
if (user_specified_port) {
// User specified port, no retry needed
this->local_hostname = local_hostname;
this->local_rpc_addr =
hostname.substr(0, colon_pos + 1) + std::to_string(local_rpc_port);
auto client_opt = mooncake::Client::Create(
this->local_hostname, metadata_server, protocol, device_name,
master_server_addr, transfer_engine);
@ -209,8 +212,10 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
std::this_thread::sleep_for(std::chrono::milliseconds(100));
continue;
}
this->local_hostname = hostname + ":" + std::to_string(port);
this->local_hostname = hostname + ":" + std::to_string(port);
this->local_rpc_addr =
hostname + ":" + std::to_string(local_rpc_port);
auto client_opt = mooncake::Client::Create(
this->local_hostname, metadata_server, protocol, device_name,
master_server_addr, transfer_engine);
@ -315,8 +320,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
}
if (enable_offload) {
auto file_storage_config = FileStorageConfig::FromEnvironment();
file_storage_ = std::make_shared<FileStorage>(client_, local_hostname,
file_storage_config);
file_storage_ = std::make_shared<FileStorage>(
file_storage_config, client_, this->local_rpc_addr);
auto init_result = file_storage_->Init();
if (!init_result) {
LOG(ERROR) << "file storage init failed with error: "
@ -324,6 +329,7 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
return init_result;
}
}
client_requester_ = std::make_shared<ClientRequester>();
return {};
}
@ -1490,6 +1496,7 @@ std::vector<tl::expected<int64_t, ErrorCode>>
RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
const std::vector<void *> &buffers,
const std::vector<size_t> &sizes) {
auto start_time = std::chrono::steady_clock::now();
// Validate preconditions
if (!client_) {
LOG(ERROR) << "Client is not initialized";
@ -1506,8 +1513,7 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
}
const size_t num_keys = keys.size();
std::vector<tl::expected<int64_t, ErrorCode>> results;
results.reserve(num_keys);
std::vector<tl::expected<int64_t, ErrorCode>> results(num_keys);
if (num_keys == 0) {
return results;
@ -1526,6 +1532,7 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
};
std::vector<ValidKeyInfo> valid_operations;
std::unordered_map<std::string, ValidKeyInfo> valid_local_disk_operations;
valid_operations.reserve(num_keys);
for (size_t i = 0; i < num_keys; ++i) {
@ -1534,7 +1541,7 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
// Handle query failures
if (!query_results[i]) {
const auto error = query_results[i].error();
results.emplace_back(tl::unexpected(error));
results[i] = tl::unexpected(error);
if (error != ErrorCode::OBJECT_NOT_FOUND &&
error != ErrorCode::REPLICA_IS_NOT_READY) {
LOG(ERROR) << "Query failed for key '" << key
@ -1547,7 +1554,7 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
auto query_result_values = query_results[i].value();
if (query_result_values.replicas.empty()) {
LOG(ERROR) << "Empty replica list for key: " << key;
results.emplace_back(tl::unexpected(ErrorCode::INVALID_REPLICA));
results[i] = tl::unexpected(ErrorCode::INVALID_REPLICA);
continue;
}
@ -1560,7 +1567,7 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
LOG(ERROR) << "Buffer too small for key '" << key
<< "': required=" << total_size
<< ", available=" << sizes[i];
results.emplace_back(tl::unexpected(ErrorCode::INVALID_PARAMS));
results[i] = tl::unexpected(ErrorCode::INVALID_PARAMS);
continue;
}
@ -1568,6 +1575,18 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
std::vector<Slice> key_slices;
allocateSlices(key_slices, replica, buffers[i]);
if (query_result_values.replicas.size() == 1 &&
query_result_values.replicas.at(0).is_local_disk_replica()) {
valid_local_disk_operations.emplace(
key,
ValidKeyInfo{.key = key,
.original_index = i,
.query_result = std::move(query_result_values),
.slices = std::move(key_slices),
.total_size = total_size});
results[i] = static_cast<int64_t>(total_size);
continue;
}
// Store operation info for batch processing
valid_operations.push_back(
{.key = key,
@ -1577,11 +1596,11 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
.total_size = total_size});
// Set success result (actual bytes transferred)
results.emplace_back(static_cast<int64_t>(total_size));
results[i] = static_cast<int64_t>(total_size);
}
// Early return if no valid operations
if (valid_operations.empty()) {
if (valid_operations.empty() && valid_local_disk_operations.empty()) {
return results;
}
@ -1598,23 +1617,66 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
batch_query_results.push_back(op.query_result);
batch_slices[op.key] = op.slices;
}
if (!valid_operations.empty()) {
// Execute batch transfer
const auto batch_get_results =
client_->BatchGet(batch_keys, batch_query_results, batch_slices);
// Execute batch transfer
const auto batch_get_results =
client_->BatchGet(batch_keys, batch_query_results, batch_slices);
// Process transfer results
for (size_t j = 0; j < batch_get_results.size(); ++j) {
const auto &op = valid_operations[j];
// Process transfer results
for (size_t j = 0; j < batch_get_results.size(); ++j) {
const auto &op = valid_operations[j];
if (!batch_get_results[j]) {
const auto error = batch_get_results[j].error();
LOG(ERROR) << "BatchGet failed for key '" << op.key
<< "': " << toString(error);
results[op.original_index] = tl::unexpected(error);
if (!batch_get_results[j]) {
const auto error = batch_get_results[j].error();
LOG(ERROR) << "BatchGet failed for key '" << op.key
<< "': " << toString(error);
results[op.original_index] = tl::unexpected(error);
}
}
}
// Prepare batch transfer data structures
std::unordered_map<std::string, std::unordered_map<std::string, Slice>>
offload_objects;
for (const auto &op_it : valid_local_disk_operations) {
const auto &replica = op_it.second.query_result.replicas.at(0);
auto [store_segment_it, _] = offload_objects.try_emplace(
replica.get_local_disk_descriptor().transport_endpoint);
store_segment_it->second.emplace(op_it.first,
op_it.second.slices.at(0));
}
size_t offload_object_count = 0;
auto start_read_store_time = std::chrono::steady_clock::now();
for (auto &offload_objects_it : offload_objects) {
offload_object_count += offload_objects_it.second.size();
auto batch_get_offload_result = batch_get_into_offload_object_internal(
offload_objects_it.first, offload_objects_it.second);
if (!batch_get_offload_result) {
LOG(ERROR) << "Batch get store object failed with error: "
<< batch_get_offload_result.error();
for (const auto &offload_object_it : offload_objects_it.second) {
results[valid_local_disk_operations.at(offload_object_it.first)
.original_index] =
tl::make_unexpected(batch_get_offload_result.error());
}
}
}
auto end_time = std::chrono::steady_clock::now();
auto elapsed_time = std::chrono::duration_cast<std::chrono::microseconds>(
end_time - start_time)
.count();
auto read_store_time =
std::chrono::duration_cast<std::chrono::microseconds>(
end_time - start_read_store_time)
.count();
LOG(INFO) << "Time taken for batch_get_into: " << elapsed_time
<< "us, read store: " << read_store_time
<< "us, with memory key count: " << valid_operations.size()
<< ", offload key count: " << offload_object_count;
return results;
}
@ -2150,5 +2212,111 @@ tl::expected<QueryTaskResponse, ErrorCode> RealClient::query_task(
const UUID &task_id) {
return client_->QueryTask(task_id);
}
tl::expected<BatchGetOffloadObjectResponse, ErrorCode>
RealClient::batch_get_offload_object(const std::vector<std::string> &keys,
const std::vector<int64_t> &sizes) {
auto result = file_storage_->BatchGet(keys, sizes);
if (!result) {
LOG(ERROR) << "Batch get offload object failed,err_code = "
<< result.error();
return tl::make_unexpected(result.error());
}
return BatchGetOffloadObjectResponse(
std::move(result.value()), client_->GetTransportEndpoint(),
file_storage_->config_.client_buffer_gc_ttl_ms);
}
tl::expected<void, ErrorCode>
RealClient::batch_get_into_offload_object_internal(
const std::string &target_rpc_service_addr,
std::unordered_map<std::string, Slice> &objects) {
auto start_time = std::chrono::steady_clock::now();
std::vector<std::string> keys;
std::vector<int64_t> sizes;
for (const auto &object_it : objects) {
keys.emplace_back(object_it.first);
sizes.emplace_back(object_it.second.size);
}
auto batchGetResp = client_requester_->batch_get_offload_object(
target_rpc_service_addr, keys, sizes);
if (!batchGetResp) {
LOG(ERROR) << "Batch get offload object failed with error: "
<< batchGetResp.error();
return tl::make_unexpected(batchGetResp.error());
}
auto result =
client_->BatchGetOffloadObject(batchGetResp->transfer_engine_addr, keys,
batchGetResp->pointers, objects);
auto end_time = std::chrono::steady_clock::now();
auto elapsed_time = static_cast<uint64_t>(
std::chrono::duration_cast<std::chrono::milliseconds>(end_time -
start_time)
.count());
LOG(INFO) << "Time taken for batch_get_into_offload_object_internal: "
<< elapsed_time
<< "ms, with target_rpc_service_addr: " << target_rpc_service_addr
<< ", key size: " << objects.size()
<< "gc ttl: " << batchGetResp->gc_ttl_ms << "ms.";
if (!result) {
LOG(ERROR) << "Batch get into offload object failed with error: "
<< result.error();
return result;
}
if (elapsed_time >= batchGetResp->gc_ttl_ms) {
return tl::make_unexpected(ErrorCode::OBJECT_HAS_LEASE);
}
return {};
}
ClientRequester::ClientRequester() {
coro_io::client_pool<coro_rpc::coro_rpc_client>::pool_config pool_conf{};
const char *value = std::getenv("MC_RPC_PROTOCOL");
if (value && std::string_view(value) == "rdma") {
pool_conf.client_config.socket_config =
coro_io::ib_socket_t::config_t{};
}
client_pools_ =
std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(
pool_conf);
}
tl::expected<BatchGetOffloadObjectResponse, ErrorCode>
ClientRequester::batch_get_offload_object(const std::string &client_addr,
const std::vector<std::string> &keys,
const std::vector<int64_t> sizes) {
auto result =
invoke_rpc<&RealClient::batch_get_offload_object,
BatchGetOffloadObjectResponse>(client_addr, keys, sizes);
if (!result) {
LOG(ERROR)
<< "Failed to invoke batch_get_offload_object, client_addr = "
<< client_addr << ", error is: " << result.error();
}
return result;
}
template <auto ServiceMethod, typename ReturnType, typename... Args>
tl::expected<ReturnType, ErrorCode> ClientRequester::invoke_rpc(
const std::string &client_addr, Args &&...args) {
auto client_pool = client_pools_->at(client_addr);
return async_simple::coro::syncAwait(
[&]() -> async_simple::coro::Lazy<tl::expected<ReturnType, ErrorCode>> {
auto ret = co_await client_pool->send_request(
[&](coro_io::client_reuse_hint,
coro_rpc::coro_rpc_client &client) {
return client.send_request<ServiceMethod>(
std::forward<Args>(args)...);
});
if (!ret.has_value()) {
LOG(ERROR) << "Dummy Client not available";
co_return tl::make_unexpected(ErrorCode::RPC_FAIL);
}
auto result = co_await std::move(ret.value());
if (!result) {
LOG(ERROR) << "RPC call failed: " << result.error().msg;
co_return tl::make_unexpected(ErrorCode::RPC_FAIL);
}
co_return result->result();
}());
}
} // namespace mooncake

4
mooncake-store/src/real_client_main.cpp
View File

@ -45,6 +45,8 @@ void RegisterClientRpcService(coro_rpc::coro_rpc_server &server,
server.register_handler<&RealClient::create_copy_task>(&real_client);
server.register_handler<&RealClient::create_move_task>(&real_client);
server.register_handler<&RealClient::query_task>(&real_client);
server.register_handler<&RealClient::batch_get_offload_object>(
&real_client);
}
} // namespace mooncake
@ -57,7 +59,7 @@ int main(int argc, char *argv[]) {
FLAGS_host, FLAGS_metadata_server, global_segment_size, 0,
FLAGS_protocol, FLAGS_device_names, FLAGS_master_server_address,
nullptr, "@mooncake_client_" + std::to_string(FLAGS_port) + ".sock",
FLAGS_enable_offload);
FLAGS_port, FLAGS_enable_offload);
if (!res) {
LOG(FATAL) << "Failed to setup client: " << toString(res.error());
return -1;

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

@ -515,6 +515,33 @@ std::optional<TransferFuture> TransferSubmitter::submit_batch(
return future;
}
std::optional<TransferFuture>
TransferSubmitter::submit_batch_get_offload_object(
const std::string& transfer_engine_addr,
const std::vector<std::string>& keys, const std::vector<uint64_t>& pointers,
const std::unordered_map<std::string, Slice>& batched_slices) {
std::optional<TransferFuture> future;
std::vector<TransferRequest> requests;
for (size_t i = 0; i < keys.size(); ++i) {
auto key = keys[i];
auto pointer = pointers[i];
SegmentHandle seg = engine_.openSegment(transfer_engine_addr);
if (seg == static_cast<uint64_t>(ERR_INVALID_ARGUMENT)) {
LOG(ERROR) << "Failed to open segment " << transfer_engine_addr;
return std::nullopt;
}
const auto& slice = batched_slices.find(key)->second;
TransferRequest request;
request.opcode = TransferRequest::READ;
request.source = static_cast<char*>(slice.ptr);
request.target_id = seg;
request.target_offset = pointer;
request.length = slice.size;
requests.emplace_back(request);
}
return submitTransfer(requests);
}
std::optional<TransferFuture> TransferSubmitter::submitMemcpyOperation(
const AllocatedBuffer::Descriptor& handle, const std::vector<Slice>& slices,
const TransferRequest::OpCode op_code) {

27
mooncake-store/tests/file_storage_test.cpp
View File

@ -49,7 +49,7 @@ class FileStorageTest : public ::testing::Test {
batch_data, buckets);
}
tl::expected<FileStorage::AllocatedBatch, ErrorCode>
tl::expected<std::shared_ptr<FileStorage::AllocatedBatch>, ErrorCode>
FileStorageAllocateBatch(FileStorage& fileStorage,
const std::vector<std::string>& keys,
const std::vector<int64_t>& sizes) {
@ -108,7 +108,7 @@ TEST_F(FileStorageTest, IsEnableOffloading) {
auto file_storage_config = FileStorageConfig::FromEnvironment();
file_storage_config.storage_filepath = data_path;
file_storage_config.local_buffer_size = 128 * 1024 * 1024;
FileStorage fileStorage1(nullptr, "localhost:9003", file_storage_config);
FileStorage fileStorage1(file_storage_config, nullptr, "localhost:9003");
ASSERT_TRUE(FileStorageBatchOffload(fileStorage1, keys, sizes, batch_data));
auto enable_offloading_result1 =
FileStorageIsEnableOffloading(fileStorage1);
@ -121,7 +121,7 @@ TEST_F(FileStorageTest, IsEnableOffloading) {
// Case 2: total_keys_limit < bucket_keys_limit => cannot offload
SetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "10");
file_storage_config.total_keys_limit = 9;
FileStorage fileStorage2(nullptr, "localhost:9003", file_storage_config);
FileStorage fileStorage2(file_storage_config, nullptr, "localhost:9003");
auto enable_offloading_result2 =
FileStorageIsEnableOffloading(fileStorage2);
ASSERT_TRUE(enable_offloading_result2 &&
@ -131,7 +131,7 @@ TEST_F(FileStorageTest, IsEnableOffloading) {
SetEnv("MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", "969");
file_storage_config.total_keys_limit = 10'000'000;
file_storage_config.total_size_limit = 100;
FileStorage fileStorage3(nullptr, "localhost:9003", file_storage_config);
FileStorage fileStorage3(file_storage_config, nullptr, "localhost:9003");
auto enable_offloading_result3 =
FileStorageIsEnableOffloading(fileStorage3);
ASSERT_TRUE(enable_offloading_result3 &&
@ -144,7 +144,7 @@ TEST_F(FileStorageTest, BatchLoad) {
std::unordered_map<std::string, std::string> batch_data;
auto file_storage_config = FileStorageConfig::FromEnvironment();
file_storage_config.storage_filepath = data_path;
FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config);
FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003");
ASSERT_TRUE(FileStorageBatchOffload(fileStorage, keys, sizes, batch_data));
std::unordered_map<std::string, Slice> batch_slice;
std::vector<BufferHandle> buff;
@ -152,7 +152,8 @@ TEST_F(FileStorageTest, BatchLoad) {
auto allocate_res = FileStorageAllocateBatch(fileStorage, keys, sizes);
ASSERT_TRUE(allocate_res);
ASSERT_TRUE(FileStorageBatchLoad(fileStorage, allocate_res.value().slices));
ASSERT_TRUE(
FileStorageBatchLoad(fileStorage, allocate_res.value()->slices));
for (auto& slice_it : batch_slice) {
std::string data(static_cast<char*>(slice_it.second.ptr),
slice_it.second.size);
@ -171,7 +172,7 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_keys_limit) {
file_storage_config.storage_filepath = data_path;
file_storage_config.scanmeta_iterator_keys_limit = 969;
SetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "10");
FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config);
FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003");
ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket(
fileStorage, offloading_objects, buckets_keys));
ASSERT_EQ(buckets_keys.size(), 3);
@ -198,7 +199,7 @@ TEST_F(FileStorageTest, GroupOffloadingKeysByBucket_bucket_size_limit) {
auto file_storage_config = FileStorageConfig::FromEnvironment();
file_storage_config.storage_filepath = data_path;
SetEnv("MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", "10");
FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config);
FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003");
ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket(
fileStorage, offloading_objects, buckets_keys));
ASSERT_EQ(buckets_keys.size(), 3);
@ -227,7 +228,7 @@ TEST_F(FileStorageTest,
file_storage_config.storage_filepath = data_path;
SetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "9");
SetEnv("MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", "496");
FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config);
FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003");
ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket(
fileStorage, offloading_objects, buckets_keys));
for (size_t i = 0; i < buckets_keys.size(); i++) {
@ -252,7 +253,7 @@ TEST_F(FileStorageTest,
std::vector<std::vector<std::string>> buckets_keys;
auto file_storage_config = FileStorageConfig::FromEnvironment();
file_storage_config.storage_filepath = data_path;
FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config);
FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003");
ASSERT_TRUE(FileStorageGroupOffloadingKeysByBucket(
fileStorage, offloading_objects, buckets_keys));
offloading_objects.clear();
@ -399,7 +400,7 @@ TEST_F(FileStorageTest, BatchLoad_WithStorageBackendAdaptor) {
auto total_path = fs::path(data_path) / file_per_key_config.fsdir;
fs::create_directories(total_path);
FileStorage fileStorage(nullptr, "localhost:9003", file_storage_config);
FileStorage fileStorage(file_storage_config, nullptr, "localhost:9003");
auto offload_res =
FileStorageBatchOffload(fileStorage, keys, sizes, batch_data);
@ -410,10 +411,10 @@ TEST_F(FileStorageTest, BatchLoad_WithStorageBackendAdaptor) {
auto batch = std::move(allocate_res.value());
auto load_res = FileStorageBatchLoad(fileStorage, batch.slices);
auto load_res = FileStorageBatchLoad(fileStorage, batch->slices);
ASSERT_TRUE(load_res) << "FileStorageBatchLoad failed";
for (const auto& it : batch.slices) {
for (const auto& it : batch->slices) {
const std::string& key = it.first;
const Slice& slice = it.second;
std::string data(static_cast<char*>(slice.ptr), slice.size);

23
mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp
View File

@ -264,12 +264,29 @@ struct Session : public std::enable_shared_from_this<Session> {
struct TcpContext {
TcpContext(short port) : acceptor(io_context) {
std::error_code ec;
asio::ip::tcp::endpoint endpoint(asio::ip::tcp::v6(), port);
acceptor.open(endpoint.protocol());
acceptor.set_option(asio::ip::v6_only(false));
acceptor.open(endpoint.protocol(), ec);
if (!ec) {
acceptor.set_option(asio::ip::v6_only(false), ec);
if (!ec) {
acceptor.set_option(
asio::ip::tcp::acceptor::reuse_address(true));
acceptor.bind(endpoint, ec);
if (!ec) {
acceptor.listen();
return;
}
}
acceptor.close();
}
LOG(ERROR) << "Failed to set up IPv6 dual-stack listener: "
<< ec.message() << " (error code: " << ec.value() << ")";
asio::ip::tcp::endpoint endpoint_v4(asio::ip::tcp::v4(), port);
acceptor.open(endpoint_v4.protocol());
acceptor.set_option(asio::ip::tcp::acceptor::reuse_address(true));
acceptor.bind(endpoint);
acceptor.bind(endpoint_v4);
acceptor.listen();
}