Add heap usage estimate to ClusterInfo

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

@@ -10,6 +10,7 @@
 
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.store.LockObtainFailedException;
+import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.index.IndexRequest;
@@ -19,6 +20,8 @@
 import org.elasticsearch.cluster.ClusterInfoServiceUtils;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.InternalClusterInfoService;
+import org.elasticsearch.cluster.ShardHeapUsage;
+import org.elasticsearch.cluster.ShardHeapUsageSupplier;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeUtils;
@@ -61,6 +64,7 @@
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.indices.recovery.RecoveryState;
+import org.elasticsearch.plugins.ClusterPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.test.DummyShardLock;
@@ -81,6 +85,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.BrokenBarrierException;
 import java.util.concurrent.CountDownLatch;
@@ -89,6 +94,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import static com.carrotsearch.randomizedtesting.RandomizedTest.randomAsciiLettersOfLength;
@@ -110,12 +116,13 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 public class IndexShardIT extends ESSingleNodeTestCase {
 
     @Override
     protected Collection<Class<? extends Plugin>> getPlugins() {
-        return pluginList(InternalSettingsPlugin.class);
+        return pluginList(InternalSettingsPlugin.class, BogusShardHeapUsagePlugin.class);
     }
 
     public void testLockTryingToDelete() throws Exception {
@@ -253,6 +260,20 @@ public void testExpectedShardSizeIsPresent() throws InterruptedException {
         assertThat(dataSetSize.get(), greaterThan(0L));
     }
 
+    public void testHeapUsageEstimateIsPresent() {
+        InternalClusterInfoService clusterInfoService = (InternalClusterInfoService) getInstanceFromNode(ClusterInfoService.class);
+        ClusterInfoServiceUtils.refresh(clusterInfoService);
+        ClusterState state = getInstanceFromNode(ClusterService.class).state();
+        Map<String, ShardHeapUsage> shardHeapUsages = clusterInfoService.getClusterInfo().getShardHeapUsages();
+        assertNotNull(shardHeapUsages);
+        assertEquals(state.nodes().size(), shardHeapUsages.size());
+        for (DiscoveryNode node : state.nodes()) {
+            assertTrue(shardHeapUsages.containsKey(node.getId()));
+            ShardHeapUsage shardHeapUsage = shardHeapUsages.get(node.getId());
+            assertThat(shardHeapUsage.estimatedFreeBytes(), lessThanOrEqualTo(shardHeapUsage.totalBytes()));
+        }
+    }
+
     public void testIndexCanChangeCustomDataPath() throws Exception {
         final String index = "test-custom-data-path";
         final Path sharedDataPath = getInstanceFromNode(Environment.class).sharedDataDir().resolve(randomAsciiLettersOfLength(10));
@@ -795,4 +816,44 @@ private static void assertAllIndicesRemovedAndDeletionCompleted(Iterable<Indices
             assertBusy(() -> assertFalse(indicesService.hasUncompletedPendingDeletes()), 1, TimeUnit.MINUTES);
         }
     }
+
+    public static class BogusShardShardHeapUsageSupplier implements ShardHeapUsageSupplier {
+
+        private final BogusShardHeapUsagePlugin plugin;
+
+        public BogusShardShardHeapUsageSupplier(BogusShardHeapUsagePlugin plugin) {
+            this.plugin = plugin;
+        }
+
+        @Override
+        public void getClusterHeapUsage(ActionListener<Map<String, ShardHeapUsage>> listener) {
+            ActionListener.completeWith(
+                listener,
+                () -> plugin.getClusterService()
+                    .state()
+                    .nodes()
+                    .stream()
+                    .collect(Collectors.toUnmodifiableMap(DiscoveryNode::getId, node -> {
+                        final long maxHeap = randomNonNegativeLong();
+                        final long usedHeap = (long) (randomFloat() * maxHeap);
+                        return new ShardHeapUsage(node.getId(), node.getName(), maxHeap, usedHeap);
+                    }))
+            );
+        }
+    }
+
+    public static class BogusShardHeapUsagePlugin extends Plugin implements ClusterPlugin {
+
+        private final SetOnce<ClusterService> clusterService = new SetOnce<>();
+
+        @Override
+        public Collection<?> createComponents(PluginServices services) {
+            clusterService.set(services.clusterService());
+            return List.of();
+        }
+
+        public ClusterService getClusterService() {
+            return clusterService.get();
+        }
+    }
 }

server/src/internalClusterTest/resources/META-INF/services/org.elasticsearch.cluster.ShardHeapUsageSupplier

@@ -0,0 +1,10 @@
+#
+# 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".
+#
+
+org.elasticsearch.index.shard.IndexShardIT$BogusShardShardHeapUsageSupplier

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

@@ -285,6 +285,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ILM_ADD_SKIP_SETTING = def(9_089_0_00);
     public static final TransportVersion ML_INFERENCE_MISTRAL_CHAT_COMPLETION_ADDED = def(9_090_0_00);
     public static final TransportVersion IDP_CUSTOM_SAML_ATTRIBUTES_ALLOW_LIST = def(9_091_0_00);
+    public static final TransportVersion HEAP_USAGE_IN_CLUSTER_INFO = def(9_092_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -57,9 +57,10 @@ public class ClusterInfo implements ChunkedToXContent, Writeable {
     final Map<ShardId, Long> shardDataSetSizes;
     final Map<NodeAndShard, String> dataPath;
     final Map<NodeAndPath, ReservedSpace> reservedSpace;
+    final Map<String, ShardHeapUsage> shardHeapUsages;
 
     protected ClusterInfo() {
-        this(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());
     }
 
     /**
@@ -71,6 +72,7 @@ protected ClusterInfo() {
      * @param shardDataSetSizes a shard id to data set size in bytes mapping per shard
      * @param dataPath the shard routing to datapath mapping
      * @param reservedSpace reserved space per shard broken down by node and data path
+     * @param shardHeapUsages shard heap usage broken down by node
      * @see #shardIdentifierFromRouting
      */
     public ClusterInfo(
@@ -79,14 +81,16 @@ public ClusterInfo(
         Map<String, Long> shardSizes,
         Map<ShardId, Long> shardDataSetSizes,
         Map<NodeAndShard, String> dataPath,
-        Map<NodeAndPath, ReservedSpace> reservedSpace
+        Map<NodeAndPath, ReservedSpace> reservedSpace,
+        Map<String, ShardHeapUsage> shardHeapUsages
     ) {
         this.leastAvailableSpaceUsage = Map.copyOf(leastAvailableSpaceUsage);
         this.mostAvailableSpaceUsage = Map.copyOf(mostAvailableSpaceUsage);
         this.shardSizes = Map.copyOf(shardSizes);
         this.shardDataSetSizes = Map.copyOf(shardDataSetSizes);
         this.dataPath = Map.copyOf(dataPath);
         this.reservedSpace = Map.copyOf(reservedSpace);
+        this.shardHeapUsages = Map.copyOf(shardHeapUsages);
     }
 
     public ClusterInfo(StreamInput in) throws IOException {
@@ -98,6 +102,11 @@ public ClusterInfo(StreamInput in) throws IOException {
             ? in.readImmutableMap(NodeAndShard::new, StreamInput::readString)
             : in.readImmutableMap(nested -> NodeAndShard.from(new ShardRouting(nested)), StreamInput::readString);
         this.reservedSpace = in.readImmutableMap(NodeAndPath::new, ReservedSpace::new);
+        if (in.getTransportVersion().onOrAfter(TransportVersions.HEAP_USAGE_IN_CLUSTER_INFO)) {
+            this.shardHeapUsages = in.readImmutableMap(ShardHeapUsage::new);
+        } else {
+            this.shardHeapUsages = Map.of();
+        }
     }
 
     @Override
@@ -112,6 +121,9 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeMap(this.dataPath, (o, k) -> createFakeShardRoutingFromNodeAndShard(k).writeTo(o), StreamOutput::writeString);
         }
         out.writeMap(this.reservedSpace);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.HEAP_USAGE_IN_CLUSTER_INFO)) {
+            out.writeMap(this.shardHeapUsages, StreamOutput::writeWriteable);
+        }
     }
 
     /**
@@ -192,9 +204,22 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
                 return builder.endObject(); // NodeAndPath
             }),
             endArray() // end "reserved_sizes"
+            // NOTE: We don't serialize shardHeapUsages at this stage, to avoid
+            // committing to API payloads until the feature is settled
         );
     }
 
+    /**
+     * Returns a node id to estimated heap usage mapping for all nodes that we have such data for.
+     * Note that these estimates should be considered minimums. They may be used to determine whether
+     * there IS NOT capacity to do something, but not to determine that there IS capacity to do something.
+     * Also note that the map may not be complete, it may contain none, or a subset of the nodes in
+     * the cluster at any time. It may also contain entries for nodes that have since left the cluster.
+     */
+    public Map<String, ShardHeapUsage> getShardHeapUsages() {
+        return shardHeapUsages;
+    }
+
     /**
      * Returns a node id to disk usage mapping for the path that has the least available space on the node.
      * Note that this does not take account of reserved space: there may be another path with less available _and unreserved_ space.

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

@@ -153,6 +153,7 @@ public ClusterInfo getClusterInfo() {
             shardSizes.toImmutableMap(),
             shardDataSetSizes,
             dataPath,
+            Map.of(),
             Map.of()
         );
     }

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

@@ -82,45 +82,55 @@ public class InternalClusterInfoService implements ClusterInfoService, ClusterSt
         Property.NodeScope
     );
 
-    private volatile boolean enabled;
+    private volatile boolean diskThresholdEnabled;
     private volatile TimeValue updateFrequency;
     private volatile TimeValue fetchTimeout;
 
     private volatile Map<String, DiskUsage> leastAvailableSpaceUsages;
     private volatile Map<String, DiskUsage> mostAvailableSpaceUsages;
     private volatile IndicesStatsSummary indicesStatsSummary;
+    private volatile Map<String, ShardHeapUsage> shardHeapUsages;
 
     private final ThreadPool threadPool;
     private final Client client;
     private final List<Consumer<ClusterInfo>> listeners = new CopyOnWriteArrayList<>();
 
     private final Object mutex = new Object();
     private final List<ActionListener<ClusterInfo>> nextRefreshListeners = new ArrayList<>();
+    private final ShardHeapUsageSupplier shardHeapUsageSupplier;
 
     private AsyncRefresh currentRefresh;
     private RefreshScheduler refreshScheduler;
 
     @SuppressWarnings("this-escape")
-    public InternalClusterInfoService(Settings settings, ClusterService clusterService, ThreadPool threadPool, Client client) {
+    public InternalClusterInfoService(
+        Settings settings,
+        ClusterService clusterService,
+        ThreadPool threadPool,
+        Client client,
+        ShardHeapUsageSupplier shardHeapUsageSupplier
+    ) {
         this.leastAvailableSpaceUsages = Map.of();
         this.mostAvailableSpaceUsages = Map.of();
+        this.shardHeapUsages = Map.of();
         this.indicesStatsSummary = IndicesStatsSummary.EMPTY;
         this.threadPool = threadPool;
         this.client = client;
+        this.shardHeapUsageSupplier = shardHeapUsageSupplier;
         this.updateFrequency = INTERNAL_CLUSTER_INFO_UPDATE_INTERVAL_SETTING.get(settings);
         this.fetchTimeout = INTERNAL_CLUSTER_INFO_TIMEOUT_SETTING.get(settings);
-        this.enabled = DiskThresholdSettings.CLUSTER_ROUTING_ALLOCATION_DISK_THRESHOLD_ENABLED_SETTING.get(settings);
+        this.diskThresholdEnabled = DiskThresholdSettings.CLUSTER_ROUTING_ALLOCATION_DISK_THRESHOLD_ENABLED_SETTING.get(settings);
         ClusterSettings clusterSettings = clusterService.getClusterSettings();
         clusterSettings.addSettingsUpdateConsumer(INTERNAL_CLUSTER_INFO_TIMEOUT_SETTING, this::setFetchTimeout);
         clusterSettings.addSettingsUpdateConsumer(INTERNAL_CLUSTER_INFO_UPDATE_INTERVAL_SETTING, this::setUpdateFrequency);
         clusterSettings.addSettingsUpdateConsumer(
             DiskThresholdSettings.CLUSTER_ROUTING_ALLOCATION_DISK_THRESHOLD_ENABLED_SETTING,
-            this::setEnabled
+            this::setDiskThresholdEnabled
         );
     }
 
-    private void setEnabled(boolean enabled) {
-        this.enabled = enabled;
+    private void setDiskThresholdEnabled(boolean diskThresholdEnabled) {
+        this.diskThresholdEnabled = diskThresholdEnabled;
     }
 
     private void setFetchTimeout(TimeValue fetchTimeout) {
@@ -169,27 +179,43 @@ private class AsyncRefresh {
         }
 
         void execute() {
-            if (enabled == false) {
-                logger.trace("skipping collecting info from cluster, notifying listeners with empty cluster info");
-                leastAvailableSpaceUsages = Map.of();
-                mostAvailableSpaceUsages = Map.of();
-                indicesStatsSummary = IndicesStatsSummary.EMPTY;
-                callListeners();
-                return;
-            }
-
             logger.trace("starting async refresh");
 
             try (var ignoredRefs = fetchRefs) {
-                try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
-                    fetchNodeStats();
+                if (diskThresholdEnabled) {
+                    try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
+                        fetchNodeStats();
+                    }
+                    try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
+                        fetchIndicesStats();
+                    }
+                } else {
+                    logger.trace("skipping collecting disk usage info from cluster, notifying listeners with empty cluster info");
+                    leastAvailableSpaceUsages = Map.of();
+                    mostAvailableSpaceUsages = Map.of();
+                    indicesStatsSummary = IndicesStatsSummary.EMPTY;
                 }
                 try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
-                    fetchIndicesStats();
+                    fetchNodesHeapUsage();
                 }
             }
         }
 
+        private void fetchNodesHeapUsage() {
+            shardHeapUsageSupplier.getClusterHeapUsage(ActionListener.releaseAfter(new ActionListener<>() {
+                @Override
+                public void onResponse(Map<String, ShardHeapUsage> currentShardHeapUsages) {
+                    shardHeapUsages = currentShardHeapUsages;
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    logger.warn("failed to fetch heap usage for nodes", e);
+                    shardHeapUsages = Map.of();
+                }
+            }, fetchRefs.acquire()));
+        }
+
         private void fetchIndicesStats() {
             final IndicesStatsRequest indicesStatsRequest = new IndicesStatsRequest();
             indicesStatsRequest.clear();
@@ -413,7 +439,8 @@ public ClusterInfo getClusterInfo() {
             indicesStatsSummary.shardSizes,
             indicesStatsSummary.shardDataSetSizes,
             indicesStatsSummary.dataPath,
-            indicesStatsSummary.reservedSpace
+            indicesStatsSummary.reservedSpace,
+            shardHeapUsages
         );
     }