Skip to content

[Bugfix][TE] Roll back failed native notification setup #4330

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
waizuichougou wants to merge 1 commit into
base: main
Choose a base branch
from
+186 −7
Open

[Bugfix][TE] Roll back failed native notification setup #4330

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {

friend class RdmaEndPointTestPeer;
friend class RdmaNotificationTestPeer;
friend class RdmaEndpointStoreNotificationTestPeer;
friend class RdmaContext;
friend class RdmaTransport;

Expand Down Expand Up @@ -153,7 +154,13 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
char *slot, const TransferMetadata::NotifyDesc &notify);
static bool decodeNotification(const char *slot, size_t bytes,
TransferMetadata::NotifyDesc &notify);
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;
Expand Down
82 changes: 75 additions & 7 deletions mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -435,22 +435,78 @@ 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<std::mutex> guard(s.mutex);
s.enabled = true;
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;
init.cap.max_send_wr = init.cap.max_recv_wr = kNotifySlots;
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;
Expand All @@ -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<char[]>(kNotifySlots * kNotifySlotBytes);
s.recv_buffer = std::make_unique<char[]>(kNotifySlots * kNotifySlotBytes);
s.send_mr =
Expand All @@ -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;
}
Expand Down
104 changes: 104 additions & 0 deletions mooncake-transfer-engine/tests/endpoint_store_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::mutex> 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<std::mutex> guard(endpoint.notify_.mutex);
endpoint.notify_.qp = reinterpret_cast<ibv_qp*>(0x1);
endpoint.notify_.send_mr = reinterpret_cast<ibv_mr*>(0x2);
endpoint.notify_.recv_mr = reinterpret_cast<ibv_mr*>(0x3);
endpoint.notify_.send_buffer = std::make_unique<char[]>(1);
endpoint.notify_.recv_buffer = std::make_unique<char[]>(1);
endpoint.notify_.enabled = true;
}

static bool enabled(const RdmaEndPoint& endpoint) {
std::lock_guard<std::mutex> guard(endpoint.notify_.mutex);
return endpoint.notify_.enabled;
}

static void clearResources(RdmaEndPoint& endpoint) {
std::lock_guard<std::mutex> 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()
Expand Down Expand Up @@ -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
Loading