[Bugfix][TransferEngine] Keep a stalled peer from wedging the handshake daemon

mooncake-transfer-engine/include/common.h

@@ -437,9 +437,15 @@ static inline ssize_t writeFully(int fd, const void *buf, size_t len) {
     size_t nbytes = len;
     while (nbytes) {
         ssize_t rc = write(fd, pos, nbytes);
-        if (rc < 0 && (errno == EAGAIN || errno == EINTR))
+        if (rc < 0 && errno == EINTR)
             continue;
-        else if (rc < 0) {
+        else if (rc < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
+            // A blocking socket only reports EAGAIN once SO_SNDTIMEO fires;
+            // retrying here would spin forever, so fail the write instead.
+            LOG(WARNING) << "Socket write timed out: expected " << len
+                         << " bytes, actual " << len - nbytes << " bytes";
+            return len - nbytes;
+        } else if (rc < 0) {
             PLOG(ERROR) << "Socket write failed";
             return rc;
         } else if (rc == 0) {

mooncake-transfer-engine/src/transfer_metadata_plugin.cpp

@@ -919,6 +919,13 @@ struct SocketHandShakePlugin : public HandShakePlugin {
                     close(conn_fd);
                     continue;
                 }
+                if (setsockopt(conn_fd, SOL_SOCKET, SO_SNDTIMEO, &timeout,
+                               sizeof(timeout))) {
+                    PLOG(ERROR)
+                        << "SocketHandShakePlugin: setsockopt(SO_SNDTIMEO)";
+                    close(conn_fd);
+                    continue;
+                }
 
                 auto peer_hostname =
                     getNetworkAddress((struct sockaddr *)&addr);

mooncake-transfer-engine/tests/transfer_metadata_test.cpp

@@ -1057,6 +1057,34 @@ TEST(HandshakeFrameTest, RejectsInvalidLength) {
     close(fds[1]);
 }
 
+TEST(HandshakeFrameTest, WriteToSilentPeerFailsWithinSendTimeout) {
+    int fds[2];
+    ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0);
+
+    // Nobody reads fds[1], so a payload far larger than the send buffer
+    // stalls the writer. With SO_SNDTIMEO the write has to fail fast instead
+    // of blocking (or spinning on EAGAIN) forever.
+    int sndbuf = 4096;
+    ASSERT_EQ(
+        setsockopt(fds[0], SOL_SOCKET, SO_SNDBUF, &sndbuf, sizeof(sndbuf)), 0);
+    struct timeval timeout;
+    timeout.tv_sec = 1;
+    timeout.tv_usec = 0;
+    ASSERT_EQ(
+        setsockopt(fds[0], SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)),
+        0);
+
+    const std::string payload(8 << 20, 'x');
+    const auto start = std::chrono::steady_clock::now();
+    EXPECT_EQ(writeString(fds[0], HandShakeRequestType::Metadata, payload),
+              ERR_SOCKET);
+    const auto elapsed = std::chrono::steady_clock::now() - start;
+    EXPECT_LT(elapsed, std::chrono::seconds(30));
+
+    close(fds[0]);
+    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