From c7cae036c3ad829b3a78fa726e288b210117ecac Mon Sep 17 00:00:00 2001
From: Stephanie Wang <swang@cs.berkeley.edu>
Date: Thu, 19 Mar 2020 18:46:16 -0700
Subject: [PATCH] [core] Only drain references for non-actor workers on
 shutdown (#7668)

* Only drain ref counter for non-actor tasks

* Don't force kill actors that have gone out of scope
---
 src/ray/core_worker/core_worker.cc  | 33 ++++++++++++++++++++++++++---
 src/ray/core_worker/task_manager.cc |  8 +++----
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc
index cd3292aee..5ef8f2018 100644
--- a/src/ray/core_worker/core_worker.cc
+++ b/src/ray/core_worker/core_worker.cc
@@ -294,7 +294,9 @@ void CoreWorker::Exit(bool intentional) {
   exiting_ = true;
   // Release the resources early in case draining takes a long time.
   RAY_CHECK_OK(local_raylet_client_->NotifyDirectCallTaskBlocked());
-  task_manager_->DrainAndShutdown([this, intentional]() {
+
+  // Callback to shutdown.
+  auto shutdown = [this, intentional]() {
     // To avoid problems, make sure shutdown is always called from the same
     // event loop each time.
     task_execution_service_.post([this, intentional]() {
@@ -303,7 +305,32 @@ void CoreWorker::Exit(bool intentional) {
       }
       Shutdown();
     });
-  });
+  };
+
+  // Callback to drain objects once all pending tasks have been drained.
+  auto drain_references_callback = [this, shutdown]() {
+    bool not_actor_task = false;
+    {
+      absl::MutexLock lock(&mutex_);
+      not_actor_task = actor_id_.IsNil();
+    }
+    if (not_actor_task) {
+      // If we are a task, then we cannot hold any object references in the
+      // heap. Therefore, any active object references are being held by other
+      // processes. Wait for these processes to release their references before
+      // we shutdown.
+      // NOTE(swang): This could still cause this worker process to stay alive
+      // forever if another process holds a reference forever.
+      reference_counter_->DrainAndShutdown(shutdown);
+    } else {
+      // If we are an actor, then we may be holding object references in the
+      // heap. Then, we should not wait to drain the object references before
+      // shutdown since this could hang.
+      shutdown();
+    }
+  };
+
+  task_manager_->DrainAndShutdown(drain_references_callback);
 }
 
 void CoreWorker::RunIOService() {
@@ -996,7 +1023,7 @@ bool CoreWorker::AddActorHandle(std::unique_ptr<ActorHandle> actor_handle,
                           << " has gone out of scope, sending message to actor "
                           << actor_id << " to do a clean exit.";
             RAY_CHECK_OK(
-                KillActor(actor_id, /*force_kill=*/true, /*no_reconstruction=*/false));
+                KillActor(actor_id, /*force_kill=*/false, /*no_reconstruction=*/false));
           }
         }));
   }
diff --git a/src/ray/core_worker/task_manager.cc b/src/ray/core_worker/task_manager.cc
index da463753a..7b369162a 100644
--- a/src/ray/core_worker/task_manager.cc
+++ b/src/ray/core_worker/task_manager.cc
@@ -86,9 +86,9 @@ void TaskManager::DrainAndShutdown(std::function<void()> shutdown) {
     }
   }
 
-  // Do not hold the lock when calling into the reference counter.
+  // Do not hold the lock when calling callbacks.
   if (!has_pending_tasks) {
-    reference_counter_->DrainAndShutdown(shutdown);
+    shutdown();
   }
 }
 
@@ -217,9 +217,9 @@ void TaskManager::ShutdownIfNeeded() {
       std::swap(shutdown_hook_, shutdown_hook);
     }
   }
-  // Do not hold the lock when calling into the reference counter.
+  // Do not hold the lock when calling callbacks.
   if (shutdown_hook != nullptr) {
-    reference_counter_->DrainAndShutdown(shutdown_hook);
+    shutdown_hook();
   }
 }