From 0e456ba5fe2c513cb75dc5d0a07b1f5774306ca3 Mon Sep 17 00:00:00 2001 From: Jeff Ma Date: Tue, 8 Sep 2026 17:58:21 +0000 Subject: [PATCH] Chunk vectored socket I/O at IOV_MAX instead of rejecting large iovec spans. WriteVExact/ReadVExact returned InvalidArgument for spans longer than IOV_MAX (1024). Per-destination-unit resharding (1eb2414) resolves a single block into thousands of small chunks, and block_transport.cc / socket_transport_adapter.cc pass the whole list in one call, so those transfers failed outright. IOV_MAX bounds one writev/readv call, not the byte stream: loop over IOV_MAX-sized sub-spans, relying on SendV/RecvV being exact within each sub-span and on TCP ordering across them. Rewrite socket_util_test.cc so it builds in this tree (it depended on //tpu_sync/transport/peregrine/src/util, which needs the absent //third_party/xxhash) and add cases at, above and far above IOV_MAX. Claude-Session: https://claude.ai/code/session_012cbF1oLqDSJhgcQS3BBhiy Signed-off-by: Jeff Ma --- tpu_sync/transport/peregrine/src/api/BUILD | 4 +- .../transport/peregrine/src/api/socket_util.h | 24 ++++++ .../peregrine/src/api/socket_util_test.cc | 73 +++++++++++++------ 3 files changed, 75 insertions(+), 26 deletions(-) diff --git a/tpu_sync/transport/peregrine/src/api/BUILD b/tpu_sync/transport/peregrine/src/api/BUILD index 2640de5af..da0b06884 100644 --- a/tpu_sync/transport/peregrine/src/api/BUILD +++ b/tpu_sync/transport/peregrine/src/api/BUILD @@ -48,10 +48,10 @@ cc_test( ":socket_util", ":transport_types", "//tpu_sync/transport/peregrine/src/internal/base:endpoint", + "//tpu_sync/transport/peregrine/src/internal/base:ipaddr", "//tpu_sync/transport/peregrine/src/internal/base:types", "//tpu_sync/transport/peregrine/src/internal/socket:socket_tcp", - "//tpu_sync/transport/peregrine/src/internal/util:test_util", - "//tpu_sync/transport/peregrine/src/util", + "//tpu_sync/transport/peregrine/src/internal/socket:socket_util", "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/synchronization", diff --git a/tpu_sync/transport/peregrine/src/api/socket_util.h b/tpu_sync/transport/peregrine/src/api/socket_util.h index c3355d3d2..0e7fe3ed5 100644 --- a/tpu_sync/transport/peregrine/src/api/socket_util.h +++ b/tpu_sync/transport/peregrine/src/api/socket_util.h @@ -17,6 +17,7 @@ #include +#include #include #include "absl/base/optimization.h" @@ -51,6 +52,18 @@ inline absl::Status WriteVExact(int fd, absl::Span iovs) { if ABSL_PREDICT_TRUE (1 <= n && n <= IOV_MAX) { return internal::TcpSocketUtil::SendV(internal::fd_t(fd), iovs); } + if (n > IOV_MAX) { + // IOV_MAX limits one syscall, not the logical byte stream. Pool + // resharding can expose many small, non-contiguous state fragments. + while (!iovs.empty()) { + const size_t count = std::min(iovs.size(), size_t{IOV_MAX}); + const absl::Status status = internal::TcpSocketUtil::SendV( + internal::fd_t(fd), iovs.subspan(0, count)); + if (!status.ok()) return status; + iovs = iovs.subspan(count); + } + return absl::OkStatus(); + } return absl::InvalidArgumentError(absl::StrCat("#iovs=", n)); } @@ -73,6 +86,17 @@ inline absl::Status ReadVExact(int fd, absl::Span iovs) { if ABSL_PREDICT_TRUE (1 <= n && n <= IOV_MAX) { return internal::TcpSocketUtil::RecvV(internal::fd_t(fd), iovs); } + if (n > IOV_MAX) { + // Receive batch boundaries need not match the sender's boundaries. + while (!iovs.empty()) { + const size_t count = std::min(iovs.size(), size_t{IOV_MAX}); + const absl::Status status = internal::TcpSocketUtil::RecvV( + internal::fd_t(fd), iovs.subspan(0, count)); + if (!status.ok()) return status; + iovs = iovs.subspan(count); + } + return absl::OkStatus(); + } return absl::InvalidArgumentError(absl::StrCat("#iovs=", n)); } diff --git a/tpu_sync/transport/peregrine/src/api/socket_util_test.cc b/tpu_sync/transport/peregrine/src/api/socket_util_test.cc index 8ba0076c8..19ac02460 100644 --- a/tpu_sync/transport/peregrine/src/api/socket_util_test.cc +++ b/tpu_sync/transport/peregrine/src/api/socket_util_test.cc @@ -32,21 +32,16 @@ #include "absl/types/span.h" #include "tpu_sync/transport/peregrine/src/api/transport_types.h" #include "tpu_sync/transport/peregrine/src/internal/base/endpoint.h" +#include "tpu_sync/transport/peregrine/src/internal/base/ipaddr.h" #include "tpu_sync/transport/peregrine/src/internal/base/types.h" #include "tpu_sync/transport/peregrine/src/internal/socket/socket_tcp.h" -#include "tpu_sync/transport/peregrine/src/internal/util/test_util.h" -#include "tpu_sync/transport/peregrine/src/util/util.h" +#include "tpu_sync/transport/peregrine/src/internal/socket/socket_util.h" namespace peregrine::testing { namespace { using ::peregrine::internal::Endpoint; using ::peregrine::internal::TcpSocket; -using ::peregrine::internal::testing::IPv4Localhost; -using ::peregrine::internal::testing::IPv6Localhost; -using ::peregrine::internal::testing::TestOnly_CreateTcpSocket; -using ::peregrine::internal::testing::TestOnly_FindFreeTcpPort; -using ::peregrine::util::RandomNonZero; using ::testing::Combine; using ::testing::Eq; using ::testing::Ne; @@ -71,11 +66,13 @@ class SocketUtilTest : public ::testing::TestWithParam { : family_(std::get<0>(GetParam())), read_iovec_(std::get<1>(GetParam())), write_iovec_(std::get<2>(GetParam())), - local_(family_ == AF_INET ? IPv4Localhost() : IPv6Localhost(), - TestOnly_FindFreeTcpPort(family_)), - peer_(local_), - listener_(TestOnly_CreateTcpSocket(family_)), - connector_(TestOnly_CreateTcpSocket(family_)) { + local_(internal::IpAddr::Create(family_ == AF_INET ? "127.0.0.1" : "::1") + .value(), + 0), + listener_(TcpSocket::Create(family_)), + connector_(TcpSocket::Create(family_)) { + CHECK_NE(listener_, nullptr); + CHECK_NE(connector_, nullptr); DCHECK(listener_->IsValid()); DCHECK(connector_->IsValid()); DCHECK(!listener_->IsConnected()); @@ -84,11 +81,12 @@ class SocketUtilTest : public ::testing::TestWithParam { } protected: + void CheckReadWrite(size_t read_iov_count, size_t write_iov_count); + const int family_; const bool read_iovec_; const bool write_iovec_; const Endpoint local_; - const Endpoint peer_; const std::unique_ptr listener_; const std::unique_ptr connector_; }; @@ -99,12 +97,15 @@ INSTANTIATE_TEST_SUITE_P(, SocketUtilTest, /*wiov=*/Values(false, true)), ToString); -TEST_P(SocketUtilTest, ReadWrite) { - // Create a big chunk of send/recv buffers with random data. +void SocketUtilTest::CheckReadWrite(size_t read_iov_count, + size_t write_iov_count) { + // Non-uniform, non-zero data detects reordering without test RNG dependencies. constexpr size_t kSize = 64UL << 20; std::vector send_buf(kSize, 0x01); std::vector recv_buf(kSize, 0x00); - RandomNonZero(absl::MakeSpan(send_buf)); + for (size_t i = 0; i < send_buf.size(); ++i) { + send_buf[i] = static_cast(1 + i % 251); + } ASSERT_THAT(recv_buf, Pointwise(Ne(), send_buf)); // First, create a server thread. @@ -122,10 +123,12 @@ TEST_P(SocketUtilTest, ReadWrite) { if (read_iovec_) { std::vector iovs; - constexpr size_t kPartial = kSize / 3; - iovs.push_back({recv_buf.data(), kPartial}); - iovs.push_back({recv_buf.data() + kPartial, kPartial}); - iovs.push_back({recv_buf.data() + kPartial * 2, kSize - kPartial * 2}); + const size_t partial = kSize / read_iov_count; + for (size_t i = 0; i < read_iov_count; ++i) { + const size_t offset = i * partial; + const size_t size = i + 1 == read_iov_count ? kSize - offset : partial; + iovs.push_back({recv_buf.data() + offset, size}); + } CHECK_OK(ReadVExact(new_socket->fd().value(), iovs)); } else { CHECK_OK(ReadExact(new_socket->fd().value(), recv_buf.data(), kSize)); @@ -135,15 +138,21 @@ TEST_P(SocketUtilTest, ReadWrite) { // Second, create a client thread. std::thread client([&]() { server_ready.WaitForNotification(); - CHECK(connector_->Connect(peer_)); + // Port zero reserves an ephemeral port, without a free-port race or + // dependencies on the optional test utility library. + const Endpoint peer = Endpoint::Create(internal::SelfAddrPort(listener_->fd())); + CHECK(connector_->Connect(peer)); DCHECK(connector_->IsBlocking()); DCHECK(connector_->IsConnected()); if (write_iovec_) { std::vector iovs; - constexpr size_t kPartial = kSize / 2; - iovs.push_back({send_buf.data(), kPartial}); - iovs.push_back({send_buf.data() + kPartial, kSize - kPartial}); + const size_t partial = kSize / write_iov_count; + for (size_t i = 0; i < write_iov_count; ++i) { + const size_t offset = i * partial; + const size_t size = i + 1 == write_iov_count ? kSize - offset : partial; + iovs.push_back({send_buf.data() + offset, size}); + } CHECK_OK(WriteVExact(connector_->fd().value(), iovs)); } else { CHECK_OK(WriteExact(connector_->fd().value(), send_buf.data(), kSize)); @@ -158,5 +167,21 @@ TEST_P(SocketUtilTest, ReadWrite) { ASSERT_THAT(recv_buf, Pointwise(Eq(), send_buf)); } +TEST_P(SocketUtilTest, ReadWrite) { CheckReadWrite(3, 2); } + +TEST_P(SocketUtilTest, ReadWriteAtIovMax) { + CheckReadWrite(IOV_MAX, IOV_MAX); +} + +TEST_P(SocketUtilTest, ReadWriteAboveIovMax) { + CheckReadWrite(IOV_MAX + 1, 2 * IOV_MAX + 3); +} + +TEST_P(SocketUtilTest, ReadWriteManyIovs) { + // Different sender/receiver batch boundaries must preserve the stream. + // This exceeds the 14,337 entries observed in hybrid-state resharding. + CheckReadWrite(16 * IOV_MAX + 3, 14 * IOV_MAX + 1); +} + } // namespace } // namespace peregrine::testing