diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 52e23bf6..f6339068 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -613,6 +613,7 @@ PYBIND11_MODULE(store, m) { return reinterpret_cast(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) { diff --git a/mooncake-store/include/client.h b/mooncake-store/include/client.h index 74369489..2f2abbf2 100644 --- a/mooncake-store/include/client.h +++ b/mooncake-store/include/client.h @@ -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& handles, std::vector& 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_list, + std::vector& handles); + // Core components TransferEngine transfer_engine_; MasterClient master_client_; + std::unique_ptr transfer_submitter_; // Client local segments struct Segment{ diff --git a/mooncake-store/include/transfer_task.h b/mooncake-store/include/transfer_task.h new file mode 100644 index 00000000..931473e2 --- /dev/null +++ b/mooncake-store/include/transfer_task.h @@ -0,0 +1,331 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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 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 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 lock(mutex_); + return result_.has_value(); + } + + void set_completed(ErrorCode error_code) { + { + std::lock_guard lock(mutex_); + assert(!result_.has_value()); + result_.emplace(error_code); + } + cv_.notify_all(); + } + + void wait_for_completion() override { + std::unique_lock 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 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 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 operations; + std::shared_ptr state; + + MemcpyTask(std::vector ops, + std::shared_ptr 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 workers_; + std::queue task_queue_; + std::mutex queue_mutex_; + std::condition_variable queue_cv_; + std::atomic 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 submit( + const std::vector& handles, + std::vector& slices, Transport::TransferRequest::OpCode op_code); + + private: + TransferEngine& engine_; + const std::string local_hostname_; + std::unique_ptr memcpy_pool_; + + /** + * @brief Select the optimal transfer strategy + */ + TransferStrategy selectStrategy( + const std::vector& handles, + const std::vector& slices) const; + + /** + * @brief Check if all handles refer to local segments + */ + bool isLocalTransfer( + const std::vector& handles) const; + + /** + * @brief Validate transfer parameters + */ + bool validateTransferParams( + const std::vector& handles, + const std::vector& slices) const; + + /** + * @brief Submit memcpy operation asynchronously + */ + std::optional submitMemcpyOperation( + const std::vector& handles, + std::vector& slices, Transport::TransferRequest::OpCode op_code); + + /** + * @brief Submit transfer engine operation asynchronously + */ + std::optional submitTransferEngineOperation( + const std::vector& handles, + std::vector& slices, Transport::TransferRequest::OpCode op_code); +}; + +} // namespace mooncake diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 08eb445d..41ddb33e 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -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 diff --git a/mooncake-store/src/client.cpp b/mooncake-store/src/client.cpp index 734a9dff..6da5377e 100644 --- a/mooncake-store/src/client.cpp +++ b/mooncake-store/src/client.cpp @@ -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(transfer_engine_, local_hostname); + return ErrorCode::OK; } @@ -196,7 +201,7 @@ std::optional> 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& object_keys, ErrorCode Client::Get(const std::string& object_key, const ObjectInfo& object_info, std::vector& 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 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 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& object_keys, BatchObjectInfo& batched_object_info, std::unordered_map>& slices) { + CHECK(transfer_submitter_) << "TransferSubmitter not initialized"; + + // Collect all transfer operations for parallel execution + std::vector> 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 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(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(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& 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& keys, std::unordered_map>& batched_slices, ReplicateConfig& config) { + CHECK(transfer_submitter_) << "TransferSubmitter not initialized"; + std::unordered_map> batched_slice_lengths; std::unordered_map 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> + 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 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& 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(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& handles, std::vector& slices, TransferRequest::OpCode op_code) { - CHECK(!handles.empty()) << "handles is empty"; - std::vector 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(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_list, + std::vector& 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 diff --git a/mooncake-store/src/transfer_task.cpp b/mooncake-store/src/transfer_task.cpp new file mode 100644 index 00000000..ea73b031 --- /dev/null +++ b/mooncake-store/src/transfer_task.cpp @@ -0,0 +1,403 @@ +#include "transfer_task.h" + +#include + +#include + +#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 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 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 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 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(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 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 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()) { + CHECK(!local_hostname_.empty()) << "Local hostname cannot be empty"; +} + +std::optional TransferSubmitter::submit( + const std::vector& handles, + std::vector& 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 TransferSubmitter::submitMemcpyOperation( + const std::vector& handles, + std::vector& slices, Transport::TransferRequest::OpCode op_code) { + auto state = std::make_shared(); + + // Create memcpy operations + std::vector 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(handle.buffer_address_); + } else { + // WRITE: from slice (local buffer) to handle (remote + // buffer) + dest = reinterpret_cast(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 TransferSubmitter::submitTransferEngineOperation( + const std::vector& handles, + std::vector& slices, Transport::TransferRequest::OpCode op_code) { + // Create transfer requests + std::vector 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(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(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( + engine_, batch_id, batch_size); + + return TransferFuture(state); +} + +TransferStrategy TransferSubmitter::selectStrategy( + const std::vector& handles, + const std::vector& 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& 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& handles, + const std::vector& 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 diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index d5bc90d8..d174aefe 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -41,4 +41,15 @@ target_link_libraries(stress_workload_test PUBLIC gtest gtest_main pthread -) \ No newline at end of file +) + +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) \ No newline at end of file diff --git a/mooncake-store/tests/client_integration_test.cpp b/mooncake-store/tests/client_integration_test.cpp index d24e20ab..473f5e4f 100644 --- a/mooncake-store/tests/client_integration_test.cpp +++ b/mooncake-store/tests/client_integration_test.cpp @@ -1,7 +1,6 @@ #include #include #include -#include #include #include @@ -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 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(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_; + static std::shared_ptr test_client_; + static std::shared_ptr 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 client_buffer_allocator_; static void* segment_ptr_; + static void* test_client_segment_ptr_; }; // Static members initialization -std::shared_ptr ClientIntegrationTest::client_ = nullptr; +std::shared_ptr ClientIntegrationTest::test_client_ = nullptr; +std::shared_ptr ClientIntegrationTest::segment_provider_client_ = + nullptr; void* ClientIntegrationTest::segment_ptr_ = nullptr; +void* ClientIntegrationTest::test_client_segment_ptr_ = nullptr; std::unique_ptr 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 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 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 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 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 keys; std::vector test_data_list; std::unordered_map> 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(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 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(end - + start) + .count() + << "us"; + + start = std::chrono::high_resolution_clock::now(); std::unordered_map> target_batched_slices; for (int i = 0; i < batch_sz; i++) { std::vector 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(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(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(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()); } } diff --git a/mooncake-store/tests/transfer_task_test.cpp b/mooncake-store/tests/transfer_task_test.cpp new file mode 100644 index 00000000..caa7c7ad --- /dev/null +++ b/mooncake-store/tests/transfer_task_test.cpp @@ -0,0 +1,168 @@ +// transfer_task_test.cpp +#include "transfer_task.h" + +#include +#include + +#include +#include +#include +#include +#include + +#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 src_data(data_size, 'A'); + std::vector 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(); + + // 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 src_data(data_size, 'X'); + std::vector dest_data(data_size, 'Y'); + + auto state = std::make_shared(); + + // Create memcpy operations + std::vector 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> src_buffers(num_ops); + std::vector> 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(); + + // Create multiple memcpy operations + std::vector 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(TransferStrategy::LOCAL_MEMCPY), 0); + EXPECT_EQ(static_cast(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(); +} diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h index 695629bb..c466e9ba 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h @@ -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 @@ -124,4 +124,4 @@ using BatchID = Transport::BatchID; } // namespace mooncake -#endif // TRANSFER_ENGINE \ No newline at end of file +#endif // RDMA_TRANSPORT_H_ \ No newline at end of file