GCS adapts to batch heartbeat table pub sub (#8346)

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -283,10 +283,7 @@ Status ServiceBasedActorInfoAccessor::AsyncGetCheckpointID(
 ServiceBasedNodeInfoAccessor::ServiceBasedNodeInfoAccessor(
     ServiceBasedGcsClient *client_impl)
     : client_impl_(client_impl),
-      resource_sub_executor_(client_impl->GetRedisGcsClient().resource_table()),
-      heartbeat_sub_executor_(client_impl->GetRedisGcsClient().heartbeat_table()),
-      heartbeat_batch_sub_executor_(
-          client_impl->GetRedisGcsClient().heartbeat_batch_table()) {}
+      resource_sub_executor_(client_impl->GetRedisGcsClient().resource_table()) {}
 
 Status ServiceBasedNodeInfoAccessor::RegisterSelf(const GcsNodeInfo &local_node_info) {
   auto node_id = ClientID::FromBinary(local_node_info.node_id());
@@ -577,12 +574,13 @@ Status ServiceBasedNodeInfoAccessor::AsyncSubscribeBatchHeartbeat(
     const StatusCallback &done) {
   RAY_LOG(DEBUG) << "Subscribing batch heartbeat.";
   RAY_CHECK(subscribe != nullptr);
-  auto on_subscribe = [subscribe](const ClientID &node_id,
-                                  const HeartbeatBatchTableData &data) {
-    subscribe(data);
+  auto on_subscribe = [subscribe](const std::string &id, const std::string &data) {
+    rpc::HeartbeatBatchTableData heartbeat_batch_table_data;
+    heartbeat_batch_table_data.ParseFromString(data);
+    subscribe(heartbeat_batch_table_data);
   };
-  auto status = heartbeat_batch_sub_executor_.AsyncSubscribeAll(ClientID::Nil(),
-                                                                on_subscribe, done);
+  auto status = client_impl_->GetGcsPubSub().Subscribe(
+      HEARTBEAT_BATCH_CHANNEL, ClientID::Nil().Hex(), on_subscribe, done);
   RAY_LOG(DEBUG) << "Finished subscribing batch heartbeat.";
   return status;
 }

src/ray/gcs/gcs_client/service_based_accessor.h

@@ -181,14 +181,6 @@ class ServiceBasedNodeInfoAccessor : public NodeInfoAccessor {
       DynamicResourceSubscriptionExecutor;
   DynamicResourceSubscriptionExecutor resource_sub_executor_;
 
-  typedef SubscriptionExecutor<ClientID, HeartbeatTableData, HeartbeatTable>
-      HeartbeatSubscriptionExecutor;
-  HeartbeatSubscriptionExecutor heartbeat_sub_executor_;
-
-  typedef SubscriptionExecutor<ClientID, HeartbeatBatchTableData, HeartbeatBatchTable>
-      HeartbeatBatchSubscriptionExecutor;
-  HeartbeatBatchSubscriptionExecutor heartbeat_batch_sub_executor_;
-
   GcsNodeInfo local_node_info_;
   ClientID local_node_id_;
 

src/ray/gcs/gcs_server/gcs_node_manager.cc

@@ -22,11 +22,13 @@ namespace gcs {
 
 GcsNodeManager::NodeFailureDetector::NodeFailureDetector(
     boost::asio::io_service &io_service, gcs::NodeInfoAccessor &node_info_accessor,
+    std::shared_ptr<gcs::GcsPubSub> gcs_pub_sub,
     std::function<void(const ClientID &)> on_node_death_callback)
     : node_info_accessor_(node_info_accessor),
       on_node_death_callback_(std::move(on_node_death_callback)),
       num_heartbeats_timeout_(RayConfig::instance().num_heartbeats_timeout()),
-      detect_timer_(io_service) {
+      detect_timer_(io_service),
+      gcs_pub_sub_(std::move(gcs_pub_sub)) {
   Tick();
 }
 
@@ -77,7 +79,12 @@ void GcsNodeManager::NodeFailureDetector::SendBatchedHeartbeat() {
     for (const auto &heartbeat : heartbeat_buffer_) {
       batch->add_batch()->CopyFrom(heartbeat.second);
     }
-    RAY_CHECK_OK(node_info_accessor_.AsyncReportBatchHeartbeat(batch, nullptr));
+
+    auto done = [this, batch](Status status) {
+      RAY_CHECK_OK(gcs_pub_sub_->Publish(HEARTBEAT_BATCH_CHANNEL, ClientID::Nil().Hex(),
+                                         batch->SerializeAsString(), nullptr));
+    };
+    RAY_CHECK_OK(node_info_accessor_.AsyncReportBatchHeartbeat(batch, done));
     heartbeat_buffer_.clear();
   }
 }
@@ -105,7 +112,7 @@ GcsNodeManager::GcsNodeManager(boost::asio::io_service &io_service,
     : node_info_accessor_(node_info_accessor),
       error_info_accessor_(error_info_accessor),
       node_failure_detector_(new NodeFailureDetector(
-          io_service, node_info_accessor,
+          io_service, node_info_accessor, gcs_pub_sub,
           [this](const ClientID &node_id) {
             if (auto node = RemoveNode(node_id, /* is_intended = */ false)) {
               node->set_state(rpc::GcsNodeInfo::DEAD);

src/ray/gcs/gcs_server/gcs_node_manager.h

@@ -132,6 +132,7 @@ class GcsNodeManager : public rpc::NodeInfoHandler {
     /// \param node_info_accessor The node info accessor.
     explicit NodeFailureDetector(
         boost::asio::io_service &io_service, gcs::NodeInfoAccessor &node_info_accessor,
+        std::shared_ptr<gcs::GcsPubSub> gcs_pub_sub,
         std::function<void(const ClientID &)> on_node_death_callback);
 
     /// Register node to this detector.
@@ -177,6 +178,8 @@ class GcsNodeManager : public rpc::NodeInfoHandler {
     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_;
+    /// A publisher for publishing gcs messages.
+    std::shared_ptr<gcs::GcsPubSub> gcs_pub_sub_;
   };
 
  private:

src/ray/gcs/pubsub/gcs_pub_sub.h

@@ -31,6 +31,7 @@ namespace gcs {
 #define OBJECT_CHANNEL "OBJECT"
 #define TASK_CHANNEL "TASK"
 #define TASK_LEASE_CHANNEL "TASK_LEASE"
+#define HEARTBEAT_BATCH_CHANNEL "HEARTBEAT_BATCH"
 
 /// \class GcsPubSub
 ///