[tpu_raiden] refactor connection pool as a separate class

tpu_raiden/transport/BUILD

@@ -32,7 +32,7 @@ cc_library(
         "//tpu_raiden/core:status_macros",
         "//tpu_raiden/core:tsl_platform_headers",
         "//tpu_raiden/transport/lib:raw_buffer_transport",
-        "//tpu_raiden/transport/lib:socket_util",
+        "//tpu_raiden/transport/lib/socket:util",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -56,7 +56,6 @@ cc_test(
     features = ["-use_header_modules"],
     deps = [
         ":block_transport",
-        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
@@ -75,7 +74,7 @@ cc_binary(
     ],
     features = ["-use_header_modules"],
     deps = [
-        "//tpu_raiden/transport/lib:socket_util",
+        "//tpu_raiden/transport/lib/socket:util",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/types:span",
     ],

tpu_raiden/transport/block_transport.cc

@@ -46,7 +46,7 @@
 #include "absl/types/span.h"
 #include "tpu_raiden/core/status_macros.h"
 #include "tpu_raiden/transport/lib/raw_buffer_transport.h"
-#include "tpu_raiden/transport/lib/socket_util.h"
+#include "tpu_raiden/transport/lib/socket/util.h"
 
 namespace tpu_raiden {
 namespace transport {
@@ -716,7 +716,7 @@ void BlockTransport::H2hWriteWorker(int stream_idx, absl::string_view peer,
                                     std::vector<absl::Status>& statuses,
                                     MajorOrder major_order, uint64_t uuid,
                                     int layer_idx, int parallelism) {
-  auto status_or_fd = BorrowConnection(peer, local_ip);
+  auto status_or_fd = conn_pool_.Borrow(peer, local_ip);
   if (!status_or_fd.ok()) {
     statuses[stream_idx] = status_or_fd.status();
     return;
@@ -725,7 +725,7 @@ void BlockTransport::H2hWriteWorker(int stream_idx, absl::string_view peer,
   const int fd = status_or_fd.value();
   bool ok_to_pool = false;
   auto fd_cleaner = absl::MakeCleanup(
-      [&] { ReturnConnection(ok_to_pool, fd, peer, local_ip); });
+      [&] { conn_pool_.Return(ok_to_pool, fd, peer, local_ip); });
 
   PacketHeader header = {};
   header.op = dst_block_ids.empty() ? 1 : 6;
@@ -842,7 +842,7 @@ 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 = BorrowConnection(peer, local_ip);
+  auto status_or_fd = conn_pool_.Borrow(peer, local_ip);
   if (!status_or_fd.ok()) {
     statuses[stream_idx] = status_or_fd.status();
     return;
@@ -851,7 +851,7 @@ void BlockTransport::H2hReadWorker(
   const int fd = status_or_fd.value();
   bool ok_to_pool = false;
   auto fd_cleaner = absl::MakeCleanup(
-      [&] { ReturnConnection(ok_to_pool, fd, peer, local_ip); });
+      [&] { conn_pool_.Return(ok_to_pool, fd, peer, local_ip); });
 
   size_t SF = block_delegate_->shard_factor();
 

tpu_raiden/transport/block_transport_test.cc

@@ -14,7 +14,6 @@
 
 #include "tpu_raiden/transport/block_transport.h"
 
-#include <algorithm>
 #include <chrono>  // NOLINT
 #include <cstddef>
 #include <cstdint>
@@ -26,7 +25,6 @@
 #include <vector>
 
 #include <gtest/gtest.h>
-#include "absl/base/thread_annotations.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
@@ -400,75 +398,6 @@ TEST(BlockTransportTest, PullSupportsBlockMajorOrder) {
             }));
 }
 
-class MockBlockTransport : public BlockTransport {
- public:
-  struct CallRecord {
-    std::string peer;
-    std::string local_ip;
-  };
-
-  MockBlockTransport(BlockTransportDelegate* delegate, int local_port,
-                     const std::vector<std::string>& local_ips)
-      : BlockTransport(delegate, local_port, local_ips) {}
-
-  absl::StatusOr<int> BorrowConnection(absl::string_view peer,
-                                       absl::string_view local_ip) override {
-    absl::MutexLock lock(mock_mu_);
-    acquire_calls_.push_back({std::string(peer), std::string(local_ip)});
-    return absl::InternalError("mock_connection_halt");
-  }
-
-  std::vector<CallRecord> acquire_calls() {
-    absl::MutexLock lock(mock_mu_);
-    return acquire_calls_;
-  }
-
- private:
-  absl::Mutex mock_mu_;
-  std::vector<CallRecord> acquire_calls_ ABSL_GUARDED_BY(mock_mu_);
-};
-
-TEST(BlockTransportTest, RoundRobinDistribution) {
-  constexpr size_t kSliceSize = 16;
-  MockDelegate delegate(kSliceSize);
-
-  std::vector<std::string> local_ips = {"10.0.0.1", "10.0.0.2"};
-  std::vector<std::string> peers = {"10.0.0.3:1234", "10.0.0.4:1234",
-                                    "10.0.0.5:1234"};
-
-  MockBlockTransport transport(&delegate, 0, local_ips);
-
-  std::vector<int> src_blocks = {0, 1, 2, 3, 4, 5};
-  auto res = transport.SyncPush(peers, src_blocks, /*dst_block_ids=*/{},
-                                /*parallelism=*/6, MajorOrder::kLayerMajor,
-                                /*uuid=*/0, /*layer_idx=*/-1);
-
-  EXPECT_FALSE(res.ok());
-  EXPECT_EQ(res.status().message(), "mock_connection_halt");
-
-  auto calls = transport.acquire_calls();
-  ASSERT_EQ(calls.size(), 6);
-
-  std::vector<MockBlockTransport::CallRecord> expected = {
-      {"10.0.0.3:1234", "10.0.0.1"}, {"10.0.0.4:1234", "10.0.0.2"},
-      {"10.0.0.5:1234", "10.0.0.1"}, {"10.0.0.3:1234", "10.0.0.2"},
-      {"10.0.0.4:1234", "10.0.0.1"}, {"10.0.0.5:1234", "10.0.0.2"}};
-
-  auto compare = [](const MockBlockTransport::CallRecord& a,
-                    const MockBlockTransport::CallRecord& b) {
-    if (a.peer != b.peer) return a.peer < b.peer;
-    return a.local_ip < b.local_ip;
-  };
-
-  std::sort(calls.begin(), calls.end(), compare);
-  std::sort(expected.begin(), expected.end(), compare);
-
-  for (size_t i = 0; i < 6; ++i) {
-    EXPECT_EQ(calls[i].peer, expected[i].peer);
-    EXPECT_EQ(calls[i].local_ip, expected[i].local_ip);
-  }
-}
-
 }  // namespace
 }  // namespace transport
 }  // namespace tpu_raiden

tpu_raiden/transport/h2h_strided_bench.cc

@@ -76,7 +76,7 @@
 
 #include "absl/status/status.h"
 #include "absl/types/span.h"
-#include "tpu_raiden/transport/lib/socket_util.h"
+#include "tpu_raiden/transport/lib/socket/util.h"
 
 namespace {
 

tpu_raiden/transport/lib/BUILD

@@ -29,14 +29,13 @@ cc_library(
     ],
     features = ["-use_header_modules"],
     deps = [
-        ":socket_util",
         "//tpu_raiden/core:status_macros",
+        "//tpu_raiden/transport/lib/conn:pool",
+        "//tpu_raiden/transport/lib/socket:util",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/cleanup",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
     ],
@@ -105,22 +104,3 @@ cc_test(
         "@com_google_googletest//:gtest_main",
     ],
 )
-
-cc_library(
-    name = "socket_util",
-    srcs = ["socket_util.cc"],
-    hdrs = ["socket_util.h"],
-    copts = [
-        "-fno-strict-aliasing",
-        "-fexceptions",
-    ],
-    features = ["-use_header_modules"],
-    deps = [
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/types:span",
-    ],
-)

tpu_raiden/transport/lib/conn/BUILD

@@ -0,0 +1,53 @@
+# Copyright 2026 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test")
+
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "pool",
+    srcs = ["pool.cc"],
+    hdrs = ["pool.h"],
+    copts = [
+        "-fno-strict-aliasing",
+        "-fexceptions",
+    ],
+    features = ["-use_header_modules"],
+    deps = [
+        "//tpu_raiden/transport/lib/socket:util",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@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/strings:string_view",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
+
+cc_test(
+    name = "pool_test",
+    srcs = ["pool_test.cc"],
+    copts = [
+        "-fno-strict-aliasing",
+        "-fexceptions",
+    ],
+    features = ["-use_header_modules"],
+    deps = [
+        ":pool",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tpu_raiden/transport/lib/conn/pool.cc

@@ -0,0 +1,103 @@
+// Copyright 2026 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tpu_raiden/transport/lib/conn/pool.h"
+
+#include <sys/poll.h>
+#include <sys/socket.h>
+
+#include <vector>
+
+#include "absl/base/optimization.h"
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "tpu_raiden/transport/lib/socket/util.h"
+
+namespace tpu_raiden::transport::lib {
+
+namespace {
+bool HasReadableData(const int fd) {
+  struct pollfd pfd = {.fd = fd, .events = POLLIN};
+  return poll(&pfd, /*nfds=*/1, /*timeout=*/0) > 0;
+}
+
+void CloseSocket(const int fd) {
+  DCHECK_GE(fd, 0);
+  ::shutdown(fd, SHUT_RDWR);
+  ::close(fd);
+}
+}  // namespace
+
+absl::StatusOr<int> ConnPool::Borrow(absl::string_view peer,
+                                     absl::string_view local_ip) {
+  const Key key = GenPoolKey(peer, local_ip);
+  {
+    absl::MutexLock lock(mu_);
+    if (stop_) {
+      return absl::FailedPreconditionError("ConnPool is closed.");
+    }
+    auto it = pool_.find(key);
+    if ABSL_PREDICT_TRUE (it != pool_.end()) {
+      Fds& fds = it->second;
+      while (!fds.empty()) {
+        const int fd = fds.back();
+        fds.pop_back();
+
+        if (HasReadableData(fd)) {
+          CloseSocket(fd);
+          continue;
+        }
+        return fd;
+      }
+    }
+  }
+  return ConnectToPeer(peer, local_ip);
+}
+
+void ConnPool::Return(bool ok, int fd, absl::string_view peer,
+                      absl::string_view local_ip) {
+  if ABSL_PREDICT_FALSE (fd < 0) {
+    return;
+  }
+
+  DCHECK_GE(fd, 0);
+  if ABSL_PREDICT_FALSE (!ok) {
+    CloseSocket(fd);
+    return;
+  }
+
+  absl::MutexLock lock(mu_);
+  if ABSL_PREDICT_FALSE (stop_) {
+    CloseSocket(fd);
+  } else {
+    const Key key = GenPoolKey(peer, local_ip);
+    pool_[key].push_back(fd);
+  }
+}
+
+void ConnPool::Close() {
+  absl::MutexLock lock(mu_);
+  stop_ = true;
+  for (auto& [_, fds] : pool_) {
+    for (const int fd : fds) {
+      CloseSocket(fd);
+    }
+  }
+  pool_.clear();
+}
+
+}  // namespace tpu_raiden::transport::lib