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 1c91aaffe9..9fdafbe3d8 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h @@ -58,6 +58,7 @@ class RdmaEndPoint : public std::enable_shared_from_this { friend class RdmaEndPointTestPeer; friend class RdmaNotificationTestPeer; + friend class RdmaEndpointStoreNotificationTestPeer; friend class RdmaContext; friend class RdmaTransport; @@ -153,7 +154,13 @@ class RdmaEndPoint : public std::enable_shared_from_this { char *slot, const TransferMetadata::NotifyDesc ¬ify); static bool decodeNotification(const char *slot, size_t bytes, TransferMetadata::NotifyDesc ¬ify); + struct NotificationCleanupOps { + int (*destroy_qp)(ibv_qp *) = nullptr; + int (*dereg_mr)(ibv_mr *) = nullptr; + }; int constructNotification(); + void rollbackNotificationConstruction(); + void rollbackNotificationConstruction(const NotificationCleanupOps &ops); int connectNotification(const ibv_gid &gid, uint32_t lid, uint32_t peer_qp, int local_gid_index); uint32_t notificationQpNum() const; 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 38e660937b..6e402555f2 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp @@ -435,6 +435,54 @@ int RdmaEndPoint::postNotificationReceive(size_t slot) { wr.num_sge = 1; return ibv_post_recv(s.qp, &wr, &bad) ? ERR_ENDPOINT : 0; } + +void RdmaEndPoint::rollbackNotificationConstruction() { + rollbackNotificationConstruction( + NotificationCleanupOps{ibv_destroy_qp, ibv_dereg_mr}); +} + +void RdmaEndPoint::rollbackNotificationConstruction( + const NotificationCleanupOps &ops) { + // constructNotification() holds notify_.mutex and has not registered the + // QP with the context yet. Best-effort cleanup keeps a failed optional + // channel from retaining verbs resources until endpoint destruction. + auto &s = notify_; + if (s.qp) { + const int ret = ops.destroy_qp(s.qp); + if (ret == 0) { + s.qp = nullptr; + } else { + LOG(ERROR) << "Failed to destroy notification QP during rollback: " + << strerror(ret); + } + } + if (s.send_mr) { + const int ret = ops.dereg_mr(s.send_mr); + if (ret == 0) { + s.send_mr = nullptr; + } else { + LOG(ERROR) << "Failed to deregister notification send MR during " + "rollback: " + << strerror(ret); + } + } + if (s.recv_mr) { + const int ret = ops.dereg_mr(s.recv_mr); + if (ret == 0) { + s.recv_mr = nullptr; + } else { + LOG(ERROR) << "Failed to deregister notification receive MR during " + "rollback: " + << strerror(ret); + } + } + if (!s.qp && !s.send_mr && !s.recv_mr) { + s.send_buffer.reset(); + s.recv_buffer.reset(); + s.enabled = false; + } +} + int RdmaEndPoint::constructNotification() { auto &s = notify_; std::lock_guard guard(s.mutex); @@ -442,7 +490,11 @@ int RdmaEndPoint::constructNotification() { s.error = 0; s.connected = s.reconnect_needed = false; s.pending = s.next_send = s.peer_qp = 0; - if (!context_.notify_cq_) return s.fail(ERR_CONTEXT); + if (!context_.notify_cq_) { + const int ret = s.fail(ERR_CONTEXT); + rollbackNotificationConstruction(); + return ret; + } ibv_qp_init_attr init{}; init.send_cq = init.recv_cq = context_.notify_cq_; init.qp_type = IBV_QPT_RC; @@ -450,7 +502,11 @@ int RdmaEndPoint::constructNotification() { init.cap.max_send_sge = init.cap.max_recv_sge = 1; init.cap.max_inline_data = globalConfig().max_inline; s.qp = ibv_create_qp(context_.pd(), &init); - if (!s.qp) return s.fail(ERR_ENDPOINT); + if (!s.qp) { + const int ret = s.fail(ERR_ENDPOINT); + rollbackNotificationConstruction(); + return ret; + } s.inline_bytes = init.cap.max_inline_data; ibv_qp_attr attr{}; attr.qp_state = IBV_QPS_INIT; @@ -459,8 +515,11 @@ int RdmaEndPoint::constructNotification() { attr.qp_access_flags = IBV_ACCESS_LOCAL_WRITE; if (ibv_modify_qp(s.qp, &attr, IBV_QP_STATE | IBV_QP_PORT | IBV_QP_PKEY_INDEX | - IBV_QP_ACCESS_FLAGS)) - return s.fail(ERR_ENDPOINT); + IBV_QP_ACCESS_FLAGS)) { + const int ret = s.fail(ERR_ENDPOINT); + rollbackNotificationConstruction(); + return ret; + } s.send_buffer = std::make_unique(kNotifySlots * kNotifySlotBytes); s.recv_buffer = std::make_unique(kNotifySlots * kNotifySlotBytes); s.send_mr = @@ -469,9 +528,18 @@ int RdmaEndPoint::constructNotification() { s.recv_mr = ibv_reg_mr(context_.pd(), s.recv_buffer.get(), kNotifySlots * kNotifySlotBytes, IBV_ACCESS_LOCAL_WRITE); - if (!s.send_mr || !s.recv_mr) return s.fail(ERR_MEMORY); - for (size_t slot = 0; slot < kNotifySlots; ++slot) - if (postNotificationReceive(slot)) return s.fail(ERR_ENDPOINT); + if (!s.send_mr || !s.recv_mr) { + const int ret = s.fail(ERR_MEMORY); + rollbackNotificationConstruction(); + return ret; + } + for (size_t slot = 0; slot < kNotifySlots; ++slot) { + if (postNotificationReceive(slot)) { + const int ret = s.fail(ERR_ENDPOINT); + rollbackNotificationConstruction(); + return ret; + } + } context_.registerNotifyQp(s.qp->qp_num, weak_from_this()); return 0; } diff --git a/mooncake-transfer-engine/tests/endpoint_store_test.cpp b/mooncake-transfer-engine/tests/endpoint_store_test.cpp index 8225a9ac93..595dc84eea 100644 --- a/mooncake-transfer-engine/tests/endpoint_store_test.cpp +++ b/mooncake-transfer-engine/tests/endpoint_store_test.cpp @@ -42,6 +42,50 @@ using namespace mooncake; +namespace mooncake { +class RdmaEndpointStoreNotificationTestPeer { + public: + using CleanupOps = RdmaEndPoint::NotificationCleanupOps; + + static void rollbackNotificationConstruction(RdmaEndPoint& endpoint, + const CleanupOps& ops) { + endpoint.rollbackNotificationConstruction(ops); + } + + static bool hasResources(const RdmaEndPoint& endpoint) { + std::lock_guard guard(endpoint.notify_.mutex); + return endpoint.notify_.qp || endpoint.notify_.send_mr || + endpoint.notify_.recv_mr || endpoint.notify_.send_buffer || + endpoint.notify_.recv_buffer; + } + + static void seedResources(RdmaEndPoint& endpoint) { + std::lock_guard guard(endpoint.notify_.mutex); + endpoint.notify_.qp = reinterpret_cast(0x1); + endpoint.notify_.send_mr = reinterpret_cast(0x2); + endpoint.notify_.recv_mr = reinterpret_cast(0x3); + endpoint.notify_.send_buffer = std::make_unique(1); + endpoint.notify_.recv_buffer = std::make_unique(1); + endpoint.notify_.enabled = true; + } + + static bool enabled(const RdmaEndPoint& endpoint) { + std::lock_guard guard(endpoint.notify_.mutex); + return endpoint.notify_.enabled; + } + + static void clearResources(RdmaEndPoint& endpoint) { + std::lock_guard guard(endpoint.notify_.mutex); + endpoint.notify_.qp = nullptr; + endpoint.notify_.send_mr = nullptr; + endpoint.notify_.recv_mr = nullptr; + endpoint.notify_.send_buffer.reset(); + endpoint.notify_.recv_buffer.reset(); + endpoint.notify_.enabled = false; + } +}; +} // namespace mooncake + namespace { // Build an RdmaEndPoint that owns zero QPs and has active_=false. construct() @@ -211,4 +255,64 @@ TEST_F(EndpointStoreTest, EXPECT_EQ(sentinel, store.getEndpointByPtr(sentinel.get())); } +struct NotificationRollbackProbe { + int destroy_qp_calls = 0; + int dereg_mr_calls = 0; + int destroy_qp_result = 0; + int dereg_mr_result = 0; + + static NotificationRollbackProbe* current; + + static int destroyQp(ibv_qp*) { + ++current->destroy_qp_calls; + return current->destroy_qp_result; + } + + static int deregMr(ibv_mr*) { + ++current->dereg_mr_calls; + return current->dereg_mr_result; + } +}; + +NotificationRollbackProbe* NotificationRollbackProbe::current = nullptr; + +TEST_F(EndpointStoreTest, NotificationConstructionRollsBackResources) { + RdmaEndPoint endpoint(*ctx_); + NotificationRollbackProbe probe; + NotificationRollbackProbe::current = &probe; + + RdmaEndpointStoreNotificationTestPeer::seedResources(endpoint); + + RdmaEndpointStoreNotificationTestPeer::rollbackNotificationConstruction( + endpoint, {NotificationRollbackProbe::destroyQp, + NotificationRollbackProbe::deregMr}); + + EXPECT_EQ(probe.destroy_qp_calls, 1); + EXPECT_EQ(probe.dereg_mr_calls, 2); + EXPECT_FALSE(RdmaEndpointStoreNotificationTestPeer::hasResources(endpoint)); + EXPECT_FALSE(RdmaEndpointStoreNotificationTestPeer::enabled(endpoint)); + + NotificationRollbackProbe::current = nullptr; +} + +TEST_F(EndpointStoreTest, NotificationConstructionRetainsFailedCleanup) { + RdmaEndPoint endpoint(*ctx_); + NotificationRollbackProbe probe; + probe.destroy_qp_result = EBUSY; + probe.dereg_mr_result = EBUSY; + NotificationRollbackProbe::current = &probe; + + RdmaEndpointStoreNotificationTestPeer::seedResources(endpoint); + + RdmaEndpointStoreNotificationTestPeer::rollbackNotificationConstruction( + endpoint, {NotificationRollbackProbe::destroyQp, + NotificationRollbackProbe::deregMr}); + + EXPECT_TRUE(RdmaEndpointStoreNotificationTestPeer::hasResources(endpoint)); + EXPECT_TRUE(RdmaEndpointStoreNotificationTestPeer::enabled(endpoint)); + + RdmaEndpointStoreNotificationTestPeer::clearResources(endpoint); + NotificationRollbackProbe::current = nullptr; +} + } // namespace