Skip to content

[Bugfix][TENT] Register notify recv slots with one MR per endpoint #3820

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

Merged
staryxchen merged 1 commit into from
Sep 2, 2026
Merged

[Bugfix][TENT] Register notify recv slots with one MR per endpoint #3820

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
17 changes: 12 additions & 5 deletions mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,12 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
void postNotifyRecv(size_t idx);
void repostAllNotifyRecvs();

static char* notifySlotPtr(char* base, size_t idx);
static bool encodeNotifyPayload(char* slot, const std::string& name,
const std::string& msg, uint32_t* out_len);
static bool decodeNotifyPayload(const char* data, size_t byte_len,
std::string* name, std::string* msg);

private:
friend class EndpointTestAccess;

Expand Down Expand Up @@ -226,13 +232,14 @@ class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
// Notification QP (one per endpoint for control plane operations)
ibv_qp* notify_qp_ = nullptr;

// Notification buffers
// Notification buffers. Send and recv each use one contiguous host buffer
// split into kNotifyMaxPendingSends slots, registered with a single MR.
static constexpr size_t kNotifyBufferSize = 65536; // 64 KB
static constexpr size_t kNotifyMaxPendingSends = 256;
std::vector<std::vector<char>> notify_recv_buffers_;
std::vector<ibv_mr*> notify_recv_mrs_; // Memory regions for recv buffers
std::vector<char> notify_send_buffer_; // Single contiguous send buffer
ibv_mr* notify_send_mr_ = nullptr; // Single MR for all send slots
std::vector<char> notify_recv_buffer_;
ibv_mr* notify_recv_mr_ = nullptr;
std::vector<char> notify_send_buffer_;
ibv_mr* notify_send_mr_ = nullptr;
// Serializes notification buffer/QP access against deconstruction.
std::mutex notify_resource_mutex_;
std::mutex notify_send_mutex_;
Expand Down
156 changes: 79 additions & 77 deletions mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <iomanip>
#include <mutex>
#include <queue>
Expand Down Expand Up @@ -177,9 +179,18 @@ int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params,
return -1;
}

// Pre-post recv buffers for notification
notify_recv_buffers_.resize(kNotifyMaxPendingSends);
notify_recv_mrs_.resize(kNotifyMaxPendingSends);
// Allocate one contiguous recv buffer, logically split into
// kNotifyMaxPendingSends slots. One MR covers every slot so high
// endpoint counts do not explode ibv_reg_mr.
notify_recv_buffer_.resize(kNotifyBufferSize * kNotifyMaxPendingSends);
notify_recv_mr_ = context_->verbs_.ibv_reg_mr_default(
context_->nativePD(), notify_recv_buffer_.data(),
notify_recv_buffer_.size(), IBV_ACCESS_LOCAL_WRITE);
if (!notify_recv_mr_) {
PLOG(ERROR) << "Failed to register notification recv buffer";
deconstruct();
return -1;
}

// Allocate one contiguous send buffer, logically split into
// kNotifyMaxPendingSends slots to avoid DMA-vs-overwrite races.
Expand All @@ -193,20 +204,6 @@ int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params,
return -1;
}

for (size_t i = 0; i < kNotifyMaxPendingSends; ++i) {
notify_recv_buffers_[i].resize(kNotifyBufferSize);

// Register memory for recv buffer
notify_recv_mrs_[i] = context_->verbs_.ibv_reg_mr_default(
context_->nativePD(), notify_recv_buffers_[i].data(),
kNotifyBufferSize, IBV_ACCESS_LOCAL_WRITE);
if (!notify_recv_mrs_[i]) {
PLOG(ERROR) << "Failed to register notification recv buffer";
deconstruct();
return -1;
}
}

return 0;
}

Expand Down Expand Up @@ -254,20 +251,15 @@ int RdmaEndPoint::deconstructUnlocked() {
// A live QP may still reference the notification MRs. Keep both MRs
// and backing buffers intact until QP destruction succeeds.
if (!notify_qp_) {
for (auto& mr : notify_recv_mrs_) {
if (!mr) continue;
if (context_->verbs_.ibv_dereg_mr(mr)) {
if (notify_recv_mr_) {
if (context_->verbs_.ibv_dereg_mr(notify_recv_mr_)) {
PLOG(ERROR) << "Failed to deregister notification recv MR";
result = -1;
} else {
mr = nullptr;
notify_recv_mr_ = nullptr;
}
}
if (std::all_of(notify_recv_mrs_.begin(), notify_recv_mrs_.end(),
[](ibv_mr* mr) { return mr == nullptr; })) {
notify_recv_mrs_.clear();
notify_recv_buffers_.clear();
}
if (!notify_recv_mr_) notify_recv_buffer_.clear();

if (notify_send_mr_) {
if (context_->verbs_.ibv_dereg_mr(notify_send_mr_)) {
Expand Down Expand Up @@ -306,8 +298,8 @@ int RdmaEndPoint::deconstructUnlocked() {
wr_depth_list_ = nullptr;
}

if (result == 0 && !notify_qp_ && notify_recv_mrs_.empty() &&
!notify_send_mr_ && qp_list_.empty()) {
if (result == 0 && !notify_qp_ && !notify_recv_mr_ && !notify_send_mr_ &&
qp_list_.empty()) {
peer_server_name_.clear();
peer_nic_name_.clear();
status_.store(EP_DESTROYED, std::memory_order_release);
Expand Down Expand Up @@ -1036,14 +1028,57 @@ int RdmaEndPoint::setupOneQP(int qp_index, const std::string& peer_gid,
return 0;
}

char* RdmaEndPoint::notifySlotPtr(char* base, size_t idx) {
return base + idx * kNotifyBufferSize;
}

bool RdmaEndPoint::encodeNotifyPayload(char* slot, const std::string& name,
const std::string& msg,
uint32_t* out_len) {
if (!slot || !out_len) return false;
if (name.size() > UINT32_MAX || msg.size() > UINT32_MAX) return false;
uint32_t name_len = static_cast<uint32_t>(name.size());
uint32_t msg_len = static_cast<uint32_t>(msg.size());
const size_t total_size =
sizeof(name_len) + name_len + sizeof(msg_len) + msg_len;
if (total_size > kNotifyBufferSize) return false;

auto* ptr = slot;
std::memcpy(ptr, &name_len, sizeof(name_len));
ptr += sizeof(name_len);
std::memcpy(ptr, name.data(), name_len);
ptr += name_len;
std::memcpy(ptr, &msg_len, sizeof(msg_len));
ptr += sizeof(msg_len);
std::memcpy(ptr, msg.data(), msg_len);
*out_len = static_cast<uint32_t>(total_size);
return true;
}

bool RdmaEndPoint::decodeNotifyPayload(const char* data, size_t byte_len,
std::string* name, std::string* msg) {
if (!data || !name || !msg || byte_len < 8) return false;
uint32_t name_len = 0;
std::memcpy(&name_len, data, sizeof(name_len));
if (name_len > byte_len - 8) return false;
uint32_t msg_len = 0;
std::memcpy(&msg_len, data + 4 + name_len, sizeof(msg_len));
if (msg_len > byte_len - 8 - name_len) return false;
name->assign(data + 4, name_len);
msg->assign(data + 4 + name_len + 4, msg_len);
return true;
}

void RdmaEndPoint::postNotifyRecv(size_t idx) {
if (idx >= notify_recv_buffers_.size() || idx >= notify_recv_mrs_.size())
if (!notify_qp_ || !notify_recv_mr_ || idx >= kNotifyMaxPendingSends ||
notify_recv_buffer_.size() < (idx + 1) * kNotifyBufferSize)
return;

ibv_sge sge = {};
sge.addr = reinterpret_cast<uint64_t>(notify_recv_buffers_[idx].data());
sge.length = notify_recv_buffers_[idx].size();
sge.lkey = notify_recv_mrs_[idx]->lkey;
sge.addr = reinterpret_cast<uint64_t>(
notifySlotPtr(notify_recv_buffer_.data(), idx));
sge.length = kNotifyBufferSize;
sge.lkey = notify_recv_mr_->lkey;

ibv_recv_wr wr = {};
wr.wr_id = idx;
Expand Down Expand Up @@ -1169,30 +1204,13 @@ bool RdmaEndPoint::sendNotification(const std::string& name,
// Pick the next send slot — flow control guarantees this slot's previous
// DMA has completed (at most kNotifyMaxPendingSends-1 in-flight).
size_t slot = notify_send_wr_id_ % kNotifyMaxPendingSends;
char* slot_ptr = notify_send_buffer_.data() + slot * kNotifyBufferSize;

// Serialize: [name_len(4)][name][msg_len(4)][msg]
if (name.size() > UINT32_MAX || msg.size() > UINT32_MAX) {
LOG(ERROR) << "Notification field exceeds uint32 limit";
return false;
}
uint32_t name_len = static_cast<uint32_t>(name.size());
uint32_t msg_len = static_cast<uint32_t>(msg.size());
size_t total_size = sizeof(name_len) + name_len + sizeof(msg_len) + msg_len;
if (total_size > kNotifyBufferSize) {
LOG(ERROR) << "Notification message too large: " << total_size;
char* slot_ptr = notifySlotPtr(notify_send_buffer_.data(), slot);
uint32_t total_size = 0;
if (!encodeNotifyPayload(slot_ptr, name, msg, &total_size)) {
LOG(ERROR) << "Failed to encode notification payload";
return false;
}

auto* ptr = slot_ptr;
std::memcpy(ptr, &name_len, sizeof(name_len));
ptr += sizeof(name_len);
std::memcpy(ptr, name.data(), name_len);
ptr += name_len;
std::memcpy(ptr, &msg_len, sizeof(msg_len));
ptr += sizeof(msg_len);
std::memcpy(ptr, msg.data(), msg_len);

// Post send
ibv_sge sge = {};
sge.addr = reinterpret_cast<uint64_t>(slot_ptr);
Expand Down Expand Up @@ -1223,7 +1241,8 @@ bool RdmaEndPoint::sendNotification(const std::string& name,

bool RdmaEndPoint::handleNotifyRecv(size_t buffer_idx, size_t byte_len) {
std::lock_guard<std::mutex> resource_guard(notify_resource_mutex_);
if (buffer_idx >= notify_recv_buffers_.size()) {
if (!notify_recv_mr_ || buffer_idx >= kNotifyMaxPendingSends ||
notify_recv_buffer_.size() < (buffer_idx + 1) * kNotifyBufferSize) {
LOG(ERROR) << "Invalid recv buffer index: " << buffer_idx;
return false;
}
Expand All @@ -1234,33 +1253,16 @@ bool RdmaEndPoint::handleNotifyRecv(size_t buffer_idx, size_t byte_len) {
return false;
}

char* data = notify_recv_buffers_[buffer_idx].data();
size_t len = byte_len;

// Deserialize: [name_len(4)][name][msg_len(4)][msg]
if (len < 8) {
LOG(ERROR) << "Invalid notification message size: " << len;
postNotifyRecv(buffer_idx);
return false;
}

uint32_t name_len = *reinterpret_cast<uint32_t*>(data);
if (name_len > len - 8) {
LOG(ERROR) << "Invalid notification message format (name too long)";
postNotifyRecv(buffer_idx);
return false;
}

std::string name(data + 4, name_len);
uint32_t msg_len = *reinterpret_cast<uint32_t*>(data + 4 + name_len);
if (msg_len > len - 8 - name_len) {
LOG(ERROR) << "Invalid notification message format (msg too long)";
char* data = notifySlotPtr(notify_recv_buffer_.data(), buffer_idx);
std::string name;
std::string msg;
if (!decodeNotifyPayload(data, byte_len, &name, &msg)) {
LOG(ERROR) << "Invalid notification message size or format: "
<< byte_len;
postNotifyRecv(buffer_idx);
return false;
}

std::string msg(data + 4 + name_len + 4, msg_len);

// Add directly to transport queue (skip endpoint queue for lower latency)
context_->transport_.addNotificationToQueue(name, msg);

Expand Down
96 changes: 96 additions & 0 deletions mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,10 @@

#include <gtest/gtest.h>

#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include "tent/common/utils/string_builder.h"
#include "tent/transport/rdma/endpoint.h"
Expand Down Expand Up @@ -45,6 +48,28 @@ class EndpointTestAccess {
static bool notifyConnected(const RdmaEndPoint& endpoint) {
return endpoint.notify_connected_.load(std::memory_order_relaxed);
}

static constexpr size_t notifyBufferSize() {
return RdmaEndPoint::kNotifyBufferSize;
}

static constexpr size_t notifyMaxPendingSends() {
return RdmaEndPoint::kNotifyMaxPendingSends;
}

static char* notifySlotPtr(char* base, size_t idx) {
return RdmaEndPoint::notifySlotPtr(base, idx);
}

static bool encodeNotifyPayload(char* slot, const std::string& name,
const std::string& msg, uint32_t* out_len) {
return RdmaEndPoint::encodeNotifyPayload(slot, name, msg, out_len);
}

static bool decodeNotifyPayload(const char* data, size_t byte_len,
std::string* name, std::string* msg) {
return RdmaEndPoint::decodeNotifyPayload(data, byte_len, name, msg);
}
};

namespace {
Expand Down Expand Up @@ -113,6 +138,77 @@ TEST(EndpointLifecycleTest, NotificationFailsWhenEndpointIsNotConnected) {
EXPECT_FALSE(endpoint.sendNotification("name", "message"));
}

TEST(EndpointLifecycleTest, NotifySlotsAreAdjacentAndNonOverlapping) {
std::vector<char> buf(EndpointTestAccess::notifyMaxPendingSends() *
EndpointTestAccess::notifyBufferSize());
char* slot0 = EndpointTestAccess::notifySlotPtr(buf.data(), 0);
char* slot1 = EndpointTestAccess::notifySlotPtr(buf.data(), 1);
char* last = EndpointTestAccess::notifySlotPtr(
buf.data(), EndpointTestAccess::notifyMaxPendingSends() - 1);
EXPECT_EQ(slot0, buf.data());
EXPECT_EQ(static_cast<size_t>(slot1 - slot0),
EndpointTestAccess::notifyBufferSize());
EXPECT_EQ(static_cast<size_t>(last - slot0),
(EndpointTestAccess::notifyMaxPendingSends() - 1) *
EndpointTestAccess::notifyBufferSize());
EXPECT_EQ(last + EndpointTestAccess::notifyBufferSize(),
buf.data() + buf.size());
}

TEST(EndpointLifecycleTest, NotifyAdjacentSlotsDoNotCrosstalk) {
// Recv used to be 256 independent vectors. After coalescing, a wrong
// offset would mix payload from a neighbor slot.
const size_t slot_size = EndpointTestAccess::notifyBufferSize();
const size_t nslots = EndpointTestAccess::notifyMaxPendingSends();
std::vector<char> buf(nslots * slot_size, '\xee');
const size_t indices[] = {0, 1, 2, 127, 254, 255};
uint32_t encoded[6] = {};
for (size_t i = 0; i < 6; ++i) {
const size_t idx = indices[i];
std::string name = "n" + std::to_string(idx);
std::string msg =
"payload-" + std::to_string(idx) + std::string(2048, 'x');
ASSERT_TRUE(EndpointTestAccess::encodeNotifyPayload(
EndpointTestAccess::notifySlotPtr(buf.data(), idx), name, msg,
&encoded[i]));
ASSERT_GT(encoded[i], 8u);
ASSERT_LE(encoded[i], slot_size);
}
for (size_t i = 0; i < 6; ++i) {
const size_t idx = indices[i];
std::string name;
std::string msg;
ASSERT_TRUE(EndpointTestAccess::decodeNotifyPayload(
EndpointTestAccess::notifySlotPtr(buf.data(), idx), encoded[i],
&name, &msg))
<< "slot " << idx;
EXPECT_EQ(name, "n" + std::to_string(idx));
EXPECT_EQ(msg,
"payload-" + std::to_string(idx) + std::string(2048, 'x'));
// Trailing bytes in the 64KiB slot stay the fill pattern, so a
// too-long read would have picked up 0xee as name_len.
ASSERT_TRUE(EndpointTestAccess::decodeNotifyPayload(
EndpointTestAccess::notifySlotPtr(buf.data(), idx), slot_size,
&name, &msg));
EXPECT_EQ(name, "n" + std::to_string(idx));
}
// Slot 3 was never written; decoding the fill pattern must fail rather
// than returning a neighbor's payload.
std::string name;
std::string msg;
EXPECT_FALSE(EndpointTestAccess::decodeNotifyPayload(
EndpointTestAccess::notifySlotPtr(buf.data(), 3), slot_size, &name,
&msg));
}

TEST(EndpointLifecycleTest, NotifyPayloadRejectedWhenLargerThanSlot) {
std::vector<char> slot(EndpointTestAccess::notifyBufferSize());
uint32_t encoded = 0;
std::string too_big(EndpointTestAccess::notifyBufferSize(), 'z');
EXPECT_FALSE(EndpointTestAccess::encodeNotifyPayload(slot.data(), "n",
too_big, &encoded));
}

TEST(EndpointLifecycleTest, NotifyLocalFaultKeepsEndpointServingData) {
// A fault confined to the notify QP must not retire the endpoint: doing so
// moves every data QP to ERR and flushes in-flight transfers.
Expand Down
Loading