Deflakey test advanced 9

python/ray/tests/test_advanced_9.py

@@ -258,7 +258,7 @@ def ready(self):
     run_string_as_driver(script.format(address=call_ray_start_2, val=2))
 
 
-@pytest.mark.skipif(sys.platform != "linux", reason="Only works on linux.")
+@pytest.mark.skipif(sys.platform == "win32", reason="Not valid on win32.")
 def test_gcs_connection_no_leak(ray_start_cluster):
     cluster = ray_start_cluster
     head_node = cluster.add_node()

src/ray/common/client_connection.h

@@ -125,6 +125,11 @@ class ServerConnection : public std::enable_shared_from_this<ServerConnection> {
 
   std::string DebugString() const;
 
+  void AsyncWaitTerminated(std::function<void()> callback) {
+    socket_.async_wait(local_stream_socket::wait_type::wait_error,
+                       [callback = std::move(callback)](auto) { callback(); });
+  }
+
  protected:
   /// A private constructor for a server connection.
   ServerConnection(local_stream_socket &&socket);

src/ray/core_worker/core_worker.cc

@@ -784,8 +784,8 @@ void CoreWorker::Exit(
          detail = std::move(detail),
          creation_task_exception_pb_bytes]() {
           rpc::DrainServerCallExecutor();
-          Disconnect(exit_type, detail, creation_task_exception_pb_bytes);
           KillChildProcs();
+          Disconnect(exit_type, detail, creation_task_exception_pb_bytes);
           Shutdown();
         },
         "CoreWorker.Shutdown");
@@ -829,10 +829,10 @@ void CoreWorker::ForceExit(const rpc::WorkerExitType exit_type,
                            const std::string &detail) {
   RAY_LOG(WARNING) << "Force exit the process. "
                    << " Details: " << detail;
-  Disconnect(exit_type, detail);
-
   KillChildProcs();
 
+  Disconnect(exit_type, detail);
+
   // NOTE(hchen): Use `QuickExit()` to force-exit this process without doing cleanup.
   // `exit()` will destruct static objects in an incorrect order, which will lead to
   // core dumps.

src/ray/gcs/gcs_server/pubsub_handler.cc

@@ -104,7 +104,6 @@ void InternalPubSubHandler::HandleGcsSubscriberCommandBatch(
   if (sender_id.empty()) {
     sender_id = request.subscriber_id();
   }
-
   auto iter = sender_to_subscribers_.find(sender_id);
   if (iter == sender_to_subscribers_.end()) {
     iter = sender_to_subscribers_.insert({sender_id, {}}).first;

src/ray/raylet/node_manager.cc

@@ -1494,9 +1494,6 @@ void NodeManager::DisconnectClient(const std::shared_ptr<ClientConnection> &clie
                                    disconnect_detail,
                                    worker->GetProcess().GetId(),
                                    creation_task_exception);
-  RAY_CHECK_OK(
-      gcs_client_->Workers().AsyncReportWorkerFailure(worker_failure_data_ptr, nullptr));
-
   if (is_worker) {
     const ActorID &actor_id = worker->GetActorId();
     const TaskID &task_id = worker->GetAssignedTaskId();
@@ -1571,9 +1568,10 @@ void NodeManager::DisconnectClient(const std::shared_ptr<ClientConnection> &clie
 
   local_task_manager_->ClearWorkerBacklog(worker->WorkerId());
   cluster_task_manager_->CancelTaskForOwner(worker->GetAssignedTaskId());
-
-  client->Close();
-
+  client->AsyncWaitTerminated([client, worker_failure_data_ptr, this] {
+    RAY_CHECK_OK(gcs_client_->Workers().AsyncReportWorkerFailure(worker_failure_data_ptr,
+                                                                 nullptr));
+  });
   // 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.