[GCS] Re-report heartbeat when gcs server restarts (#10040)

* Retry to send failed heartbeat when light heartbeat enalbed

* Re-report heartbeat when gcs server restarts

* remove is_pubsub_server_restarted

* add lock per comment

* minor change, name related

src/ray/gcs/accessor.h

@@ -551,6 +551,9 @@ class NodeInfoAccessor {
       const std::shared_ptr<rpc::HeartbeatTableData> &data_ptr,
       const StatusCallback &callback) = 0;
 
+  /// Resend heartbeat when GCS restarts from a failure.
+  virtual void AsyncReReportHeartbeat() = 0;
+
   /// Subscribe to the heartbeat of each node from GCS.
   ///
   /// \param subscribe Callback that will be called each time when heartbeat is updated.

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -709,10 +709,11 @@ Status ServiceBasedNodeInfoAccessor::AsyncSubscribeToResources(
 Status ServiceBasedNodeInfoAccessor::AsyncReportHeartbeat(
     const std::shared_ptr<rpc::HeartbeatTableData> &data_ptr,
     const StatusCallback &callback) {
-  rpc::ReportHeartbeatRequest request;
-  request.mutable_heartbeat()->CopyFrom(*data_ptr);
+  absl::MutexLock lock(&mutex_);
+  cached_heartbeat_.mutable_heartbeat()->CopyFrom(*data_ptr);
   client_impl_->GetGcsRpcClient().ReportHeartbeat(
-      request, [callback](const Status &status, const rpc::ReportHeartbeatReply &reply) {
+      cached_heartbeat_,
+      [callback](const Status &status, const rpc::ReportHeartbeatReply &reply) {
         if (callback) {
           callback(status);
         }
@@ -720,6 +721,16 @@ Status ServiceBasedNodeInfoAccessor::AsyncReportHeartbeat(
   return Status::OK();
 }
 
+void ServiceBasedNodeInfoAccessor::AsyncReReportHeartbeat() {
+  absl::MutexLock lock(&mutex_);
+  if (cached_heartbeat_.has_heartbeat()) {
+    RAY_LOG(INFO) << "Rereport heartbeat.";
+    client_impl_->GetGcsRpcClient().ReportHeartbeat(
+        cached_heartbeat_,
+        [](const Status &status, const rpc::ReportHeartbeatReply &reply) {});
+  }
+}
+
 Status ServiceBasedNodeInfoAccessor::AsyncSubscribeHeartbeat(
     const SubscribeCallback<ClientID, rpc::HeartbeatTableData> &subscribe,
     const StatusCallback &done) {

src/ray/gcs/gcs_client/service_based_accessor.h

@@ -18,6 +18,7 @@
 #include "ray/gcs/accessor.h"
 #include "ray/gcs/subscription_executor.h"
 #include "ray/util/sequencer.h"
+#include "src/ray/protobuf/gcs_service.pb.h"
 
 namespace ray {
 namespace gcs {
@@ -193,6 +194,8 @@ class ServiceBasedNodeInfoAccessor : public NodeInfoAccessor {
   Status AsyncReportHeartbeat(const std::shared_ptr<rpc::HeartbeatTableData> &data_ptr,
                               const StatusCallback &callback) override;
 
+  void AsyncReReportHeartbeat() override;
+
   Status AsyncSubscribeHeartbeat(
       const SubscribeCallback<ClientID, rpc::HeartbeatTableData> &subscribe,
       const StatusCallback &done) override;
@@ -225,6 +228,13 @@ class ServiceBasedNodeInfoAccessor : public NodeInfoAccessor {
   /// server restarts from a failure.
   FetchDataOperation fetch_node_data_operation_;
 
+  // Mutex to protect the cached_heartbeat_ field.
+  absl::Mutex mutex_;
+
+  /// Save the heartbeat data, so we can resend it again when GCS server restarts from a
+  /// failure.
+  rpc::ReportHeartbeatRequest cached_heartbeat_ GUARDED_BY(mutex_);
+
   void HandleNotification(const GcsNodeInfo &node_info);
 
   ServiceBasedGcsClient *client_impl_;

src/ray/gcs/gcs_client/service_based_gcs_client.cc

@@ -175,6 +175,8 @@ void ServiceBasedGcsClient::GcsServiceFailureDetected(rpc::GcsServiceFailureType
     // because we use the same Redis server for both GCS storage and pub-sub. So the
     // following flag is alway false.
     resubscribe_func_(false);
+    // Resend heartbeat after reconnected, needed by resource view in GCS.
+    node_accessor_->AsyncReReportHeartbeat();
     break;
   default:
     RAY_LOG(FATAL) << "Unsupported failure type: " << type;

src/ray/gcs/redis_accessor.cc

@@ -660,6 +660,8 @@ Status RedisNodeInfoAccessor::AsyncReportHeartbeat(
   return heartbeat_table.Add(JobID::Nil(), node_id, data_ptr, on_done);
 }
 
+void RedisNodeInfoAccessor::AsyncReReportHeartbeat() {}
+
 Status RedisNodeInfoAccessor::AsyncSubscribeHeartbeat(
     const SubscribeCallback<ClientID, HeartbeatTableData> &subscribe,
     const StatusCallback &done) {

src/ray/gcs/redis_accessor.h

@@ -362,6 +362,8 @@ class RedisNodeInfoAccessor : public NodeInfoAccessor {
   Status AsyncReportHeartbeat(const std::shared_ptr<HeartbeatTableData> &data_ptr,
                               const StatusCallback &callback) override;
 
+  void AsyncReReportHeartbeat() override;
+
   Status AsyncSubscribeHeartbeat(
       const SubscribeCallback<ClientID, HeartbeatTableData> &subscribe,
       const StatusCallback &done) override;

src/ray/raylet/node_manager.cc

@@ -464,9 +464,8 @@ void NodeManager::Heartbeat() {
     should_local_gc_ = false;
   }
 
-  ray::Status status = gcs_client_->Nodes().AsyncReportHeartbeat(heartbeat_data,
-                                                                 /*done*/ nullptr);
-  RAY_CHECK_OK_PREPEND(status, "Heartbeat failed");
+  RAY_CHECK_OK(
+      gcs_client_->Nodes().AsyncReportHeartbeat(heartbeat_data, /*done*/ nullptr));
 
   if (debug_dump_period_ > 0 &&
       static_cast<int64_t>(now_ms - last_debug_dump_at_ms_) > debug_dump_period_) {