diff --git a/tpu_raiden/transport/BUILD b/tpu_raiden/transport/BUILD index 7a846034..2092abaa 100644 --- a/tpu_raiden/transport/BUILD +++ b/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", ], diff --git a/tpu_raiden/transport/block_transport.cc b/tpu_raiden/transport/block_transport.cc index 3c126be5..1c059641 100644 --- a/tpu_raiden/transport/block_transport.cc +++ b/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& 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& explicit_dst_ptrs, std::vector& 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(); diff --git a/tpu_raiden/transport/block_transport_test.cc b/tpu_raiden/transport/block_transport_test.cc index b79ba330..aa22b08e 100644 --- a/tpu_raiden/transport/block_transport_test.cc +++ b/tpu_raiden/transport/block_transport_test.cc @@ -14,7 +14,6 @@ #include "tpu_raiden/transport/block_transport.h" -#include #include // NOLINT #include #include @@ -26,7 +25,6 @@ #include #include -#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& local_ips) - : BlockTransport(delegate, local_port, local_ips) {} - - absl::StatusOr 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 acquire_calls() { - absl::MutexLock lock(mock_mu_); - return acquire_calls_; - } - - private: - absl::Mutex mock_mu_; - std::vector acquire_calls_ ABSL_GUARDED_BY(mock_mu_); -}; - -TEST(BlockTransportTest, RoundRobinDistribution) { - constexpr size_t kSliceSize = 16; - MockDelegate delegate(kSliceSize); - - std::vector local_ips = {"10.0.0.1", "10.0.0.2"}; - std::vector peers = {"10.0.0.3:1234", "10.0.0.4:1234", - "10.0.0.5:1234"}; - - MockBlockTransport transport(&delegate, 0, local_ips); - - std::vector 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 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 diff --git a/tpu_raiden/transport/h2h_strided_bench.cc b/tpu_raiden/transport/h2h_strided_bench.cc index 47dbc870..f2898544 100644 --- a/tpu_raiden/transport/h2h_strided_bench.cc +++ b/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 { diff --git a/tpu_raiden/transport/lib/BUILD b/tpu_raiden/transport/lib/BUILD index 090916c7..62a351c2 100644 --- a/tpu_raiden/transport/lib/BUILD +++ b/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", - ], -) diff --git a/tpu_raiden/transport/lib/conn/BUILD b/tpu_raiden/transport/lib/conn/BUILD new file mode 100644 index 00000000..aea49246 --- /dev/null +++ b/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", + ], +) diff --git a/tpu_raiden/transport/lib/conn/pool.cc b/tpu_raiden/transport/lib/conn/pool.cc new file mode 100644 index 00000000..20e2b569 --- /dev/null +++ b/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 +#include + +#include + +#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 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 diff --git a/tpu_raiden/transport/lib/conn/pool.h b/tpu_raiden/transport/lib/conn/pool.h new file mode 100644 index 00000000..65349d45 --- /dev/null +++ b/tpu_raiden/transport/lib/conn/pool.h @@ -0,0 +1,73 @@ +// 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. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TRANSPORT_LIB_CONN_POOL_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TRANSPORT_LIB_CONN_POOL_H_ + +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/flat_hash_map.h" +#include "absl/log/check.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" + +namespace tpu_raiden::transport::lib { + +// This class manages a pool of TCP connections between multiple pairs of +// local and remote peers. +// It is thread-safe. +class ConnPool { + public: + // Constructor. + ConnPool() : stop_(false) {} + + // Destructor. + ~ConnPool() { DCHECK(stop_ && pool_.empty()); } + + // Closes all connections. + void Close(); + + // Borrows a connection from the pool. If no connection is available, creates + // a new one. Returns the socket descriptor of the connection if successful. + // Otherwise returns an error status. + absl::StatusOr Borrow(absl::string_view peer, + absl::string_view local_ip = ""); + + // Returns a connection to the pool if ok is true. Otherwise, closes the + // connection. + void Return(bool ok, int fd, absl::string_view peer, + absl::string_view local_ip = ""); + + private: + using Key = std::string; + using Fds = std::vector; + + // Generates a pool key from the local/peer ip address pair. + static Key GenPoolKey(absl::string_view peer, absl::string_view local_ip) { + return absl::StrCat(local_ip, "->", peer); + } + + private: + absl::Mutex mu_; + bool stop_ ABSL_GUARDED_BY(mu_); + absl::flat_hash_map pool_ ABSL_GUARDED_BY(mu_); +}; + +} // namespace tpu_raiden::transport::lib + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TRANSPORT_LIB_CONN_POOL_H_ diff --git a/tpu_raiden/transport/lib/conn/pool_test.cc b/tpu_raiden/transport/lib/conn/pool_test.cc new file mode 100644 index 00000000..293d3eba --- /dev/null +++ b/tpu_raiden/transport/lib/conn/pool_test.cc @@ -0,0 +1,80 @@ +// 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 + +#include +#include + +namespace tpu_raiden::transport::lib::testing { +namespace { + +TEST(ConnPoolTest, Basic) { + ConnPool pool; + const std::string peer = "localhost:12345"; + + const auto fd_or = pool.Borrow(peer); + ASSERT_OK(fd_or) << fd_or.status().message(); + + const int fd = fd_or.value(); + pool.Return(/*ok=*/true, fd, peer); + + pool.Close(); +} + +TEST(ConnPoolTest, MultiIpPoolingIsolation) { + const std::string peer = "127.0.0.1:12345"; + ConnPool pool; + + // 1. Borrow connection with local_ip = "127.0.0.1" + const auto fd1_or = pool.Borrow(peer, "127.0.0.1"); + ASSERT_OK(fd1_or) << fd1_or.status().message(); + const int fd1 = fd1_or.value(); + + // Return it. It should be pooled under "127.0.0.1->peer". + pool.Return(/*ok=*/true, fd1, peer, "127.0.0.1"); + + // 2. Borrow connection with local_ip = "127.0.0.2" + // This should NOT reuse fd1 because it's a different local IP. + const auto fd2_or = pool.Borrow(peer, "127.0.0.2"); + ASSERT_OK(fd2_or) << fd2_or.status().message(); + const int fd2 = fd2_or.value(); + EXPECT_NE(fd1, fd2); + + // Return it. It should be pooled under "127.0.0.2->peer". + pool.Return(/*ok=*/true, fd2, peer, "127.0.0.2"); + + // 3. Borrow connection with local_ip = "127.0.0.1" again. + // This SHOULD reuse fd1. + const auto fd3_or = pool.Borrow(peer, "127.0.0.1"); + ASSERT_OK(fd3_or) << fd3_or.status().message(); + const int fd3 = fd3_or.value(); + EXPECT_EQ(fd3, fd1); + pool.Return(/*ok=*/true, fd3, peer, "127.0.0.1"); + + // 4. Borrow connection with local_ip = "127.0.0.2" again. + // This SHOULD reuse fd2. + const auto fd4_or = pool.Borrow(peer, "127.0.0.2"); + ASSERT_OK(fd4_or) << fd4_or.status().message(); + const int fd4 = fd4_or.value(); + EXPECT_EQ(fd4, fd2); + pool.Return(/*ok=*/true, fd4, peer, "127.0.0.2"); + + pool.Close(); +} + +} // namespace +} // namespace tpu_raiden::transport::lib::testing diff --git a/tpu_raiden/transport/lib/raw_buffer_transport.cc b/tpu_raiden/transport/lib/raw_buffer_transport.cc index a4b6adbb..1e8d0e42 100644 --- a/tpu_raiden/transport/lib/raw_buffer_transport.cc +++ b/tpu_raiden/transport/lib/raw_buffer_transport.cc @@ -31,7 +31,6 @@ #include #include #include // NOLINT -#include #include #include "absl/cleanup/cleanup.h" @@ -41,7 +40,8 @@ #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "tpu_raiden/core/status_macros.h" -#include "tpu_raiden/transport/lib/socket_util.h" +#include "tpu_raiden/transport/lib/conn/pool.h" +#include "tpu_raiden/transport/lib/socket/util.h" namespace tpu_raiden::transport::lib { @@ -123,7 +123,7 @@ RawBufferTransport::RawBufferTransport( RawBufferTransport::~RawBufferTransport() { stopping_ = true; - ClosePooledConnections(); + conn_pool_.Close(); if (server_fd_ >= 0) { shutdown(server_fd_, SHUT_RDWR); close(server_fd_); @@ -146,69 +146,6 @@ RawBufferTransport::~RawBufferTransport() { } } -static std::string GetPoolKey(absl::string_view peer, - absl::string_view local_ip) { - return absl::StrCat(local_ip, "->", peer); -} - -absl::StatusOr RawBufferTransport::BorrowConnection( - absl::string_view peer, absl::string_view local_ip) { - { - absl::MutexLock lock(pool_mu_); - std::string key = GetPoolKey(peer, local_ip); - if (auto it = conn_pool_.find(key); it != conn_pool_.end()) { - while (!it->second.empty()) { - int fd = it->second.back(); - it->second.pop_back(); - - struct pollfd pfd; - pfd.fd = fd; - pfd.events = POLLIN; - if (poll(&pfd, 1, 0) > 0) { - shutdown(fd, SHUT_RDWR); - close(fd); - continue; - } - return fd; - } - } - } - return ConnectToPeer(peer, local_ip); -} - -void RawBufferTransport::ReturnConnection(bool ok, int fd, - absl::string_view peer, - absl::string_view local_ip) { - if (fd < 0) return; - - if (!ok) { - shutdown(fd, SHUT_RDWR); - close(fd); - return; - } - - absl::MutexLock lock( pool_mu_ ); - if (stopping_) { - shutdown(fd, SHUT_RDWR); - close(fd); - return; - } - - const std::string key = GetPoolKey(peer, local_ip); - conn_pool_[key].push_back(fd); -} - -void RawBufferTransport::ClosePooledConnections() { - absl::MutexLock lock( pool_mu_ ); - for (auto& entry : conn_pool_) { - for (int fd : entry.second) { - shutdown(fd, SHUT_RDWR); - close(fd); - } - } - conn_pool_.clear(); -} - absl::Status RawBufferTransport::ProcessSingleRequest(int client_fd) { PacketHeader header = {}; RETURN_IF_ERROR(ReadExact(client_fd, &header, sizeof(header))); @@ -337,10 +274,10 @@ absl::Status RawBufferTransport::PullBuffer( ", Size: ", size_bytes, ", Shard Host Size: ", host_size)); } - ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer)); + ASSIGN_OR_RETURN(const int fd, conn_pool_.Borrow(peer)); bool ok_to_pool = false; auto fd_cleaner = - absl::MakeCleanup([&] { ReturnConnection(ok_to_pool, fd, peer); }); + absl::MakeCleanup([&] { conn_pool_.Return(ok_to_pool, fd, peer); }); PacketHeader header = {}; header.op = 3; @@ -367,10 +304,10 @@ absl::Status RawBufferTransport::PushBuffer( "Destination peer address cannot be empty"); } - ASSIGN_OR_RETURN(const int fd, BorrowConnection(peer)); + ASSIGN_OR_RETURN(const int fd, conn_pool_.Borrow(peer)); bool ok_to_pool = false; auto fd_cleaner = - absl::MakeCleanup([&] { ReturnConnection(ok_to_pool, fd, peer); }); + absl::MakeCleanup([&] { conn_pool_.Return(ok_to_pool, fd, peer); }); PacketHeader header = {}; header.op = 5; diff --git a/tpu_raiden/transport/lib/raw_buffer_transport.h b/tpu_raiden/transport/lib/raw_buffer_transport.h index 817c8b42..6b565fd3 100644 --- a/tpu_raiden/transport/lib/raw_buffer_transport.h +++ b/tpu_raiden/transport/lib/raw_buffer_transport.h @@ -25,11 +25,10 @@ #include #include "absl/base/thread_annotations.h" -#include "absl/container/flat_hash_map.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/conn/pool.h" namespace tpu_raiden::transport::lib { @@ -88,12 +87,6 @@ class RawBufferTransport { const std::string& bound_ip() const { return bound_ip_; } protected: - virtual absl::StatusOr BorrowConnection(absl::string_view peer, - absl::string_view local_ip = ""); - virtual void ReturnConnection(bool ok, int fd, absl::string_view peer, - absl::string_view local_ip = ""); - void ClosePooledConnections(); - virtual absl::Status ProcessSingleRequest(int client_fd); virtual absl::Status HandleCustomRequest(int client_fd, const PacketHeader& header); @@ -111,9 +104,7 @@ class RawBufferTransport { absl::Mutex mu_; std::vector active_client_fds_ ABSL_GUARDED_BY(mu_); - absl::Mutex pool_mu_; - absl::flat_hash_map> conn_pool_ - ABSL_GUARDED_BY(pool_mu_); + ConnPool conn_pool_; std::thread listener_thread_; std::vector worker_threads_; diff --git a/tpu_raiden/transport/lib/raw_buffer_transport_test.cc b/tpu_raiden/transport/lib/raw_buffer_transport_test.cc index 89eab14f..4d88e3ee 100644 --- a/tpu_raiden/transport/lib/raw_buffer_transport_test.cc +++ b/tpu_raiden/transport/lib/raw_buffer_transport_test.cc @@ -204,59 +204,5 @@ TEST(RawBufferTransportTest, RejectsOutOfBounds) { EXPECT_FALSE(pull_res.ok()) << pull_res.message(); } -class TestRawBufferTransport : public RawBufferTransport { - public: - using RawBufferTransport::BorrowConnection; - using RawBufferTransport::RawBufferTransport; - using RawBufferTransport::ReturnConnection; -}; - -TEST(RawBufferTransportTest, MultiIpPoolingIsolation) { - // Set up src/dst buffers. - constexpr size_t size = 1024; - RawMockDelegate src(size); - RawMockDelegate dst(size); - - // Create two transports. - RawBufferTransport src_transport(&src, 0); - TestRawBufferTransport dst_transport(&dst, 0); - std::this_thread::sleep_for(std::chrono::milliseconds(50)); - - // 1. Borrow connection with local_ip = "127.0.0.1" - const std::string src_addr = GetIpPort(src_transport); - const auto fd1_or = dst_transport.BorrowConnection(src_addr, "127.0.0.1"); - ASSERT_OK(fd1_or) << fd1_or.status().message(); - const int fd1 = fd1_or.value(); - - // Return it. It should be pooled under "127.0.0.1->peer1". - dst_transport.ReturnConnection(/*ok=*/true, fd1, src_addr, "127.0.0.1"); - - // 2. Borrow connection with local_ip = "127.0.0.2" - // This should NOT reuse fd1 because it's a different local IP. - const auto fd2_or = dst_transport.BorrowConnection(src_addr, "127.0.0.2"); - ASSERT_OK(fd2_or) << fd2_or.status().message(); - const int fd2 = fd2_or.value(); - EXPECT_NE(fd1, fd2); - - // Return it. It should be pooled under "127.0.0.2->peer1". - dst_transport.ReturnConnection(/*ok=*/true, fd2, src_addr, "127.0.0.2"); - - // 3. Borrow connection with local_ip = "127.0.0.1" again. - // This SHOULD reuse fd1. - const auto fd3_or = dst_transport.BorrowConnection(src_addr, "127.0.0.1"); - ASSERT_OK(fd3_or) << fd3_or.status().message(); - const int fd3 = fd3_or.value(); - EXPECT_EQ(fd1, fd3); - dst_transport.ReturnConnection(/*ok=*/true, fd3, src_addr, "127.0.0.1"); - - // 4. Borrow connection with local_ip = "127.0.0.2" again. - // This SHOULD reuse fd2. - const auto fd4_or = dst_transport.BorrowConnection(src_addr, "127.0.0.2"); - ASSERT_OK(fd4_or) << fd4_or.status().message(); - const int fd4 = fd4_or.value(); - EXPECT_EQ(fd2, fd4); - dst_transport.ReturnConnection(/*ok=*/true, fd4, src_addr, "127.0.0.2"); -} - } // namespace } // namespace tpu_raiden::transport::lib::testing diff --git a/tpu_raiden/transport/lib/socket/BUILD b/tpu_raiden/transport/lib/socket/BUILD new file mode 100644 index 00000000..ea525658 --- /dev/null +++ b/tpu_raiden/transport/lib/socket/BUILD @@ -0,0 +1,36 @@ +# 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") + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "util", + srcs = ["util.cc"], + hdrs = ["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", + ], +) diff --git a/tpu_raiden/transport/lib/socket_util.cc b/tpu_raiden/transport/lib/socket/util.cc similarity index 99% rename from tpu_raiden/transport/lib/socket_util.cc rename to tpu_raiden/transport/lib/socket/util.cc index e2b5ab4a..6a1ff437 100644 --- a/tpu_raiden/transport/lib/socket_util.cc +++ b/tpu_raiden/transport/lib/socket/util.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "tpu_raiden/transport/lib/socket_util.h" +#include "tpu_raiden/transport/lib/socket/util.h" #include #include diff --git a/tpu_raiden/transport/lib/socket_util.h b/tpu_raiden/transport/lib/socket/util.h similarity index 100% rename from tpu_raiden/transport/lib/socket_util.h rename to tpu_raiden/transport/lib/socket/util.h