[GCS] Add node failure detector (#8119)

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -296,23 +296,30 @@ Status ServiceBasedNodeInfoAccessor::RegisterSelf(const GcsNodeInfo &local_node_
   RAY_CHECK(local_node_info.state() == GcsNodeInfo::ALIVE);
   rpc::RegisterNodeRequest request;
   request.mutable_node_info()->CopyFrom(local_node_info);
-  client_impl_->GetGcsRpcClient().RegisterNode(
-      request, [this, node_id, &local_node_info](const Status &status,
-                                                 const rpc::RegisterNodeReply &reply) {
-        if (status.ok()) {
-          local_node_info_.CopyFrom(local_node_info);
-          local_node_id_ = ClientID::FromBinary(local_node_info.node_id());
-        }
-        RAY_LOG(DEBUG) << "Finished registering node info, status = " << status
-                       << ", node id = " << node_id;
-      });
+
+  auto operation = [this, request, local_node_info,
+                    node_id](SequencerDoneCallback done_callback) {
+    client_impl_->GetGcsRpcClient().RegisterNode(
+        request, [this, node_id, local_node_info, done_callback](
+                     const Status &status, const rpc::RegisterNodeReply &reply) {
+          if (status.ok()) {
+            local_node_info_.CopyFrom(local_node_info);
+            local_node_id_ = ClientID::FromBinary(local_node_info.node_id());
+          }
+          RAY_LOG(DEBUG) << "Finished registering node info, status = " << status
+                         << ", node id = " << node_id;
+          done_callback();
+        });
+  };
+
+  sequencer_.Post(node_id, operation);
   return Status::OK();
 }
 
 Status ServiceBasedNodeInfoAccessor::UnregisterSelf() {
   RAY_CHECK(!local_node_id_.IsNil()) << "This node is disconnected.";
   ClientID node_id = ClientID::FromBinary(local_node_info_.node_id());
-  RAY_LOG(DEBUG) << "Unregistering node info, node id = " << node_id;
+  RAY_LOG(INFO) << "Unregistering node info, node id = " << node_id;
   rpc::UnregisterNodeRequest request;
   request.set_node_id(local_node_info_.node_id());
   client_impl_->GetGcsRpcClient().UnregisterNode(
@@ -322,8 +329,8 @@ Status ServiceBasedNodeInfoAccessor::UnregisterSelf() {
           local_node_info_.set_state(GcsNodeInfo::DEAD);
           local_node_id_ = ClientID::Nil();
         }
-        RAY_LOG(DEBUG) << "Finished unregistering node info, status = " << status
-                       << ", node id = " << node_id;
+        RAY_LOG(INFO) << "Finished unregistering node info, status = " << status
+                      << ", node id = " << node_id;
       });
   return Status::OK();
 }

src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc

@@ -630,8 +630,12 @@ TEST_F(ServiceBasedGcsClientTest, TestNodeResources) {
   };
   ASSERT_TRUE(SubscribeToResources(on_subscribe));
 
+  // Register node.
+  auto node_info = Mocker::GenNodeInfo();
+  RAY_CHECK(RegisterNode(*node_info));
+
   // Update resources of node in GCS.
-  ClientID node_id = ClientID::FromRandom();
+  ClientID node_id = ClientID::FromBinary(node_info->node_id());
   gcs::NodeInfoAccessor::ResourceMap resource_map;
   std::string key = "CPU";
   auto resource = std::make_shared<rpc::ResourceTableData>();
@@ -656,8 +660,12 @@ TEST_F(ServiceBasedGcsClientTest, TestNodeHeartbeat) {
       };
   ASSERT_TRUE(SubscribeBatchHeartbeat(on_subscribe));
 
+  // Register node.
+  auto node_info = Mocker::GenNodeInfo();
+  RAY_CHECK(RegisterNode(*node_info));
+
   // Report heartbeat of a node to GCS.
-  ClientID node_id = ClientID::FromRandom();
+  ClientID node_id = ClientID::FromBinary(node_info->node_id());
   auto heartbeat = std::make_shared<rpc::HeartbeatTableData>();
   heartbeat->set_client_id(node_id.Binary());
   ASSERT_TRUE(ReportHeartbeat(heartbeat));

src/ray/gcs/gcs_server/gcs_node_manager.cc

@@ -19,100 +19,59 @@
 
 namespace ray {
 namespace gcs {
-GcsNodeManager::GcsNodeManager(boost::asio::io_service &io_service,
-                               gcs::NodeInfoAccessor &node_info_accessor,
-                               gcs::ErrorInfoAccessor &error_info_accessor)
+
+GcsNodeManager::NodeFailureDetector::NodeFailureDetector(
+    boost::asio::io_service &io_service, gcs::NodeInfoAccessor &node_info_accessor,
+    std::function<void(const ClientID &)> on_node_death_callback)
     : node_info_accessor_(node_info_accessor),
-      error_info_accessor_(error_info_accessor),
+      on_node_death_callback_(std::move(on_node_death_callback)),
       num_heartbeats_timeout_(RayConfig::instance().num_heartbeats_timeout()),
-      heartbeat_timer_(io_service) {
-  Start();
+      detect_timer_(io_service) {
+  Tick();
 }
 
-void GcsNodeManager::HandleHeartbeat(const ClientID &node_id,
-                                     const rpc::HeartbeatTableData &heartbeat_data) {
-  heartbeats_[node_id] = num_heartbeats_timeout_;
+void GcsNodeManager::NodeFailureDetector::AddNode(const ray::ClientID &node_id) {
+  heartbeats_.emplace(node_id, num_heartbeats_timeout_);
+}
+
+void GcsNodeManager::NodeFailureDetector::HandleHeartbeat(
+    const ClientID &node_id, const rpc::HeartbeatTableData &heartbeat_data) {
+  auto iter = heartbeats_.find(node_id);
+  if (iter == heartbeats_.end()) {
+    // Ignore this heartbeat as the node is not registered.
+    // TODO(Shanly): Maybe we should reply the raylet with an error. So the raylet can
+    // crash itself as soon as possible.
+    return;
+  }
+
+  iter->second = num_heartbeats_timeout_;
   heartbeat_buffer_[node_id] = heartbeat_data;
 }
 
-void GcsNodeManager::Start() {
-  RAY_LOG(INFO) << "Starting gcs node manager.";
-  const auto lookup_callback =
-      [this](Status status, const std::vector<rpc::GcsNodeInfo> &node_info_list) {
-        for (const auto &node_info : node_info_list) {
-          if (node_info.state() != rpc::GcsNodeInfo::DEAD) {
-            // If there're any existing alive clients in client table, add them to
-            // our `heartbeats_` cache. Thus, if they died before monitor starts,
-            // we can also detect their death.
-            // Use `emplace` instead of `operator []` because we just want to add this
-            // client to `heartbeats_` only if it has not yet received heartbeat event.
-            // Besides, it is not necessary to add an empty `HeartbeatTableData`
-            // to `heartbeat_buffer_` as it doesn't make sense to broadcast an empty
-            // message to the cluster and it's ok to add it when actually receive
-            // its heartbeat event.
-            heartbeats_.emplace(ClientID::FromBinary(node_info.node_id()),
-                                num_heartbeats_timeout_);
-          }
-        }
-        Tick();
-      };
-  RAY_CHECK_OK(node_info_accessor_.AsyncGetAll(lookup_callback));
-}
-
 /// A periodic timer that checks for timed out clients.
-void GcsNodeManager::Tick() {
+void GcsNodeManager::NodeFailureDetector::Tick() {
   DetectDeadNodes();
   SendBatchedHeartbeat();
   ScheduleTick();
 }
 
-void GcsNodeManager::DetectDeadNodes() {
+void GcsNodeManager::NodeFailureDetector::DetectDeadNodes() {
   for (auto it = heartbeats_.begin(); it != heartbeats_.end();) {
     auto current = it++;
     current->second = current->second - 1;
     if (current->second == 0) {
-      if (dead_nodes_.count(current->first) == 0) {
-        auto node_id = current->first;
-        RAY_LOG(WARNING) << "Node timed out: " << node_id;
-        auto lookup_callback = [this, node_id](
-                                   Status status,
-                                   const std::vector<rpc::GcsNodeInfo> &all_node) {
-          RAY_CHECK_OK(status);
-          bool marked = false;
-          for (const auto &node : all_node) {
-            if (node_id.Binary() == node.node_id() &&
-                node.state() == rpc::GcsNodeInfo::DEAD) {
-              // The node has been marked dead by itself.
-              marked = true;
-              break;
-            }
-          }
-          if (!marked) {
-            RemoveNode(node_id);
-            RAY_CHECK_OK(node_info_accessor_.AsyncUnregister(node_id, nullptr));
-            // Broadcast a warning to all of the drivers indicating that the node
-            // has been marked as dead.
-            // TODO(rkn): Define this constant somewhere else.
-            std::string type = "node_removed";
-            std::ostringstream error_message;
-            error_message << "The node with node id " << node_id
-                          << " has been marked dead because the monitor"
-                          << " has missed too many heartbeats from it.";
-            auto error_data_ptr =
-                gcs::CreateErrorTableData(type, error_message.str(), current_time_ms());
-            RAY_CHECK_OK(
-                error_info_accessor_.AsyncReportJobError(error_data_ptr, nullptr));
-          }
-        };
-        RAY_CHECK_OK(node_info_accessor_.AsyncGetAll(lookup_callback));
-        dead_nodes_.insert(node_id);
-      }
+      auto node_id = current->first;
+      RAY_LOG(WARNING) << "Node timed out: " << node_id;
       heartbeats_.erase(current);
+      heartbeat_buffer_.erase(node_id);
+      if (on_node_death_callback_) {
+        on_node_death_callback_(node_id);
+      }
     }
   }
 }
 
-void GcsNodeManager::SendBatchedHeartbeat() {
+void GcsNodeManager::NodeFailureDetector::SendBatchedHeartbeat() {
   if (!heartbeat_buffer_.empty()) {
     auto batch = std::make_shared<rpc::HeartbeatBatchTableData>();
     for (const auto &heartbeat : heartbeat_buffer_) {
@@ -123,13 +82,13 @@ void GcsNodeManager::SendBatchedHeartbeat() {
   }
 }
 
-void GcsNodeManager::ScheduleTick() {
+void GcsNodeManager::NodeFailureDetector::ScheduleTick() {
   auto heartbeat_period = boost::posix_time::milliseconds(
       RayConfig::instance().raylet_heartbeat_timeout_milliseconds());
-  heartbeat_timer_.expires_from_now(heartbeat_period);
-  heartbeat_timer_.async_wait([this](const boost::system::error_code &error) {
+  detect_timer_.expires_from_now(heartbeat_period);
+  detect_timer_.async_wait([this](const boost::system::error_code &error) {
     if (error == boost::system::errc::operation_canceled) {
-      // `operation_canceled` is set when `heartbeat_timer_` is canceled or destroyed.
+      // `operation_canceled` is set when `detect_timer_` is canceled or destroyed.
       // The Monitor lifetime may be short than the object who use it. (e.g. gcs_server)
       return;
     }
@@ -138,6 +97,156 @@ void GcsNodeManager::ScheduleTick() {
   });
 }
 
+//////////////////////////////////////////////////////////////////////////////////////////
+GcsNodeManager::GcsNodeManager(boost::asio::io_service &io_service,
+                               gcs::NodeInfoAccessor &node_info_accessor,
+                               gcs::ErrorInfoAccessor &error_info_accessor)
+    : node_info_accessor_(node_info_accessor),
+      error_info_accessor_(error_info_accessor),
+      node_failure_detector_(new NodeFailureDetector(
+          io_service, node_info_accessor, [this](const ClientID &node_id) {
+            if (auto node = RemoveNode(node_id, /* is_intended = */ false)) {
+              node->set_state(rpc::GcsNodeInfo::DEAD);
+              RAY_CHECK(dead_nodes_.emplace(node_id, node).second);
+              RAY_CHECK_OK(node_info_accessor_.AsyncUnregister(node_id, nullptr));
+              // TODO(Shanly): Remove node resources from resource table.
+            }
+          })) {
+  // TODO(Shanly): Load node info list from storage synchronously.
+  // TODO(Shanly): Load cluster resources from storage synchronously.
+}
+
+void GcsNodeManager::HandleRegisterNode(const rpc::RegisterNodeRequest &request,
+                                        rpc::RegisterNodeReply *reply,
+                                        rpc::SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.node_info().node_id());
+  RAY_LOG(INFO) << "Registering node info, node id = " << node_id;
+  AddNode(std::make_shared<rpc::GcsNodeInfo>(request.node_info()));
+  auto on_done = [node_id, reply, send_reply_callback](Status status) {
+    RAY_CHECK_OK(status);
+    RAY_LOG(INFO) << "Finished registering node info, node id = " << node_id;
+    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
+  };
+  RAY_CHECK_OK(node_info_accessor_.AsyncRegister(request.node_info(), on_done));
+}
+
+void GcsNodeManager::HandleUnregisterNode(const rpc::UnregisterNodeRequest &request,
+                                          rpc::UnregisterNodeReply *reply,
+                                          rpc::SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.node_id());
+  RAY_LOG(INFO) << "Unregistering node info, node id = " << node_id;
+  auto on_done = [node_id, request, reply, send_reply_callback](Status status) {
+    RAY_CHECK_OK(status);
+    RAY_LOG(INFO) << "Finished unregistering node info, node id = " << node_id;
+    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
+  };
+  if (auto node = RemoveNode(node_id, /* is_intended = */ true)) {
+    node->set_state(rpc::GcsNodeInfo::DEAD);
+    RAY_CHECK(dead_nodes_.emplace(node_id, node).second);
+    RAY_CHECK_OK(node_info_accessor_.AsyncUnregister(node_id, on_done));
+    // TODO(Shanly): Remove node resources from resource table.
+  }
+}
+
+void GcsNodeManager::HandleGetAllNodeInfo(const rpc::GetAllNodeInfoRequest &request,
+                                          rpc::GetAllNodeInfoReply *reply,
+                                          rpc::SendReplyCallback send_reply_callback) {
+  RAY_LOG(DEBUG) << "Getting all nodes info.";
+  for (const auto &entry : alive_nodes_) {
+    reply->add_node_info_list()->CopyFrom(*entry.second);
+  }
+  for (const auto &entry : dead_nodes_) {
+    reply->add_node_info_list()->CopyFrom(*entry.second);
+  }
+  GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
+  RAY_LOG(DEBUG) << "Finished getting all node info.";
+}
+
+void GcsNodeManager::HandleReportHeartbeat(const rpc::ReportHeartbeatRequest &request,
+                                           rpc::ReportHeartbeatReply *reply,
+                                           rpc::SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.heartbeat().client_id());
+  RAY_LOG(DEBUG) << "Reporting heartbeat, node id = " << node_id;
+  auto heartbeat_data = std::make_shared<rpc::HeartbeatTableData>();
+  heartbeat_data->CopyFrom(request.heartbeat());
+  node_failure_detector_->HandleHeartbeat(node_id, *heartbeat_data);
+  GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
+  // TODO(Shanly): Remove it later.
+  // The heartbeat data is reported here because some python unit tests rely on the
+  // heartbeat data in redis.
+  RAY_CHECK_OK(node_info_accessor_.AsyncReportHeartbeat(heartbeat_data, nullptr));
+  RAY_LOG(DEBUG) << "Finished reporting heartbeat, node id = " << node_id;
+}
+
+void GcsNodeManager::HandleGetResources(const rpc::GetResourcesRequest &request,
+                                        rpc::GetResourcesReply *reply,
+                                        rpc::SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.node_id());
+  RAY_LOG(DEBUG) << "Getting node resources, node id = " << node_id;
+  auto iter = cluster_resources_.find(node_id);
+  if (iter != cluster_resources_.end()) {
+    for (auto &resource : iter->second) {
+      (*reply->mutable_resources())[resource.first] = *resource.second;
+    }
+  }
+  GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
+  RAY_LOG(DEBUG) << "Finished getting node resources, node id = " << node_id;
+}
+
+void GcsNodeManager::HandleUpdateResources(const rpc::UpdateResourcesRequest &request,
+                                           rpc::UpdateResourcesReply *reply,
+                                           rpc::SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.node_id());
+  RAY_LOG(DEBUG) << "Updating resources, node id = " << node_id;
+  auto iter = cluster_resources_.find(node_id);
+  if (iter != cluster_resources_.end()) {
+    auto to_be_updated_resources = std::make_shared<gcs::NodeInfoAccessor::ResourceMap>();
+    for (auto resource : request.resources()) {
+      (*to_be_updated_resources)[resource.first] =
+          std::make_shared<rpc::ResourceTableData>(resource.second);
+    }
+    for (auto &entry : *to_be_updated_resources) {
+      iter->second[entry.first] = entry.second;
+    }
+    auto on_done = [node_id, to_be_updated_resources, reply,
+                    send_reply_callback](Status status) {
+      RAY_CHECK_OK(status);
+      GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
+      RAY_LOG(DEBUG) << "Finished updating resources, node id = " << node_id;
+    };
+
+    RAY_CHECK_OK(node_info_accessor_.AsyncUpdateResources(
+        node_id, *to_be_updated_resources, on_done));
+  } else {
+    GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::Invalid("Node is not exist."));
+    RAY_LOG(ERROR) << "Failed to update resources as node " << node_id
+                   << " is not registered.";
+  }
+}
+
+void GcsNodeManager::HandleDeleteResources(const rpc::DeleteResourcesRequest &request,
+                                           rpc::DeleteResourcesReply *reply,
+                                           rpc::SendReplyCallback send_reply_callback) {
+  ClientID node_id = ClientID::FromBinary(request.node_id());
+  RAY_LOG(DEBUG) << "Deleting node resources, node id = " << node_id;
+  auto resource_names = VectorFromProtobuf(request.resource_name_list());
+  auto iter = cluster_resources_.find(node_id);
+  if (iter != cluster_resources_.end()) {
+    for (auto &resource_name : resource_names) {
+      iter->second.erase(resource_name);
+    }
+    auto on_done = [reply, send_reply_callback](Status status) {
+      RAY_CHECK_OK(status);
+      GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
+    };
+    RAY_CHECK_OK(
+        node_info_accessor_.AsyncDeleteResources(node_id, resource_names, on_done));
+  } else {
+    GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
+    RAY_LOG(DEBUG) << "Finished deleting node resources, node id = " << node_id;
+  }
+}
+
 std::shared_ptr<rpc::GcsNodeInfo> GcsNodeManager::GetNode(
     const ray::ClientID &node_id) const {
   auto iter = alive_nodes_.find(node_id);
@@ -148,31 +257,53 @@ std::shared_ptr<rpc::GcsNodeInfo> GcsNodeManager::GetNode(
   return iter->second;
 }
 
-const absl::flat_hash_map<ClientID, std::shared_ptr<rpc::GcsNodeInfo>>
-    &GcsNodeManager::GetAllAliveNodes() const {
-  return alive_nodes_;
-}
-
 void GcsNodeManager::AddNode(std::shared_ptr<rpc::GcsNodeInfo> node) {
   auto node_id = ClientID::FromBinary(node->node_id());
   auto iter = alive_nodes_.find(node_id);
   if (iter == alive_nodes_.end()) {
     alive_nodes_.emplace(node_id, node);
+    // Add an empty resources for this node.
+    RAY_CHECK(
+        cluster_resources_.emplace(node_id, gcs::NodeInfoAccessor::ResourceMap()).second);
+    // Register this node to the `node_failure_detector_` which will start monitoring it.
+    node_failure_detector_->AddNode(node_id);
+    // Notify all listeners.
     for (auto &listener : node_added_listeners_) {
       listener(node);
     }
   }
 }
 
-void GcsNodeManager::RemoveNode(const ray::ClientID &node_id) {
+std::shared_ptr<rpc::GcsNodeInfo> GcsNodeManager::RemoveNode(
+    const ray::ClientID &node_id, bool is_intended /*= false*/) {
+  std::shared_ptr<rpc::GcsNodeInfo> removed_node;
   auto iter = alive_nodes_.find(node_id);
   if (iter != alive_nodes_.end()) {
-    auto node = std::move(iter->second);
+    removed_node = std::move(iter->second);
+    // Remove from alive nodes.
     alive_nodes_.erase(iter);
+    // Remove from cluster resources.
+    RAY_CHECK(cluster_resources_.erase(node_id) != 0);
+    if (!is_intended) {
+      // Broadcast a warning to all of the drivers indicating that the node
+      // has been marked as dead.
+      // TODO(rkn): Define this constant somewhere else.
+      std::string type = "node_removed";
+      std::ostringstream error_message;
+      error_message << "The node with node id " << node_id
+                    << " has been marked dead because the detector"
+                    << " has missed too many heartbeats from it.";
+      auto error_data_ptr =
+          gcs::CreateErrorTableData(type, error_message.str(), current_time_ms());
+      RAY_CHECK_OK(error_info_accessor_.AsyncReportJobError(error_data_ptr, nullptr));
+    }
+
+    // Notify all listeners.
     for (auto &listener : node_removed_listeners_) {
-      listener(node);
+      listener(removed_node);
     }
   }
+  return removed_node;
 }
 
 }  // namespace gcs

src/ray/gcs/gcs_server/gcs_node_manager.h

@@ -19,14 +19,17 @@
 #include <ray/gcs/accessor.h>
 #include <ray/protobuf/gcs.pb.h>
 #include <ray/rpc/client_call.h>
+#include <ray/rpc/gcs_server/gcs_rpc_server.h>
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 
 namespace ray {
 namespace gcs {
-/// GcsNodeManager is responsible for managing and monitoring nodes.
+
+/// GcsNodeManager is responsible for managing and monitoring nodes as well as handing
+/// node and resource related rpc requests.
 /// This class is not thread-safe.
-class GcsNodeManager {
+class GcsNodeManager : public rpc::NodeInfoHandler {
  public:
   /// Create a GcsNodeManager.
   ///
@@ -38,6 +41,41 @@ class GcsNodeManager {
                           gcs::NodeInfoAccessor &node_info_accessor,
                           gcs::ErrorInfoAccessor &error_info_accessor);
 
+  /// Handle register rpc request come from raylet.
+  void HandleRegisterNode(const rpc::RegisterNodeRequest &request,
+                          rpc::RegisterNodeReply *reply,
+                          rpc::SendReplyCallback send_reply_callback) override;
+
+  /// Handle unregister rpc request come from raylet.
+  void HandleUnregisterNode(const rpc::UnregisterNodeRequest &request,
+                            rpc::UnregisterNodeReply *reply,
+                            rpc::SendReplyCallback send_reply_callback) override;
+
+  /// Handle get all node info rpc request.
+  void HandleGetAllNodeInfo(const rpc::GetAllNodeInfoRequest &request,
+                            rpc::GetAllNodeInfoReply *reply,
+                            rpc::SendReplyCallback send_reply_callback) override;
+
+  /// Handle heartbeat rpc come from raylet.
+  void HandleReportHeartbeat(const rpc::ReportHeartbeatRequest &request,
+                             rpc::ReportHeartbeatReply *reply,
+                             rpc::SendReplyCallback send_reply_callback) override;
+
+  /// Handle get resource rpc request.
+  void HandleGetResources(const rpc::GetResourcesRequest &request,
+                          rpc::GetResourcesReply *reply,
+                          rpc::SendReplyCallback send_reply_callback) override;
+
+  /// Handle update resource rpc request.
+  void HandleUpdateResources(const rpc::UpdateResourcesRequest &request,
+                             rpc::UpdateResourcesReply *reply,
+                             rpc::SendReplyCallback send_reply_callback) override;
+
+  /// Handle delete resource rpc request.
+  void HandleDeleteResources(const rpc::DeleteResourcesRequest &request,
+                             rpc::DeleteResourcesReply *reply,
+                             rpc::SendReplyCallback send_reply_callback) override;
+
   /// Add an alive node.
   ///
   /// \param node The info of the node to be added.
@@ -46,7 +84,10 @@ class GcsNodeManager {
   /// Remove from alive nodes.
   ///
   /// \param node_id The ID of the node to be removed.
-  void RemoveNode(const ClientID &node_id);
+  /// \param is_intended False if this is triggered by `node_failure_detector_`, else
+  /// True.
+  std::shared_ptr<rpc::GcsNodeInfo> RemoveNode(const ClientID &node_id,
+                                               bool is_intended = false);
 
   /// Get alive node by ID.
   ///
@@ -58,7 +99,9 @@ class GcsNodeManager {
   ///
   /// \return all alive nodes.
   const absl::flat_hash_map<ClientID, std::shared_ptr<rpc::GcsNodeInfo>>
-      &GetAllAliveNodes() const;
+      &GetAllAliveNodes() const {
+    return alive_nodes_;
+  }
 
   /// Add listener to monitor the remove action of nodes.
   ///
@@ -78,51 +121,75 @@ class GcsNodeManager {
     node_added_listeners_.emplace_back(std::move(listener));
   }
 
-  /// Handle a heartbeat from a Raylet.
-  ///
-  /// \param node_id The client ID of the Raylet that sent the heartbeat.
-  /// \param heartbeat_data The heartbeat sent by the client.
-  void HandleHeartbeat(const ClientID &node_id,
-                       const rpc::HeartbeatTableData &heartbeat_data);
-
  protected:
-  /// Listen for heartbeats from Raylets and mark Raylets
-  /// that do not send a heartbeat within a given period as dead.
-  void Start();
+  class NodeFailureDetector {
+   public:
+    /// Create a NodeFailureDetector.
+    ///
+    /// \param io_service The event loop to run the monitor on.
+    /// \param node_info_accessor The node info accessor.
+    explicit NodeFailureDetector(
+        boost::asio::io_service &io_service, gcs::NodeInfoAccessor &node_info_accessor,
+        std::function<void(const ClientID &)> on_node_death_callback);
 
-  /// A periodic timer that fires on every heartbeat period. Raylets that have
-  /// not sent a heartbeat within the last num_heartbeats_timeout ticks will be
-  /// marked as dead in the client table.
-  void Tick();
+    /// Register node to this detector.
+    /// Only if the node has registered, its heartbeat data will be accepted.
+    ///
+    /// \param node_id ID of the node to be registered.
+    void AddNode(const ClientID &node_id);
 
-  /// Check that if any raylet is inactive due to no heartbeat for a period of time.
-  /// If found any, mark it as dead.
-  void DetectDeadNodes();
+    /// Handle a heartbeat from a Raylet.
+    ///
+    /// \param node_id The client ID of the Raylet that sent the heartbeat.
+    /// \param heartbeat_data The heartbeat sent by the client.
+    void HandleHeartbeat(const ClientID &node_id,
+                         const rpc::HeartbeatTableData &heartbeat_data);
 
-  /// Send any buffered heartbeats as a single publish.
-  void SendBatchedHeartbeat();
+   protected:
+    /// A periodic timer that fires on every heartbeat period. Raylets that have
+    /// not sent a heartbeat within the last num_heartbeats_timeout ticks will be
+    /// marked as dead in the client table.
+    void Tick();
 
-  /// Schedule another tick after a short time.
-  void ScheduleTick();
+    /// Check that if any raylet is inactive due to no heartbeat for a period of time.
+    /// If found any, mark it as dead.
+    void DetectDeadNodes();
+
+    /// Send any buffered heartbeats as a single publish.
+    void SendBatchedHeartbeat();
+
+    /// Schedule another tick after a short time.
+    void ScheduleTick();
+
+   protected:
+    /// Node info accessor.
+    gcs::NodeInfoAccessor &node_info_accessor_;
+    /// The callback of node death.
+    std::function<void(const ClientID &)> on_node_death_callback_;
+    /// The number of heartbeats that can be missed before a node is removed.
+    int64_t num_heartbeats_timeout_;
+    /// A timer that ticks every heartbeat_timeout_ms_ milliseconds.
+    boost::asio::deadline_timer detect_timer_;
+    /// For each Raylet that we receive a heartbeat from, the number of ticks
+    /// that may pass before the Raylet will be declared dead.
+    absl::flat_hash_map<ClientID, int64_t> heartbeats_;
+    /// A buffer containing heartbeats received from node managers in the last tick.
+    absl::flat_hash_map<ClientID, rpc::HeartbeatTableData> heartbeat_buffer_;
+  };
 
  private:
-  /// Alive nodes.
-  absl::flat_hash_map<ClientID, std::shared_ptr<rpc::GcsNodeInfo>> alive_nodes_;
   /// Node info accessor.
   gcs::NodeInfoAccessor &node_info_accessor_;
   /// Error info accessor.
   gcs::ErrorInfoAccessor &error_info_accessor_;
-  /// The number of heartbeats that can be missed before a node is removed.
-  int64_t num_heartbeats_timeout_;
-  /// A timer that ticks every heartbeat_timeout_ms_ milliseconds.
-  boost::asio::deadline_timer heartbeat_timer_;
-  /// For each Raylet that we receive a heartbeat from, the number of ticks
-  /// that may pass before the Raylet will be declared dead.
-  absl::flat_hash_map<ClientID, int64_t> heartbeats_;
-  /// The Raylets that have been marked as dead in gcs.
-  absl::flat_hash_set<ClientID> dead_nodes_;
-  /// A buffer containing heartbeats received from node managers in the last tick.
-  absl::flat_hash_map<ClientID, rpc::HeartbeatTableData> heartbeat_buffer_;
+  /// Detector to detect the failure of node.
+  std::unique_ptr<NodeFailureDetector> node_failure_detector_;
+  /// Alive nodes.
+  absl::flat_hash_map<ClientID, std::shared_ptr<rpc::GcsNodeInfo>> alive_nodes_;
+  /// Dead nodes.
+  absl::flat_hash_map<ClientID, std::shared_ptr<rpc::GcsNodeInfo>> dead_nodes_;
+  /// Cluster resources.
+  absl::flat_hash_map<ClientID, gcs::NodeInfoAccessor::ResourceMap> cluster_resources_;
   /// Listeners which monitors the addition of nodes.
   std::vector<std::function<void(std::shared_ptr<rpc::GcsNodeInfo>)>>
       node_added_listeners_;

src/ray/gcs/gcs_server/gcs_server.cc

@@ -18,7 +18,6 @@
 #include "gcs_actor_manager.h"
 #include "gcs_node_manager.h"
 #include "job_info_handler_impl.h"
-#include "node_info_handler_impl.h"
 #include "object_info_handler_impl.h"
 #include "ray/common/network_util.h"
 #include "ray/common/ray_config.h"
@@ -64,9 +63,8 @@ void GcsServer::Start() {
       new rpc::ActorInfoGrpcService(main_service_, *actor_info_handler_));
   rpc_server_.RegisterService(*actor_info_service_);
 
-  node_info_handler_ = InitNodeInfoHandler();
   node_info_service_.reset(
-      new rpc::NodeInfoGrpcService(main_service_, *node_info_handler_));
+      new rpc::NodeInfoGrpcService(main_service_, *gcs_node_manager_));
   rpc_server_.RegisterService(*node_info_service_);
 
   object_info_handler_ = InitObjectInfoHandler();
@@ -149,11 +147,6 @@ std::unique_ptr<rpc::ActorInfoHandler> GcsServer::InitActorInfoHandler() {
       new rpc::DefaultActorInfoHandler(*redis_gcs_client_, *gcs_actor_manager_));
 }
 
-std::unique_ptr<rpc::NodeInfoHandler> GcsServer::InitNodeInfoHandler() {
-  return std::unique_ptr<rpc::DefaultNodeInfoHandler>(
-      new rpc::DefaultNodeInfoHandler(*redis_gcs_client_, *gcs_node_manager_));
-}
-
 std::unique_ptr<rpc::ObjectInfoHandler> GcsServer::InitObjectInfoHandler() {
   return std::unique_ptr<rpc::DefaultObjectInfoHandler>(
       new rpc::DefaultObjectInfoHandler(*redis_gcs_client_));

src/ray/gcs/gcs_server/gcs_server.h

@@ -83,9 +83,6 @@ class GcsServer {
   /// The actor info handler
   virtual std::unique_ptr<rpc::ActorInfoHandler> InitActorInfoHandler();
 
-  /// The node info handler
-  virtual std::unique_ptr<rpc::NodeInfoHandler> InitNodeInfoHandler();
-
   /// The object info handler
   virtual std::unique_ptr<rpc::ObjectInfoHandler> InitObjectInfoHandler();
 
@@ -128,7 +125,6 @@ class GcsServer {
   std::unique_ptr<rpc::ActorInfoHandler> actor_info_handler_;
   std::unique_ptr<rpc::ActorInfoGrpcService> actor_info_service_;
   /// Node info handler and service
-  std::unique_ptr<rpc::NodeInfoHandler> node_info_handler_;
   std::unique_ptr<rpc::NodeInfoGrpcService> node_info_service_;
   /// Object info handler and service
   std::unique_ptr<rpc::ObjectInfoHandler> object_info_handler_;

src/ray/gcs/gcs_server/node_info_handler_impl.cc

@@ -1,193 +0,0 @@
-// Copyright 2017 The Ray Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//  http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#include "node_info_handler_impl.h"
-#include "ray/util/logging.h"
-
-namespace ray {
-namespace rpc {
-
-void DefaultNodeInfoHandler::HandleRegisterNode(
-    const rpc::RegisterNodeRequest &request, rpc::RegisterNodeReply *reply,
-    rpc::SendReplyCallback send_reply_callback) {
-  ClientID node_id = ClientID::FromBinary(request.node_info().node_id());
-  RAY_LOG(DEBUG) << "Registering node info, node id = " << node_id;
-  gcs_node_manager_.AddNode(std::make_shared<rpc::GcsNodeInfo>(request.node_info()));
-  auto on_done = [node_id, reply, send_reply_callback](Status status) {
-    if (!status.ok()) {
-      RAY_LOG(ERROR) << "Failed to register node info: " << status.ToString()
-                     << ", node id = " << node_id;
-    }
-    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
-  };
-
-  Status status = gcs_client_.Nodes().AsyncRegister(request.node_info(), on_done);
-  if (!status.ok()) {
-    on_done(status);
-  }
-  RAY_LOG(DEBUG) << "Finished registering node info, node id = " << node_id;
-}
-
-void DefaultNodeInfoHandler::HandleUnregisterNode(
-    const rpc::UnregisterNodeRequest &request, rpc::UnregisterNodeReply *reply,
-    rpc::SendReplyCallback send_reply_callback) {
-  ClientID node_id = ClientID::FromBinary(request.node_id());
-  RAY_LOG(DEBUG) << "Unregistering node info, node id = " << node_id;
-  gcs_node_manager_.RemoveNode(node_id);
-
-  auto on_done = [node_id, reply, send_reply_callback](Status status) {
-    if (!status.ok()) {
-      RAY_LOG(ERROR) << "Failed to unregister node info: " << status.ToString()
-                     << ", node id = " << node_id;
-    }
-    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
-  };
-
-  Status status = gcs_client_.Nodes().AsyncUnregister(node_id, on_done);
-  if (!status.ok()) {
-    on_done(status);
-  }
-  RAY_LOG(DEBUG) << "Finished unregistering node info, node id = " << node_id;
-}
-
-void DefaultNodeInfoHandler::HandleGetAllNodeInfo(
-    const rpc::GetAllNodeInfoRequest &request, rpc::GetAllNodeInfoReply *reply,
-    rpc::SendReplyCallback send_reply_callback) {
-  RAY_LOG(DEBUG) << "Getting all nodes info.";
-  auto on_done = [reply, send_reply_callback](
-                     Status status, const std::vector<rpc::GcsNodeInfo> &result) {
-    if (status.ok()) {
-      for (const rpc::GcsNodeInfo &node_info : result) {
-        reply->add_node_info_list()->CopyFrom(node_info);
-      }
-    } else {
-      RAY_LOG(ERROR) << "Failed to get all nodes info: " << status.ToString();
-    }
-    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
-  };
-
-  Status status = gcs_client_.Nodes().AsyncGetAll(on_done);
-  if (!status.ok()) {
-    on_done(status, std::vector<rpc::GcsNodeInfo>());
-  }
-  RAY_LOG(DEBUG) << "Finished getting all node info.";
-}
-
-void DefaultNodeInfoHandler::HandleReportHeartbeat(
-    const ReportHeartbeatRequest &request, ReportHeartbeatReply *reply,
-    SendReplyCallback send_reply_callback) {
-  ClientID node_id = ClientID::FromBinary(request.heartbeat().client_id());
-  RAY_LOG(DEBUG) << "Reporting heartbeat, node id = " << node_id;
-  auto on_done = [node_id, reply, send_reply_callback](Status status) {
-    if (!status.ok()) {
-      RAY_LOG(ERROR) << "Failed to report heartbeat: " << status.ToString()
-                     << ", node id = " << node_id;
-    }
-    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
-  };
-
-  auto heartbeat_data = std::make_shared<rpc::HeartbeatTableData>();
-  heartbeat_data->CopyFrom(request.heartbeat());
-  gcs_node_manager_.HandleHeartbeat(node_id, *heartbeat_data);
-
-  Status status = gcs_client_.Nodes().AsyncReportHeartbeat(heartbeat_data, on_done);
-  if (!status.ok()) {
-    on_done(status);
-  }
-  RAY_LOG(DEBUG) << "Finished reporting heartbeat, node id = " << node_id;
-}
-
-void DefaultNodeInfoHandler::HandleGetResources(const GetResourcesRequest &request,
-                                                GetResourcesReply *reply,
-                                                SendReplyCallback send_reply_callback) {
-  ClientID node_id = ClientID::FromBinary(request.node_id());
-  RAY_LOG(DEBUG) << "Getting node resources, node id = " << node_id;
-
-  auto on_done = [node_id, reply, send_reply_callback](
-                     Status status,
-                     const boost::optional<gcs::NodeInfoAccessor::ResourceMap> &result) {
-    if (status.ok()) {
-      if (result) {
-        for (auto &resource : *result) {
-          (*reply->mutable_resources())[resource.first] = *resource.second;
-        }
-      }
-    } else {
-      RAY_LOG(ERROR) << "Failed to get node resources: " << status.ToString()
-                     << ", node id = " << node_id;
-    }
-    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
-  };
-
-  Status status = gcs_client_.Nodes().AsyncGetResources(node_id, on_done);
-  if (!status.ok()) {
-    on_done(status, boost::none);
-  }
-
-  RAY_LOG(DEBUG) << "Finished getting node resources, node id = " << node_id;
-}
-
-void DefaultNodeInfoHandler::HandleUpdateResources(
-    const UpdateResourcesRequest &request, UpdateResourcesReply *reply,
-    SendReplyCallback send_reply_callback) {
-  ClientID node_id = ClientID::FromBinary(request.node_id());
-  RAY_LOG(DEBUG) << "Updating node resources, node id = " << node_id;
-
-  gcs::NodeInfoAccessor::ResourceMap resources;
-  for (auto resource : request.resources()) {
-    resources[resource.first] = std::make_shared<rpc::ResourceTableData>(resource.second);
-  }
-
-  auto on_done = [node_id, reply, send_reply_callback](Status status) {
-    if (!status.ok()) {
-      RAY_LOG(ERROR) << "Failed to update node resources: " << status.ToString()
-                     << ", node id = " << node_id;
-    }
-    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
-  };
-
-  Status status = gcs_client_.Nodes().AsyncUpdateResources(node_id, resources, on_done);
-  if (!status.ok()) {
-    on_done(status);
-  }
-
-  RAY_LOG(DEBUG) << "Finished updating node resources, node id = " << node_id;
-}
-
-void DefaultNodeInfoHandler::HandleDeleteResources(
-    const DeleteResourcesRequest &request, DeleteResourcesReply *reply,
-    SendReplyCallback send_reply_callback) {
-  ClientID node_id = ClientID::FromBinary(request.node_id());
-  auto resource_names = VectorFromProtobuf(request.resource_name_list());
-  RAY_LOG(DEBUG) << "Deleting node resources, node id = " << node_id;
-
-  auto on_done = [node_id, reply, send_reply_callback](Status status) {
-    if (!status.ok()) {
-      RAY_LOG(ERROR) << "Failed to delete node resources: " << status.ToString()
-                     << ", node id = " << node_id;
-    }
-    GCS_RPC_SEND_REPLY(send_reply_callback, reply, status);
-  };
-
-  Status status =
-      gcs_client_.Nodes().AsyncDeleteResources(node_id, resource_names, on_done);
-  if (!status.ok()) {
-    on_done(status);
-  }
-
-  RAY_LOG(DEBUG) << "Finished deleting node resources, node id = " << node_id;
-}
-
-}  // namespace rpc
-}  // namespace ray

src/ray/gcs/gcs_server/node_info_handler_impl.h

@@ -1,67 +0,0 @@
-// Copyright 2017 The Ray Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//  http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#ifndef RAY_GCS_NODE_INFO_HANDLER_IMPL_H
-#define RAY_GCS_NODE_INFO_HANDLER_IMPL_H
-
-#include "gcs_node_manager.h"
-#include "ray/gcs/redis_gcs_client.h"
-#include "ray/rpc/gcs_server/gcs_rpc_server.h"
-
-namespace ray {
-
-namespace rpc {
-
-/// This implementation class of `NodeInfoHandler`.
-class DefaultNodeInfoHandler : public rpc::NodeInfoHandler {
- public:
-  explicit DefaultNodeInfoHandler(gcs::RedisGcsClient &gcs_client,
-                                  gcs::GcsNodeManager &gcs_node_manager)
-      : gcs_client_(gcs_client), gcs_node_manager_(gcs_node_manager) {}
-
-  void HandleRegisterNode(const RegisterNodeRequest &request, RegisterNodeReply *reply,
-                          SendReplyCallback send_reply_callback) override;
-
-  void HandleUnregisterNode(const UnregisterNodeRequest &request,
-                            UnregisterNodeReply *reply,
-                            SendReplyCallback send_reply_callback) override;
-
-  void HandleGetAllNodeInfo(const GetAllNodeInfoRequest &request,
-                            GetAllNodeInfoReply *reply,
-                            SendReplyCallback send_reply_callback) override;
-
-  void HandleReportHeartbeat(const ReportHeartbeatRequest &request,
-                             ReportHeartbeatReply *reply,
-                             SendReplyCallback send_reply_callback) override;
-
-  void HandleGetResources(const GetResourcesRequest &request, GetResourcesReply *reply,
-                          SendReplyCallback send_reply_callback) override;
-
-  void HandleUpdateResources(const UpdateResourcesRequest &request,
-                             UpdateResourcesReply *reply,
-                             SendReplyCallback send_reply_callback) override;
-
-  void HandleDeleteResources(const DeleteResourcesRequest &request,
-                             DeleteResourcesReply *reply,
-                             SendReplyCallback send_reply_callback) override;
-
- private:
-  gcs::RedisGcsClient &gcs_client_;
-  gcs::GcsNodeManager &gcs_node_manager_;
-};
-
-}  // namespace rpc
-}  // namespace ray
-
-#endif  // RAY_GCS_NODE_INFO_HANDLER_IMPL_H

src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc

@@ -476,15 +476,6 @@ TEST_F(GcsServerTest, TestNodeInfo) {
   report_heartbeat_request.mutable_heartbeat()->set_client_id(gcs_node_info->node_id());
   ASSERT_TRUE(ReportHeartbeat(report_heartbeat_request));
 
-  // Unregister node info
-  rpc::UnregisterNodeRequest unregister_node_info_request;
-  unregister_node_info_request.set_node_id(gcs_node_info->node_id());
-  ASSERT_TRUE(UnregisterNode(unregister_node_info_request));
-  node_info_list = GetAllNodeInfo();
-  ASSERT_TRUE(node_info_list.size() == 1);
-  ASSERT_TRUE(node_info_list[0].state() ==
-              rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_DEAD);
-
   // Update node resources
   rpc::UpdateResourcesRequest update_resources_request;
   update_resources_request.set_node_id(gcs_node_info->node_id());
@@ -503,6 +494,15 @@ TEST_F(GcsServerTest, TestNodeInfo) {
   ASSERT_TRUE(DeleteResources(delete_resources_request));
   resources = GetResources(gcs_node_info->node_id());
   ASSERT_TRUE(resources.empty());
+
+  // Unregister node info
+  rpc::UnregisterNodeRequest unregister_node_info_request;
+  unregister_node_info_request.set_node_id(gcs_node_info->node_id());
+  ASSERT_TRUE(UnregisterNode(unregister_node_info_request));
+  node_info_list = GetAllNodeInfo();
+  ASSERT_TRUE(node_info_list.size() == 1);
+  ASSERT_TRUE(node_info_list[0].state() ==
+              rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_DEAD);
 }
 
 TEST_F(GcsServerTest, TestObjectInfo) {