[tpu_raiden]Encapsulate connection borrowing and returning in RawBufferTransport

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<size_t>(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<uint8_t*>& explicit_dst_ptrs,
     std::vector<absl::Status>& 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();

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",

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;

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<const std::string> 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<int> 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