Reset factory eagerly so that WaitForQuiesce() is sure to drop all refs and

pschuh · Google-ML-Automation · commit bb760b047bdb · 2026-01-15T12:11:40.000-08:00
track refs for SocketServer::Connection.

PiperOrigin-RevId: 856761720
diff --git a/xla/python/transfer/BUILD b/xla/python/transfer/BUILD
@@ -231,7 +231,6 @@ cc_library(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
diff --git a/xla/python/transfer/socket-server.cc b/xla/python/transfer/socket-server.cc
@@ -59,15 +59,21 @@ class SocketServer::SocketNetworkState : public SocketFdPacketState {
     connections_->list.push_back(this);
     connection_it_ = --connections_->list.end();
   }
-  explicit SocketNetworkState(std::shared_ptr<PullTable> table,
+  explicit SocketNetworkState(std::shared_ptr<ConnectionList> connections,
+                              std::shared_ptr<PullTable> table,
                               std::shared_ptr<BulkTransportFactory> factory,
                               const SocketAddress& addr)
       : table_(std::move(table)),
         factory_(std::move(factory)),
-        remote_addr_(addr) {
+        remote_addr_(addr),
+        connections_(std::move(connections)) {
+    absl::MutexLock l(connections_->mu);
+    connections_->list.push_back(this);
+    connection_it_ = --connections_->list.end();
   }
 
   ~SocketNetworkState() override {
+    factory_.reset();
     if (connections_) {
       absl::MutexLock l(connections_->mu);
       connections_->list.erase(connection_it_);
@@ -550,8 +556,8 @@ SocketServer::~SocketServer() {
 
 tsl::RCReference<SocketServer::Connection> SocketServer::Connect(
     const SocketAddress& other_addr) {
-  auto* local_ =
-      new SocketNetworkState(pull_table_, bulk_transport_factory_, other_addr);
+  auto* local_ = new SocketNetworkState(connections_, pull_table_,
+                                        bulk_transport_factory_, other_addr);
   local_->StartBulkTransporting();
   local_->IncRef();
   local_->StartConnect();
