Add heartbeat methods to gcs server node info handler (#6647)

src/ray/gcs/gcs_server/node_info_handler_impl.cc

@@ -61,5 +61,54 @@ void DefaultNodeInfoHandler::HandleGetAllNodeInfo(
   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, send_reply_callback](Status status) {
+    if (!status.ok()) {
+      RAY_LOG(ERROR) << "Failed to report heartbeat: " << status.ToString()
+                     << ", node id = " << node_id;
+    }
+    send_reply_callback(status, nullptr, nullptr);
+  };
+
+  auto heartbeat_data = std::make_shared<rpc::HeartbeatTableData>();
+  heartbeat_data->CopyFrom(request.heartbeat());
+  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::HandleReportBatchHeartbeat(
+    const ReportBatchHeartbeatRequest &request, ReportBatchHeartbeatReply *reply,
+    SendReplyCallback send_reply_callback) {
+  RAY_LOG(DEBUG) << "Reporting batch heartbeat, batch size = "
+                 << request.heartbeat_batch().batch_size();
+
+  auto on_done = [&request, send_reply_callback](Status status) {
+    if (!status.ok()) {
+      RAY_LOG(ERROR) << "Failed to report batch heartbeat: " << status.ToString()
+                     << ", batch size = " << request.heartbeat_batch().batch_size();
+    }
+    send_reply_callback(status, nullptr, nullptr);
+  };
+
+  auto heartbeat_batch_data = std::make_shared<rpc::HeartbeatBatchTableData>();
+  heartbeat_batch_data->CopyFrom(request.heartbeat_batch());
+  Status status =
+      gcs_client_.Nodes().AsyncReportBatchHeartbeat(heartbeat_batch_data, on_done);
+  if (!status.ok()) {
+    on_done(status);
+  }
+
+  RAY_LOG(DEBUG) << "Finished reporting batch heartbeat, batch size = "
+                 << request.heartbeat_batch().batch_size();
+}
+
 }  // namespace rpc
 }  // namespace ray

src/ray/gcs/gcs_server/node_info_handler_impl.h

@@ -24,6 +24,14 @@ class DefaultNodeInfoHandler : public rpc::NodeInfoHandler {
                             GetAllNodeInfoReply *reply,
                             SendReplyCallback send_reply_callback) override;
 
+  void HandleReportHeartbeat(const ReportHeartbeatRequest &request,
+                             ReportHeartbeatReply *reply,
+                             SendReplyCallback send_reply_callback) override;
+
+  void HandleReportBatchHeartbeat(const ReportBatchHeartbeatRequest &request,
+                                  ReportBatchHeartbeatReply *reply,
+                                  SendReplyCallback send_reply_callback) override;
+
  private:
   gcs::RedisGcsClient &gcs_client_;
 };

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

@@ -144,6 +144,27 @@ class GcsServerTest : public RedisServiceManagerForTest {
     return node_info_list;
   }
 
+  bool ReportHeartbeat(const rpc::ReportHeartbeatRequest &request) {
+    std::promise<bool> promise;
+    client_->ReportHeartbeat(request, [&promise](const Status &status,
+                                                 const rpc::ReportHeartbeatReply &reply) {
+      RAY_CHECK_OK(status);
+      promise.set_value(true);
+    });
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
+  bool ReportBatchHeartbeat(const rpc::ReportBatchHeartbeatRequest &request) {
+    std::promise<bool> promise;
+    client_->ReportBatchHeartbeat(
+        request,
+        [&promise](const Status &status, const rpc::ReportBatchHeartbeatReply &reply) {
+          RAY_CHECK_OK(status);
+          promise.set_value(true);
+        });
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool AddObjectLocation(const rpc::AddObjectLocationRequest &request) {
     std::promise<bool> promise;
     client_->AddObjectLocation(
@@ -291,6 +312,15 @@ TEST_F(GcsServerTest, TestNodeInfo) {
   ASSERT_TRUE(node_info_list[0].state() ==
               rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_ALIVE);
 
+  // Report heartbeat
+  rpc::ReportHeartbeatRequest report_heartbeat_request;
+  report_heartbeat_request.mutable_heartbeat()->set_client_id(node_id.Binary());
+  ASSERT_TRUE(ReportHeartbeat(report_heartbeat_request));
+  rpc::ReportBatchHeartbeatRequest report_batch_heartbeat_request;
+  report_batch_heartbeat_request.mutable_heartbeat_batch()->add_batch()->set_client_id(
+      node_id.Binary());
+  ASSERT_TRUE(ReportBatchHeartbeat(report_batch_heartbeat_request));
+
   // Unregister node info
   rpc::UnregisterNodeRequest unregister_node_info_request;
   unregister_node_info_request.set_node_id(node_id.Binary());

src/ray/protobuf/gcs_service.proto

@@ -89,6 +89,20 @@ message GetAllNodeInfoReply {
   repeated GcsNodeInfo node_info_list = 1;
 }
 
+message ReportHeartbeatRequest {
+  HeartbeatTableData heartbeat = 1;
+}
+
+message ReportHeartbeatReply {
+}
+
+message ReportBatchHeartbeatRequest {
+  HeartbeatBatchTableData heartbeat_batch = 1;
+}
+
+message ReportBatchHeartbeatReply {
+}
+
 // Service for node info access.
 service NodeInfoGcsService {
   // Register a node to GCS Service.
@@ -97,6 +111,11 @@ service NodeInfoGcsService {
   rpc UnregisterNode(UnregisterNodeRequest) returns (UnregisterNodeReply);
   // Get information of all nodes from GCS Service.
   rpc GetAllNodeInfo(GetAllNodeInfoRequest) returns (GetAllNodeInfoReply);
+  // Report heartbeat of a node to GCS Service.
+  rpc ReportHeartbeat(ReportHeartbeatRequest) returns (ReportHeartbeatReply);
+  // Report batch heartbeat to GCS Service.
+  rpc ReportBatchHeartbeat(ReportBatchHeartbeatRequest)
+      returns (ReportBatchHeartbeatReply);
 }
 
 message GetObjectLocationsRequest {

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -76,6 +76,30 @@ class GcsRpcClient {
   VOID_RPC_CLIENT_METHOD(ObjectInfoGcsService, RemoveObjectLocation, request, callback,
                          object_info_grpc_client_)
 
+  /// Report heartbeat of a node to GCS Service.
+  ///
+  /// \param request The request message.
+  /// \param callback The callback function that handles reply from server.
+  void ReportHeartbeat(const ReportHeartbeatRequest &request,
+                       const ClientCallback<ReportHeartbeatReply> &callback) {
+    client_call_manager_
+        .CreateCall<NodeInfoGcsService, ReportHeartbeatRequest, ReportHeartbeatReply>(
+            *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncReportHeartbeat,
+            request, callback);
+  }
+
+  /// Report batch heartbeat to GCS Service.
+  ///
+  /// \param request The request message.
+  /// \param callback The callback function that handles reply from server.
+  void ReportBatchHeartbeat(const ReportBatchHeartbeatRequest &request,
+                            const ClientCallback<ReportBatchHeartbeatReply> &callback) {
+    client_call_manager_.CreateCall<NodeInfoGcsService, ReportBatchHeartbeatRequest,
+                                    ReportBatchHeartbeatReply>(
+        *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncReportBatchHeartbeat,
+        request, callback);
+  }
+
  private:
   /// The gRPC-generated stub.
   std::unique_ptr<GrpcClient<JobInfoGcsService>> job_info_grpc_client_;

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -146,6 +146,14 @@ class NodeInfoHandler {
   virtual void HandleGetAllNodeInfo(const GetAllNodeInfoRequest &request,
                                     GetAllNodeInfoReply *reply,
                                     SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleReportHeartbeat(const ReportHeartbeatRequest &request,
+                                     ReportHeartbeatReply *reply,
+                                     SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleReportBatchHeartbeat(const ReportBatchHeartbeatRequest &request,
+                                          ReportBatchHeartbeatReply *reply,
+                                          SendReplyCallback send_reply_callback) = 0;
 };
 
 /// The `GrpcService` for `NodeInfoGcsService`.
@@ -168,6 +176,8 @@ class NodeInfoGrpcService : public GrpcService {
     NODE_INFO_SERVICE_RPC_HANDLER(RegisterNode, 1);
     NODE_INFO_SERVICE_RPC_HANDLER(UnregisterNode, 1);
     NODE_INFO_SERVICE_RPC_HANDLER(GetAllNodeInfo, 1);
+    NODE_INFO_SERVICE_RPC_HANDLER(ReportHeartbeat, 1);
+    NODE_INFO_SERVICE_RPC_HANDLER(ReportBatchHeartbeat, 1);
   }
 
  private: