Revert "Revert "Deflakey test advanced 9""

src/ray/common/client_connection.cc

@@ -183,6 +183,41 @@ void ServerConnection::ReadBufferAsync(
         });
   }
 }
+void ServerConnection::AsyncWaitTerminated(std::function<void()> callback) {
+  // Different platforms have different behavior for sockets in asio
+  // Wait for error and read until the end.
+  if (status_ == ConnectionStatus::RUNNING) {
+    socket_.async_wait(local_stream_socket::wait_type::wait_error,
+                       [this, callback](auto) {
+                         if (status_ != ConnectionStatus::TERMINATED) {
+                           callback();
+                           // Close the connection so it'll cancel all other operations
+                           Close();
+                         }
+                       });
+    status_ = ConnectionStatus::TERMINATING;
+  }
+
+  if (status_ == ConnectionStatus::TERMINATING) {
+    // This is only used for buffer. Never read. So race condition is
+    // not an issue here.
+    static char buffer[1024];
+    auto read_cb = [this, cb = callback](boost::system::error_code ec,
+                                         std::size_t /* length */) mutable {
+      if (ec) {
+        if (status_ != ConnectionStatus::TERMINATED) {
+          cb();
+          // Close the connection so it'll cancel all other operations
+          Close();
+        }
+      } else {
+        AsyncWaitTerminated(std::move(cb));
+      }
+    };
+    socket_.async_read_some(boost::asio::mutable_buffer(buffer, sizeof(buffer)),
+                            std::move(read_cb));
+  }
+}
 
 ray::Status ServerConnection::WriteMessage(int64_t type,
                                            int64_t length,

src/ray/common/client_connection.h

@@ -111,6 +111,7 @@ class ServerConnection : public std::enable_shared_from_this<ServerConnection> {
   void Close() {
     boost::system::error_code ec;
     socket_.close(ec);
+    status_ = ConnectionStatus::TERMINATED;
   }
 
   /// Get the native handle of the socket.
@@ -125,11 +126,7 @@ class ServerConnection : public std::enable_shared_from_this<ServerConnection> {
 
   std::string DebugString() const;
 
-  void AsyncWaitTerminated(std::function<void()> callback) {
-    // Async wait until the connection is disconnected.
-    socket_.async_wait(local_stream_socket::wait_type::wait_error,
-                       [callback = std::move(callback)](auto) { callback(); });
-  }
+  void AsyncWaitTerminated(std::function<void()> callback);
 
  protected:
   /// A private constructor for a server connection.
@@ -144,6 +141,11 @@ class ServerConnection : public std::enable_shared_from_this<ServerConnection> {
     std::function<void(const ray::Status &)> handler;
   };
 
+  enum struct ConnectionStatus { RUNNING = 0, TERMINATING, TERMINATED };
+
+  /// Whether the connection is terminating.
+  ConnectionStatus status_ = ConnectionStatus::RUNNING;
+
   /// The socket connection to the server.
   local_stream_socket socket_;
 

src/ray/raylet/node_manager.cc

@@ -1561,23 +1561,10 @@ void NodeManager::DisconnectClient(const std::shared_ptr<ClientConnection> &clie
 
   local_task_manager_->ClearWorkerBacklog(worker->WorkerId());
   cluster_task_manager_->CancelTaskForOwner(worker->GetAssignedTaskId());
-#ifdef _WIN32
-  // On Windows, when the worker is killed, client async wait won't get notified
-  // somehow.
-  RAY_CHECK_OK(
-      gcs_client_->Workers().AsyncReportWorkerFailure(worker_failure_data_ptr, nullptr));
-  client->Close();
-#else
-  // ReportWorkerFailure should happen after the worker exit completely.
-  // A better way is to monitor the pid exit. But that needs Process.h
-  // support async operation.
-  // Here we monitor the socket to achieve similar result.
-  // When the worker exited, the pid will be disconnected (local stream socket).
   client->AsyncWaitTerminated([client, worker_failure_data_ptr, this] {
     RAY_CHECK_OK(gcs_client_->Workers().AsyncReportWorkerFailure(worker_failure_data_ptr,
                                                                  nullptr));
   });
-#endif
   // TODO(rkn): Tell the object manager that this client has disconnected so
   // that it can clean up the wait requests for this client. Currently I think
   // these can be leaked.