Fix worker exit cleanup (#6450)

* working but ugly

* comments

* proper but hanging in grpc server destructor

* grpc server shutdown deadline

* fix disconnect

* lint

* shutdown_only in test

* replace shutdown

src/ray/core_worker/core_worker.cc

@@ -70,9 +70,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
                        const std::string &log_dir, const std::string &node_ip_address,
                        int node_manager_port,
                        const TaskExecutionCallback &task_execution_callback,
-                       std::function<Status()> check_signals,
-                       const std::function<void()> exit_handler,
-                       bool ref_counting_enabled)
+                       std::function<Status()> check_signals, bool ref_counting_enabled)
     : worker_type_(worker_type),
       language_(language),
       log_dir_(log_dir),
@@ -119,13 +117,13 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
     RAY_CHECK(task_execution_callback_ != nullptr);
     auto execute_task = std::bind(&CoreWorker::ExecuteTask, this, std::placeholders::_1,
                                   std::placeholders::_2, std::placeholders::_3);
+    auto exit = std::bind(&CoreWorker::Shutdown, this);
     raylet_task_receiver_ =
         std::unique_ptr<CoreWorkerRayletTaskReceiver>(new CoreWorkerRayletTaskReceiver(
-            worker_context_.GetWorkerID(), local_raylet_client_, execute_task,
-            exit_handler));
+            worker_context_.GetWorkerID(), local_raylet_client_, execute_task, exit));
     direct_task_receiver_ =
         std::unique_ptr<CoreWorkerDirectTaskReceiver>(new CoreWorkerDirectTaskReceiver(
-            worker_context_, task_execution_service_, execute_task, exit_handler));
+            worker_context_, task_execution_service_, execute_task, exit));
   }
 
   // Start RPC server after all the task receivers are properly initialized.
@@ -239,26 +237,25 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
 }
 
 CoreWorker::~CoreWorker() {
-  Shutdown();
+  io_service_.stop();
   io_thread_.join();
+  if (log_dir_ != "") {
+    RayLog::ShutDownRayLog();
+  }
 }
 
 void CoreWorker::Shutdown() {
-  if (!shutdown_) {
-    shutdown_ = true;
-    io_service_.stop();
-    if (worker_type_ == WorkerType::WORKER) {
-      task_execution_service_.stop();
-    }
-    if (log_dir_ != "") {
-      RayLog::ShutDownRayLog();
-    }
+  io_service_.stop();
+  if (worker_type_ == WorkerType::WORKER) {
+    task_execution_service_.stop();
   }
 }
 
 void CoreWorker::Disconnect() {
   io_service_.stop();
-  gcs_client_->Disconnect();
+  if (gcs_client_) {
+    gcs_client_->Disconnect();
+  }
   if (local_raylet_client_) {
     RAY_IGNORE_EXPR(local_raylet_client_->Disconnect());
   }

src/ray/core_worker/core_worker.h

@@ -67,8 +67,6 @@ class CoreWorker {
   /// \param[in] check_signals Language worker function to check for signals and handle
   ///            them. If the function returns anything but StatusOK, any long-running
   ///            operations in the core worker will short circuit and return that status.
-  /// \param[in] exit_handler Language worker function to orderly shutdown the worker.
-  ///            We guarantee this will be run on the main thread of the worker.
   /// \param[in] ref_counting_enabled Whether to enable object ref counting.
   ///
   /// NOTE(zhijunfu): the constructor would throw if a failure happens.
@@ -78,7 +76,6 @@ class CoreWorker {
              const std::string &log_dir, const std::string &node_ip_address,
              int node_manager_port, const TaskExecutionCallback &task_execution_callback,
              std::function<Status()> check_signals = nullptr,
-             std::function<void()> exit_handler = nullptr,
              bool ref_counting_enabled = false);
 
   ~CoreWorker();

src/ray/core_worker/transport/direct_actor_transport.h

@@ -431,7 +431,10 @@ class CoreWorkerDirectTaskReceiver {
 
   ~CoreWorkerDirectTaskReceiver() {
     fiber_shutdown_event_.Notify();
-    fiber_runner_thread_.join();
+    // Only join the fiber thread if it was spawned in the first place.
+    if (fiber_runner_thread_.joinable()) {
+      fiber_runner_thread_.join();
+    }
   }
 
   /// Initialize this receiver. This must be called prior to use.

src/ray/rpc/grpc_server.h

@@ -1,12 +1,12 @@
 #ifndef RAY_RPC_GRPC_SERVER_H
 #define RAY_RPC_GRPC_SERVER_H
 
+#include <grpcpp/grpcpp.h>
+
+#include <boost/asio.hpp>
 #include <thread>
 #include <utility>
 
-#include <grpcpp/grpcpp.h>
-#include <boost/asio.hpp>
-
 #include "ray/common/status.h"
 #include "ray/rpc/server_call.h"
 
@@ -42,7 +42,9 @@ class GrpcServer {
   // Shutdown this server
   void Shutdown() {
     if (!is_closed_) {
-      server_->Shutdown();
+      // Shutdown the server with an immediate deadline.
+      // TODO(edoakes): do we want to do this in all cases?
+      server_->Shutdown(gpr_now(GPR_CLOCK_REALTIME));
       for (const auto &cq : cqs_) {
         cq->Shutdown();
       }