diff --git a/tpu_sync/common/BUILD b/tpu_sync/common/BUILD index 4283f121c..b04d05871 100644 --- a/tpu_sync/common/BUILD +++ b/tpu_sync/common/BUILD @@ -48,6 +48,28 @@ cc_test( ], ) +cc_library( + name = "accept_backoff", + srcs = ["accept_backoff.cc"], + hdrs = ["accept_backoff.h"], + visibility = ["//visibility:public"], + deps = [ + "@com_google_absl//absl/log", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/time", + ], +) + +cc_test( + name = "accept_backoff_test", + srcs = ["accept_backoff_test.cc"], + deps = [ + ":accept_backoff", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "detached_thread_group", srcs = ["detached_thread_group.cc"], diff --git a/tpu_sync/common/accept_backoff.cc b/tpu_sync/common/accept_backoff.cc new file mode 100644 index 000000000..f57d385f6 --- /dev/null +++ b/tpu_sync/common/accept_backoff.cc @@ -0,0 +1,91 @@ +// 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_sync/common/accept_backoff.h" + +#include +#include +#include + +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" + +namespace tpu_raiden { +namespace { + +constexpr absl::Duration kInitialBackoff = absl::Milliseconds(10); +constexpr absl::Duration kMaxBackoff = absl::Milliseconds(500); + +} // namespace + +AcceptErrorAction ClassifyAcceptError(int err) { + switch (err) { + // A signal arrived, or the peer went away before we accepted it. Both are + // routine; backing off would penalise healthy peers that reconnect often. + case EINTR: + case ECONNABORTED: + case EPROTO: + return AcceptErrorAction::kRetry; + // Out of descriptors, buffers or memory. + case EMFILE: + case ENFILE: + case ENOBUFS: + case ENOMEM: + return AcceptErrorAction::kBackoff; + // The descriptor is closed, or is no longer a listening socket. On Linux + // EINVAL is what shutdown() on a listening socket reports, which is how + // teardown unblocks accept(). + case EBADF: + case EINVAL: + case ENOTSOCK: + case EOPNOTSUPP: + return AcceptErrorAction::kStop; + default: + // Unknown, so assume it can repeat. Backing off is the safe default: + // the worst case is a slower retry, not a spin. + return AcceptErrorAction::kBackoff; + } +} + +absl::Duration NextAcceptBackoff(absl::Duration current) { + if (current <= absl::ZeroDuration()) return kInitialBackoff; + return std::min(2 * current, kMaxBackoff); +} + +AcceptBackoff::AcceptBackoff(absl::string_view name) : name_(name) {} + +bool AcceptBackoff::OnError(int err) { + switch (ClassifyAcceptError(err)) { + case AcceptErrorAction::kRetry: + return true; + case AcceptErrorAction::kBackoff: + delay_ = NextAcceptBackoff(delay_); + LOG_EVERY_N_SEC(WARNING, 1) + << name_ << " accept() failed: " << std::strerror(err) + << "; retrying in " << delay_; + absl::SleepFor(delay_); + return true; + case AcceptErrorAction::kStop: + LOG(ERROR) << name_ << " stopped accepting connections: " + << std::strerror(err); + return false; + } + return false; +} + +void AcceptBackoff::OnSuccess() { delay_ = absl::ZeroDuration(); } + +} // namespace tpu_raiden diff --git a/tpu_sync/common/accept_backoff.h b/tpu_sync/common/accept_backoff.h new file mode 100644 index 000000000..e5fef6b6c --- /dev/null +++ b/tpu_sync/common/accept_backoff.h @@ -0,0 +1,72 @@ +// 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_SYNC_COMMON_ACCEPT_BACKOFF_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_COMMON_ACCEPT_BACKOFF_H_ + +#include + +#include "absl/strings/string_view.h" +#include "absl/time/time.h" + +namespace tpu_raiden { + +// What an accept loop should do after accept(), or the poll() that guards it, +// reports an error. +enum class AcceptErrorAction { + // Routine, and the listener is healthy: retry straight away. + kRetry, + // The process is out of a resource. The connection that could not be + // accepted stays queued, so an immediate retry fails the same way and the + // loop spins at 100% CPU, starving the very threads that would free the + // resource. Sleep first. + kBackoff, + // The listening descriptor is not usable any more: leave the loop. + kStop, +}; + +AcceptErrorAction ClassifyAcceptError(int err); + +// Delay to wait after |current| (pass zero for the first backoff). Capped, so +// a stop request is never delayed by more than one step. +absl::Duration NextAcceptBackoff(absl::Duration current); + +// Backoff state for one accept loop. Not thread-safe: each loop owns one on +// its own stack. +class AcceptBackoff { + public: + // |name| identifies the listener in log messages. + explicit AcceptBackoff(absl::string_view name); + + // Handles a failed accept(). Sleeps for errors worth retrying slowly, and + // returns false if the caller must leave its accept loop. + // + // Callers check their own stopping flag before calling this, so a kStop + // error here always means the descriptor died unexpectedly and is logged as + // an error rather than as routine shutdown. + bool OnError(int err); + + // Resets the delay. Call after a successful accept(). + void OnSuccess(); + + absl::Duration delay_for_testing() const { return delay_; } + + private: + const std::string name_; + absl::Duration delay_ = absl::ZeroDuration(); +}; + +} // namespace tpu_raiden + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_COMMON_ACCEPT_BACKOFF_H_ diff --git a/tpu_sync/common/accept_backoff_test.cc b/tpu_sync/common/accept_backoff_test.cc new file mode 100644 index 000000000..007a70e75 --- /dev/null +++ b/tpu_sync/common/accept_backoff_test.cc @@ -0,0 +1,90 @@ +// 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_sync/common/accept_backoff.h" + +#include + +#include +#include "absl/time/time.h" + +namespace tpu_raiden { +namespace { + +TEST(ClassifyAcceptErrorTest, RoutineErrorsRetryImmediately) { + EXPECT_EQ(ClassifyAcceptError(EINTR), AcceptErrorAction::kRetry); + EXPECT_EQ(ClassifyAcceptError(ECONNABORTED), AcceptErrorAction::kRetry); +} + +TEST(ClassifyAcceptErrorTest, ResourceExhaustionBacksOff) { + EXPECT_EQ(ClassifyAcceptError(EMFILE), AcceptErrorAction::kBackoff); + EXPECT_EQ(ClassifyAcceptError(ENFILE), AcceptErrorAction::kBackoff); + EXPECT_EQ(ClassifyAcceptError(ENOBUFS), AcceptErrorAction::kBackoff); + EXPECT_EQ(ClassifyAcceptError(ENOMEM), AcceptErrorAction::kBackoff); +} + +TEST(ClassifyAcceptErrorTest, UnusableDescriptorStops) { + EXPECT_EQ(ClassifyAcceptError(EBADF), AcceptErrorAction::kStop); + EXPECT_EQ(ClassifyAcceptError(ENOTSOCK), AcceptErrorAction::kStop); + // What shutdown() on a listening socket reports, so teardown relies on it. + EXPECT_EQ(ClassifyAcceptError(EINVAL), AcceptErrorAction::kStop); +} + +TEST(ClassifyAcceptErrorTest, UnknownErrorsBackOffRatherThanSpin) { + EXPECT_EQ(ClassifyAcceptError(EIO), AcceptErrorAction::kBackoff); +} + +TEST(NextAcceptBackoffTest, GrowsFromZero) { + const absl::Duration first = NextAcceptBackoff(absl::ZeroDuration()); + EXPECT_GT(first, absl::ZeroDuration()); + EXPECT_GT(NextAcceptBackoff(first), first); +} + +TEST(NextAcceptBackoffTest, CapsLowEnoughNotToDelayTeardown) { + absl::Duration delay = absl::ZeroDuration(); + for (int i = 0; i < 30; ++i) { + delay = NextAcceptBackoff(delay); + } + EXPECT_EQ(delay, NextAcceptBackoff(delay)); + EXPECT_LE(delay, absl::Seconds(1)); +} + +TEST(AcceptBackoffTest, KeepsLoopRunningForRoutineErrors) { + AcceptBackoff backoff("test"); + EXPECT_TRUE(backoff.OnError(EINTR)); + // Routine errors must not slow the loop down. + EXPECT_EQ(backoff.delay_for_testing(), absl::ZeroDuration()); +} + +TEST(AcceptBackoffTest, LeavesLoopWhenDescriptorIsUnusable) { + AcceptBackoff backoff("test"); + EXPECT_FALSE(backoff.OnError(EBADF)); +} + +TEST(AcceptBackoffTest, GrowsDelayWhileFailingAndResetsOnSuccess) { + AcceptBackoff backoff("test"); + + EXPECT_TRUE(backoff.OnError(EMFILE)); + const absl::Duration first = backoff.delay_for_testing(); + EXPECT_GT(first, absl::ZeroDuration()); + + EXPECT_TRUE(backoff.OnError(EMFILE)); + EXPECT_GT(backoff.delay_for_testing(), first); + + backoff.OnSuccess(); + EXPECT_EQ(backoff.delay_for_testing(), absl::ZeroDuration()); +} + +} // namespace +} // namespace tpu_raiden diff --git a/tpu_sync/kv_cache/BUILD b/tpu_sync/kv_cache/BUILD index 88cfeaca3..b88cf5d19 100644 --- a/tpu_sync/kv_cache/BUILD +++ b/tpu_sync/kv_cache/BUILD @@ -833,6 +833,7 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":kv_cache_manager_base", + "//tpu_sync/common:accept_backoff", "//tpu_sync/common:detached_thread_group", "//tpu_sync/rpc:raiden_service_cc_proto", "@com_google_absl//absl/log", diff --git a/tpu_sync/kv_cache/kv_cache_listener.cc b/tpu_sync/kv_cache/kv_cache_listener.cc index ea75971f1..ab28b1c36 100644 --- a/tpu_sync/kv_cache/kv_cache_listener.cc +++ b/tpu_sync/kv_cache/kv_cache_listener.cc @@ -30,6 +30,7 @@ #include "absl/log/log.h" #include "absl/status/status.h" +#include "tpu_sync/common/accept_backoff.h" #include "tpu_sync/kv_cache/kv_cache_manager_base.h" #include "tpu_sync/rpc/raiden_service.pb.h" @@ -137,15 +138,19 @@ KVCacheListener::~KVCacheListener() { } void KVCacheListener::ListenerLoop() { + AcceptBackoff backoff("KVCacheListener"); while (!stopping_) { sockaddr_in6 client_addr{}; socklen_t client_len = sizeof(client_addr); int client_fd = accept( server_fd_, reinterpret_cast(&client_addr), &client_len); if (client_fd < 0) { + const int err = errno; if (stopping_) break; + if (!backoff.OnError(err)) break; continue; } + backoff.OnSuccess(); connection_threads_.Spawn( [this, client_fd] { ConnectionWorker(client_fd); }); diff --git a/tpu_sync/kv_cache/reshard/BUILD b/tpu_sync/kv_cache/reshard/BUILD index 6aad0ba93..fbb6e58ed 100644 --- a/tpu_sync/kv_cache/reshard/BUILD +++ b/tpu_sync/kv_cache/reshard/BUILD @@ -44,6 +44,7 @@ cc_library( copts = RESHARD_COPTS, features = ["-use_header_modules"], deps = [ + "//tpu_sync/common:accept_backoff", "//tpu_sync/common:detached_thread_group", "@com_google_absl//absl/log", "@com_google_absl//absl/status", diff --git a/tpu_sync/kv_cache/reshard/framed_rpc.cc b/tpu_sync/kv_cache/reshard/framed_rpc.cc index 208b68a37..9255e7c26 100644 --- a/tpu_sync/kv_cache/reshard/framed_rpc.cc +++ b/tpu_sync/kv_cache/reshard/framed_rpc.cc @@ -34,6 +34,7 @@ #include "absl/strings/string_view.h" #include "absl/time/clock.h" #include "absl/time/time.h" +#include "tpu_sync/common/accept_backoff.h" namespace tpu_raiden { namespace kv_cache { @@ -252,15 +253,19 @@ void FramedServer::Stop() { } void FramedServer::AcceptLoop() { + AcceptBackoff backoff("FramedServer"); while (!stopping_) { sockaddr_storage client_addr{}; socklen_t client_len = sizeof(client_addr); int client_fd = accept( server_fd_, reinterpret_cast(&client_addr), &client_len); if (client_fd < 0) { + const int err = errno; if (stopping_) break; + if (!backoff.OnError(err)) break; continue; } + backoff.OnSuccess(); if (stopping_) { close(client_fd); break; diff --git a/tpu_sync/transport/lib/BUILD b/tpu_sync/transport/lib/BUILD index c16adcf7b..29d59814b 100644 --- a/tpu_sync/transport/lib/BUILD +++ b/tpu_sync/transport/lib/BUILD @@ -176,6 +176,7 @@ cc_library( ":histogram", ":raw_buffer_transport_delegate", ":transport_adapter", + "//tpu_sync/common:accept_backoff", "//tpu_sync/common:detached_thread_group", "//tpu_sync/core:numa_thread_pool", "//tpu_sync/transport:buffer_push_task", diff --git a/tpu_sync/transport/lib/raw_buffer_transport.cc b/tpu_sync/transport/lib/raw_buffer_transport.cc index 2e69327c9..1f1b50142 100644 --- a/tpu_sync/transport/lib/raw_buffer_transport.cc +++ b/tpu_sync/transport/lib/raw_buffer_transport.cc @@ -51,6 +51,7 @@ #include "absl/synchronization/blocking_counter.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" +#include "tpu_sync/common/accept_backoff.h" #include "tpu_sync/core/numa_thread_pool.h" #include "tpu_sync/transport/buffer_push_task.h" #include "tpu_sync/transport/lib/transport_adapter.h" @@ -469,14 +470,20 @@ RawBufferTransport::RegisterPspPeer(uint32_t client_spi, } void RawBufferTransport::ListenerLoop() { + AcceptBackoff backoff("RawBufferTransport"); while (!stopping_) { DCHECK(IsValidSocket(server_fd_)); struct pollfd pfd; pfd.fd = server_fd_; pfd.events = POLLIN; int ret = poll(&pfd, 1, 50); - if (ret <= 0) { + if (ret == 0) continue; // Idle tick: nothing waiting to be accepted. + if (ret < 0) { + const int err = errno; if (stopping_) break; + // poll() reports the same descriptor and resource errors as accept(), + // so the two share one policy. + if (!backoff.OnError(err)) break; continue; } @@ -485,9 +492,12 @@ void RawBufferTransport::ListenerLoop() { int client_fd = accept( server_fd_, reinterpret_cast(&client_addr), &clilen); if (client_fd < 0) { + const int err = errno; if (stopping_) break; + if (!backoff.OnError(err)) break; continue; } + backoff.OnSuccess(); if (require_psp_tcp_ && !PspEnabled(client_fd)) { close(client_fd); LOG_EVERY_N_SEC(ERROR, 1) diff --git a/tpu_sync/weight_sync/BUILD b/tpu_sync/weight_sync/BUILD index 1a52a9188..208e99bed 100644 --- a/tpu_sync/weight_sync/BUILD +++ b/tpu_sync/weight_sync/BUILD @@ -56,6 +56,7 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":tiling_utils", + "//tpu_sync/common:accept_backoff", "//tpu_sync/common:detached_thread_group", "//tpu_sync/common:trace", "//tpu_sync/core:host_memory_allocator", diff --git a/tpu_sync/weight_sync/weight_synchronizer_listener.cc b/tpu_sync/weight_sync/weight_synchronizer_listener.cc index b33e39a74..2cf807d71 100644 --- a/tpu_sync/weight_sync/weight_synchronizer_listener.cc +++ b/tpu_sync/weight_sync/weight_synchronizer_listener.cc @@ -30,6 +30,7 @@ #include "absl/container/flat_hash_map.h" #include "absl/log/log.h" +#include "tpu_sync/common/accept_backoff.h" #include "tpu_sync/rpc/raiden_service.pb.h" #include "tpu_sync/weight_sync/weight_synchronizer_base.h" @@ -100,6 +101,7 @@ WeightSynchronizerListener::~WeightSynchronizerListener() { } void WeightSynchronizerListener::ListenerLoop() { + AcceptBackoff backoff("WeightSynchronizerListener"); while (!stopping_) { sockaddr_in6 client_addr{}; socklen_t client_len = sizeof(client_addr); @@ -107,9 +109,12 @@ void WeightSynchronizerListener::ListenerLoop() { accept(server_fd_.load(), reinterpret_cast(&client_addr), &client_len); if (client_fd < 0) { + const int err = errno; if (stopping_) break; + if (!backoff.OnError(err)) break; continue; } + backoff.OnSuccess(); connection_threads_.Spawn( [this, client_fd] { ConnectionWorker(client_fd); });