Repository: ray-project/ray
File: src/ray/gcs/gcs_node_manager.h (C/C++)

Code before:
  absl::flat_hash_map<NodeID, std::shared_ptr<const rpc::GcsNodeInfo>> alive_nodes_
      ABSL_GUARDED_BY(mutex_);
  /// Draining nodes.
  /// This map is used to store the nodes which have received the drain request.
  /// Invariant: its keys should always be a subset of the keys of `alive_nodes_`,
  /// and entry in it should be removed whenever a node is removed from `alive_nodes_`.
  absl::flat_hash_map<NodeID, std::shared_ptr<const rpc::autoscaler::DrainNodeRequest>>
      draining_nodes_ ABSL_GUARDED_BY(mutex_);
  /// Dead nodes.
  absl::flat_hash_map<NodeID, std::shared_ptr<const rpc::GcsNodeInfo>> dead_nodes_
      ABSL_GUARDED_BY(mutex_);
  /// The nodes are sorted according to the timestamp, and the oldest is at the head of
  /// the deque.
  std::deque<std::pair<NodeID, int64_t>> sorted_dead_node_list_ ABSL_GUARDED_BY(mutex_);

  /// Listeners which monitors the addition of nodes.
  std::vector<Postable<void(std::shared_ptr<const rpc::GcsNodeInfo>)>>
      node_added_listeners_ ABSL_GUARDED_BY(mutex_);

  /// Listeners which monitors the removal of nodes.
  std::vector<Postable<void(std::shared_ptr<const rpc::GcsNodeInfo>)>>
      node_removed_listeners_ ABSL_GUARDED_BY(mutex_);

  /// Listeners which monitors when nodes are set to draining.
  /// Uses std::function (not Postable) for synchronous invocation to ensure
  /// scheduler sees draining state before HandleDrainNode returns.
  std::vector<std::function<void(const NodeID &, bool, int64_t)>> node_draining_listeners_
      ABSL_GUARDED_BY(mutex_);

  /// A publisher for publishing gcs messages.
  pubsub::GcsPublisher *gcs_publisher_;
  /// Storage for GCS tables.
  GcsTableStorage *gcs_table_storage_;
  instrumented_io_context &io_context_;
  /// Raylet client pool.
  rpc::RayletClientPool *raylet_client_pool_;
  /// Cluster ID to be shared with clients when connecting.
  const ClusterID cluster_id_;
  /// Class lock for node manager
  mutable absl::Mutex mutex_;

  observability::RayEventRecorderInterface &ray_event_recorder_;
  std::string session_name_;

  // Debug info.
  enum CountType {
    REGISTER_NODE_REQUEST = 0,
    DRAIN_NODE_REQUEST = 1,
    GET_ALL_NODE_INFO_REQUEST = 2,
    CountType_MAX = 3,
  };

Review comment (suggestion):
```suggestion
  /// Listeners which monitor when nodes are set to draining.
```

Code after:
  absl::flat_hash_map<NodeID, std::shared_ptr<const rpc::GcsNodeInfo>> alive_nodes_
      ABSL_GUARDED_BY(mutex_);
  /// Draining nodes.
  /// This map is used to store the nodes which have received the drain request.
  /// Invariant: its keys should always be a subset of the keys of `alive_nodes_`,
  /// and entry in it should be removed whenever a node is removed from `alive_nodes_`.
  absl::flat_hash_map<NodeID, std::shared_ptr<const rpc::autoscaler::DrainNodeRequest>>
      draining_nodes_ ABSL_GUARDED_BY(mutex_);
  /// Dead nodes.
  absl::flat_hash_map<NodeID, std::shared_ptr<const rpc::GcsNodeInfo>> dead_nodes_
      ABSL_GUARDED_BY(mutex_);
  /// The nodes are sorted according to the timestamp, and the oldest is at the head of
  /// the deque.
  std::deque<std::pair<NodeID, int64_t>> sorted_dead_node_list_ ABSL_GUARDED_BY(mutex_);

  /// Listeners which monitors the addition of nodes.
  std::vector<Postable<void(std::shared_ptr<const rpc::GcsNodeInfo>)>>
      node_added_listeners_ ABSL_GUARDED_BY(mutex_);

  /// Listeners which monitors the removal of nodes.
  std::vector<Postable<void(std::shared_ptr<const rpc::GcsNodeInfo>)>>
      node_removed_listeners_ ABSL_GUARDED_BY(mutex_);

  /// Listeners which monitor when nodes are set to draining.
  std::vector<std::function<void(const NodeID &, bool, int64_t)>> node_draining_listeners_
      ABSL_GUARDED_BY(mutex_);

  /// A publisher for publishing gcs messages.
  pubsub::GcsPublisher *gcs_publisher_;
  /// Storage for GCS tables.
  GcsTableStorage *gcs_table_storage_;
  instrumented_io_context &io_context_;
  /// Raylet client pool.
  rpc::RayletClientPool *raylet_client_pool_;
  /// Cluster ID to be shared with clients when connecting.
  const ClusterID cluster_id_;
  /// Class lock for node manager
  mutable absl::Mutex mutex_;

  observability::RayEventRecorderInterface &ray_event_recorder_;
  std::string session_name_;

  // Debug info.
  enum CountType {
    REGISTER_NODE_REQUEST = 0,
    DRAIN_NODE_REQUEST = 1,
    GET_ALL_NODE_INFO_REQUEST = 2,
    CountType_MAX = 3,
  };
  std::atomic<uint64_t> counts_[CountType::CountType_MAX] = {0};