Expose node heap size in cluster info (#134436)

This is needed for a new Stateless decider that limits concurrent
recoveries based on node heap size.

Relates ES-12554

Author: pxsalehi

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

@@ -410,6 +410,19 @@ public void testShardWriteLoadsArePresent() {
         }
     }
 
+    public void testMaxHeapPerNodeIsPresent() {
+        InternalClusterInfoService clusterInfoService = (InternalClusterInfoService) getInstanceFromNode(ClusterInfoService.class);
+        ClusterInfoServiceUtils.refresh(clusterInfoService);
+        Map<String, ByteSizeValue> maxHeapSizePerNode = clusterInfoService.getClusterInfo().getMaxHeapSizePerNode();
+        assertNotNull(maxHeapSizePerNode);
+        ClusterState state = getInstanceFromNode(ClusterService.class).state();
+        assertEquals(state.nodes().size(), maxHeapSizePerNode.size());
+        for (DiscoveryNode node : state.nodes()) {
+            assertTrue(maxHeapSizePerNode.containsKey(node.getId()));
+            assertThat(maxHeapSizePerNode.get(node.getId()), greaterThan(ByteSizeValue.ZERO));
+        }
+    }
+
     public void testIndexCanChangeCustomDataPath() throws Exception {
         final String index = "test-custom-data-path";
         final Path sharedDataPath = getInstanceFromNode(Environment.class).sharedDataDir().resolve(randomAsciiLettersOfLength(10));

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -324,6 +324,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INFERENCE_API_EIS_DIAGNOSTICS = def(9_156_0_00);
     public static final TransportVersion ML_INFERENCE_ENDPOINT_CACHE = def(9_157_0_00);
     public static final TransportVersion INDEX_SOURCE = def(9_158_0_00);
+    public static final TransportVersion MAX_HEAP_SIZE_PER_NODE_IN_CLUSTER_INFO = def(9_159_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -63,9 +63,11 @@ public class ClusterInfo implements ChunkedToXContent, Writeable {
     final Map<String, EstimatedHeapUsage> estimatedHeapUsages;
     final Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools;
     final Map<ShardId, Double> shardWriteLoads;
+    // max heap size per node ID
+    final Map<String, ByteSizeValue> maxHeapSizePerNode;
 
     protected ClusterInfo() {
-        this(Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of());
+        this(Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of());
     }
 
     /**
@@ -80,6 +82,7 @@ protected ClusterInfo() {
      * @param estimatedHeapUsages estimated heap usage broken down by node
      * @param nodeUsageStatsForThreadPools node-level usage stats (operational load) broken down by node
      * @see #shardIdentifierFromRouting
+     * @param maxHeapSizePerNode node id to max heap size
      */
     public ClusterInfo(
         Map<String, DiskUsage> leastAvailableSpaceUsage,
@@ -90,7 +93,8 @@ public ClusterInfo(
         Map<NodeAndPath, ReservedSpace> reservedSpace,
         Map<String, EstimatedHeapUsage> estimatedHeapUsages,
         Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools,
-        Map<ShardId, Double> shardWriteLoads
+        Map<ShardId, Double> shardWriteLoads,
+        Map<String, ByteSizeValue> maxHeapSizePerNode
     ) {
         this.leastAvailableSpaceUsage = Map.copyOf(leastAvailableSpaceUsage);
         this.mostAvailableSpaceUsage = Map.copyOf(mostAvailableSpaceUsage);
@@ -101,6 +105,7 @@ public ClusterInfo(
         this.estimatedHeapUsages = Map.copyOf(estimatedHeapUsages);
         this.nodeUsageStatsForThreadPools = Map.copyOf(nodeUsageStatsForThreadPools);
         this.shardWriteLoads = Map.copyOf(shardWriteLoads);
+        this.maxHeapSizePerNode = Map.copyOf(maxHeapSizePerNode);
     }
 
     public ClusterInfo(StreamInput in) throws IOException {
@@ -125,6 +130,11 @@ public ClusterInfo(StreamInput in) throws IOException {
         } else {
             this.shardWriteLoads = Map.of();
         }
+        if (in.getTransportVersion().onOrAfter(TransportVersions.MAX_HEAP_SIZE_PER_NODE_IN_CLUSTER_INFO)) {
+            this.maxHeapSizePerNode = in.readImmutableMap(ByteSizeValue::readFrom);
+        } else {
+            this.maxHeapSizePerNode = Map.of();
+        }
     }
 
     @Override
@@ -144,6 +154,9 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().supports(SHARD_WRITE_LOAD_IN_CLUSTER_INFO)) {
             out.writeMap(this.shardWriteLoads, StreamOutput::writeWriteable, StreamOutput::writeDouble);
         }
+        if (out.getTransportVersion().onOrAfter(TransportVersions.MAX_HEAP_SIZE_PER_NODE_IN_CLUSTER_INFO)) {
+            out.writeMap(this.maxHeapSizePerNode, StreamOutput::writeWriteable);
+        }
     }
 
     /**
@@ -224,8 +237,8 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
                 return builder.endObject(); // NodeAndPath
             }),
             endArray() // end "reserved_sizes"
-            // NOTE: We don't serialize estimatedHeapUsages/nodeUsageStatsForThreadPools/shardWriteLoads at this stage, to avoid
-            // committing to API payloads until the features are settled
+            // NOTE: We don't serialize estimatedHeapUsages/nodeUsageStatsForThreadPools/shardWriteLoads/maxHeapSizePerNode at this stage,
+            // to avoid committing to API payloads until the features are settled
         );
     }
 
@@ -326,6 +339,10 @@ public ReservedSpace getReservedSpace(String nodeId, String dataPath) {
         return result == null ? ReservedSpace.EMPTY : result;
     }
 
+    public Map<String, ByteSizeValue> getMaxHeapSizePerNode() {
+        return this.maxHeapSizePerNode;
+    }
+
     /**
      * Method that incorporates the ShardId for the shard into a string that
      * includes a 'p' or 'r' depending on whether the shard is a primary.
@@ -351,7 +368,8 @@ public boolean equals(Object o) {
             && reservedSpace.equals(that.reservedSpace)
             && estimatedHeapUsages.equals(that.estimatedHeapUsages)
             && nodeUsageStatsForThreadPools.equals(that.nodeUsageStatsForThreadPools)
-            && shardWriteLoads.equals(that.shardWriteLoads);
+            && shardWriteLoads.equals(that.shardWriteLoads)
+            && maxHeapSizePerNode.equals(that.maxHeapSizePerNode);
     }
 
     @Override
@@ -365,7 +383,8 @@ public int hashCode() {
             reservedSpace,
             estimatedHeapUsages,
             nodeUsageStatsForThreadPools,
-            shardWriteLoads
+            shardWriteLoads,
+            maxHeapSizePerNode
         );
     }
 
@@ -489,6 +508,7 @@ public static class Builder {
         private Map<String, EstimatedHeapUsage> estimatedHeapUsages = Map.of();
         private Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools = Map.of();
         private Map<ShardId, Double> shardWriteLoads = Map.of();
+        private Map<String, ByteSizeValue> maxHeapSizePerNode = Map.of();
 
         public ClusterInfo build() {
             return new ClusterInfo(
@@ -500,7 +520,8 @@ public ClusterInfo build() {
                 reservedSpace,
                 estimatedHeapUsages,
                 nodeUsageStatsForThreadPools,
-                shardWriteLoads
+                shardWriteLoads,
+                maxHeapSizePerNode
             );
         }
 
@@ -548,5 +569,10 @@ public Builder shardWriteLoads(Map<ShardId, Double> shardWriteLoads) {
             this.shardWriteLoads = shardWriteLoads;
             return this;
         }
+
+        public Builder maxHeapSizePerNode(Map<String, ByteSizeValue> maxHeapSizePerNode) {
+            this.maxHeapSizePerNode = maxHeapSizePerNode;
+            return this;
+        }
     }
 }

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

@@ -13,6 +13,7 @@
 import org.elasticsearch.cluster.routing.ShardMovementWriteLoadSimulator;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.CopyOnFirstWriteMap;
 import org.elasticsearch.index.shard.ShardId;
 
@@ -35,6 +36,7 @@ public class ClusterInfoSimulator {
     private final Map<ShardId, Long> shardDataSetSizes;
     private final Map<NodeAndShard, String> dataPath;
     private final Map<String, EstimatedHeapUsage> estimatedHeapUsages;
+    private final Map<String, ByteSizeValue> maxHeapSizePerNode;
     private final ShardMovementWriteLoadSimulator shardMovementWriteLoadSimulator;
 
     public ClusterInfoSimulator(RoutingAllocation allocation) {
@@ -45,6 +47,7 @@ public ClusterInfoSimulator(RoutingAllocation allocation) {
         this.shardDataSetSizes = Map.copyOf(allocation.clusterInfo().shardDataSetSizes);
         this.dataPath = Map.copyOf(allocation.clusterInfo().dataPath);
         this.estimatedHeapUsages = allocation.clusterInfo().getEstimatedHeapUsages();
+        this.maxHeapSizePerNode = Map.copyOf(allocation.clusterInfo().maxHeapSizePerNode);
         this.shardMovementWriteLoadSimulator = new ShardMovementWriteLoadSimulator(allocation);
     }
 
@@ -162,7 +165,8 @@ public ClusterInfo getClusterInfo() {
             Map.of(),
             estimatedHeapUsages,
             shardMovementWriteLoadSimulator.simulatedNodeUsageStatsForThreadPools(),
-            allocation.clusterInfo().getShardWriteLoads()
+            allocation.clusterInfo().getShardWriteLoads(),
+            maxHeapSizePerNode
         );
     }
 }

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

@@ -539,7 +539,8 @@ private boolean shouldRefresh() {
     public ClusterInfo getClusterInfo() {
         final IndicesStatsSummary indicesStatsSummary = this.indicesStatsSummary; // single volatile read
         final Map<String, EstimatedHeapUsage> estimatedHeapUsages = new HashMap<>();
-        maxHeapPerNode.forEach((nodeId, maxHeapSize) -> {
+        final var currentMaxHeapPerNode = this.maxHeapPerNode; // Make sure we use a consistent view
+        currentMaxHeapPerNode.forEach((nodeId, maxHeapSize) -> {
             final Long estimatedHeapUsage = estimatedHeapUsagePerNode.get(nodeId);
             if (estimatedHeapUsage != null) {
                 estimatedHeapUsages.put(nodeId, new EstimatedHeapUsage(nodeId, maxHeapSize.getBytes(), estimatedHeapUsage));
@@ -554,7 +555,8 @@ public ClusterInfo getClusterInfo() {
             indicesStatsSummary.reservedSpace,
             estimatedHeapUsages,
             nodeThreadPoolUsageStatsPerNode,
-            indicesStatsSummary.shardWriteLoads()
+            indicesStatsSummary.shardWriteLoads(),
+            currentMaxHeapPerNode
         );
     }
 

server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ThrottlingAllocationDecider.java

@@ -210,7 +210,7 @@ public Decision canForceAllocateDuringReplace(ShardRouting shardRouting, Routing
      *
      * This method returns the corresponding initializing shard that would be allocated to this node.
      */
-    private static ShardRouting initializingShard(ShardRouting shardRouting, String currentNodeId) {
+    public static ShardRouting initializingShard(ShardRouting shardRouting, String currentNodeId) {
         final ShardRouting initializingShard;
         if (shardRouting.unassigned()) {
             initializingShard = shardRouting.initialize(currentNodeId, null, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE);

server/src/test/java/org/elasticsearch/cluster/ClusterInfoTests.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.cluster;
 
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
@@ -45,7 +46,8 @@ public static ClusterInfo randomClusterInfo() {
             randomReservedSpace(),
             randomNodeHeapUsage(),
             randomNodeUsageStatsForThreadPools(),
-            randomShardWriteLoad()
+            randomShardWriteLoad(),
+            randomMaxHeapSizes()
         );
     }
 
@@ -58,6 +60,15 @@ private static Map<ShardId, Double> randomShardWriteLoad() {
         return builder;
     }
 
+    private static Map<String, ByteSizeValue> randomMaxHeapSizes() {
+        int numEntries = randomIntBetween(0, 128);
+        Map<String, ByteSizeValue> nodeMaxHeapSizes = new HashMap<>(numEntries);
+        for (int i = 0; i < numEntries; i++) {
+            nodeMaxHeapSizes.put(randomAlphaOfLength(32), randomByteSizeValue());
+        }
+        return nodeMaxHeapSizes;
+    }
+
     private static Map<String, EstimatedHeapUsage> randomNodeHeapUsage() {
         int numEntries = randomIntBetween(0, 128);
         Map<String, EstimatedHeapUsage> nodeHeapUsage = new HashMap<>(numEntries);

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceReconcilerTests.java

@@ -622,6 +622,7 @@ public void testUnassignedAllocationPredictsDiskUsage() {
             ImmutableOpenMap.of(),
             ImmutableOpenMap.of(),
             ImmutableOpenMap.of(),
+            ImmutableOpenMap.of(),
             ImmutableOpenMap.of()
         );
 

server/src/test/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDeciderTests.java

@@ -1415,6 +1415,7 @@ static class DevNullClusterInfo extends ClusterInfo {
                 reservedSpace,
                 Map.of(),
                 Map.of(),
+                Map.of(),
                 Map.of()
             );
         }

server/src/test/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDeciderUnitTests.java

@@ -111,6 +111,7 @@ public void testCanAllocateUsesMaxAvailableSpace() {
             Map.of(),
             Map.of(),
             Map.of(),
+            Map.of(),
             Map.of()
         );
         RoutingAllocation allocation = new RoutingAllocation(
@@ -185,6 +186,7 @@ private void doTestCannotAllocateDueToLackOfDiskResources(boolean testMaxHeadroo
             Map.of(),
             Map.of(),
             Map.of(),
+            Map.of(),
             Map.of()
         );
         RoutingAllocation allocation = new RoutingAllocation(
@@ -333,6 +335,7 @@ private void doTestCanRemainUsesLeastAvailableSpace(boolean testMaxHeadroom) {
             Map.of(),
             Map.of(),
             Map.of(),
+            Map.of(),
             Map.of()
         );
         RoutingAllocation allocation = new RoutingAllocation(