Fetch node and shard heap metrics in a single call (#158087)

- Exposes method to fetch node and shard heap metrics in one call using one snapshot of memory metrics
- Adds tests
- Closes #4520

Author: surya-estc

server/src/internalClusterTest/java/org/elasticsearch/index/shard/IndexShardIT.java

@@ -22,6 +22,7 @@
 import org.elasticsearch.cluster.ClusterInfoServiceUtils;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.EstimatedHeapUsageCollector;
+import org.elasticsearch.cluster.EstimatedHeapUsageStats;
 import org.elasticsearch.cluster.InternalClusterInfoService;
 import org.elasticsearch.cluster.NodeHeapEstimates;
 import org.elasticsearch.cluster.NodeHeapMetrics;
@@ -1008,29 +1009,19 @@ public BogusEstimatedEstimatedHeapUsageCollector(BogusEstimatedHeapUsagePlugin p
         }
 
         @Override
-        public void collectClusterHeapUsage(ActionListener<Map<String, NodeHeapEstimates>> listener) {
+        public void collectEstimatedHeapUsage(ActionListener<EstimatedHeapUsageStats> listener) {
             final long totalHeapUsageBytes = randomNonNegativeLong();
-            ActionListener.completeWith(
-                listener,
-                () -> plugin.getClusterService()
-                    .state()
-                    .nodes()
+            ActionListener.completeWith(listener, () -> {
+                ClusterState state = plugin.getClusterService().state();
+                Map<String, NodeHeapEstimates> nodeHeapEstimates = state.nodes()
                     .stream()
                     .collect(
                         Collectors.toUnmodifiableMap(
                             DiscoveryNode::getId,
                             node -> new NodeHeapEstimates(totalHeapUsageBytes, randomLongBetween(0, totalHeapUsageBytes))
                         )
-                    )
-            );
-        }
-
-        @Override
-        public void collectShardHeapUsage(ActionListener<ShardHeapUsageEstimates> listener) {
-            ActionListener.completeWith(listener, () -> {
-                var perShard = plugin.getClusterService()
-                    .state()
-                    .getRoutingNodes()
+                    );
+                var perShard = state.getRoutingNodes()
                     .stream()
                     .map(node -> node.started())
                     .flatMap(nodeIt -> StreamSupport.stream(nodeIt.spliterator(), false))
@@ -1040,7 +1031,10 @@ public void collectShardHeapUsage(ActionListener<ShardHeapUsageEstimates> listen
                             shardRouting -> new ShardAndIndexHeapUsage(randomShardHeapUsage(), randomIndexHeapUsage())
                         )
                     );
-                return new ShardHeapUsageEstimates(perShard, new ShardAndIndexHeapUsage(randomShardHeapUsage(), randomIndexHeapUsage()));
+                return new EstimatedHeapUsageStats(
+                    nodeHeapEstimates,
+                    new ShardHeapUsageEstimates(perShard, new ShardAndIndexHeapUsage(randomShardHeapUsage(), randomIndexHeapUsage()))
+                );
             });
         }
     }

server/src/main/java/org/elasticsearch/cluster/EstimatedHeapUsageCollector.java

@@ -11,12 +11,8 @@
 
 import org.elasticsearch.action.ActionListener;
 
-import java.util.Map;
-
 /**
- * Collect the estimated heap usage for each node in the cluster.
- * <p>
- * Results are returned as a map of node ID to estimated heap usage in bytes
+ * Collects node and shard heap usage estimates for {@link ClusterInfo}.
  *
  * @see NodeHeapMetrics
  */
@@ -25,29 +21,12 @@ public interface EstimatedHeapUsageCollector {
     /**
      * This will be used when there is no EstimatedHeapUsageCollector available
      */
-    EstimatedHeapUsageCollector EMPTY = new EstimatedHeapUsageCollector() {
-        @Override
-        public void collectClusterHeapUsage(ActionListener<Map<String, NodeHeapEstimates>> listener) {
-            listener.onResponse(Map.of());
-        }
-
-        @Override
-        public void collectShardHeapUsage(ActionListener<ShardHeapUsageEstimates> listener) {
-            listener.onResponse(ShardHeapUsageEstimates.empty());
-        }
-    };
-
-    /**
-     * Collect the estimated heap usage for every node in the cluster
-     *
-     * @param listener The listener which will receive the results
-     */
-    void collectClusterHeapUsage(ActionListener<Map<String, NodeHeapEstimates>> listener);
+    EstimatedHeapUsageCollector EMPTY = listener -> listener.onResponse(EstimatedHeapUsageStats.EMPTY);
 
     /**
-     * Collects the estimated heap usage for every shard in the cluster.
+     * Collects node and shard heap usage estimates from a single source snapshot.
      *
      * @param listener The listener which will receive the results
      */
-    void collectShardHeapUsage(ActionListener<ShardHeapUsageEstimates> listener);
+    void collectEstimatedHeapUsage(ActionListener<EstimatedHeapUsageStats> listener);
 }

server/src/main/java/org/elasticsearch/cluster/EstimatedHeapUsageStats.java

@@ -0,0 +1,25 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.cluster;
+
+import java.util.Map;
+
+/**
+ * Node heap usage estimates and individual shard heap usage estimates collected from the same source snapshot.
+ */
+public record EstimatedHeapUsageStats(Map<String, NodeHeapEstimates> nodeHeapEstimates, ShardHeapUsageEstimates shardHeapUsageEstimates) {
+
+    public static final EstimatedHeapUsageStats EMPTY = new EstimatedHeapUsageStats(Map.of(), ShardHeapUsageEstimates.empty());
+
+    public EstimatedHeapUsageStats {
+        nodeHeapEstimates = Map.copyOf(nodeHeapEstimates);
+        assert shardHeapUsageEstimates != null;
+    }
+}

server/src/main/java/org/elasticsearch/cluster/InternalClusterInfoService.java

@@ -376,28 +376,17 @@ public void onFailure(Exception e) {
         }
 
         private void fetchEstimatedHeapUsage() {
-            estimatedHeapUsageCollector.collectClusterHeapUsage(ActionListener.releaseAfter(new ActionListener<>() {
+            estimatedHeapUsageCollector.collectEstimatedHeapUsage(ActionListener.releaseAfter(new ActionListener<>() {
                 @Override
-                public void onResponse(Map<String, NodeHeapEstimates> currentNodeHeapEstimates) {
-                    nodeHeapEstimates = currentNodeHeapEstimates;
+                public void onResponse(EstimatedHeapUsageStats estimatedHeapUsageStats) {
+                    nodeHeapEstimates = estimatedHeapUsageStats.nodeHeapEstimates();
+                    estimatedShardHeapUsageEstimates = estimatedHeapUsageStats.shardHeapUsageEstimates();
                 }
 
                 @Override
                 public void onFailure(Exception e) {
-                    logger.warn("failed to fetch heap usage for nodes", e);
+                    logger.warn("failed to fetch node heap estimates and shard heap usage estimates", e);
                     nodeHeapEstimates = Map.of();
-                }
-            }, fetchRefs.acquire()));
-
-            estimatedHeapUsageCollector.collectShardHeapUsage(ActionListener.releaseAfter(new ActionListener<>() {
-                @Override
-                public void onResponse(ShardHeapUsageEstimates currentEstimatedHeapUsages) {
-                    estimatedShardHeapUsageEstimates = currentEstimatedHeapUsages;
-                }
-
-                @Override
-                public void onFailure(Exception e) {
-                    logger.warn("failed to fetch heap usage for shards", e);
                     estimatedShardHeapUsageEstimates = ShardHeapUsageEstimates.empty();
                 }
             }, fetchRefs.acquire()));