forked from mooncake-track/Mooncake
feat(client): Abstract client-side data transmission for async and batch optimization (#455)
* feat(client): add transfer submitter for optimized data transfer Signed-off-by: Jinyang Su <751080330@qq.com> * feat(store): implement async memcpy task execution with worker pool Add `MemcpyWorkerPool` to manage asynchronous execution of memcpy tasks. Refactor `BatchGet` and `BatchPut` methods for parallel execution and enhance logging for better traceability. * Squashed commit of the following: commit38c435fcc6Author: Feng Ren <alogfans@users.noreply.github.com> Date: Wed Jun 11 16:50:29 2025 +0800 Revert "[TransferEngine] Fix minor bugs in NVLink transport and benchmark (#468)" (#469) This reverts commitffaad6aa18. commit41b1df7954Author: ykwd <oneday117@qq.com> Date: Wed Jun 11 16:37:05 2025 +0800 [Store] Add initial support for master high availability failover (#451) * A temp version. Better to continue development after merging the latest main branch * Temp version to merge the latest main branch * Allow optional use HA mode, in default use non-HA mode. Fix a minor metrics bug. * Refactor the etcd_helper * refactor ha_helper * Add some unit tests. Refactor the code * Update cmakelists: build etcd_wrapper in default * Fix ci problems. Compile etcd wrapper only when use_etcd or with_store are set. * Update python config relating to mooncake-store client * make some blocking etcd helper function cancellable. bug fix: add string name of new errors that will be used in tostring. * Refactor etcd related code * Bug fix * Add basic masterviewhelper unit tests * In ci flow, install and start etcd to run HA feature unit test. * Fix a ci bug * Reuse master_server_address parameter and remove enable_ha parameter. * Format the code. Fix a minor bug. * Handle the error case: the coro server may fail to start or return internal error. commitffaad6aa18Author: Feng Ren <alogfans@users.noreply.github.com> Date: Wed Jun 11 16:02:41 2025 +0800 [TransferEngine] Fix minor bugs in NVLink transport and benchmark (#468) * [TransferEngine] Fix compilation bug in NVLink xport * [TransferEngine] Fix minor bugs in nvlink benchmark Signed-off-by: Jinyang Su <751080330@qq.com> --------- Signed-off-by: Jinyang Su <751080330@qq.com>
This commit is contained in:
JinYan Su
GitHub
parent
883a5a734e
commit
1a53701e3d
|
|
@ -613,6 +613,7 @@ PYBIND11_MODULE(store, m) {
|
|||
return reinterpret_cast<uintptr_t>(self.ptr());
|
||||
})
|
||||
.def("size", &SliceBuffer::size)
|
||||
.def("__len__", &SliceBuffer::size)
|
||||
.def_buffer([](SliceBuffer &self) -> py::buffer_info {
|
||||
// SliceBuffer now always contains contiguous memory
|
||||
if (self.size() > 0) {
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@
|
|||
#include "master_client.h"
|
||||
#include "rpc_service.h"
|
||||
#include "transfer_engine.h"
|
||||
#include "transfer_task.h"
|
||||
#include "types.h"
|
||||
#include "ha_helper.h"
|
||||
|
||||
|
|
@ -215,9 +216,22 @@ class Client {
|
|||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
std::vector<Slice>& slices);
|
||||
|
||||
/**
|
||||
* @brief Find the first complete replica from a replica list
|
||||
* @param replica_list List of replicas to search through
|
||||
* @param handles Output vector to store the buffer handles of the found
|
||||
* replica
|
||||
* @return ErrorCode::OK if found, ErrorCode::INVALID_REPLICA if no complete
|
||||
* replica
|
||||
*/
|
||||
ErrorCode FindFirstCompleteReplica(
|
||||
const std::vector<Replica::Descriptor>& replica_list,
|
||||
std::vector<AllocatedBuffer::Descriptor>& handles);
|
||||
|
||||
// Core components
|
||||
TransferEngine transfer_engine_;
|
||||
MasterClient master_client_;
|
||||
std::unique_ptr<TransferSubmitter> transfer_submitter_;
|
||||
|
||||
// Client local segments
|
||||
struct Segment{
|
||||
|
|
|
|||
|
|
@ -0,0 +1,331 @@
|
|||
#pragma once
|
||||
|
||||
#include <atomic>
|
||||
#include <condition_variable>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <optional>
|
||||
#include <ostream>
|
||||
#include <queue>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
#include "transfer_engine.h"
|
||||
#include "transport/transport.h"
|
||||
#include "types.h"
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
/**
|
||||
* @brief Transfer strategy enumeration
|
||||
*/
|
||||
enum class TransferStrategy {
|
||||
LOCAL_MEMCPY = 0, // Local memory copy using memcpy
|
||||
TRANSFER_ENGINE = 1 // Remote transfer using transfer engine
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Stream operator for TransferStrategy
|
||||
*/
|
||||
inline std::ostream& operator<<(std::ostream& os,
|
||||
const TransferStrategy& strategy) noexcept {
|
||||
switch (strategy) {
|
||||
case TransferStrategy::LOCAL_MEMCPY:
|
||||
return os << "LOCAL_MEMCPY";
|
||||
case TransferStrategy::TRANSFER_ENGINE:
|
||||
return os << "TRANSFER_ENGINE";
|
||||
default:
|
||||
return os << "UNKNOWN";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Abstract base class for operation state management
|
||||
*
|
||||
* This class encapsulates the common state and behavior for async transfer
|
||||
* operations. Derived classes implement strategy-specific waiting logic.
|
||||
*/
|
||||
class OperationState {
|
||||
public:
|
||||
OperationState() = default;
|
||||
virtual ~OperationState() = default;
|
||||
|
||||
// Non-copyable, non-movable
|
||||
OperationState(const OperationState&) = delete;
|
||||
OperationState& operator=(const OperationState&) = delete;
|
||||
OperationState(OperationState&&) = delete;
|
||||
OperationState& operator=(OperationState&&) = delete;
|
||||
|
||||
/**
|
||||
* @brief Check if the operation has completed
|
||||
*/
|
||||
virtual bool is_completed() = 0;
|
||||
|
||||
/**
|
||||
* @brief Get the operation result. Make sure to call is_completed() first.
|
||||
*/
|
||||
ErrorCode get_result() const { // lock mutex
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
assert(result_.has_value() &&
|
||||
"get_result() called on an incomplete or failed-to-set "
|
||||
"operation state.");
|
||||
return result_.value_or(ErrorCode::INVALID_PARAMS);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get the transfer strategy
|
||||
*/
|
||||
virtual TransferStrategy get_strategy() const = 0;
|
||||
|
||||
/**
|
||||
* @brief Wait for the operation to complete (strategy-specific
|
||||
* implementation)
|
||||
*/
|
||||
virtual void wait_for_completion() = 0;
|
||||
|
||||
protected:
|
||||
std::optional<ErrorCode> result_ = std::nullopt;
|
||||
mutable std::mutex mutex_;
|
||||
std::condition_variable cv_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Operation state for local memcpy transfers
|
||||
*/
|
||||
class MemcpyOperationState : public OperationState {
|
||||
public:
|
||||
bool is_completed() override {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
return result_.has_value();
|
||||
}
|
||||
|
||||
void set_completed(ErrorCode error_code) {
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
assert(!result_.has_value());
|
||||
result_.emplace(error_code);
|
||||
}
|
||||
cv_.notify_all();
|
||||
}
|
||||
|
||||
void wait_for_completion() override {
|
||||
std::unique_lock<std::mutex> lock(mutex_);
|
||||
cv_.wait(lock, [this] { return result_.has_value(); });
|
||||
}
|
||||
|
||||
TransferStrategy get_strategy() const override {
|
||||
return TransferStrategy::LOCAL_MEMCPY;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Operation state for transfer engine operations
|
||||
*/
|
||||
class TransferEngineOperationState : public OperationState {
|
||||
public:
|
||||
TransferEngineOperationState(TransferEngine& engine, BatchID batch_id,
|
||||
size_t batch_size)
|
||||
: engine_(engine), batch_id_(batch_id), batch_size_(batch_size) {
|
||||
CHECK(batch_id_ != Transport::INVALID_BATCH_ID)
|
||||
<< "Invalid batch ID for transfer engine operation";
|
||||
}
|
||||
|
||||
~TransferEngineOperationState() { engine_.freeBatchID(batch_id_); }
|
||||
|
||||
bool is_completed() override;
|
||||
|
||||
void wait_for_completion() override;
|
||||
|
||||
TransferStrategy get_strategy() const override {
|
||||
return TransferStrategy::TRANSFER_ENGINE;
|
||||
}
|
||||
|
||||
private:
|
||||
/**
|
||||
* @brief Check the current completion status of the task, make sure to lock
|
||||
* the mutex before calling this function.
|
||||
* Updates the internal state and returns true if the task is completed.
|
||||
*/
|
||||
void check_task_status();
|
||||
|
||||
void set_result_internal(ErrorCode error_code);
|
||||
|
||||
TransferEngine& engine_;
|
||||
BatchID batch_id_;
|
||||
size_t batch_size_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Represents the future result of an asynchronous transfer operation
|
||||
*
|
||||
* This class provides a std::future-like interface for asynchronous transfer
|
||||
* operations. Users can check completion status, wait for results, or get the
|
||||
* final error code.
|
||||
*/
|
||||
class TransferFuture {
|
||||
public:
|
||||
explicit TransferFuture(std::shared_ptr<OperationState> state);
|
||||
|
||||
// Non-copyable but movable
|
||||
TransferFuture(const TransferFuture&) = delete;
|
||||
TransferFuture& operator=(const TransferFuture&) = delete;
|
||||
TransferFuture(TransferFuture&&) = default;
|
||||
TransferFuture& operator=(TransferFuture&&) = default;
|
||||
|
||||
/**
|
||||
* @brief Check if the operation has completed (non-blocking)
|
||||
* @return true if the operation is finished, false otherwise
|
||||
*/
|
||||
bool isReady() const;
|
||||
|
||||
/**
|
||||
* @brief Wait for the operation to complete (blocking)
|
||||
* @return ErrorCode indicating success or failure
|
||||
*/
|
||||
ErrorCode wait();
|
||||
|
||||
/**
|
||||
* @brief Get the result, waiting if necessary (blocking)
|
||||
* @return ErrorCode indicating success or failure
|
||||
*/
|
||||
ErrorCode get();
|
||||
|
||||
/**
|
||||
* @brief Get the transfer strategy used by this operation
|
||||
* @return TransferStrategy enum value
|
||||
*/
|
||||
TransferStrategy strategy() const;
|
||||
|
||||
private:
|
||||
std::shared_ptr<OperationState> state_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Memory copy operation descriptor
|
||||
*/
|
||||
struct MemcpyOperation {
|
||||
void* dest;
|
||||
const void* src;
|
||||
size_t size;
|
||||
|
||||
MemcpyOperation(void* d, const void* s, size_t sz)
|
||||
: dest(d), src(s), size(sz) {}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Memcpy task for async execution
|
||||
*/
|
||||
struct MemcpyTask {
|
||||
std::vector<MemcpyOperation> operations;
|
||||
std::shared_ptr<MemcpyOperationState> state;
|
||||
|
||||
MemcpyTask(std::vector<MemcpyOperation> ops,
|
||||
std::shared_ptr<MemcpyOperationState> s)
|
||||
: operations(std::move(ops)), state(std::move(s)) {}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Thread pool for asynchronous memcpy operations
|
||||
*
|
||||
* This class manages a single worker thread that executes memcpy operations
|
||||
* asynchronously.
|
||||
*/
|
||||
class MemcpyWorkerPool {
|
||||
public:
|
||||
explicit MemcpyWorkerPool();
|
||||
~MemcpyWorkerPool();
|
||||
|
||||
// Non-copyable, non-movable
|
||||
MemcpyWorkerPool(const MemcpyWorkerPool&) = delete;
|
||||
MemcpyWorkerPool& operator=(const MemcpyWorkerPool&) = delete;
|
||||
MemcpyWorkerPool(MemcpyWorkerPool&&) = delete;
|
||||
MemcpyWorkerPool& operator=(MemcpyWorkerPool&&) = delete;
|
||||
|
||||
/**
|
||||
* @brief Submit a memcpy task for async execution
|
||||
* @param task The memcpy task to execute
|
||||
*/
|
||||
void submitTask(MemcpyTask task);
|
||||
|
||||
private:
|
||||
void workerThread();
|
||||
|
||||
std::vector<std::thread> workers_;
|
||||
std::queue<MemcpyTask> task_queue_;
|
||||
std::mutex queue_mutex_;
|
||||
std::condition_variable queue_cv_;
|
||||
std::atomic<bool> shutdown_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Submitter class for asynchronous transfer operations
|
||||
*
|
||||
* This class analyzes transfer requirements, selects optimal strategies, and
|
||||
* immediately submits operations returning TransferFuture objects for result
|
||||
* tracking.
|
||||
*/
|
||||
class TransferSubmitter {
|
||||
public:
|
||||
explicit TransferSubmitter(TransferEngine& engine,
|
||||
const std::string& local_hostname);
|
||||
|
||||
/**
|
||||
* @brief Submit an asynchronous transfer operation
|
||||
*
|
||||
* Analyzes the transfer requirements, selects the optimal strategy,
|
||||
* and immediately submits the operation. Returns a TransferFuture
|
||||
* that can be used to track completion and get results.
|
||||
*
|
||||
* @param handles Buffer descriptors for the transfer
|
||||
* @param slices Memory slices for the transfer
|
||||
* @param op_code Transfer operation (READ/WRITE)
|
||||
* @return TransferFuture representing the async operation, or nullopt on
|
||||
* failure
|
||||
*/
|
||||
std::optional<TransferFuture> submit(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
std::vector<Slice>& slices, Transport::TransferRequest::OpCode op_code);
|
||||
|
||||
private:
|
||||
TransferEngine& engine_;
|
||||
const std::string local_hostname_;
|
||||
std::unique_ptr<MemcpyWorkerPool> memcpy_pool_;
|
||||
|
||||
/**
|
||||
* @brief Select the optimal transfer strategy
|
||||
*/
|
||||
TransferStrategy selectStrategy(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
const std::vector<Slice>& slices) const;
|
||||
|
||||
/**
|
||||
* @brief Check if all handles refer to local segments
|
||||
*/
|
||||
bool isLocalTransfer(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles) const;
|
||||
|
||||
/**
|
||||
* @brief Validate transfer parameters
|
||||
*/
|
||||
bool validateTransferParams(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
const std::vector<Slice>& slices) const;
|
||||
|
||||
/**
|
||||
* @brief Submit memcpy operation asynchronously
|
||||
*/
|
||||
std::optional<TransferFuture> submitMemcpyOperation(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
std::vector<Slice>& slices, Transport::TransferRequest::OpCode op_code);
|
||||
|
||||
/**
|
||||
* @brief Submit transfer engine operation asynchronously
|
||||
*/
|
||||
std::optional<TransferFuture> submitTransferEngineOperation(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
std::vector<Slice>& slices, Transport::TransferRequest::OpCode op_code);
|
||||
};
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
@ -11,6 +11,9 @@ set(MOONCAKE_STORE_SOURCES
|
|||
master_metric_manager.cpp
|
||||
etcd_helper.cpp
|
||||
ha_helper.cpp
|
||||
transfer_task.cpp
|
||||
etcd_helper.cpp
|
||||
ha_helper.cpp
|
||||
)
|
||||
|
||||
# The cache_allocator library
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@
|
|||
|
||||
#include "rpc_service.h"
|
||||
#include "transfer_engine.h"
|
||||
#include "transfer_task.h"
|
||||
#include "transport/transport.h"
|
||||
#include "types.h"
|
||||
|
||||
|
|
@ -179,6 +180,10 @@ ErrorCode Client::InitTransferEngine(const std::string& local_hostname,
|
|||
}
|
||||
CHECK(transport) << "Failed to install transport";
|
||||
|
||||
// Initialize TransferSubmitter after transfer engine is ready
|
||||
transfer_submitter_ =
|
||||
std::make_unique<TransferSubmitter>(transfer_engine_, local_hostname);
|
||||
|
||||
return ErrorCode::OK;
|
||||
}
|
||||
|
||||
|
|
@ -196,7 +201,7 @@ std::optional<std::shared_ptr<Client>> Client::Create(
|
|||
|
||||
// Initialize transfer engine
|
||||
err = client->InitTransferEngine(local_hostname, metadata_connstring,
|
||||
protocol, protocol_args);
|
||||
protocol, protocol_args);
|
||||
if (err != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to initialize transfer engine";
|
||||
return std::nullopt;
|
||||
|
|
@ -270,51 +275,34 @@ ErrorCode Client::BatchQuery(const std::vector<std::string>& object_keys,
|
|||
ErrorCode Client::Get(const std::string& object_key,
|
||||
const ObjectInfo& object_info,
|
||||
std::vector<Slice>& slices) {
|
||||
// Get the first complete replica
|
||||
for (size_t i = 0; i < object_info.replica_list.size(); ++i) {
|
||||
if (object_info.replica_list[i].status == ReplicaStatus::COMPLETE) {
|
||||
const auto& replica = object_info.replica_list[i];
|
||||
|
||||
std::vector<AllocatedBuffer::Descriptor> handles;
|
||||
for (const auto& handle : replica.buffer_descriptors) {
|
||||
VLOG(1) << "handle: segment_name=" << handle.segment_name_
|
||||
<< " buffer=" << handle.buffer_address_
|
||||
<< " size=" << handle.size_;
|
||||
if (handle.status_ != BufStatus::COMPLETE) {
|
||||
LOG(ERROR) << "incomplete_handle_found segment_name="
|
||||
<< handle.segment_name_;
|
||||
return ErrorCode::INVALID_PARAMS;
|
||||
}
|
||||
handles.push_back(handle);
|
||||
}
|
||||
// Fast path: if segment is on local host and we have single slice
|
||||
// and handle, use memcpy instead of going through the transfer
|
||||
// engine to improve performance and save bandwidth
|
||||
if (slices.size() == 1 && handles.size() == 1 &&
|
||||
handles[0].size_ == slices[0].size &&
|
||||
handles[0].segment_name_ == this->local_hostname_) {
|
||||
VLOG(1) << "Using fast path (memcpy) for local transfer";
|
||||
memcpy(slices[0].ptr, (char*)handles[0].buffer_address_,
|
||||
handles[0].size_);
|
||||
return ErrorCode::OK;
|
||||
}
|
||||
|
||||
if (TransferRead(handles, slices) != ErrorCode::OK) {
|
||||
LOG(ERROR) << "transfer_read_failed key=" << object_key;
|
||||
return ErrorCode::INVALID_PARAMS;
|
||||
}
|
||||
return ErrorCode::OK;
|
||||
// Find the first complete replica
|
||||
std::vector<AllocatedBuffer::Descriptor> handles;
|
||||
ErrorCode err = FindFirstCompleteReplica(object_info.replica_list, handles);
|
||||
if (err != ErrorCode::OK) {
|
||||
if (err == ErrorCode::INVALID_REPLICA) {
|
||||
LOG(ERROR) << "no_complete_replicas_found key=" << object_key;
|
||||
}
|
||||
return err;
|
||||
}
|
||||
|
||||
LOG(ERROR) << "no_complete_replicas_found key=" << object_key;
|
||||
return ErrorCode::INVALID_REPLICA;
|
||||
if (TransferRead(handles, slices) != ErrorCode::OK) {
|
||||
LOG(ERROR) << "transfer_read_failed key=" << object_key;
|
||||
return ErrorCode::INVALID_PARAMS;
|
||||
}
|
||||
return ErrorCode::OK;
|
||||
}
|
||||
|
||||
ErrorCode Client::BatchGet(
|
||||
const std::vector<std::string>& object_keys,
|
||||
BatchObjectInfo& batched_object_info,
|
||||
std::unordered_map<std::string, std::vector<Slice>>& slices) {
|
||||
CHECK(transfer_submitter_) << "TransferSubmitter not initialized";
|
||||
|
||||
// Collect all transfer operations for parallel execution
|
||||
std::vector<std::pair<std::string, TransferFuture>> pending_transfers;
|
||||
pending_transfers.reserve(object_keys.size());
|
||||
|
||||
// Submit all transfers in parallel
|
||||
for (const auto& key : object_keys) {
|
||||
auto object_info_it = batched_object_info.batch_replica_list.find(key);
|
||||
auto slices_it = slices.find(key);
|
||||
|
|
@ -325,15 +313,48 @@ ErrorCode Client::BatchGet(
|
|||
return ErrorCode::INVALID_PARAMS;
|
||||
}
|
||||
|
||||
ObjectInfo object_info;
|
||||
object_info.replica_list = object_info_it->second;
|
||||
object_info.error_code = ErrorCode::OK;
|
||||
if (Get(key, object_info, slices_it->second) != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to get key: " << key;
|
||||
// Find the first complete replica for this key
|
||||
const auto& replica_list = object_info_it->second;
|
||||
std::vector<AllocatedBuffer::Descriptor> handles;
|
||||
ErrorCode err = FindFirstCompleteReplica(replica_list, handles);
|
||||
if (err != ErrorCode::OK) {
|
||||
if (err == ErrorCode::INVALID_REPLICA) {
|
||||
LOG(ERROR) << "no_complete_replicas_found key=" << key;
|
||||
}
|
||||
slices.clear();
|
||||
return ErrorCode::INVALID_PARAMS;
|
||||
return err;
|
||||
}
|
||||
|
||||
// Submit transfer operation asynchronously
|
||||
auto future = transfer_submitter_->submit(handles, slices_it->second,
|
||||
TransferRequest::READ);
|
||||
if (!future) {
|
||||
LOG(ERROR) << "Failed to submit transfer operation for key: "
|
||||
<< key;
|
||||
slices.clear();
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
|
||||
VLOG(1) << "Submitted transfer for key " << key
|
||||
<< " using strategy: " << static_cast<int>(future->strategy());
|
||||
|
||||
pending_transfers.emplace_back(key, std::move(*future));
|
||||
}
|
||||
|
||||
// Wait for all transfers to complete
|
||||
for (auto& [key, future] : pending_transfers) {
|
||||
ErrorCode result = future.get();
|
||||
if (result != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Transfer failed for key: " << key
|
||||
<< " with error: " << static_cast<int>(result);
|
||||
slices.clear();
|
||||
return result;
|
||||
}
|
||||
VLOG(1) << "Transfer completed successfully for key: " << key;
|
||||
}
|
||||
|
||||
VLOG(1) << "BatchGet completed successfully for " << object_keys.size()
|
||||
<< " keys";
|
||||
return ErrorCode::OK;
|
||||
}
|
||||
|
||||
|
|
@ -368,26 +389,15 @@ ErrorCode Client::Put(const ObjectKey& key, std::vector<Slice>& slices,
|
|||
handles.push_back(handle);
|
||||
}
|
||||
|
||||
// Fast path: if segment is on local host and we have single slice and
|
||||
// handle, use memcpy instead of going through the transfer engine
|
||||
if (slices.size() == 1 && handles.size() == 1 &&
|
||||
handles[0].size_ == slices[0].size &&
|
||||
handles[0].segment_name_ == this->local_hostname_) {
|
||||
VLOG(1) << "Using fast path (memcpy) for local transfer";
|
||||
memcpy((char*)handles[0].buffer_address_, slices[0].ptr,
|
||||
handles[0].size_);
|
||||
} else {
|
||||
// Normal path: use transfer engine
|
||||
ErrorCode transfer_err = TransferWrite(handles, slices);
|
||||
if (transfer_err != ErrorCode::OK) {
|
||||
// Revoke put operation
|
||||
auto revoke_err = master_client_.PutRevoke(key);
|
||||
if (revoke_err.error_code != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to revoke put operation";
|
||||
return revoke_err.error_code;
|
||||
}
|
||||
return transfer_err;
|
||||
ErrorCode transfer_err = TransferWrite(handles, slices);
|
||||
if (transfer_err != ErrorCode::OK) {
|
||||
// Revoke put operation
|
||||
auto revoke_err = master_client_.PutRevoke(key);
|
||||
if (revoke_err.error_code != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to revoke put operation";
|
||||
return revoke_err.error_code;
|
||||
}
|
||||
return transfer_err;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -404,6 +414,8 @@ ErrorCode Client::BatchPut(
|
|||
const std::vector<ObjectKey>& keys,
|
||||
std::unordered_map<std::string, std::vector<Slice>>& batched_slices,
|
||||
ReplicateConfig& config) {
|
||||
CHECK(transfer_submitter_) << "TransferSubmitter not initialized";
|
||||
|
||||
std::unordered_map<std::string, std::vector<size_t>> batched_slice_lengths;
|
||||
std::unordered_map<std::string, size_t> batched_value_lengths;
|
||||
for (const auto& key : keys) {
|
||||
|
|
@ -433,7 +445,11 @@ ErrorCode Client::BatchPut(
|
|||
return err;
|
||||
}
|
||||
|
||||
// Transfer data using allocated handles from all replicas
|
||||
// Collect all transfer operations for parallel execution
|
||||
std::vector<std::tuple<std::string, size_t, TransferFuture>>
|
||||
pending_transfers;
|
||||
|
||||
// Submit all transfers in parallel
|
||||
for (const auto& key : keys) {
|
||||
const auto& slices_it = batched_slices.find(key);
|
||||
if (slices_it == batched_slices.end()) {
|
||||
|
|
@ -445,7 +461,10 @@ ErrorCode Client::BatchPut(
|
|||
LOG(ERROR) << "Cannot find replica_list for key: " << key;
|
||||
return ErrorCode::INVALID_PARAMS;
|
||||
}
|
||||
for (auto& replica : replica_list->second) {
|
||||
|
||||
for (size_t replica_idx = 0; replica_idx < replica_list->second.size();
|
||||
++replica_idx) {
|
||||
const auto& replica = replica_list->second[replica_idx];
|
||||
std::vector<AllocatedBuffer::Descriptor> handles;
|
||||
for (const auto& handle : replica.buffer_descriptors) {
|
||||
CHECK(handle.buffer_address_ != 0)
|
||||
|
|
@ -453,37 +472,58 @@ ErrorCode Client::BatchPut(
|
|||
handles.push_back(handle);
|
||||
}
|
||||
|
||||
// Fast path: if segment is on local host and we have single slice
|
||||
// and handle, use memcpy instead of going through the transfer
|
||||
// engine
|
||||
std::vector<Slice>& slices = slices_it->second;
|
||||
if (slices.size() == 1 && handles.size() == 1 &&
|
||||
handles[0].size_ == slices[0].size &&
|
||||
handles[0].segment_name_ == this->local_hostname_) {
|
||||
VLOG(1) << "Using fast path (memcpy) for local transfer";
|
||||
memcpy((char*)handles[0].buffer_address_, slices[0].ptr,
|
||||
handles[0].size_);
|
||||
} else {
|
||||
// Normal path: use transfer engine
|
||||
ErrorCode transfer_err = TransferWrite(handles, slices);
|
||||
if (transfer_err != ErrorCode::OK) {
|
||||
// Revoke put operation
|
||||
auto revoke_err = master_client_.BatchPutRevoke(keys);
|
||||
if (revoke_err.error_code != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to revoke put operation";
|
||||
return revoke_err.error_code;
|
||||
}
|
||||
return transfer_err;
|
||||
// Submit transfer operation asynchronously
|
||||
auto future = transfer_submitter_->submit(
|
||||
handles, slices_it->second, TransferRequest::WRITE);
|
||||
if (!future) {
|
||||
LOG(ERROR) << "Failed to submit transfer operation for key: "
|
||||
<< key << " replica: " << replica_idx;
|
||||
// Revoke put operation
|
||||
auto revoke_err = master_client_.BatchPutRevoke(keys);
|
||||
if (revoke_err.error_code != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to revoke put operation";
|
||||
return revoke_err.error_code;
|
||||
}
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
|
||||
VLOG(1) << "Submitted transfer for key " << key << " replica "
|
||||
<< replica_idx << " using strategy: "
|
||||
<< static_cast<int>(future->strategy());
|
||||
|
||||
pending_transfers.emplace_back(key, replica_idx,
|
||||
std::move(*future));
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for all transfers to complete
|
||||
for (auto& [key, replica_idx, future] : pending_transfers) {
|
||||
ErrorCode result = future.get();
|
||||
if (result != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Transfer failed for key: " << key
|
||||
<< " replica: " << replica_idx
|
||||
<< " with error: " << result;
|
||||
// Revoke put operation
|
||||
auto revoke_err = master_client_.BatchPutRevoke(keys);
|
||||
if (revoke_err.error_code != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to revoke put operation";
|
||||
return revoke_err.error_code;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
VLOG(1) << "Transfer completed successfully for key: " << key
|
||||
<< " replica: " << replica_idx;
|
||||
}
|
||||
|
||||
// End put operation
|
||||
err = master_client_.BatchPutEnd(keys).error_code;
|
||||
if (err != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to end put operation: " << err;
|
||||
return err;
|
||||
}
|
||||
|
||||
VLOG(1) << "BatchPut completed successfully for " << keys.size()
|
||||
<< " keys with " << pending_transfers.size() << " total transfers";
|
||||
return ErrorCode::OK;
|
||||
}
|
||||
|
||||
|
|
@ -591,99 +631,17 @@ ErrorCode Client::IsExist(const std::string& key) {
|
|||
ErrorCode Client::TransferData(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
std::vector<Slice>& slices, TransferRequest::OpCode op_code) {
|
||||
CHECK(!handles.empty()) << "handles is empty";
|
||||
std::vector<TransferRequest> transfer_tasks;
|
||||
if (handles.size() > slices.size()) {
|
||||
LOG(ERROR) << "invalid_partition_count handles_size=" << handles.size()
|
||||
<< " slices_size=" << slices.size();
|
||||
CHECK(transfer_submitter_) << "TransferSubmitter not initialized";
|
||||
|
||||
auto future = transfer_submitter_->submit(handles, slices, op_code);
|
||||
if (!future) {
|
||||
LOG(ERROR) << "Failed to submit transfer operation";
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
|
||||
for (uint64_t idx = 0; idx < handles.size(); ++idx) {
|
||||
auto& handle = handles[idx];
|
||||
auto& slice = slices[idx];
|
||||
if (handle.size_ > slice.size) {
|
||||
LOG(ERROR)
|
||||
<< "Size of replica partition more than provided buffers";
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
Transport::SegmentHandle seg =
|
||||
transfer_engine_.openSegment(handle.segment_name_);
|
||||
if (seg == (uint64_t)ERR_INVALID_ARGUMENT) {
|
||||
LOG(ERROR) << "Failed to open segment " << handle.segment_name_;
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
TransferRequest request;
|
||||
request.opcode = op_code;
|
||||
request.source = static_cast<char*>(slice.ptr);
|
||||
request.target_id = seg;
|
||||
request.target_offset = handle.buffer_address_;
|
||||
request.length = handle.size_;
|
||||
transfer_tasks.push_back(request);
|
||||
}
|
||||
VLOG(1) << "Using transfer strategy: " << future->strategy();
|
||||
|
||||
const size_t batch_size = transfer_tasks.size();
|
||||
BatchID batch_id = transfer_engine_.allocateBatchID(batch_size);
|
||||
if (batch_id == Transport::INVALID_BATCH_ID) {
|
||||
LOG(ERROR) << "Failed to allocate batch ID";
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
|
||||
Status s = transfer_engine_.submitTransfer(batch_id, transfer_tasks);
|
||||
if (!s.ok()) {
|
||||
LOG(ERROR) << "Failed to submit all transfers, error code is "
|
||||
<< s.code();
|
||||
transfer_engine_.freeBatchID(batch_id);
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
|
||||
bool has_err = false;
|
||||
bool all_ready = true;
|
||||
uint32_t try_num = 0;
|
||||
const uint32_t max_try_num = 3;
|
||||
int64_t start_ts = getCurrentTimeInNano();
|
||||
const static int64_t kOneSecondInNano = 1000 * 1000 * 1000;
|
||||
|
||||
while (try_num < max_try_num) {
|
||||
has_err = false;
|
||||
all_ready = true;
|
||||
if (getCurrentTimeInNano() - start_ts > 60 * kOneSecondInNano) {
|
||||
LOG(ERROR) << "Failed to complete transfers after 60 seconds";
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
for (size_t i = 0; i < batch_size; ++i) {
|
||||
TransferStatus status;
|
||||
s = transfer_engine_.getTransferStatus(batch_id, i, status);
|
||||
if (!s.ok()) {
|
||||
LOG(ERROR) << "Transfer " << i
|
||||
<< " error, error_code=" << s.code();
|
||||
transfer_engine_.freeBatchID(batch_id);
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
if (status.s != TransferStatusEnum::COMPLETED) all_ready = false;
|
||||
if (status.s == TransferStatusEnum::FAILED) {
|
||||
LOG(ERROR) << "Transfer failed for task" << i;
|
||||
has_err = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (has_err) {
|
||||
LOG(WARNING) << "Transfer incomplete, retrying... (attempt "
|
||||
<< try_num + 1 << "/" << max_try_num << ")";
|
||||
++try_num;
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(10));
|
||||
}
|
||||
|
||||
if (all_ready) break;
|
||||
}
|
||||
|
||||
if (!all_ready) {
|
||||
LOG(ERROR) << "transfer_incomplete max_attempts=" << max_try_num;
|
||||
return ErrorCode::TRANSFER_FAIL;
|
||||
}
|
||||
|
||||
transfer_engine_.freeBatchID(batch_id);
|
||||
return ErrorCode::OK;
|
||||
return future->get();
|
||||
}
|
||||
|
||||
ErrorCode Client::TransferWrite(
|
||||
|
|
@ -807,4 +765,21 @@ void Client::PingThreadFunc(int current_version) {
|
|||
}
|
||||
}
|
||||
|
||||
ErrorCode Client::FindFirstCompleteReplica(
|
||||
const std::vector<Replica::Descriptor>& replica_list,
|
||||
std::vector<AllocatedBuffer::Descriptor>& handles) {
|
||||
handles.clear();
|
||||
|
||||
// Find the first complete replica
|
||||
for (size_t i = 0; i < replica_list.size(); ++i) {
|
||||
if (replica_list[i].status == ReplicaStatus::COMPLETE) {
|
||||
handles = replica_list[i].buffer_descriptors;
|
||||
return ErrorCode::OK;
|
||||
}
|
||||
}
|
||||
|
||||
// No complete replica found
|
||||
return ErrorCode::INVALID_REPLICA;
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
|
|||
|
|
@ -0,0 +1,403 @@
|
|||
#include "transfer_task.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#include "utils.h"
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
// ============================================================================
|
||||
// MemcpyWorkerPool Implementation
|
||||
// ============================================================================
|
||||
// Since memcpy is bound by memory bandwidth, we only need one worker thread.
|
||||
constexpr int kDefaultMemcpyWorkers = 1;
|
||||
|
||||
MemcpyWorkerPool::MemcpyWorkerPool() : shutdown_(false) {
|
||||
VLOG(1) << "Creating MemcpyWorkerPool with " << kDefaultMemcpyWorkers
|
||||
<< " workers";
|
||||
|
||||
// Start worker threads
|
||||
workers_.reserve(kDefaultMemcpyWorkers);
|
||||
for (int i = 0; i < kDefaultMemcpyWorkers; ++i) {
|
||||
workers_.emplace_back(&MemcpyWorkerPool::workerThread, this);
|
||||
}
|
||||
}
|
||||
|
||||
MemcpyWorkerPool::~MemcpyWorkerPool() {
|
||||
// Signal shutdown
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(queue_mutex_);
|
||||
shutdown_.store(true);
|
||||
}
|
||||
queue_cv_.notify_all();
|
||||
|
||||
// Wait for all workers to finish
|
||||
for (auto& worker : workers_) {
|
||||
if (worker.joinable()) {
|
||||
worker.join();
|
||||
}
|
||||
}
|
||||
|
||||
VLOG(1) << "MemcpyWorkerPool destroyed";
|
||||
}
|
||||
|
||||
void MemcpyWorkerPool::submitTask(MemcpyTask task) {
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(queue_mutex_);
|
||||
if (shutdown_.load()) {
|
||||
LOG(WARNING)
|
||||
<< "Attempting to submit task to shutdown MemcpyWorkerPool";
|
||||
task.state->set_completed(ErrorCode::TRANSFER_FAIL);
|
||||
return;
|
||||
}
|
||||
task_queue_.push(std::move(task));
|
||||
}
|
||||
queue_cv_.notify_one();
|
||||
}
|
||||
|
||||
void MemcpyWorkerPool::workerThread() {
|
||||
VLOG(2) << "MemcpyWorkerPool worker thread started";
|
||||
|
||||
while (true) {
|
||||
MemcpyTask task({}, nullptr);
|
||||
|
||||
// Wait for task or shutdown signal
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(queue_mutex_);
|
||||
queue_cv_.wait(lock, [this] {
|
||||
return shutdown_.load() || !task_queue_.empty();
|
||||
});
|
||||
|
||||
if (shutdown_.load() && task_queue_.empty()) {
|
||||
break;
|
||||
}
|
||||
|
||||
if (!task_queue_.empty()) {
|
||||
task = std::move(task_queue_.front());
|
||||
task_queue_.pop();
|
||||
}
|
||||
}
|
||||
|
||||
// Execute the task if we have one
|
||||
if (task.state) {
|
||||
try {
|
||||
for (const auto& op : task.operations) {
|
||||
std::memcpy(op.dest, op.src, op.size);
|
||||
}
|
||||
|
||||
VLOG(2) << "Memcpy task completed successfully with "
|
||||
<< task.operations.size() << " operations";
|
||||
task.state->set_completed(ErrorCode::OK);
|
||||
} catch (const std::exception& e) {
|
||||
LOG(ERROR) << "Exception during async memcpy: " << e.what();
|
||||
task.state->set_completed(ErrorCode::TRANSFER_FAIL);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
VLOG(2) << "MemcpyWorkerPool worker thread exiting";
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// TransferEngineOperationState Implementation
|
||||
// ============================================================================
|
||||
|
||||
bool TransferEngineOperationState::is_completed() {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
if (result_.has_value()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
check_task_status();
|
||||
return result_.has_value();
|
||||
}
|
||||
|
||||
void TransferEngineOperationState::check_task_status() {
|
||||
// Check all transfers in the batch
|
||||
bool all_completed = true;
|
||||
bool has_failure = false;
|
||||
|
||||
for (size_t i = 0; i < batch_size_; ++i) {
|
||||
TransferStatus status;
|
||||
Status s = engine_.getTransferStatus(batch_id_, i, status);
|
||||
if (!s.ok()) {
|
||||
LOG(ERROR) << "Failed to get transfer status for task " << i
|
||||
<< " with error " << s.message();
|
||||
set_result_internal(ErrorCode::TRANSFER_FAIL);
|
||||
return;
|
||||
}
|
||||
|
||||
switch (status.s) {
|
||||
case TransferStatusEnum::COMPLETED:
|
||||
// This transfer is done, continue checking others
|
||||
break;
|
||||
case TransferStatusEnum::FAILED:
|
||||
case TransferStatusEnum::CANCELED:
|
||||
case TransferStatusEnum::INVALID:
|
||||
LOG(ERROR) << "Transfer failed for task " << i
|
||||
<< " with status " << static_cast<int>(status.s);
|
||||
has_failure = true;
|
||||
break;
|
||||
default:
|
||||
// Transfer is still pending (PENDING, RUNNING, etc.)
|
||||
all_completed = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (has_failure) {
|
||||
set_result_internal(ErrorCode::TRANSFER_FAIL);
|
||||
return;
|
||||
}
|
||||
|
||||
if (all_completed) {
|
||||
set_result_internal(ErrorCode::OK);
|
||||
return;
|
||||
}
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
void TransferEngineOperationState::set_result_internal(ErrorCode error_code) {
|
||||
assert(!result_.has_value() && "Result should only be set once.");
|
||||
result_.emplace(error_code);
|
||||
|
||||
cv_.notify_all();
|
||||
}
|
||||
|
||||
void TransferEngineOperationState::wait_for_completion() {
|
||||
if (is_completed()) {
|
||||
return;
|
||||
}
|
||||
|
||||
VLOG(1) << "Starting transfer engine polling for batch " << batch_id_;
|
||||
constexpr int64_t timeout_seconds = 60;
|
||||
constexpr int64_t kOneSecondInNano = 1000 * 1000 * 1000;
|
||||
|
||||
const int64_t start_ts = getCurrentTimeInNano();
|
||||
|
||||
while (true) {
|
||||
if (getCurrentTimeInNano() - start_ts >
|
||||
timeout_seconds * kOneSecondInNano) {
|
||||
LOG(ERROR) << "Failed to complete transfers after "
|
||||
<< timeout_seconds << " seconds";
|
||||
set_result_internal(ErrorCode::TRANSFER_FAIL);
|
||||
return;
|
||||
}
|
||||
|
||||
std::unique_lock<std::mutex> lock(mutex_);
|
||||
check_task_status();
|
||||
if (result_.has_value()) {
|
||||
VLOG(1) << "Transfer engine operation completed successfully";
|
||||
break;
|
||||
}
|
||||
// Continue polling
|
||||
VLOG(1) << "Transfer engine operation still pending";
|
||||
}
|
||||
|
||||
VLOG(1) << "Transfer engine operation completed successfully";
|
||||
set_result_internal(ErrorCode::OK);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// TransferFuture Implementation
|
||||
// ============================================================================
|
||||
|
||||
TransferFuture::TransferFuture(std::shared_ptr<OperationState> state)
|
||||
: state_(std::move(state)) {
|
||||
CHECK(state_) << "TransferFuture requires valid state";
|
||||
}
|
||||
|
||||
bool TransferFuture::isReady() const { return state_->is_completed(); }
|
||||
|
||||
ErrorCode TransferFuture::wait() {
|
||||
if (!isReady()) {
|
||||
state_->wait_for_completion();
|
||||
}
|
||||
return state_->get_result();
|
||||
}
|
||||
|
||||
ErrorCode TransferFuture::get() { return wait(); }
|
||||
|
||||
TransferStrategy TransferFuture::strategy() const {
|
||||
return state_->get_strategy();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// TransferSubmitter Implementation
|
||||
// ============================================================================
|
||||
|
||||
TransferSubmitter::TransferSubmitter(TransferEngine& engine,
|
||||
const std::string& local_hostname)
|
||||
: engine_(engine),
|
||||
local_hostname_(local_hostname),
|
||||
memcpy_pool_(std::make_unique<MemcpyWorkerPool>()) {
|
||||
CHECK(!local_hostname_.empty()) << "Local hostname cannot be empty";
|
||||
}
|
||||
|
||||
std::optional<TransferFuture> TransferSubmitter::submit(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
std::vector<Slice>& slices, Transport::TransferRequest::OpCode op_code) {
|
||||
if (!validateTransferParams(handles, slices)) {
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
TransferStrategy strategy = selectStrategy(handles, slices);
|
||||
|
||||
switch (strategy) {
|
||||
case TransferStrategy::LOCAL_MEMCPY:
|
||||
return submitMemcpyOperation(handles, slices, op_code);
|
||||
case TransferStrategy::TRANSFER_ENGINE:
|
||||
return submitTransferEngineOperation(handles, slices, op_code);
|
||||
default:
|
||||
LOG(ERROR) << "Unknown transfer strategy: " << strategy;
|
||||
return std::nullopt;
|
||||
}
|
||||
}
|
||||
|
||||
std::optional<TransferFuture> TransferSubmitter::submitMemcpyOperation(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
std::vector<Slice>& slices, Transport::TransferRequest::OpCode op_code) {
|
||||
auto state = std::make_shared<MemcpyOperationState>();
|
||||
|
||||
// Create memcpy operations
|
||||
std::vector<MemcpyOperation> operations;
|
||||
operations.reserve(handles.size());
|
||||
|
||||
for (size_t i = 0; i < handles.size(); ++i) {
|
||||
const auto& handle = handles[i];
|
||||
const auto& slice = slices[i];
|
||||
|
||||
void* dest;
|
||||
const void* src;
|
||||
|
||||
if (op_code == Transport::TransferRequest::READ) {
|
||||
// READ: from handle (remote buffer) to slice (local
|
||||
// buffer)
|
||||
dest = slice.ptr;
|
||||
src = reinterpret_cast<const void*>(handle.buffer_address_);
|
||||
} else {
|
||||
// WRITE: from slice (local buffer) to handle (remote
|
||||
// buffer)
|
||||
dest = reinterpret_cast<void*>(handle.buffer_address_);
|
||||
src = slice.ptr;
|
||||
}
|
||||
|
||||
operations.emplace_back(dest, src, handle.size_);
|
||||
}
|
||||
|
||||
// Submit memcpy operations to worker pool for async execution
|
||||
MemcpyTask task(std::move(operations), state);
|
||||
memcpy_pool_->submitTask(std::move(task));
|
||||
|
||||
VLOG(1) << "Memcpy transfer submitted to worker pool with "
|
||||
<< handles.size() << " operations";
|
||||
|
||||
return TransferFuture(state);
|
||||
}
|
||||
|
||||
std::optional<TransferFuture> TransferSubmitter::submitTransferEngineOperation(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
std::vector<Slice>& slices, Transport::TransferRequest::OpCode op_code) {
|
||||
// Create transfer requests
|
||||
std::vector<Transport::TransferRequest> requests;
|
||||
requests.reserve(handles.size());
|
||||
|
||||
for (size_t i = 0; i < handles.size(); ++i) {
|
||||
const auto& handle = handles[i];
|
||||
const auto& slice = slices[i];
|
||||
|
||||
Transport::SegmentHandle seg =
|
||||
engine_.openSegment(handle.segment_name_);
|
||||
if (seg == static_cast<uint64_t>(ERR_INVALID_ARGUMENT)) {
|
||||
LOG(ERROR) << "Failed to open segment " << handle.segment_name_;
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
Transport::TransferRequest request;
|
||||
request.opcode = op_code;
|
||||
request.source = static_cast<char*>(slice.ptr);
|
||||
request.target_id = seg;
|
||||
request.target_offset = handle.buffer_address_;
|
||||
request.length = handle.size_;
|
||||
|
||||
requests.emplace_back(request);
|
||||
}
|
||||
|
||||
// Allocate batch ID
|
||||
const size_t batch_size = requests.size();
|
||||
BatchID batch_id = engine_.allocateBatchID(batch_size);
|
||||
if (batch_id == Transport::INVALID_BATCH_ID) {
|
||||
LOG(ERROR) << "Failed to allocate batch ID";
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
// Submit transfer
|
||||
Status s = engine_.submitTransfer(batch_id, requests);
|
||||
if (!s.ok()) {
|
||||
LOG(ERROR) << "Failed to submit all transfers, error code is "
|
||||
<< s.code();
|
||||
// Note: batch_id will be freed by TransferEngineOperationState
|
||||
// destructor if we create the state object, otherwise we need to free
|
||||
// it here
|
||||
engine_.freeBatchID(batch_id);
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
// Create state with transfer engine context - no polling thread
|
||||
// needed
|
||||
auto state = std::make_shared<TransferEngineOperationState>(
|
||||
engine_, batch_id, batch_size);
|
||||
|
||||
return TransferFuture(state);
|
||||
}
|
||||
|
||||
TransferStrategy TransferSubmitter::selectStrategy(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
const std::vector<Slice>& slices) const {
|
||||
// Check conditions for local memcpy optimization
|
||||
if (isLocalTransfer(handles)) {
|
||||
return TransferStrategy::LOCAL_MEMCPY;
|
||||
}
|
||||
|
||||
return TransferStrategy::TRANSFER_ENGINE;
|
||||
}
|
||||
|
||||
bool TransferSubmitter::isLocalTransfer(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles) const {
|
||||
return std::all_of(handles.begin(), handles.end(),
|
||||
[this](const auto& handle) {
|
||||
return handle.segment_name_ == local_hostname_;
|
||||
});
|
||||
}
|
||||
|
||||
bool TransferSubmitter::validateTransferParams(
|
||||
const std::vector<AllocatedBuffer::Descriptor>& handles,
|
||||
const std::vector<Slice>& slices) const {
|
||||
if (handles.empty()) {
|
||||
LOG(ERROR) << "handles is empty";
|
||||
return false;
|
||||
}
|
||||
|
||||
if (handles.size() > slices.size()) {
|
||||
LOG(ERROR) << "invalid_partition_count handles_size=" << handles.size()
|
||||
<< " slices_size=" << slices.size();
|
||||
return false;
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < handles.size(); ++i) {
|
||||
if (handles[i].size_ != slices[i].size) {
|
||||
LOG(ERROR) << "Size of replica partition " << i << " ("
|
||||
<< handles[i].size_
|
||||
<< ") does not match provided buffer (" << slices[i].size
|
||||
<< ")";
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
|
@ -41,4 +41,15 @@ target_link_libraries(stress_workload_test PUBLIC
|
|||
gtest
|
||||
gtest_main
|
||||
pthread
|
||||
)
|
||||
)
|
||||
|
||||
add_executable(transfer_task_test transfer_task_test.cpp)
|
||||
target_link_libraries(transfer_task_test PUBLIC
|
||||
mooncake_store
|
||||
cachelib_memory_allocator
|
||||
glog
|
||||
gtest
|
||||
gtest_main
|
||||
pthread
|
||||
)
|
||||
add_test(NAME transfer_task_test COMMAND transfer_task_test)
|
||||
|
|
@ -1,7 +1,6 @@
|
|||
#include <gflags/gflags.h>
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <numa.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
|
|
@ -16,7 +15,7 @@
|
|||
DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp");
|
||||
DEFINE_string(device_name, "ibp6s0",
|
||||
"Device name to use, valid if protocol=rdma");
|
||||
DEFINE_string(transfer_engine_metadata_url, "http://127.0.0.1:8090/metadata",
|
||||
DEFINE_string(transfer_engine_metadata_url, "localhost:2379",
|
||||
"Metadata connection string for transfer engine");
|
||||
DEFINE_uint64(default_kv_lease_ttl, mooncake::DEFAULT_DEFAULT_KV_LEASE_TTL,
|
||||
"Default lease time for kv objects, must be set to the "
|
||||
|
|
@ -27,12 +26,29 @@ namespace testing {
|
|||
|
||||
class ClientIntegrationTest : public ::testing::Test {
|
||||
protected:
|
||||
static std::shared_ptr<Client> CreateClient(const std::string& host_name) {
|
||||
void** args =
|
||||
(FLAGS_protocol == "rdma") ? rdma_args(FLAGS_device_name) : nullptr;
|
||||
|
||||
auto client_opt = Client::Create(
|
||||
host_name, // Local hostname
|
||||
FLAGS_transfer_engine_metadata_url, // Metadata connection string
|
||||
FLAGS_protocol, args,
|
||||
"localhost:50051" // Master server 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;
|
||||
}
|
||||
|
||||
static void SetUpTestSuite() {
|
||||
// Initialize glog
|
||||
google::InitGoogleLogging("ClientIntegrationTest");
|
||||
|
||||
// Set VLOG level to 1 for detailed logs
|
||||
google::SetVLOGLevel("*", 1);
|
||||
FLAGS_logtostderr = 1;
|
||||
|
||||
// Override flags from environment variables if present
|
||||
|
|
@ -45,77 +61,103 @@ class ClientIntegrationTest : public ::testing::Test {
|
|||
<< ", Device name: " << FLAGS_device_name
|
||||
<< ", Metadata URL: " << FLAGS_transfer_engine_metadata_url;
|
||||
|
||||
InitializeClient();
|
||||
InitializeClients();
|
||||
InitializeSegment();
|
||||
}
|
||||
|
||||
static void TearDownTestSuite() {
|
||||
CleanupSegment();
|
||||
CleanupClient();
|
||||
CleanupClients();
|
||||
google::ShutdownGoogleLogging();
|
||||
}
|
||||
|
||||
static void InitializeSegment() {
|
||||
const size_t ram_buffer_size = 1024 * 1024 * 1024; // 1GB
|
||||
const size_t ram_buffer_size = 512 * 1024 * 1024; // 512 MB
|
||||
segment_ptr_ = allocate_buffer_allocator_memory(ram_buffer_size);
|
||||
LOG_ASSERT(segment_ptr_);
|
||||
ErrorCode rc = client_->MountSegment("localhost:17812", segment_ptr_,
|
||||
ram_buffer_size);
|
||||
ErrorCode rc = segment_provider_client_->MountSegment(
|
||||
"localhost:17812", segment_ptr_, ram_buffer_size);
|
||||
if (rc != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to mount segment: " << toString(rc);
|
||||
}
|
||||
LOG(INFO) << "Segment mounted successfully";
|
||||
}
|
||||
|
||||
static void InitializeClient() {
|
||||
void** args =
|
||||
(FLAGS_protocol == "rdma") ? rdma_args(FLAGS_device_name) : nullptr;
|
||||
static void InitializeClients() {
|
||||
// This client is used for testing purposes.
|
||||
test_client_ = CreateClient("localhost:17813");
|
||||
ASSERT_TRUE(test_client_ != nullptr);
|
||||
|
||||
auto client_opt = Client::Create(
|
||||
"localhost:17812", // Local hostname
|
||||
FLAGS_transfer_engine_metadata_url, // Metadata connection string
|
||||
FLAGS_protocol, args,
|
||||
"localhost:50051" // Master server address
|
||||
);
|
||||
|
||||
ASSERT_TRUE(client_opt.has_value()) << "Failed to create client";
|
||||
client_ = *client_opt;
|
||||
// This client is used to provide segments.
|
||||
segment_provider_client_ = CreateClient("localhost:17812");
|
||||
ASSERT_TRUE(segment_provider_client_ != nullptr);
|
||||
|
||||
client_buffer_allocator_ =
|
||||
std::make_unique<SimpleAllocator>(128 * 1024 * 1024);
|
||||
ErrorCode error_code = client_->RegisterLocalMemory(
|
||||
ErrorCode error_code = test_client_->RegisterLocalMemory(
|
||||
client_buffer_allocator_->getBase(), 128 * 1024 * 1024, "cpu:0",
|
||||
false, false);
|
||||
if (error_code != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to allocate transfer buffer: "
|
||||
<< toString(error_code);
|
||||
}
|
||||
|
||||
// Mount segment for test_client_ as well
|
||||
const size_t test_client_ram_buffer_size = 512 * 1024 * 1024; // 512 MB
|
||||
test_client_segment_ptr_ =
|
||||
allocate_buffer_allocator_memory(test_client_ram_buffer_size);
|
||||
LOG_ASSERT(test_client_segment_ptr_);
|
||||
ErrorCode rc = test_client_->MountSegment("localhost:17813",
|
||||
test_client_segment_ptr_,
|
||||
test_client_ram_buffer_size);
|
||||
if (rc != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to mount segment for test_client_: "
|
||||
<< toString(rc);
|
||||
}
|
||||
LOG(INFO) << "Test client segment mounted successfully";
|
||||
}
|
||||
|
||||
static void CleanupClient() {
|
||||
if (client_) {
|
||||
client_.reset(); // Release the client
|
||||
static void CleanupClients() {
|
||||
// Unmount test client segment first
|
||||
if (test_client_ && test_client_segment_ptr_) {
|
||||
if (test_client_->UnmountSegment("localhost:17813",
|
||||
test_client_segment_ptr_) !=
|
||||
ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to unmount test client segment";
|
||||
}
|
||||
}
|
||||
|
||||
if (test_client_) {
|
||||
test_client_.reset();
|
||||
}
|
||||
if (segment_provider_client_) {
|
||||
segment_provider_client_.reset();
|
||||
}
|
||||
}
|
||||
|
||||
static void CleanupSegment() {
|
||||
if (client_->UnmountSegment("localhost:17812", segment_ptr_) !=
|
||||
ErrorCode::OK) {
|
||||
if (segment_provider_client_->UnmountSegment(
|
||||
"localhost:17812", segment_ptr_) != ErrorCode::OK) {
|
||||
LOG(ERROR) << "Failed to unmount segment";
|
||||
}
|
||||
}
|
||||
|
||||
static std::shared_ptr<Client> client_;
|
||||
static std::shared_ptr<Client> test_client_;
|
||||
static std::shared_ptr<Client> segment_provider_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 void* test_client_segment_ptr_;
|
||||
};
|
||||
|
||||
// Static members initialization
|
||||
std::shared_ptr<Client> ClientIntegrationTest::client_ = nullptr;
|
||||
std::shared_ptr<Client> ClientIntegrationTest::test_client_ = nullptr;
|
||||
std::shared_ptr<Client> ClientIntegrationTest::segment_provider_client_ =
|
||||
nullptr;
|
||||
void* ClientIntegrationTest::segment_ptr_ = nullptr;
|
||||
void* ClientIntegrationTest::test_client_segment_ptr_ = nullptr;
|
||||
std::unique_ptr<SimpleAllocator>
|
||||
ClientIntegrationTest::client_buffer_allocator_ = nullptr;
|
||||
|
||||
|
|
@ -133,14 +175,14 @@ TEST_F(ClientIntegrationTest, BasicPutGetOperations) {
|
|||
// Test Put operation
|
||||
ReplicateConfig config;
|
||||
config.replica_num = 1;
|
||||
ASSERT_EQ(client_->Put(key, slices, config), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Put(key, slices, config), ErrorCode::OK);
|
||||
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
|
||||
ErrorCode error_code = client_->Get(key, slices);
|
||||
ErrorCode error_code = test_client_->Get(key, slices);
|
||||
ASSERT_EQ(error_code, ErrorCode::OK);
|
||||
ASSERT_EQ(slices.size(), 1);
|
||||
ASSERT_EQ(slices[0].size, test_data.size());
|
||||
|
|
@ -153,10 +195,10 @@ TEST_F(ClientIntegrationTest, BasicPutGetOperations) {
|
|||
memcpy(buffer, test_data.data(), test_data.size());
|
||||
slices.clear();
|
||||
slices.emplace_back(Slice{buffer, test_data.size()});
|
||||
ASSERT_EQ(client_->Put(key, slices, config), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Put(key, slices, config), ErrorCode::OK);
|
||||
std::this_thread::sleep_for(
|
||||
std::chrono::milliseconds(FLAGS_default_kv_lease_ttl));
|
||||
ASSERT_EQ(client_->Remove(key), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Remove(key), ErrorCode::OK);
|
||||
client_buffer_allocator_->deallocate(buffer, test_data.size());
|
||||
}
|
||||
|
||||
|
|
@ -172,17 +214,17 @@ TEST_F(ClientIntegrationTest, RemoveOperation) {
|
|||
slices.emplace_back(Slice{buffer, test_data.size()});
|
||||
ReplicateConfig config;
|
||||
config.replica_num = 1;
|
||||
ASSERT_EQ(client_->Put(key, slices, config), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Put(key, slices, config), ErrorCode::OK);
|
||||
client_buffer_allocator_->deallocate(buffer, test_data.size());
|
||||
|
||||
// Remove the data
|
||||
ASSERT_EQ(client_->Remove(key), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Remove(key), ErrorCode::OK);
|
||||
|
||||
// Try to get the removed data - should fail
|
||||
buffer = client_buffer_allocator_->allocate(test_data.size());
|
||||
slices.clear();
|
||||
slices.emplace_back(Slice{buffer, test_data.size()});
|
||||
ErrorCode error_code = client_->Get(key, slices);
|
||||
ErrorCode error_code = test_client_->Get(key, slices);
|
||||
ASSERT_NE(error_code, ErrorCode::OK);
|
||||
client_buffer_allocator_->deallocate(buffer, test_data.size());
|
||||
}
|
||||
|
|
@ -205,7 +247,7 @@ TEST_F(ClientIntegrationTest, LocalPreferredAllocationTest) {
|
|||
// compatibility issues in the future.
|
||||
config.preferred_segment = "localhost:17812"; // Local segment
|
||||
|
||||
ASSERT_EQ(client_->Put(key, slices, config), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Put(key, slices, config), ErrorCode::OK);
|
||||
client_buffer_allocator_->deallocate(buffer, test_data.size());
|
||||
|
||||
// Verify data through Get operation
|
||||
|
|
@ -214,14 +256,14 @@ TEST_F(ClientIntegrationTest, LocalPreferredAllocationTest) {
|
|||
slices.emplace_back(Slice{buffer, test_data.size()});
|
||||
|
||||
Client::ObjectInfo objectinfo;
|
||||
ErrorCode error_code = client_->Query(key, objectinfo);
|
||||
ErrorCode error_code = test_client_->Query(key, objectinfo);
|
||||
ASSERT_EQ(error_code, ErrorCode::OK);
|
||||
ASSERT_EQ(objectinfo.replica_list.size(), 1);
|
||||
ASSERT_EQ(objectinfo.replica_list[0].buffer_descriptors.size(), 1);
|
||||
ASSERT_EQ(objectinfo.replica_list[0].buffer_descriptors[0].segment_name_,
|
||||
"localhost:17812");
|
||||
|
||||
error_code = client_->Get(key, objectinfo, slices);
|
||||
error_code = test_client_->Get(key, objectinfo, slices);
|
||||
ASSERT_EQ(error_code, ErrorCode::OK);
|
||||
ASSERT_EQ(slices.size(), 1);
|
||||
ASSERT_EQ(slices[0].size, test_data.size());
|
||||
|
|
@ -231,7 +273,7 @@ TEST_F(ClientIntegrationTest, LocalPreferredAllocationTest) {
|
|||
// Clean up
|
||||
std::this_thread::sleep_for(
|
||||
std::chrono::milliseconds(FLAGS_default_kv_lease_ttl));
|
||||
ASSERT_EQ(client_->Remove(key), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Remove(key), ErrorCode::OK);
|
||||
}
|
||||
|
||||
// Test heavy workload operations
|
||||
|
|
@ -254,14 +296,14 @@ TEST_F(ClientIntegrationTest, DISABLED_AllocateTest) {
|
|||
memcpy(buffer, large_data.data(), data_size);
|
||||
std::vector<Slice> put_slices;
|
||||
put_slices.emplace_back(Slice{buffer, data_size});
|
||||
ErrorCode error_code = client_->Put(key, put_slices, config);
|
||||
ErrorCode error_code = test_client_->Put(key, put_slices, config);
|
||||
if (error_code != ErrorCode::OK) break;
|
||||
client_buffer_allocator_->deallocate(buffer, data_size);
|
||||
// Get and verify data
|
||||
buffer = client_buffer_allocator_->allocate(data_size);
|
||||
std::vector<Slice> get_slices;
|
||||
get_slices.emplace_back(Slice{buffer, data_size});
|
||||
error_code = client_->Get(key, get_slices);
|
||||
error_code = test_client_->Get(key, get_slices);
|
||||
ASSERT_EQ(error_code, ErrorCode::OK);
|
||||
ASSERT_EQ(get_slices[0].size, data_size);
|
||||
|
||||
|
|
@ -276,7 +318,7 @@ TEST_F(ClientIntegrationTest, DISABLED_AllocateTest) {
|
|||
std::vector<Slice> failed_slices;
|
||||
failed_slices.emplace_back(Slice{failed_buffer, data_size});
|
||||
memcpy(failed_buffer, large_data.data(), data_size);
|
||||
ASSERT_NE(client_->Put(allocate_failed_key, failed_slices, config),
|
||||
ASSERT_NE(test_client_->Put(allocate_failed_key, failed_slices, config),
|
||||
ErrorCode::OK);
|
||||
client_buffer_allocator_->deallocate(failed_buffer, data_size);
|
||||
|
||||
|
|
@ -288,10 +330,10 @@ TEST_F(ClientIntegrationTest, DISABLED_AllocateTest) {
|
|||
std::vector<Slice> success_slices;
|
||||
success_slices.emplace_back(Slice{success_buffer, data_size});
|
||||
memcpy(success_buffer, large_data.data(), data_size);
|
||||
ASSERT_EQ(client_->Put(allocate_failed_key, success_slices, config),
|
||||
ASSERT_EQ(test_client_->Put(allocate_failed_key, success_slices, config),
|
||||
ErrorCode::OK);
|
||||
client_buffer_allocator_->deallocate(success_buffer, data_size);
|
||||
ASSERT_EQ(client_->Remove(allocate_failed_key), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Remove(allocate_failed_key), ErrorCode::OK);
|
||||
}
|
||||
|
||||
// Test large allocation operations
|
||||
|
|
@ -320,7 +362,7 @@ TEST_F(ClientIntegrationTest, LargeAllocateTest) {
|
|||
}
|
||||
|
||||
// Put operation
|
||||
ASSERT_EQ(client_->Put(key, slices, config), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Put(key, slices, config), ErrorCode::OK);
|
||||
|
||||
// Clear buffers before Get
|
||||
for (size_t i = 0; i < kNumBuffers; ++i) {
|
||||
|
|
@ -328,7 +370,7 @@ TEST_F(ClientIntegrationTest, LargeAllocateTest) {
|
|||
}
|
||||
|
||||
// Get operation
|
||||
ErrorCode error_code = client_->Get(key, slices);
|
||||
ErrorCode error_code = test_client_->Get(key, slices);
|
||||
ASSERT_EQ(error_code, ErrorCode::OK);
|
||||
|
||||
// Verify data and deallocate buffers
|
||||
|
|
@ -345,12 +387,12 @@ TEST_F(ClientIntegrationTest, LargeAllocateTest) {
|
|||
// Remove the key
|
||||
std::this_thread::sleep_for(
|
||||
std::chrono::milliseconds(FLAGS_default_kv_lease_ttl));
|
||||
ASSERT_EQ(client_->Remove(key), ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->Remove(key), ErrorCode::OK);
|
||||
}
|
||||
|
||||
// Test batch Put/Get operations through the client
|
||||
TEST_F(ClientIntegrationTest, BatchPutGetOperations) {
|
||||
int batch_sz = 10;
|
||||
int batch_sz = 100;
|
||||
std::vector<std::string> keys;
|
||||
std::vector<std::string> test_data_list;
|
||||
std::unordered_map<std::string, std::vector<Slice>> batched_slices;
|
||||
|
|
@ -370,9 +412,34 @@ TEST_F(ClientIntegrationTest, BatchPutGetOperations) {
|
|||
// Test Batch Put operation
|
||||
ReplicateConfig config;
|
||||
config.replica_num = 1;
|
||||
ASSERT_EQ(client_->BatchPut(keys, batched_slices, config), ErrorCode::OK);
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
ASSERT_EQ(test_client_->BatchPut(keys, batched_slices, config),
|
||||
ErrorCode::OK);
|
||||
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";
|
||||
|
||||
// Allocate slice buffers for GetBatch
|
||||
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()});
|
||||
ASSERT_EQ(test_client_->Get(keys[i], slices), ErrorCode::OK);
|
||||
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;
|
||||
|
|
@ -382,32 +449,23 @@ TEST_F(ClientIntegrationTest, BatchPutGetOperations) {
|
|||
Slice{target_buffer, test_data_list[i].size()});
|
||||
target_batched_slices.emplace(keys[i], target_slices);
|
||||
}
|
||||
|
||||
ErrorCode error_code = client_->BatchGet(keys, target_batched_slices);
|
||||
ASSERT_EQ(error_code, ErrorCode::OK);
|
||||
for (int i = 0; i < batch_sz; i++) {
|
||||
const auto& slices = target_batched_slices[keys[i]];
|
||||
memcmp(static_cast<char*>(slices[0].ptr), test_data_list[i].data(),
|
||||
test_data_list[i].size());
|
||||
std::cout << "Key: " << keys[i] << std::endl;
|
||||
for (const auto& slice : slices) {
|
||||
std::cout << "Value: "
|
||||
<< std::string(static_cast<char*>(slices[0].ptr),
|
||||
slice.size)
|
||||
<< std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
error_code = client_->BatchGet({keys[0], keys[0]}, target_batched_slices);
|
||||
ASSERT_NE(error_code, ErrorCode::OK);
|
||||
ASSERT_EQ(test_client_->BatchGet(keys, target_batched_slices),
|
||||
ErrorCode::OK);
|
||||
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++) {
|
||||
for (auto& slice : target_batched_slices[keys[i]]) {
|
||||
client_buffer_allocator_->deallocate(slice.ptr, slice.size);
|
||||
}
|
||||
for (auto& slice : batched_slices[keys[i]]) {
|
||||
client_buffer_allocator_->deallocate(slice.ptr, slice.size);
|
||||
}
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,168 @@
|
|||
// transfer_task_test.cpp
|
||||
#include "transfer_task.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <chrono>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
#include "types.h"
|
||||
|
||||
namespace mooncake {
|
||||
|
||||
// Test fixture for TransferTask tests
|
||||
// TODO: Currently, this test does not cover TransferSubmitter and
|
||||
// TransferEngine integration. Will add more tests in the future.
|
||||
class TransferTaskTest : public ::testing::Test {
|
||||
protected:
|
||||
void SetUp() override {
|
||||
// Initialize glog for logging
|
||||
google::InitGoogleLogging("TransferTaskTest");
|
||||
FLAGS_logtostderr = 1; // Output logs to stderr
|
||||
}
|
||||
|
||||
void TearDown() override {
|
||||
// Cleanup glog
|
||||
google::ShutdownGoogleLogging();
|
||||
}
|
||||
};
|
||||
|
||||
// Test basic MemcpyOperation functionality
|
||||
TEST_F(TransferTaskTest, MemcpyOperationBasic) {
|
||||
const size_t data_size = 1024;
|
||||
std::vector<char> src_data(data_size, 'A');
|
||||
std::vector<char> dest_data(data_size, 'B');
|
||||
|
||||
// Create memcpy operation
|
||||
MemcpyOperation op(dest_data.data(), src_data.data(), data_size);
|
||||
|
||||
// Verify operation parameters
|
||||
EXPECT_EQ(op.dest, dest_data.data());
|
||||
EXPECT_EQ(op.src, src_data.data());
|
||||
EXPECT_EQ(op.size, data_size);
|
||||
|
||||
// Perform memcpy manually to test
|
||||
std::memcpy(op.dest, op.src, op.size);
|
||||
|
||||
// Verify data was copied correctly
|
||||
EXPECT_EQ(dest_data, src_data);
|
||||
for (size_t i = 0; i < data_size; ++i) {
|
||||
EXPECT_EQ(dest_data[i], 'A');
|
||||
}
|
||||
}
|
||||
|
||||
// Test MemcpyOperationState functionality
|
||||
TEST_F(TransferTaskTest, MemcpyOperationState) {
|
||||
auto state = std::make_shared<MemcpyOperationState>();
|
||||
|
||||
// Initially not completed
|
||||
EXPECT_FALSE(state->is_completed());
|
||||
EXPECT_EQ(state->get_strategy(), TransferStrategy::LOCAL_MEMCPY);
|
||||
|
||||
// Set completed with success
|
||||
state->set_completed(ErrorCode::OK);
|
||||
EXPECT_TRUE(state->is_completed());
|
||||
EXPECT_EQ(state->get_result(), ErrorCode::OK);
|
||||
}
|
||||
|
||||
// Test MemcpyWorkerPool basic functionality
|
||||
TEST_F(TransferTaskTest, MemcpyWorkerPoolBasic) {
|
||||
MemcpyWorkerPool pool;
|
||||
|
||||
const size_t data_size = 512;
|
||||
std::vector<char> src_data(data_size, 'X');
|
||||
std::vector<char> dest_data(data_size, 'Y');
|
||||
|
||||
auto state = std::make_shared<MemcpyOperationState>();
|
||||
|
||||
// Create memcpy operations
|
||||
std::vector<MemcpyOperation> operations;
|
||||
operations.emplace_back(dest_data.data(), src_data.data(), data_size);
|
||||
|
||||
// Create and submit task
|
||||
MemcpyTask task(std::move(operations), state);
|
||||
pool.submitTask(std::move(task));
|
||||
|
||||
// Wait for completion
|
||||
state->wait_for_completion();
|
||||
|
||||
// Verify completion and result
|
||||
EXPECT_TRUE(state->is_completed());
|
||||
EXPECT_EQ(state->get_result(), ErrorCode::OK);
|
||||
|
||||
// Verify data was copied correctly
|
||||
for (size_t i = 0; i < data_size; ++i) {
|
||||
EXPECT_EQ(dest_data[i], 'X');
|
||||
}
|
||||
}
|
||||
|
||||
// Test multiple memcpy operations in one task
|
||||
TEST_F(TransferTaskTest, MemcpyWorkerPoolMultipleOperations) {
|
||||
MemcpyWorkerPool pool;
|
||||
|
||||
const size_t num_ops = 3;
|
||||
const size_t data_size = 256;
|
||||
|
||||
std::vector<std::vector<char>> src_buffers(num_ops);
|
||||
std::vector<std::vector<char>> dest_buffers(num_ops);
|
||||
|
||||
// Initialize source buffers with different patterns
|
||||
for (size_t i = 0; i < num_ops; ++i) {
|
||||
src_buffers[i].resize(data_size, 'A' + i);
|
||||
dest_buffers[i].resize(data_size, 'Z');
|
||||
}
|
||||
|
||||
auto state = std::make_shared<MemcpyOperationState>();
|
||||
|
||||
// Create multiple memcpy operations
|
||||
std::vector<MemcpyOperation> operations;
|
||||
for (size_t i = 0; i < num_ops; ++i) {
|
||||
operations.emplace_back(dest_buffers[i].data(), src_buffers[i].data(),
|
||||
data_size);
|
||||
}
|
||||
|
||||
// Create and submit task
|
||||
MemcpyTask task(std::move(operations), state);
|
||||
pool.submitTask(std::move(task));
|
||||
|
||||
// Wait for completion
|
||||
state->wait_for_completion();
|
||||
|
||||
// Verify completion and result
|
||||
EXPECT_TRUE(state->is_completed());
|
||||
EXPECT_EQ(state->get_result(), ErrorCode::OK);
|
||||
|
||||
// Verify all data was copied correctly
|
||||
for (size_t i = 0; i < num_ops; ++i) {
|
||||
for (size_t j = 0; j < data_size; ++j) {
|
||||
EXPECT_EQ(dest_buffers[i][j], 'A' + i);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Test TransferStrategy enum and stream operator
|
||||
TEST_F(TransferTaskTest, TransferStrategyEnum) {
|
||||
// Test enum values
|
||||
EXPECT_EQ(static_cast<int>(TransferStrategy::LOCAL_MEMCPY), 0);
|
||||
EXPECT_EQ(static_cast<int>(TransferStrategy::TRANSFER_ENGINE), 1);
|
||||
|
||||
// Test stream operator
|
||||
std::ostringstream oss;
|
||||
oss << TransferStrategy::LOCAL_MEMCPY;
|
||||
EXPECT_EQ(oss.str(), "LOCAL_MEMCPY");
|
||||
|
||||
oss.str("");
|
||||
oss << TransferStrategy::TRANSFER_ENGINE;
|
||||
EXPECT_EQ(oss.str(), "TRANSFER_ENGINE");
|
||||
}
|
||||
|
||||
} // namespace mooncake
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
::testing::InitGoogleTest(&argc, argv);
|
||||
return RUN_ALL_TESTS();
|
||||
}
|
||||
|
|
@ -12,8 +12,8 @@
|
|||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#ifndef TRANSFER_ENGINE
|
||||
#define TRANSFER_ENGINE
|
||||
#ifndef RDMA_TRANSPORT_H_
|
||||
#define RDMA_TRANSPORT_H_
|
||||
|
||||
#include <infiniband/verbs.h>
|
||||
|
||||
|
|
@ -124,4 +124,4 @@ using BatchID = Transport::BatchID;
|
|||
|
||||
} // namespace mooncake
|
||||
|
||||
#endif // TRANSFER_ENGINE
|
||||
#endif // RDMA_TRANSPORT_H_
|
||||
Loading…
Reference in New Issue