[TE] Fix simultaneous open handshake in RdmaEndpoint (#1733)
* [TE] Fix simultaneous open handshake in RdmaEndpoint * Keep the same logic for ERDMA. * apply gemini-code-assist's suggestion. * Fix endpoint reinitialization. * Add disconnect and waiting with back-off. * Add eRDMA Endpoint Re-establishment Test * Include the test in cmake * Address reviewer comments * Better log message.
This commit is contained in:
Zhanhao Cao
GitHub
parent
41d40dabd7
commit
4b3d44f39f
|
|
@ -41,6 +41,7 @@ class RdmaEndPoint {
|
|||
enum Status {
|
||||
INITIALIZING,
|
||||
UNCONNECTED,
|
||||
CONNECTING,
|
||||
CONNECTED,
|
||||
};
|
||||
|
||||
|
|
@ -53,6 +54,7 @@ class RdmaEndPoint {
|
|||
size_t max_wr = 256, size_t max_inline = 64);
|
||||
|
||||
private:
|
||||
int reconstruct();
|
||||
int deconstruct();
|
||||
|
||||
public:
|
||||
|
|
@ -98,7 +100,28 @@ class RdmaEndPoint {
|
|||
int destroyQP();
|
||||
|
||||
private:
|
||||
void disconnectUnlocked();
|
||||
int disconnectUnlocked();
|
||||
|
||||
// Resets the connection.
|
||||
//
|
||||
// The main difference between this function and `disconnectUnlocked`
|
||||
// is that it will reconstruct QPs when `CONFIG_ERDMA` is defined.
|
||||
// Without `CONFIG_ERDMA`, it is essentially the same as
|
||||
// `disconnectUnlocked` but with additional logging.
|
||||
//
|
||||
// This serves as a workaround for Aliyun eRDMA devices (i.e., once a QP is
|
||||
// transitioned to the RTS state, it cannot be reset to RTS again directly).
|
||||
// For more details:
|
||||
// https://github.com/kvcache-ai/Mooncake/pull/1733#discussion_r2992088663
|
||||
//
|
||||
// In practice:
|
||||
// - Call `resetConnection` if the QPs' state may have transitioned to RTS.
|
||||
// - Call `disconnectUnlocked` otherwise.
|
||||
//
|
||||
// This is mainly used in `setupConnectionsByActive` or
|
||||
// `setupConnectionsByPassive`. It is NOT invoked in the normal execution
|
||||
// flow, so a `reason` argument is passed for internal logging purposes.
|
||||
int resetConnection(const std::string &reason);
|
||||
|
||||
public:
|
||||
const std::string toString() const;
|
||||
|
|
@ -125,6 +148,12 @@ class RdmaEndPoint {
|
|||
std::string *reply_msg = nullptr);
|
||||
|
||||
private:
|
||||
static constexpr uint64_t kWaitExistingHandshakeTimeoutNano =
|
||||
10 * 1000000000ull; // 10 seconds
|
||||
static constexpr uint32_t kWaitExistingHandshakeSpinCount = 500;
|
||||
static constexpr uint32_t kWaitExistingHandshakeInitialSleepUs = 50;
|
||||
static constexpr uint32_t kWaitExistingHandshakeMaxSleepUs = 2000;
|
||||
|
||||
RdmaContext &context_;
|
||||
std::atomic<Status> status_;
|
||||
|
||||
|
|
@ -132,9 +161,12 @@ class RdmaEndPoint {
|
|||
std::vector<ibv_qp *> qp_list_;
|
||||
|
||||
std::string peer_nic_path_;
|
||||
std::vector<uint32_t> peer_qp_num_list_;
|
||||
|
||||
volatile int *wr_depth_list_;
|
||||
int max_wr_depth_;
|
||||
size_t max_sge_per_wr_;
|
||||
size_t max_inline_bytes_;
|
||||
|
||||
volatile bool active_;
|
||||
volatile int *cq_outstanding_;
|
||||
|
|
|
|||
|
|
@ -18,7 +18,10 @@
|
|||
|
||||
#include <cassert>
|
||||
#include <cstddef>
|
||||
#include <chrono>
|
||||
#include <thread>
|
||||
|
||||
#include "common.h"
|
||||
#include "config.h"
|
||||
|
||||
namespace mooncake {
|
||||
|
|
@ -48,6 +51,9 @@ int RdmaEndPoint::construct(ibv_cq *cq, size_t num_qp_list,
|
|||
cq_outstanding_ = (volatile int *)cq->cq_context;
|
||||
|
||||
max_wr_depth_ = (int)max_wr_depth;
|
||||
max_sge_per_wr_ = max_sge_per_wr;
|
||||
max_inline_bytes_ = max_inline_bytes;
|
||||
|
||||
wr_depth_list_ = new volatile int[num_qp_list];
|
||||
if (!wr_depth_list_) {
|
||||
LOG(ERROR) << "Failed to allocate memory for work request depth list";
|
||||
|
|
@ -76,6 +82,35 @@ int RdmaEndPoint::construct(ibv_cq *cq, size_t num_qp_list,
|
|||
return 0;
|
||||
}
|
||||
|
||||
int RdmaEndPoint::reconstruct() {
|
||||
// Save original construction parameters
|
||||
size_t num_qp = qp_list_.size();
|
||||
auto max_wr_depth = max_wr_depth_;
|
||||
auto max_sge_per_wr = max_sge_per_wr_;
|
||||
auto max_inline_bytes = max_inline_bytes_;
|
||||
|
||||
// Deconstruct and reconstruct to get fresh QPs (same as delete+create)
|
||||
int ret = deconstruct();
|
||||
if (ret) {
|
||||
LOG(ERROR) << "Failed to deconstruct endpoint: " << ret;
|
||||
return ret;
|
||||
}
|
||||
|
||||
// Get CQ from context for reconstruction
|
||||
ibv_cq *cq = context_.cq();
|
||||
if (!cq) {
|
||||
LOG(ERROR) << "No CQ available for endpoint reconstruction";
|
||||
return ERR_ENDPOINT;
|
||||
}
|
||||
|
||||
// Reconstruct with same parameters as original construction
|
||||
status_.store(INITIALIZING, std::memory_order_relaxed);
|
||||
active_ = true;
|
||||
|
||||
return construct(cq, num_qp, max_sge_per_wr, max_wr_depth,
|
||||
max_inline_bytes);
|
||||
}
|
||||
|
||||
int RdmaEndPoint::deconstruct() {
|
||||
for (size_t i = 0; i < qp_list_.size(); ++i) {
|
||||
if (ibv_destroy_qp(qp_list_[i])) {
|
||||
|
|
@ -95,6 +130,7 @@ int RdmaEndPoint::deconstruct() {
|
|||
}
|
||||
}
|
||||
qp_list_.clear();
|
||||
peer_qp_num_list_.clear();
|
||||
delete[] wr_depth_list_;
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -111,44 +147,134 @@ void RdmaEndPoint::setPeerNicPath(const std::string &peer_nic_path) {
|
|||
}
|
||||
|
||||
int RdmaEndPoint::setupConnectionsByActive() {
|
||||
RWSpinlock::WriteGuard guard(lock_);
|
||||
if (connected()) {
|
||||
LOG(INFO) << "Connection has been established";
|
||||
return 0;
|
||||
}
|
||||
|
||||
// loopback mode
|
||||
if (context_.nicPath() == peer_nic_path_) {
|
||||
auto segment_desc =
|
||||
context_.engine().meta()->getSegmentDescByID(LOCAL_SEGMENT_ID);
|
||||
if (segment_desc) {
|
||||
for (auto &nic : segment_desc->devices)
|
||||
if (nic.name == context_.deviceName())
|
||||
return doSetupConnection(nic.gid, nic.lid, qpNum());
|
||||
}
|
||||
LOG(ERROR) << "Peer NIC " << context_.deviceName()
|
||||
<< " not found in localhost";
|
||||
return ERR_DEVICE_NOT_FOUND;
|
||||
}
|
||||
|
||||
HandShakeDesc local_desc, peer_desc;
|
||||
local_desc.local_nic_path = context_.nicPath();
|
||||
local_desc.peer_nic_path = peer_nic_path_;
|
||||
local_desc.qp_num = qpNum();
|
||||
std::string peer_server_name, peer_nic_name;
|
||||
bool do_rpc = false;
|
||||
|
||||
auto peer_server_name = getServerNameFromNicPath(peer_nic_path_);
|
||||
auto peer_nic_name = getNicNameFromNicPath(peer_nic_path_);
|
||||
if (peer_server_name.empty() || peer_nic_name.empty()) {
|
||||
LOG(ERROR) << "Parse peer nic path failed: " << peer_nic_path_;
|
||||
return ERR_INVALID_ARGUMENT;
|
||||
{
|
||||
RWSpinlock::WriteGuard guard(lock_);
|
||||
if (connected()) {
|
||||
LOG(INFO) << "Connection has been established";
|
||||
return 0;
|
||||
}
|
||||
|
||||
// loopback mode
|
||||
if (context_.nicPath() == peer_nic_path_) {
|
||||
auto segment_desc =
|
||||
context_.engine().meta()->getSegmentDescByID(LOCAL_SEGMENT_ID);
|
||||
if (segment_desc) {
|
||||
for (auto &nic : segment_desc->devices)
|
||||
if (nic.name == context_.deviceName())
|
||||
return doSetupConnection(nic.gid, nic.lid, qpNum());
|
||||
}
|
||||
LOG(ERROR) << "Peer NIC " << context_.deviceName()
|
||||
<< " not found in localhost";
|
||||
return ERR_DEVICE_NOT_FOUND;
|
||||
}
|
||||
|
||||
// Only proceed with RPC if we are the first to transition from
|
||||
// UNCONNECTED. This prevents duplicate concurrent handshake attempts
|
||||
// from the same endpoint.
|
||||
auto current_status = status_.load(std::memory_order_relaxed);
|
||||
if (current_status == UNCONNECTED) {
|
||||
status_.store(CONNECTING, std::memory_order_relaxed);
|
||||
do_rpc = true;
|
||||
|
||||
peer_server_name = getServerNameFromNicPath(peer_nic_path_);
|
||||
peer_nic_name = getNicNameFromNicPath(peer_nic_path_);
|
||||
if (peer_server_name.empty() || peer_nic_name.empty()) {
|
||||
LOG(ERROR) << "Parse peer nic path failed: " << peer_nic_path_;
|
||||
disconnectUnlocked();
|
||||
return ERR_INVALID_ARGUMENT;
|
||||
}
|
||||
|
||||
local_desc.local_nic_path = context_.nicPath();
|
||||
local_desc.peer_nic_path = peer_nic_path_;
|
||||
local_desc.qp_num = qpNum();
|
||||
}
|
||||
}
|
||||
|
||||
if (!do_rpc) {
|
||||
LOG(INFO) << "Another thread is already performing the endpoint "
|
||||
"handshake, waiting for it to complete";
|
||||
uint64_t start_time = getCurrentTimeInNano();
|
||||
uint32_t spin_count = 0;
|
||||
uint32_t sleep_us = kWaitExistingHandshakeInitialSleepUs;
|
||||
while (status_.load(std::memory_order_acquire) == CONNECTING) {
|
||||
if (spin_count < kWaitExistingHandshakeSpinCount) {
|
||||
PAUSE();
|
||||
} else {
|
||||
std::this_thread::sleep_for(
|
||||
std::chrono::microseconds(sleep_us));
|
||||
uint32_t next = sleep_us * 2;
|
||||
sleep_us = next > kWaitExistingHandshakeMaxSleepUs
|
||||
? kWaitExistingHandshakeMaxSleepUs
|
||||
: next;
|
||||
}
|
||||
++spin_count;
|
||||
// Prevent infinite wait with a timeout
|
||||
if (getCurrentTimeInNano() - start_time >
|
||||
kWaitExistingHandshakeTimeoutNano) {
|
||||
// Timeout while waiting for another thread's handshake.
|
||||
// The QP state on this endpoint may have changed; therefore,
|
||||
// reset the connection so that subsequent callers can retry.
|
||||
RWSpinlock::WriteGuard write_guard(lock_);
|
||||
resetConnection("wait existing handshake timeout");
|
||||
return ERR_ENDPOINT;
|
||||
}
|
||||
}
|
||||
RWSpinlock::ReadGuard guard(lock_);
|
||||
return connected() ? 0 : ERR_ENDPOINT;
|
||||
}
|
||||
|
||||
// Perform the RPC without holding the lock to avoid deadlock and allow
|
||||
// "simultaneous open" handshake handling.
|
||||
int rc = context_.engine().sendHandshake(peer_server_name, local_desc,
|
||||
peer_desc);
|
||||
if (rc) return rc;
|
||||
|
||||
// We should check the RPC return code before comparing `peer_qp_num_list_`
|
||||
// with `peer_desc.qp_num`, since a failed RPC may result in an
|
||||
// invalid `peer_desc.qp_num`.
|
||||
//
|
||||
// If the RPC is failed, even if the state is CONNECTED, (which means
|
||||
// it is handled by setupConnectionsByPassive in another thread during the
|
||||
// RPC, or "simultaneous open"), we should resetConnection to be safe.
|
||||
// Because we're not sure whether the peer needs a connection
|
||||
// re-establishment. (We don't know `peer_desc.qp_num`)
|
||||
if (rc) {
|
||||
RWSpinlock::WriteGuard write_guard(lock_);
|
||||
resetConnection("handshake RPC failure");
|
||||
return rc;
|
||||
}
|
||||
|
||||
// Re-acquire lock after RPC to finalize state transition
|
||||
RWSpinlock::WriteGuard guard(lock_);
|
||||
|
||||
// Handle simultaneous open: if the peer initiates a connection during our
|
||||
// RPC and it is passively established in setupConnectionsByPassive, simply
|
||||
// reuse the existing endpoint.
|
||||
if (connected()) {
|
||||
if (peer_qp_num_list_ == peer_desc.qp_num) {
|
||||
LOG(INFO) << "Received same peer QP numbers, reusing connection.";
|
||||
return 0;
|
||||
}
|
||||
|
||||
// This mismatch scenario should be rare. It may occur when a peer
|
||||
// first sends us an Active RPC and establishes a connection,
|
||||
// then restarts, and eventually accepts and responds to our
|
||||
// Active RPC.
|
||||
LOG(WARNING) << "Peer QP list mismatch on connected endpoint, "
|
||||
"re-establishing connection: "
|
||||
<< toString();
|
||||
|
||||
int ret = resetConnection("re-establishing connection (active)");
|
||||
if (ret) return ret;
|
||||
}
|
||||
|
||||
if (!peer_desc.reply_msg.empty()) {
|
||||
LOG(ERROR) << "Reject the handshake request by peer "
|
||||
LOG(ERROR) << "Rejected handshake request by peer "
|
||||
<< local_desc.peer_nic_path;
|
||||
disconnectUnlocked();
|
||||
return ERR_REJECT_HANDSHAKE;
|
||||
}
|
||||
|
||||
|
|
@ -160,18 +286,26 @@ int RdmaEndPoint::setupConnectionsByActive() {
|
|||
<< ", local.peer_nic_path: " << local_desc.peer_nic_path
|
||||
<< ", peer.local_nic_path: " << peer_desc.local_nic_path
|
||||
<< ", peer.peer_nic_path: " << peer_desc.peer_nic_path;
|
||||
disconnectUnlocked();
|
||||
return ERR_REJECT_HANDSHAKE;
|
||||
}
|
||||
|
||||
auto segment_desc =
|
||||
context_.engine().meta()->getSegmentDescByName(peer_server_name);
|
||||
if (segment_desc) {
|
||||
for (auto &nic : segment_desc->devices)
|
||||
if (nic.name == peer_nic_name)
|
||||
return doSetupConnection(nic.gid, nic.lid, peer_desc.qp_num);
|
||||
for (auto &nic : segment_desc->devices) {
|
||||
if (nic.name == peer_nic_name) {
|
||||
int ret = doSetupConnection(nic.gid, nic.lid, peer_desc.qp_num);
|
||||
if (ret != 0) {
|
||||
resetConnection("failed connection setup (active)");
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
}
|
||||
}
|
||||
LOG(ERROR) << "Peer NIC " << peer_nic_name << " not found in "
|
||||
<< peer_server_name;
|
||||
disconnectUnlocked();
|
||||
return ERR_DEVICE_NOT_FOUND;
|
||||
}
|
||||
|
||||
|
|
@ -179,10 +313,30 @@ int RdmaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc,
|
|||
HandShakeDesc &local_desc) {
|
||||
RWSpinlock::WriteGuard guard(lock_);
|
||||
if (connected()) {
|
||||
// If already connected with the same peer QP info, return success
|
||||
if (peer_qp_num_list_ == peer_desc.qp_num) {
|
||||
local_desc.local_nic_path = context_.nicPath();
|
||||
local_desc.peer_nic_path = peer_nic_path_;
|
||||
local_desc.qp_num = qpNum();
|
||||
LOG(INFO) << "Received same peer QP numbers, reusing connection.";
|
||||
return 0;
|
||||
}
|
||||
// Different peer (e.g., peer restarted)
|
||||
LOG(WARNING) << "Re-establish connection: " << toString();
|
||||
disconnectUnlocked();
|
||||
|
||||
int ret = resetConnection("re-establishing connection (passive)");
|
||||
if (ret) return ret;
|
||||
}
|
||||
|
||||
// At this point, the state can only be UNCONNECTED or CONNECTING.
|
||||
// Even if the state is CONNECTING, we can still safely proceed to
|
||||
// establish the connection on this same endpoint. Because we're holding
|
||||
// the lock, even if there are already Active RPCs sent to the same
|
||||
// peer nic path by setupConnectionsByActive on another thread, it will
|
||||
// be blocked after the RPC return. Once the lock is released,
|
||||
// they will simply observe the CONNECTED state and safely reuse the QP.
|
||||
// This inherently handles simultaneous open.
|
||||
|
||||
if (peer_desc.peer_nic_path != context_.nicPath() ||
|
||||
peer_desc.local_nic_path != peer_nic_path_) {
|
||||
local_desc.reply_msg =
|
||||
|
|
@ -209,10 +363,16 @@ int RdmaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc,
|
|||
auto segment_desc =
|
||||
context_.engine().meta()->getSegmentDescByName(peer_server_name);
|
||||
if (segment_desc) {
|
||||
for (auto &nic : segment_desc->devices)
|
||||
if (nic.name == peer_nic_name)
|
||||
return doSetupConnection(nic.gid, nic.lid, peer_desc.qp_num,
|
||||
&local_desc.reply_msg);
|
||||
for (auto &nic : segment_desc->devices) {
|
||||
if (nic.name == peer_nic_name) {
|
||||
int ret = doSetupConnection(nic.gid, nic.lid, peer_desc.qp_num,
|
||||
&local_desc.reply_msg);
|
||||
if (ret != 0) {
|
||||
resetConnection("failed connection setup (passive)");
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
}
|
||||
}
|
||||
local_desc.reply_msg =
|
||||
"Peer nic not found in that server: " + peer_nic_path_;
|
||||
|
|
@ -225,13 +385,20 @@ void RdmaEndPoint::disconnect() {
|
|||
disconnectUnlocked();
|
||||
}
|
||||
|
||||
void RdmaEndPoint::disconnectUnlocked() {
|
||||
int RdmaEndPoint::disconnectUnlocked() {
|
||||
auto curr_status = status_.load(std::memory_order_acquire);
|
||||
if (curr_status != CONNECTED && curr_status != CONNECTING) return 0;
|
||||
|
||||
ibv_qp_attr attr;
|
||||
memset(&attr, 0, sizeof(attr));
|
||||
attr.qp_state = IBV_QPS_RESET;
|
||||
int ret = 0;
|
||||
for (size_t i = 0; i < qp_list_.size(); ++i) {
|
||||
int ret = ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE);
|
||||
if (ret) PLOG(ERROR) << "Failed to modify QP to RESET";
|
||||
int curr_ret = ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE);
|
||||
if (curr_ret) {
|
||||
PLOG(ERROR) << "Failed to modify QP to RESET";
|
||||
ret = ERR_ENDPOINT;
|
||||
}
|
||||
// After resetting QP, the wr_depth_list_ won't change
|
||||
bool displayed = false;
|
||||
if (wr_depth_list_[i] != 0) {
|
||||
|
|
@ -244,7 +411,29 @@ void RdmaEndPoint::disconnectUnlocked() {
|
|||
wr_depth_list_[i] = 0;
|
||||
}
|
||||
}
|
||||
peer_qp_num_list_.clear();
|
||||
status_.store(UNCONNECTED, std::memory_order_release);
|
||||
return ret;
|
||||
}
|
||||
|
||||
int RdmaEndPoint::resetConnection(const std::string &reason) {
|
||||
auto curr_status = status_.load(std::memory_order_acquire);
|
||||
if (curr_status != CONNECTING && curr_status != CONNECTED) return 0;
|
||||
|
||||
#ifdef CONFIG_ERDMA
|
||||
int ret = reconstruct();
|
||||
#else
|
||||
int ret = disconnectUnlocked();
|
||||
#endif
|
||||
|
||||
if (ret) {
|
||||
LOG(ERROR) << "Failed to reset the endpoint (triggered by: " << reason
|
||||
<< "): error=" << ret;
|
||||
} else {
|
||||
LOG(INFO) << "Successfully reset the endpoint (triggered by: " << reason
|
||||
<< ").";
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
const std::string RdmaEndPoint::toString() const {
|
||||
|
|
@ -346,6 +535,7 @@ int RdmaEndPoint::doSetupConnection(const std::string &peer_gid,
|
|||
if (ret) return ret;
|
||||
}
|
||||
|
||||
peer_qp_num_list_ = std::move(peer_qp_num_list);
|
||||
status_.store(CONNECTED, std::memory_order_relaxed);
|
||||
return 0;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -634,9 +634,7 @@ int RdmaTransport::onSetupRdmaConnections(const HandShakeDesc &peer_desc,
|
|||
}
|
||||
if (!context) return ERR_INVALID_ARGUMENT;
|
||||
|
||||
#ifdef CONFIG_ERDMA
|
||||
if (context->deleteEndpoint(peer_desc.local_nic_path)) return ERR_ENDPOINT;
|
||||
#endif
|
||||
// Use existing endpoint or create new one.
|
||||
auto endpoint = context->endpoint(peer_desc.local_nic_path);
|
||||
if (!endpoint) return ERR_ENDPOINT;
|
||||
return endpoint->setupConnectionsByPassive(peer_desc, local_desc);
|
||||
|
|
|
|||
|
|
@ -24,6 +24,11 @@ add_executable(rdma_loopback_test ${WORKSPACE}/rdma_loopback_test.cpp)
|
|||
target_link_libraries(rdma_loopback_test PUBLIC transfer_engine gtest gtest_main )
|
||||
# add_test(NAME rdma_loopback_test COMMAND rdma_loopback_test)
|
||||
|
||||
# This test verifies endpoint re-establishment in RDMATransport.
|
||||
# Intended for manual testing only.
|
||||
add_executable(rdma_endpoint_reestablish_test ${WORKSPACE}/rdma_endpoint_reestablish_test.cpp)
|
||||
target_link_libraries(rdma_endpoint_reestablish_test PUBLIC transfer_engine gtest gtest_main )
|
||||
|
||||
if (USE_CXL)
|
||||
add_executable(cxl_transport_test ${WORKSPACE}/cxl_transport_test.cpp)
|
||||
target_link_libraries(cxl_transport_test PUBLIC transfer_engine gtest gtest_main )
|
||||
|
|
|
|||
|
|
@ -0,0 +1,260 @@
|
|||
// Copyright 2026 KVCache.AI
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
/*
|
||||
* RDMA Endpoint Re-establishment Test
|
||||
*
|
||||
* Purpose:
|
||||
* This test verifies that TE correctly handles endpoint re-establish
|
||||
* during simulated Initiator restarts.
|
||||
*
|
||||
* How to run:
|
||||
* 1. Start etcd:
|
||||
* etcd --listen-client-urls http://127.0.0.1:18222 \
|
||||
* --advertise-client-urls http://127.0.0.1:18222
|
||||
*
|
||||
* 2. Run test
|
||||
* sudo env MC_METADATA_SERVER=127.0.0.1:18222 \
|
||||
* MC_TARGET_SERVER_NAME=127.0.0.1:12345 \
|
||||
* MC_INITIATOR_SERVER_NAME=127.0.0.1:12346 \
|
||||
* MC_TARGET_DEVICE_NAME=erdma_0 MC_INITIATOR_DEVICE_NAME=erdma_1 \
|
||||
* ./build/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test
|
||||
*/
|
||||
|
||||
#include <gflags/gflags.h>
|
||||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <numa.h>
|
||||
#include <sys/time.h>
|
||||
#include <unistd.h>
|
||||
|
||||
#include <cstdlib>
|
||||
#include <fstream>
|
||||
#include <iomanip>
|
||||
#include <memory>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
|
||||
#include "transfer_engine.h"
|
||||
#include "transport/transport.h"
|
||||
#include "common.h"
|
||||
|
||||
using namespace mooncake;
|
||||
|
||||
// Size of the pre-registered memory.
|
||||
constexpr size_t kRAMBufSize = 256ull << 24; // 256 MB
|
||||
|
||||
// Actual data payload size for RDMA Read/Write.
|
||||
constexpr size_t kDataLength = 16ull << 24; // 16MB
|
||||
|
||||
std::string formatDeviceNames(const std::string &device_names) {
|
||||
std::stringstream ss(device_names);
|
||||
std::string item;
|
||||
std::vector<std::string> tokens;
|
||||
while (getline(ss, item, ',')) {
|
||||
tokens.push_back(item);
|
||||
}
|
||||
|
||||
std::string formatted;
|
||||
for (size_t i = 0; i < tokens.size(); ++i) {
|
||||
formatted += "\"" + tokens[i] + "\"";
|
||||
if (i < tokens.size() - 1) {
|
||||
formatted += ",";
|
||||
}
|
||||
}
|
||||
return formatted;
|
||||
}
|
||||
|
||||
std::string makeNicPriorityMatrix(const std::string &device_name) {
|
||||
auto formatted_devices = formatDeviceNames(device_name);
|
||||
return "{\"cpu:0\": [[" + formatted_devices +
|
||||
"],[]], "
|
||||
" \"cpu:1\": [[" +
|
||||
formatted_devices + "],[]]}";
|
||||
}
|
||||
|
||||
void wait_for_transfer(TransferEngine *engine, BatchID batch_id,
|
||||
const std::string &op_name) {
|
||||
bool completed = false;
|
||||
TransferStatus status;
|
||||
while (!completed) {
|
||||
Status s = engine->getTransferStatus(batch_id, 0, status);
|
||||
EXPECT_EQ(s, Status::OK());
|
||||
if (status.s == TransferStatusEnum::COMPLETED) {
|
||||
completed = true;
|
||||
} else if (status.s == TransferStatusEnum::FAILED) {
|
||||
FAIL() << op_name << " FAILED";
|
||||
}
|
||||
}
|
||||
Status s = engine->freeBatchID(batch_id);
|
||||
EXPECT_EQ(s, Status::OK());
|
||||
}
|
||||
|
||||
struct TEContext {
|
||||
std::unique_ptr<TransferEngine> engine_{};
|
||||
uint8_t *local_addr_{};
|
||||
bool segment_opened_{false};
|
||||
SegmentHandle segment_handle_{};
|
||||
uint64_t remote_base_{};
|
||||
|
||||
TEContext(const std::string &local_server_name,
|
||||
const std::string &metadata_server, const std::string &segment_id,
|
||||
const std::string &device_name) {
|
||||
engine_ = std::make_unique<TransferEngine>(false);
|
||||
auto hostname_port = parseHostNameWithPort(local_server_name);
|
||||
engine_->init(metadata_server, local_server_name, hostname_port.first,
|
||||
hostname_port.second);
|
||||
|
||||
auto nic_priority_matrix = makeNicPriorityMatrix(device_name);
|
||||
void *args[2] = {const_cast<char *>(nic_priority_matrix.c_str()),
|
||||
nullptr};
|
||||
|
||||
Transport *xport = engine_->installTransport("rdma", args);
|
||||
LOG_ASSERT(xport);
|
||||
|
||||
local_addr_ = static_cast<uint8_t *>(numa_alloc_onnode(kRAMBufSize, 0));
|
||||
memset(local_addr_, 0, kDataLength);
|
||||
|
||||
int rc =
|
||||
engine_->registerLocalMemory(local_addr_, kRAMBufSize, "cpu:0");
|
||||
LOG_ASSERT(!rc);
|
||||
|
||||
if (!segment_id.empty()) {
|
||||
segment_opened_ = true;
|
||||
LOG(INFO) << "Opening segment " << segment_id << "...";
|
||||
segment_handle_ = engine_->openSegment(segment_id);
|
||||
auto segment_desc =
|
||||
engine_->getMetadata()->getSegmentDescByID(segment_handle_);
|
||||
remote_base_ = (uint64_t)segment_desc->buffers[0].addr;
|
||||
}
|
||||
}
|
||||
|
||||
~TEContext() {
|
||||
engine_->unregisterLocalMemory(local_addr_);
|
||||
numa_free(local_addr_, kRAMBufSize);
|
||||
if (segment_opened_) engine_->closeSegment(segment_handle_);
|
||||
}
|
||||
};
|
||||
|
||||
class RDMAEndpointReestablishTest : public ::testing::Test {
|
||||
protected:
|
||||
void SetUp() override {
|
||||
google::InitGoogleLogging("RDMAEndpointReestablishTest");
|
||||
FLAGS_logtostderr = true;
|
||||
|
||||
const char *env = std::getenv("MC_METADATA_SERVER");
|
||||
metadata_server = env ? env : "127.0.0.1:18222";
|
||||
LOG(INFO) << "metadata_server: " << metadata_server;
|
||||
|
||||
env = std::getenv("MC_TARGET_SERVER_NAME");
|
||||
target_server_name = env ? env : "127.0.0.1:12345";
|
||||
LOG(INFO) << "target_server_name: " << target_server_name;
|
||||
|
||||
env = std::getenv("MC_INITIATOR_SERVER_NAME");
|
||||
initiator_server_name = env ? env : "127.0.0.1:12346";
|
||||
LOG(INFO) << "initiator_server_name: " << initiator_server_name;
|
||||
|
||||
env = std::getenv("MC_TARGET_DEVICE_NAME");
|
||||
target_device_name = env ? env : "erdma_0";
|
||||
LOG(INFO) << "target_device_name: " << target_device_name;
|
||||
|
||||
env = std::getenv("MC_INITIATOR_DEVICE_NAME");
|
||||
initiator_device_name = env ? env : "erdma_1";
|
||||
LOG(INFO) << "initiator_device_name: " << initiator_device_name;
|
||||
}
|
||||
|
||||
void TearDown() override { google::ShutdownGoogleLogging(); }
|
||||
|
||||
std::string metadata_server;
|
||||
std::string target_server_name;
|
||||
std::string initiator_server_name;
|
||||
std::string target_device_name;
|
||||
std::string initiator_device_name;
|
||||
};
|
||||
|
||||
TEST_F(RDMAEndpointReestablishTest, EndpointReestablish) {
|
||||
// 1. Setup Target (will stay alive until test function exits)
|
||||
LOG(INFO) << "========== Setting up Target ==========";
|
||||
TEContext target_ctx(target_server_name, metadata_server, "",
|
||||
target_device_name);
|
||||
LOG(INFO) << "Target is up. Waiting for RDMA connections and operations...";
|
||||
|
||||
// 2. Phase 1: Initiator Start, Connect & Write
|
||||
LOG(INFO) << "========== Phase 1: Start, Connect & Write ==========";
|
||||
{
|
||||
TEContext init_ctx(initiator_server_name, metadata_server,
|
||||
target_server_name, initiator_device_name);
|
||||
|
||||
// Fill buffer with test pattern
|
||||
for (size_t i = 0; i < kDataLength; ++i) {
|
||||
init_ctx.local_addr_[i] = static_cast<uint8_t>(i % 256);
|
||||
}
|
||||
|
||||
LOG(INFO) << "Writing " << kDataLength << " bytes to Target...";
|
||||
auto batch_id = init_ctx.engine_->allocateBatchID(1);
|
||||
TransferRequest entry;
|
||||
entry.opcode = TransferRequest::WRITE;
|
||||
entry.length = kDataLength;
|
||||
entry.source = init_ctx.local_addr_;
|
||||
entry.target_id = init_ctx.segment_handle_;
|
||||
entry.target_offset = init_ctx.remote_base_;
|
||||
|
||||
Status s = init_ctx.engine_->submitTransfer(batch_id, {entry});
|
||||
ASSERT_EQ(s, Status::OK());
|
||||
|
||||
wait_for_transfer(init_ctx.engine_.get(), batch_id, "WRITE");
|
||||
|
||||
LOG(INFO) << "Phase 1: Write Completed. Tearing down connection...";
|
||||
}
|
||||
|
||||
LOG(INFO) << "Simulating Initiator Crash/Restart... Waiting 2 seconds.";
|
||||
sleep(2);
|
||||
|
||||
// 3. Phase 2: Restart, Re-establish Endpoint & Read
|
||||
LOG(INFO) << "========== Phase 2: Restart, Re-establish Endpoint & Read "
|
||||
"==========";
|
||||
{
|
||||
TEContext init_ctx(initiator_server_name, metadata_server,
|
||||
target_server_name, initiator_device_name);
|
||||
|
||||
LOG(INFO) << "Reading data back over new Endpoint...";
|
||||
auto batch_id = init_ctx.engine_->allocateBatchID(1);
|
||||
TransferRequest entry;
|
||||
entry.opcode = TransferRequest::READ;
|
||||
entry.length = kDataLength;
|
||||
entry.source = init_ctx.local_addr_;
|
||||
entry.target_id = init_ctx.segment_handle_;
|
||||
entry.target_offset = init_ctx.remote_base_;
|
||||
|
||||
Status s = init_ctx.engine_->submitTransfer(batch_id, {entry});
|
||||
ASSERT_EQ(s, Status::OK());
|
||||
|
||||
wait_for_transfer(init_ctx.engine_.get(), batch_id, "READ");
|
||||
|
||||
bool ok = true;
|
||||
for (size_t i = 0; i < kDataLength; ++i) {
|
||||
if (init_ctx.local_addr_[i] != static_cast<uint8_t>(i % 256)) {
|
||||
ok = false;
|
||||
LOG(ERROR) << "Data mismatch at offset " << i << ", expected "
|
||||
<< (i % 256) << ", got "
|
||||
<< (int)init_ctx.local_addr_[i];
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
ASSERT_TRUE(ok) << "Endpoint Reconstruction Verification Failed!";
|
||||
LOG(INFO)
|
||||
<< ">>> ENDPOINT RECONSTRUCTION VERIFICATION: \033[32mSUCCESS\033";
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue