diff --git a/mooncake-transfer-engine/include/common.h b/mooncake-transfer-engine/include/common.h index 883b51cedf..aaf1911c9b 100644 --- a/mooncake-transfer-engine/include/common.h +++ b/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) { diff --git a/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp b/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp index 95d4cd05a9..ac1f3f66fe 100644 --- a/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp +++ b/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); diff --git a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp index 2ebb1eefc0..e8592dc93c 100644 --- a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp +++ b/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