Add shard heap usage to ClusterInfo

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

@@ -25,6 +25,7 @@
 import org.elasticsearch.cluster.EstimatedHeapUsageCollector;
 import org.elasticsearch.cluster.InternalClusterInfoService;
 import org.elasticsearch.cluster.NodeUsageStatsForThreadPools;
+import org.elasticsearch.cluster.ShardAndIndexHeapUsage;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.ProjectId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
@@ -105,6 +106,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 import static com.carrotsearch.randomizedtesting.RandomizedTest.randomAsciiLettersOfLength;
 import static java.util.Collections.emptySet;
@@ -304,6 +306,49 @@ public void testHeapUsageEstimateIsPresent() {
         }
     }
 
+    public void testShardHeapUsagesArePresent() {
+        // Create some indices so we can later expect a precise number of (random) shard-level heap usage reports.
+        final int numIndices = randomIntBetween(1, 5);
+        final int numShards = randomIntBetween(1, 3);
+        final String indexPrefix = randomIdentifier();
+        IntStream.range(0, numIndices).forEach(i -> {
+            final String indexName = indexPrefix + "_" + i;
+            createIndex(indexName, Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, numShards).build());
+        });
+
+        InternalClusterInfoService clusterInfoService = (InternalClusterInfoService) getInstanceFromNode(ClusterInfoService.class);
+        ClusterInfoServiceUtils.refresh(clusterInfoService);
+
+        Map<ShardId, ShardAndIndexHeapUsage> estimatedShardHeapUsages = clusterInfoService.getClusterInfo().getEstimatedShardHeapUsages();
+        assertNotNull(estimatedShardHeapUsages);
+        // No shard heap usage is reported because it is not yet enabled.
+        assertTrue(estimatedShardHeapUsages.isEmpty());
+
+        // Enable collection of heap usages for ClusterInfo.
+        updateClusterSettings(
+            Settings.builder()
+                .put(InternalClusterInfoService.CLUSTER_ROUTING_ALLOCATION_ESTIMATED_HEAP_THRESHOLD_DECIDER_ENABLED.getKey(), true)
+                .build()
+        );
+
+        try {
+            ClusterInfoServiceUtils.refresh(clusterInfoService);
+            estimatedShardHeapUsages = clusterInfoService.getClusterInfo().getEstimatedShardHeapUsages();
+            assertNotNull(estimatedShardHeapUsages);
+            assertEquals(estimatedShardHeapUsages.size(), numIndices * numShards);
+            for (var entry : estimatedShardHeapUsages.entrySet()) {
+                assertThat(entry.getValue().shardHeapUsageBytes(), greaterThanOrEqualTo(0L));
+                assertThat(entry.getValue().indexHeapUsageBytes(), greaterThanOrEqualTo(0L));
+            }
+        } finally {
+            updateClusterSettings(
+                Settings.builder()
+                    .putNull(InternalClusterInfoService.CLUSTER_ROUTING_ALLOCATION_ESTIMATED_HEAP_THRESHOLD_DECIDER_ENABLED.getKey())
+                    .build()
+            );
+        }
+    }
+
     public void testNodeWriteLoadsArePresent() {
         InternalClusterInfoService clusterInfoService = (InternalClusterInfoService) getInstanceFromNode(ClusterInfoService.class);
 
@@ -970,6 +1015,25 @@ public void collectClusterHeapUsage(ActionListener<Map<String, Long>> listener)
                     .collect(Collectors.toUnmodifiableMap(DiscoveryNode::getId, node -> randomNonNegativeLong()))
             );
         }
+
+        @Override
+        public void collectShardHeapUsage(ActionListener<Map<ShardId, ShardAndIndexHeapUsage>> listener) {
+            ActionListener.completeWith(
+                listener,
+                () -> plugin.getClusterService()
+                    .state()
+                    .getRoutingNodes()
+                    .stream()
+                    .map(node -> node.started())
+                    .flatMap(nodeIt -> StreamSupport.stream(nodeIt.spliterator(), false))
+                    .collect(
+                        Collectors.toUnmodifiableMap(
+                            shardRouting -> shardRouting.shardId(),
+                            shardRouting -> new ShardAndIndexHeapUsage(randomNonNegativeLong(), randomNonNegativeLong())
+                        )
+                    )
+            );
+        }
     }
 
     public static class BogusEstimatedHeapUsagePlugin extends Plugin implements ClusterPlugin {

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

@@ -63,6 +63,7 @@ public class ClusterInfo implements ChunkedToXContent, Writeable, ExpectedShardS
     private static final TransportVersion NODES_WRITE_LOAD_HOTSPOTTING_IN_CLUSTER_INFO = TransportVersion.fromName(
         "nodes_write_load_hotspotting_in_cluster_info"
     );
+    private static final TransportVersion SHARD_HEAP_USAGE_IN_CLUSTER_INFO = TransportVersion.fromName("shard_heap_usage_in_cluster_info");
 
     private final Map<String, DiskUsage> leastAvailableSpaceUsage;
     private final Map<String, DiskUsage> mostAvailableSpaceUsage;
@@ -71,6 +72,7 @@ public class ClusterInfo implements ChunkedToXContent, Writeable, ExpectedShardS
     final Map<NodeAndShard, String> dataPath;
     final Map<NodeAndPath, ReservedSpace> reservedSpace;
     final Map<String, EstimatedHeapUsage> estimatedHeapUsages;
+    final Map<ShardId, ShardAndIndexHeapUsage> estimatedShardHeapUsages;
     final Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools;
     final Map<ShardId, Double> shardWriteLoads;
     // max heap size per node ID
@@ -79,7 +81,7 @@ public class ClusterInfo implements ChunkedToXContent, Writeable, ExpectedShardS
     private final Map<ShardId, Set<String>> shardToNodeIds;
 
     protected ClusterInfo() {
-        this(Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Set.of());
+        this(Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Set.of());
     }
 
     /**
@@ -92,6 +94,7 @@ protected ClusterInfo() {
      * @param dataPath the shard routing to datapath mapping
      * @param reservedSpace reserved space per shard broken down by node and data path
      * @param estimatedHeapUsages estimated heap usage broken down by node
+     * @param estimatedShardHeapUsages estimated heap usage broken down by {@link ShardId}
      * @param nodeUsageStatsForThreadPools node-level usage stats (operational load) broken down by node
      * @see #shardIdentifierFromRouting
      * @param maxHeapSizePerNode node id to max heap size
@@ -105,6 +108,7 @@ public ClusterInfo(
         Map<NodeAndShard, String> dataPath,
         Map<NodeAndPath, ReservedSpace> reservedSpace,
         Map<String, EstimatedHeapUsage> estimatedHeapUsages,
+        Map<ShardId, ShardAndIndexHeapUsage> estimatedShardHeapUsages,
         Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools,
         Map<ShardId, Double> shardWriteLoads,
         Map<String, ByteSizeValue> maxHeapSizePerNode,
@@ -118,6 +122,7 @@ public ClusterInfo(
             dataPath,
             reservedSpace,
             estimatedHeapUsages,
+            estimatedShardHeapUsages,
             nodeUsageStatsForThreadPools,
             shardWriteLoads,
             maxHeapSizePerNode,
@@ -134,6 +139,7 @@ private ClusterInfo(
         Map<NodeAndShard, String> dataPath,
         Map<NodeAndPath, ReservedSpace> reservedSpace,
         Map<String, EstimatedHeapUsage> estimatedHeapUsages,
+        Map<ShardId, ShardAndIndexHeapUsage> estimatedShardHeapUsages,
         Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools,
         Map<ShardId, Double> shardWriteLoads,
         Map<String, ByteSizeValue> maxHeapSizePerNode,
@@ -147,6 +153,7 @@ private ClusterInfo(
         this.dataPath = Map.copyOf(dataPath);
         this.reservedSpace = Map.copyOf(reservedSpace);
         this.estimatedHeapUsages = Map.copyOf(estimatedHeapUsages);
+        this.estimatedShardHeapUsages = Map.copyOf(estimatedShardHeapUsages);
         this.nodeUsageStatsForThreadPools = Map.copyOf(nodeUsageStatsForThreadPools);
         this.shardWriteLoads = Map.copyOf(shardWriteLoads);
         this.maxHeapSizePerNode = Map.copyOf(maxHeapSizePerNode);
@@ -186,6 +193,11 @@ public ClusterInfo(StreamInput in) throws IOException {
         } else {
             this.nodeIdsWriteLoadHotspotting = Set.of();
         }
+        if (in.getTransportVersion().supports(SHARD_HEAP_USAGE_IN_CLUSTER_INFO)) {
+            this.estimatedShardHeapUsages = in.readImmutableMap(ShardId::new, ShardAndIndexHeapUsage::new);
+        } else {
+            this.estimatedShardHeapUsages = Map.of();
+        }
         this.shardToNodeIds = computeShardToNodeIds(dataPath);
     }
 
@@ -195,6 +207,7 @@ ClusterInfo updateWith(
         Map<String, Long> shardSizes,
         Map<NodeAndPath, ReservedSpace> reservedSpace,
         Map<String, EstimatedHeapUsage> estimatedHeapUsages,
+        Map<ShardId, ShardAndIndexHeapUsage> estimatedShardHeapUsages,
         Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools
     ) {
         return new ClusterInfo(
@@ -205,6 +218,7 @@ ClusterInfo updateWith(
             dataPath,
             reservedSpace,
             estimatedHeapUsages,
+            estimatedShardHeapUsages,
             nodeUsageStatsForThreadPools,
             shardWriteLoads,
             maxHeapSizePerNode,
@@ -253,6 +267,9 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().supports(NODES_WRITE_LOAD_HOTSPOTTING_IN_CLUSTER_INFO)) {
             out.writeStringCollection(this.nodeIdsWriteLoadHotspotting);
         }
+        if (out.getTransportVersion().supports(SHARD_HEAP_USAGE_IN_CLUSTER_INFO)) {
+            out.writeMap(this.estimatedShardHeapUsages);
+        }
     }
 
     /**
@@ -334,7 +351,8 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
             }),
             endArray() // end "reserved_sizes"
             // NOTE: We don't serialize estimatedHeapUsages/nodeUsageStatsForThreadPools/shardWriteLoads/maxHeapSizePerNode/
-            // nodeIdsWriteLoadHotspotting at this stage, to avoid committing to API payloads until the features are settled
+            // nodeIdsWriteLoadHotspotting/estimatedShardHeapUsages at this stage, to avoid committing to API payloads until
+            // the features are settled
         );
     }
 
@@ -349,6 +367,10 @@ public Map<String, EstimatedHeapUsage> getEstimatedHeapUsages() {
         return estimatedHeapUsages;
     }
 
+    public Map<ShardId, ShardAndIndexHeapUsage> getEstimatedShardHeapUsages() {
+        return estimatedShardHeapUsages;
+    }
+
     /**
      * Returns a map containing thread pool usage stats for each node, keyed by node ID.
      */
@@ -455,6 +477,7 @@ public boolean equals(Object o) {
             && dataPath.equals(that.dataPath)
             && reservedSpace.equals(that.reservedSpace)
             && estimatedHeapUsages.equals(that.estimatedHeapUsages)
+            && estimatedShardHeapUsages.equals(that.estimatedShardHeapUsages)
             && nodeUsageStatsForThreadPools.equals(that.nodeUsageStatsForThreadPools)
             && shardWriteLoads.equals(that.shardWriteLoads)
             && maxHeapSizePerNode.equals(that.maxHeapSizePerNode)
@@ -471,6 +494,7 @@ public int hashCode() {
             dataPath,
             reservedSpace,
             estimatedHeapUsages,
+            estimatedShardHeapUsages,
             nodeUsageStatsForThreadPools,
             shardWriteLoads,
             maxHeapSizePerNode,
@@ -596,6 +620,7 @@ public static class Builder {
         private Map<NodeAndShard, String> dataPath = Map.of();
         private Map<NodeAndPath, ReservedSpace> reservedSpace = Map.of();
         private Map<String, EstimatedHeapUsage> estimatedHeapUsages = Map.of();
+        private Map<ShardId, ShardAndIndexHeapUsage> estimatedShardHeapUsages = Map.of();
         private Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools = Map.of();
         private Map<ShardId, Double> shardWriteLoads = Map.of();
         private Map<String, ByteSizeValue> maxHeapSizePerNode = Map.of();
@@ -610,6 +635,7 @@ public ClusterInfo build() {
                 dataPath,
                 reservedSpace,
                 estimatedHeapUsages,
+                estimatedShardHeapUsages,
                 nodeUsageStatsForThreadPools,
                 shardWriteLoads,
                 maxHeapSizePerNode,
@@ -652,6 +678,11 @@ public Builder estimatedHeapUsages(Map<String, EstimatedHeapUsage> estimatedHeap
             return this;
         }
 
+        public Builder estimatedShardHeapUsages(Map<ShardId, ShardAndIndexHeapUsage> estimatedShardHeapUsages) {
+            this.estimatedShardHeapUsages = estimatedShardHeapUsages;
+            return this;
+        }
+
         public Builder nodeUsageStatsForThreadPools(Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools) {
             this.nodeUsageStatsForThreadPools = nodeUsageStatsForThreadPools;
             return this;

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

@@ -39,6 +39,7 @@ public class ClusterInfoSimulator {
     private final Map<String, DiskUsage> mostAvailableSpaceUsage;
     private final CopyOnFirstWriteMap<String, Long> shardSizes;
     private final Map<String, EstimatedHeapUsage> estimatedHeapUsages;
+    private final Map<ShardId, ShardAndIndexHeapUsage> estimatedShardHeapUsages;
     private final ShardMovementWriteLoadSimulator shardMovementWriteLoadSimulator;
 
     public ClusterInfoSimulator(RoutingAllocation allocation) {
@@ -47,6 +48,7 @@ public ClusterInfoSimulator(RoutingAllocation allocation) {
         this.mostAvailableSpaceUsage = getAdjustedDiskSpace(allocation, allocation.clusterInfo().getNodeMostAvailableDiskUsages());
         this.shardSizes = new CopyOnFirstWriteMap<>(allocation.clusterInfo().shardSizes);
         this.estimatedHeapUsages = allocation.clusterInfo().getEstimatedHeapUsages();
+        this.estimatedShardHeapUsages = allocation.clusterInfo().getEstimatedShardHeapUsages();
         this.shardMovementWriteLoadSimulator = new ShardMovementWriteLoadSimulator(allocation);
     }
 
@@ -192,6 +194,7 @@ public ClusterInfo getClusterInfo() {
                 shardSizes.toImmutableMap(),
                 Map.of(),
                 estimatedHeapUsages,
+                estimatedShardHeapUsages,
                 shardMovementWriteLoadSimulator.simulatedNodeUsageStatsForThreadPools()
             );
     }

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

@@ -10,6 +10,7 @@
 package org.elasticsearch.cluster;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.index.shard.ShardId;
 
 import java.util.Map;
 
@@ -25,12 +26,29 @@ public interface EstimatedHeapUsageCollector {
     /**
      * This will be used when there is no EstimatedHeapUsageCollector available
      */
-    EstimatedHeapUsageCollector EMPTY = listener -> listener.onResponse(Map.of());
+    EstimatedHeapUsageCollector EMPTY = new EstimatedHeapUsageCollector() {
+        @Override
+        public void collectClusterHeapUsage(ActionListener<Map<String, Long>> listener) {
+            listener.onResponse(Map.of());
+        }
+
+        @Override
+        public void collectShardHeapUsage(ActionListener<Map<ShardId, ShardAndIndexHeapUsage>> listener) {
+            listener.onResponse(Map.of());
+        }
+    };
 
     /**
      * 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, Long>> listener);
+
+    /**
+     * Collects the estimated heap usage for every shard in the cluster.
+     *
+     * @param listener The listener which will receive the results
+     */
+    void collectShardHeapUsage(ActionListener<Map<ShardId, ShardAndIndexHeapUsage>> listener);
 }

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

@@ -215,6 +215,7 @@ private class AsyncRefresh {
         private volatile Map<String, DiskUsage> mostAvailableSpaceUsages;
         private volatile Map<String, ByteSizeValue> maxHeapPerNode;
         private volatile Map<String, Long> estimatedHeapUsagePerNode;
+        private volatile Map<ShardId, ShardAndIndexHeapUsage> estimatedHeapUsagePerShard;
         private volatile Map<String, NodeUsageStatsForThreadPools> nodeThreadPoolUsageStatsPerNode;
         private volatile IndicesStatsSummary indicesStatsSummary;
 
@@ -231,7 +232,7 @@ void execute() {
             try (var ignoredRefs = fetchRefs) {
                 maybeFetchIndicesStats(diskThresholdEnabled || writeLoadConstraintEnabled.atLeastLowThresholdEnabled());
                 maybeFetchNodeStats(diskThresholdEnabled || estimatedHeapThresholdEnabled);
-                maybeFetchNodesEstimatedHeapUsage(estimatedHeapThresholdEnabled);
+                maybeFetchEstimatedHeapUsage(estimatedHeapThresholdEnabled);
                 fetchNodesUsageStatsForThreadPools();
             }
         }
@@ -260,14 +261,15 @@ private void maybeFetchNodeStats(boolean shouldFetch) {
             }
         }
 
-        private void maybeFetchNodesEstimatedHeapUsage(boolean shouldFetch) {
+        private void maybeFetchEstimatedHeapUsage(boolean shouldFetch) {
             if (shouldFetch) {
                 try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
-                    fetchNodesEstimatedHeapUsage();
+                    fetchEstimatedHeapUsage();
                 }
             } else {
                 logger.trace("skipping collecting estimated heap usage from cluster, notifying listeners with empty estimated heap usage");
                 estimatedHeapUsagePerNode = Map.of();
+                estimatedHeapUsagePerShard = Map.of();
             }
         }
 
@@ -292,7 +294,7 @@ public void onFailure(Exception e) {
             }
         }
 
-        private void fetchNodesEstimatedHeapUsage() {
+        private void fetchEstimatedHeapUsage() {
             estimatedHeapUsageCollector.collectClusterHeapUsage(ActionListener.releaseAfter(new ActionListener<>() {
                 @Override
                 public void onResponse(Map<String, Long> currentEstimatedHeapUsages) {
@@ -305,6 +307,19 @@ public void onFailure(Exception e) {
                     estimatedHeapUsagePerNode = Map.of();
                 }
             }, fetchRefs.acquire()));
+
+            estimatedHeapUsageCollector.collectShardHeapUsage(ActionListener.releaseAfter(new ActionListener<>() {
+                @Override
+                public void onResponse(Map<ShardId, ShardAndIndexHeapUsage> currentEstimatedHeapUsages) {
+                    estimatedHeapUsagePerShard = currentEstimatedHeapUsages;
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    logger.warn("failed to fetch heap usage for shards", e);
+                    estimatedHeapUsagePerShard = Map.of();
+                }
+            }, fetchRefs.acquire()));
         }
 
         private void fetchIndicesStats() {
@@ -497,6 +512,7 @@ private ClusterInfo updateAndGetCurrentClusterInfo() {
                 indicesStatsSummary.dataPath,
                 indicesStatsSummary.reservedSpace,
                 estimatedHeapUsages,
+                estimatedHeapUsagePerShard,
                 nodeThreadPoolUsageStatsPerNode,
                 indicesStatsSummary.shardWriteLoads(),
                 maxHeapPerNode,