From 57644b8d6666fe885e80712d0be4aa0c8d99a64c Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 12 Feb 2026 11:06:44 -0800 Subject: [PATCH] Improve tsl::net utilities PiperOrigin-RevId: 869297974 --- tsl/platform/BUILD | 5 ++- tsl/platform/net.h | 18 ++++++++ tsl/platform/net_test.cc | 95 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 1 deletion(-) diff --git a/tsl/platform/BUILD b/tsl/platform/BUILD index d03395f32..83b7a4c84 100644 --- a/tsl/platform/BUILD +++ b/tsl/platform/BUILD @@ -668,7 +668,10 @@ filegroup( cc_library( name = "net", textual_hdrs = ["net.h"], - deps = tf_windows_aware_platform_deps("net") + ["@com_google_absl//absl/base:core_headers"], + deps = tf_windows_aware_platform_deps("net") + [ + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/strings", + ], ) cc_library( diff --git a/tsl/platform/net.h b/tsl/platform/net.h index 7ae03696f..e42ec6b91 100644 --- a/tsl/platform/net.h +++ b/tsl/platform/net.h @@ -16,10 +16,28 @@ limitations under the License. #ifndef TENSORFLOW_TSL_PLATFORM_NET_H_ #define TENSORFLOW_TSL_PLATFORM_NET_H_ +#include + #include "absl/base/macros.h" +#include "absl/strings/str_cat.h" namespace tsl { namespace net { +// Checks whether the given port is available for binding to a TCP or UDP +// socket. If the port is available, returns true. Otherwise, returns false. If +// error is not null, sets error to a string describing the error. +bool IsPortAvailable(int* port, bool is_tcp, std::string* error); + +inline bool IsPortAvailable(int port, bool is_tcp, std::string* error) { + if (port <= 0) { + if (error != nullptr) { + *error = + absl::StrCat("Invalid port number: ", port, ". Port must be > 0."); + } + return false; + } + return IsPortAvailable(&port, is_tcp, error); +} // Return a port number that is not currently bound to any TCP or UDP port. // On success returns the assigned port number. Otherwise returns -1. int PickUnusedPort(); diff --git a/tsl/platform/net_test.cc b/tsl/platform/net_test.cc index 138a9e83a..25a82b62e 100644 --- a/tsl/platform/net_test.cc +++ b/tsl/platform/net_test.cc @@ -15,11 +15,88 @@ limitations under the License. #include "tsl/platform/net.h" +#if !defined(PLATFORM_WINDOWS) +#include +#include +#include +#include +#include +#endif // PLATFORM_WINDOWS + #include "xla/tsl/platform/logging.h" #include "xla/tsl/platform/test.h" namespace tsl { namespace net { +#if !defined(PLATFORM_WINDOWS) +namespace { + +bool CanBindAndListen(int port, int family, int socktype) { + int protocol = (socktype == SOCK_STREAM) ? IPPROTO_TCP : 0; + int sock = socket(family, socktype, protocol); + if (sock < 0) { + // If we can't create a socket for this family, maybe it's not supported. + if (errno == EAFNOSUPPORT || errno == EPROTONOSUPPORT) { + LOG(INFO) << "Skipping bind test for family " << family + << ": socket() failed with errno=" << errno; + return true; + } + LOG(ERROR) << "socket(" << family << ", " << socktype << ", " << protocol + << ") failed: " << strerror(errno); + return false; + } + + struct sockaddr_storage addr_storage; + socklen_t addr_len; + if (family == AF_INET6) { + struct sockaddr_in6* addr = + reinterpret_cast(&addr_storage); + addr_len = sizeof(*addr); + memset(addr, 0, addr_len); + addr->sin6_family = AF_INET6; + addr->sin6_addr = in6addr_any; + addr->sin6_port = htons(port); + } else { + struct sockaddr_in* addr = + reinterpret_cast(&addr_storage); + addr_len = sizeof(*addr); + memset(addr, 0, addr_len); + addr->sin_family = AF_INET; + addr->sin_addr.s_addr = INADDR_ANY; + addr->sin_port = htons(port); + } + + // Use SO_REUSEADDR like IsPortAvailable does. + int one = 1; + if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) < 0) { + LOG(ERROR) << "setsockopt(SO_REUSEADDR) failed: " << strerror(errno); + close(sock); + return false; + } + + if (bind(sock, reinterpret_cast(&addr_storage), addr_len) != + 0) { + LOG(ERROR) << "bind(" << port << ") failed for family " << family << ": " + << strerror(errno); + close(sock); + return false; + } + + if (socktype == SOCK_STREAM) { + if (listen(sock, 1) != 0) { + LOG(ERROR) << "listen(" << port << ") failed for family " << family + << ": " << strerror(errno); + close(sock); + return false; + } + } + + close(sock); + return true; +} + +} // namespace +#endif // PLATFORM_WINDOWS TEST(Net, PickUnusedPortOrDie) { int port0 = PickUnusedPortOrDie(); @@ -29,8 +106,22 @@ TEST(Net, PickUnusedPortOrDie) { CHECK_GE(port1, 0); CHECK_LT(port1, 65536); CHECK_NE(port0, port1); + RecycleUnusedPort(port0); + RecycleUnusedPort(port1); } +#if !defined(PLATFORM_WINDOWS) +TEST(Net, PickedPortIsBindable) { + int port = PickUnusedPortOrDie(); + ASSERT_GT(port, 0); + EXPECT_TRUE(CanBindAndListen(port, AF_INET, SOCK_STREAM)); + EXPECT_TRUE(CanBindAndListen(port, AF_INET6, SOCK_STREAM)); + EXPECT_TRUE(CanBindAndListen(port, AF_INET, SOCK_DGRAM)); + EXPECT_TRUE(CanBindAndListen(port, AF_INET6, SOCK_DGRAM)); + RecycleUnusedPort(port); +} +#endif // PLATFORM_WINDOWS + TEST(Net, RecycleUnusedPort) { for (int i = 0; i < 1000; ++i) { int port0 = PickUnusedPortOrDie(); @@ -49,5 +140,9 @@ TEST(Net, RecycleUnusedPortTwiceShallFail) { EXPECT_DEATH(RecycleUnusedPort(port0), ""); } +TEST(Net, RecycleUnusedPortNegativeShallFail) { + EXPECT_DEATH(RecycleUnusedPort(-1), ""); +} + } // namespace net } // namespace tsl