From 80e939dd5c48d98e3f95da9fa1f444b0ea2edda3 Mon Sep 17 00:00:00 2001 From: Googler Date: Thu, 20 Aug 2026 10:50:56 -0700 Subject: [PATCH] [tpu_raiden]Encapsulate connection borrowing and returning in RawBufferTransport PiperOrigin-RevId: 967922647 --- tpu_sync/transport/block_transport.cc | 20 +++++++++---------- tpu_sync/transport/lib/BUILD | 1 + .../transport/lib/raw_buffer_transport.cc | 12 +++++------ tpu_sync/transport/lib/raw_buffer_transport.h | 14 +++++++++++-- 4 files changed, 29 insertions(+), 18 deletions(-) diff --git a/tpu_sync/transport/block_transport.cc b/tpu_sync/transport/block_transport.cc index 6260e801..06153e52 100644 --- a/tpu_sync/transport/block_transport.cc +++ b/tpu_sync/transport/block_transport.cc @@ -1015,15 +1015,15 @@ absl::Status BlockTransport::ProcessSocketPush( const uint8_t major_order = first.major_order; const size_t block_count = static_cast(count_or_size); - auto status_or_fd = raw_transport_.conn_pool().Borrow(peer, local_ip); - if (!status_or_fd.ok()) { - return status_or_fd.status(); + auto borrowed_fd = raw_transport_.BorrowConnection(peer, local_ip); + if (!borrowed_fd.ok()) { + return borrowed_fd.status(); } - const int fd = status_or_fd.value(); + const int fd = borrowed_fd.value(); bool ok_to_pool = false; auto fd_cleaner = absl::MakeCleanup([&] { - raw_transport_.conn_pool().Return(ok_to_pool, fd, peer, local_ip); + raw_transport_.ReturnConnection(ok_to_pool, fd, peer, local_ip); }); lib::ChunkHeader header = {}; @@ -1154,16 +1154,16 @@ void BlockTransport::H2hReadWorker( const std::vector& explicit_dst_ptrs, std::vector& statuses, MajorOrder major_order, BlockReceivedCallback on_block_received, uint64_t uuid) { - auto status_or_fd = raw_transport_.conn_pool().Borrow(peer, local_ip); - if (!status_or_fd.ok()) { - statuses[stream_idx] = status_or_fd.status(); + auto borrowed_fd = raw_transport_.BorrowConnection(peer, local_ip); + if (!borrowed_fd.ok()) { + statuses[stream_idx] = borrowed_fd.status(); return; } - const int fd = status_or_fd.value(); + const int fd = borrowed_fd.value(); bool ok_to_pool = false; auto fd_cleaner = absl::MakeCleanup([&] { - raw_transport_.conn_pool().Return(ok_to_pool, fd, peer, local_ip); + raw_transport_.ReturnConnection(ok_to_pool, fd, peer, local_ip); }); size_t SF = block_delegate_->shard_factor(); diff --git a/tpu_sync/transport/lib/BUILD b/tpu_sync/transport/lib/BUILD index 02460982..8619059d 100644 --- a/tpu_sync/transport/lib/BUILD +++ b/tpu_sync/transport/lib/BUILD @@ -105,6 +105,7 @@ cc_library( "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", diff --git a/tpu_sync/transport/lib/raw_buffer_transport.cc b/tpu_sync/transport/lib/raw_buffer_transport.cc index e800f067..8f421e59 100644 --- a/tpu_sync/transport/lib/raw_buffer_transport.cc +++ b/tpu_sync/transport/lib/raw_buffer_transport.cc @@ -477,10 +477,10 @@ absl::Status RawBufferTransport::PullBuffer( ", Size: ", size_bytes, ", Shard Host Size: ", host_size)); } - ASSIGN_OR_RETURN(const int fd, conn_pool_.Borrow(peer)); + ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer)); bool ok_to_pool = false; auto fd_cleaner = - absl::MakeCleanup([&] { conn_pool_.Return(ok_to_pool, fd, peer); }); + absl::MakeCleanup([&] { ReturnConnection(ok_to_pool, fd, peer); }); ChunkHeader header = {}; header.version = 1; @@ -565,10 +565,10 @@ absl::Status RawBufferTransport::ProcessSocketBufferPush( "Destination peer address cannot be empty"); } - ASSIGN_OR_RETURN(const int fd, conn_pool_.Borrow(peer)); + ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer)); bool ok_to_pool = false; auto fd_cleaner = - absl::MakeCleanup([&] { conn_pool_.Return(ok_to_pool, fd, peer); }); + absl::MakeCleanup([&] { ReturnConnection(ok_to_pool, fd, peer); }); const uint8_t opcode = request.socket_opcode; const uint64_t uuid = request.uuid; @@ -722,10 +722,10 @@ absl::Status RawBufferTransport::PushBatch( "Destination peer address cannot be empty"); } - ASSIGN_OR_RETURN(const int fd, conn_pool_.Borrow(peer)); + ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer)); bool ok_to_pool = false; auto fd_cleaner = - absl::MakeCleanup([&] { conn_pool_.Return(ok_to_pool, fd, peer); }); + absl::MakeCleanup([&] { ReturnConnection(ok_to_pool, fd, peer); }); ChunkHeader header = {}; header.version = 1; diff --git a/tpu_sync/transport/lib/raw_buffer_transport.h b/tpu_sync/transport/lib/raw_buffer_transport.h index 643405a1..77a4f943 100644 --- a/tpu_sync/transport/lib/raw_buffer_transport.h +++ b/tpu_sync/transport/lib/raw_buffer_transport.h @@ -30,6 +30,7 @@ #include "absl/container/flat_hash_set.h" #include "absl/functional/any_invocable.h" #include "absl/status/status.h" +#include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" @@ -75,8 +76,17 @@ class RawBufferTransport final { // Return the local IP addresses. absl::Span local_ips() const { return local_ips_; } - // Return the connection pool that manages the sockets that connect to peers. - ConnPool& conn_pool() { return conn_pool_; } + // Borrows a connection from the connection pool. + absl::StatusOr BorrowConnection(absl::string_view peer, + absl::string_view local_ip = "") { + return conn_pool_.Borrow(peer, local_ip); + } + + // Returns a connection to the connection pool. + void ReturnConnection(bool ok_to_pool, int fd, absl::string_view peer, + absl::string_view local_ip = "") { + conn_pool_.Return(ok_to_pool, fd, peer, local_ip); + } // Synchronously pulls a buffer identified by `buffer_id` from the remote // `peer`, by sending out a `kOpBufferPull ChunkHeader` and then receiving