diff --git a/mooncake-transfer-engine/include/common.h b/mooncake-transfer-engine/include/common.h index 5118874c3f..883b51cedf 100644 --- a/mooncake-transfer-engine/include/common.h +++ b/mooncake-transfer-engine/include/common.h @@ -585,6 +585,37 @@ static inline size_t getHandshakeMaxLength() { return max_length; } +// Constants for the handshake daemon's pending-notify bound. +constexpr size_t kDefaultHandshakeMaxNotifyEntries = 1024; +constexpr size_t kMaxHandshakeMaxNotifyEntries = 1ULL << 20; + +// Load the pending-notify bound from the environment variable. +static inline size_t loadHandshakeMaxNotifyEntries() { + const char *env = std::getenv("MC_HANDSHAKE_MAX_NOTIFY_ENTRIES"); + if (env != nullptr) { + std::string_view env_sv(env); + size_t val = 0; + auto [ptr, ec] = + std::from_chars(env_sv.data(), env_sv.data() + env_sv.size(), val); + if (ec == std::errc() && ptr == env_sv.data() + env_sv.size() && + val >= 1 && val <= kMaxHandshakeMaxNotifyEntries) { + LOG(INFO) << "MC_HANDSHAKE_MAX_NOTIFY_ENTRIES set to " << val; + return val; + } + LOG(WARNING) << "Invalid MC_HANDSHAKE_MAX_NOTIFY_ENTRIES value: " << env + << ", valid range: 1 to " << kMaxHandshakeMaxNotifyEntries + << ", using default " << kDefaultHandshakeMaxNotifyEntries; + } + return kDefaultHandshakeMaxNotifyEntries; +} + +// Get the pending-notify bound. +// Loaded once from MC_HANDSHAKE_MAX_NOTIFY_ENTRIES at first call. +static inline size_t getHandshakeMaxNotifyEntries() { + static size_t max_entries = loadHandshakeMaxNotifyEntries(); + return max_entries; +} + static inline std::pair readString(int fd) { HandShakeRequestType type = HandShakeRequestType::Invalid; diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index a8f87e63fb..c12d94d317 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -402,6 +402,17 @@ int TransferMetadata::receivePeerNotify(const Json::Value &peer_json, RWSpinlock::WriteGuard guard(notify_lock_); TransferMetadata::NotifyDesc peer_notify, local_reply; TransferNotifyUtil::decode(peer_json, peer_notify); + const size_t max_entries = getHandshakeMaxNotifyEntries(); + if (notifys.size() >= max_entries) { + // The daemon accepts notify frames from any peer, so an unbounded + // vector lets a flood grow RSS without limit (#4445). Evict the + // oldest undrained entry: a flood degrades pending notifications + // instead of the process. + notifys.erase(notifys.begin()); + LOG_EVERY_N(WARNING, 1000) + << "pending notify queue is full (" << max_entries + << " entries); evicting the oldest entry"; + } notifys.push_back(peer_notify); // reply local_reply.name = ""; diff --git a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp index 1002f3d9df..2ebb1eefc0 100644 --- a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp +++ b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp @@ -1057,6 +1057,59 @@ TEST(HandshakeFrameTest, RejectsInvalidLength) { close(fds[1]); } +// Every notify frame the daemon accepts lands in an in-memory vector that +// used to have no bound, so any peer that can reach the RPC port could grow +// the host's RSS forever (#4445). The queue now caps at +// getHandshakeMaxNotifyEntries() and evicts the oldest undrained entry once +// full. Flood cap + 64 frames over the real socket, then drain. +TEST(HandshakeNotifyTest, BoundedQueueEvictsOldestUnderFlood) { + TransferMetadata server(P2PHANDSHAKE); + int sockfd = -1; + const uint16_t port = findAvailableTcpPort(sockfd); + ASSERT_NE(port, 0); + + TransferMetadata::RpcMetaDesc rpc{}; + rpc.ip_or_host_name = "127.0.0.1"; + rpc.rpc_port = port; + rpc.sockfd = sockfd; + ASSERT_EQ(server.addRpcMetaEntry("notify-bounded-test", rpc), 0); + + const size_t cap = getHandshakeMaxNotifyEntries(); + const size_t total = cap + 64; + for (size_t i = 0; i < total; ++i) { + int fd = socket(AF_INET, SOCK_STREAM, 0); + ASSERT_NE(fd, -1); + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + ASSERT_EQ( + connect(fd, reinterpret_cast(&addr), sizeof(addr)), 0); + Json::Value body; + body["name"] = "atk"; + body["notify_msg"] = "m" + std::to_string(i); + ASSERT_EQ(writeString(fd, HandShakeRequestType::Notify, + Json::FastWriter{}.write(body)), + 0); + // The daemon answers each accepted frame before half-closing, so + // reading the reply orders frame i before frame i+1. Read the whole + // frame: a bare read() can return after the header alone, and closing + // while the daemon is still writing the payload resets the socket + // and kills the process with SIGPIPE on its next write. + auto [rtype, rjson] = readString(fd); + ASSERT_EQ(rtype, HandShakeRequestType::Notify); + ASSERT_FALSE(rjson.empty()); + close(fd); + } + + std::vector drained; + ASSERT_EQ(server.getNotifies(drained), 0); + ASSERT_EQ(drained.size(), cap); + // The oldest 64 frames were evicted; the survivors keep arrival order. + EXPECT_EQ(drained.front().notify_msg, "m64"); + EXPECT_EQ(drained.back().notify_msg, "m" + std::to_string(total - 1)); +} + namespace { // Hardware-free in-memory MetadataStoragePlugin. get() returns false for keys