diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h index 5cc3e94b..042d8661 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h @@ -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_; @@ -132,9 +161,12 @@ class RdmaEndPoint { std::vector qp_list_; std::string peer_nic_path_; + std::vector 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_; diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp index ba7a2fa5..fcc715c5 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp @@ -18,7 +18,10 @@ #include #include +#include +#include +#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; } diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp index 1a0aa636..80d7995d 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp @@ -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); diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 2dc89be5..a8dab571 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -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 ) diff --git a/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp b/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp new file mode 100644 index 00000000..3498b9bf --- /dev/null +++ b/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp @@ -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 +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#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 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 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(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(nic_priority_matrix.c_str()), + nullptr}; + + Transport *xport = engine_->installTransport("rdma", args); + LOG_ASSERT(xport); + + local_addr_ = static_cast(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(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(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"; + } +}