From fa98694dd045b3ea424483efdedc6f291143273a Mon Sep 17 00:00:00 2001
From: Zhijun Fu <37800433+zhijunfu@users.noreply.github.com>
Date: Thu, 5 Dec 2019 11:08:58 +0800
Subject: [PATCH] Fix raylet crash during cluster shutdown (#6272)

---
 src/ray/common/constants.h                         |  3 +++
 src/ray/gcs/tables.cc                              |  2 ++
 src/ray/gcs/tables.h                               |  4 ++++
 .../object_store_notification_manager.cc           | 11 ++++++++---
 src/ray/raylet/node_manager.cc                     | 14 +++++++++++---
 5 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/ray/common/constants.h b/src/ray/common/constants.h
index e117d4bac..2abc68898 100644
--- a/src/ray/common/constants.h
+++ b/src/ray/common/constants.h
@@ -16,6 +16,9 @@ constexpr int kObjectIdIndexSize = 32;
 static_assert(kObjectIdIndexSize % CHAR_BIT == 0,
               "ObjectID prefix not a multiple of bytes");
 
+/// Raylet exit code on plasma store socket error.
+constexpr int kRayletStoreErrorExitCode = 100;
+
 /// Prefix for the object table keys in redis.
 constexpr char kObjectTablePrefix[] = "ObjectTable";
 /// Prefix for the task table keys in redis.
diff --git a/src/ray/gcs/tables.cc b/src/ray/gcs/tables.cc
index a7569d979..c891dbeef 100644
--- a/src/ray/gcs/tables.cc
+++ b/src/ray/gcs/tables.cc
@@ -673,6 +673,8 @@ Status ClientTable::Disconnect(const DisconnectCallback &callback) {
   return Status::OK();
 }
 
+bool ClientTable::IsDisconnected() const { return disconnected_; }
+
 ray::Status ClientTable::MarkDisconnected(const ClientID &dead_node_id) {
   auto node_info = std::make_shared<GcsNodeInfo>();
   node_info->set_node_id(dead_node_id.Binary());
diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h
index 511ba4fd4..f02bc453c 100644
--- a/src/ray/gcs/tables.h
+++ b/src/ray/gcs/tables.h
@@ -860,6 +860,10 @@ class ClientTable : public Log<ClientID, GcsNodeInfo> {
   /// \return Status
   ray::Status Disconnect(const DisconnectCallback &callback = nullptr);
 
+  /// Whether the client is disconnected from the GCS.
+  /// \return Whether the client is disconnected.
+  bool IsDisconnected() const;
+
   /// Mark a different client as disconnected. The client ID should never be
   /// reused for a new client.
   ///
diff --git a/src/ray/object_manager/object_store_notification_manager.cc b/src/ray/object_manager/object_store_notification_manager.cc
index 3b29ab2d5..151c4f548 100644
--- a/src/ray/object_manager/object_store_notification_manager.cc
+++ b/src/ray/object_manager/object_store_notification_manager.cc
@@ -43,9 +43,14 @@ void ObjectStoreNotificationManager::ProcessStoreLength(
     const boost::system::error_code &error) {
   notification_.resize(length_);
   if (error) {
-    RAY_LOG(FATAL)
-        << "Problem communicating with the object store from raylet, check logs or "
-        << "dmesg for previous errors: " << boost_to_ray_status(error).ToString();
+    // When shutting down a cluster, it's possible that the plasma store is killed
+    // earlier than raylet, in this case we don't want raylet to crash, we instead
+    // log an error message and exit.
+    RAY_LOG(ERROR) << "Failed to process store length: "
+                   << boost_to_ray_status(error).ToString()
+                   << ", most likely plasma store is down, raylet will exit";
+    // Exit raylet process.
+    _exit(kRayletStoreErrorExitCode);
   }
   boost::asio::async_read(
       socket_, boost::asio::buffer(notification_),
diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc
index e72b6a98e..fae949196 100644
--- a/src/ray/raylet/node_manager.cc
+++ b/src/ray/raylet/node_manager.cc
@@ -518,9 +518,17 @@ void NodeManager::ClientRemoved(const GcsNodeInfo &node_info) {
   const ClientID client_id = ClientID::FromBinary(node_info.node_id());
   RAY_LOG(DEBUG) << "[ClientRemoved] Received callback from client id " << client_id;
 
-  RAY_CHECK(client_id != gcs_client_->client_table().GetLocalClientId())
-      << "Exiting because this node manager has mistakenly been marked dead by the "
-      << "monitor.";
+  if (!gcs_client_->client_table().IsDisconnected()) {
+    // We could receive a notification for our own death when we disconnect from client
+    // table after receiving a 'SIGTERM' signal, in that case we disconnect from gcs
+    // client table and then do some cleanup in the disconnect callback, and it's possible
+    // that we receive the notification in between, for more details refer to the SIGTERM
+    // handler in main.cc. In this case check for intentional disconnection and rule it
+    // out.
+    RAY_CHECK(client_id != gcs_client_->client_table().GetLocalClientId())
+        << "Exiting because this node manager has mistakenly been marked dead by the "
+        << "monitor.";
+  }
 
   // Below, when we remove client_id from all of these data structures, we could
   // check that it is actually removed, or log a warning otherwise, but that may