[Bugfix][TransferEngine] Bound the handshake daemon's pending-notify queue

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<HandShakeRequestType, std::string> readString(int fd) {
     HandShakeRequestType type = HandShakeRequestType::Invalid;
 

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 = "";

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<sockaddr *>(&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<TransferMetadata::NotifyDesc> 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