Add heap usage estimate to ClusterInfo

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

@@ -273,6 +273,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INFERENCE_CUSTOM_SERVICE_ADDED = def(9_084_0_00);
     public static final TransportVersion ESQL_LIMIT_ROW_SIZE = def(9_085_0_00);
     public static final TransportVersion ESQL_REGEX_MATCH_WITH_CASE_INSENSITIVITY = def(9_086_0_00);
+    public static final TransportVersion HEAP_USAGE_IN_CLUSTER_INFO = def(9_087_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -36,7 +36,7 @@
 import static org.elasticsearch.cluster.routing.ShardRouting.newUnassigned;
 import static org.elasticsearch.cluster.routing.UnassignedInfo.Reason.REINITIALIZED;
 import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.chunk;
-import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.endArray;
+import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.endObject;
 import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.startObject;
 
 /**
@@ -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, HeapUsage> nodesHeapUsage;
 
     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());
     }
 
     /**
@@ -79,14 +80,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, HeapUsage> nodesHeapUsage
     ) {
         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.nodesHeapUsage = Map.copyOf(nodesHeapUsage);
     }
 
     public ClusterInfo(StreamInput in) throws IOException {
@@ -98,6 +101,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.nodesHeapUsage = in.readImmutableMap(HeapUsage::new);
+        } else {
+            this.nodesHeapUsage = Map.of();
+        }
     }
 
     @Override
@@ -112,6 +120,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.nodesHeapUsage, StreamOutput::writeWriteable);
+        }
     }
 
     /**
@@ -191,7 +202,17 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
                 }
                 return builder.endObject(); // NodeAndPath
             }),
-            endArray() // end "reserved_sizes"
+            chunk(
+                (builder, p) -> builder.endArray() // end "reserved_sizes"
+                    .startObject("heap_usage")
+            ),
+            Iterators.map(nodesHeapUsage.entrySet().iterator(), c -> (builder, p) -> {
+                builder.startObject(c.getKey());
+                c.getValue().toShortXContent(builder);
+                builder.endObject();
+                return builder;
+            }),
+            endObject() // end "heap_usage"
         );
     }
 

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/HeapUsage.java

@@ -0,0 +1,70 @@
+/*
+ * 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 org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.xcontent.ToXContentFragment;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+
+/**
+ * Record representing the heap usage for a single cluster node
+ */
+public record HeapUsage(String nodeId, String nodeName, long totalBytes, long freeBytes) implements ToXContentFragment, Writeable {
+
+    public HeapUsage(StreamInput in) throws IOException {
+        this(in.readString(), in.readString(), in.readVLong(), in.readVLong());
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeString(this.nodeId);
+        out.writeString(this.nodeName);
+        out.writeVLong(this.totalBytes);
+        out.writeVLong(this.freeBytes);
+    }
+
+    public XContentBuilder toShortXContent(XContentBuilder builder) throws IOException {
+        builder.field("node_name", this.nodeName);
+        builder.humanReadableField("total_heap_bytes", "total", ByteSizeValue.ofBytes(this.totalBytes));
+        builder.humanReadableField("used_heap_bytes", "used", ByteSizeValue.ofBytes(this.usedBytes()));
+        builder.humanReadableField("free_heap_bytes", "free", ByteSizeValue.ofBytes(this.freeBytes));
+        builder.field("free_heap_percent", truncatePercent(this.freeHeapAsPercentage()));
+        builder.field("used_heap_percent", truncatePercent(this.usedHeapAsPercentage()));
+        return builder;
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.field("node_id", this.nodeId);
+        toShortXContent(builder);
+        return builder;
+    }
+
+    public double freeHeapAsPercentage() {
+        return 100.0 * freeBytes / (double) totalBytes;
+    }
+
+    public double usedHeapAsPercentage() {
+        return 100.0 - freeHeapAsPercentage();
+    }
+
+    public long usedBytes() {
+        return totalBytes - freeBytes;
+    }
+
+    private static double truncatePercent(double pct) {
+        return Math.round(pct * 10.0) / 10.0;
+    }
+}

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

@@ -0,0 +1,29 @@
+/*
+ * 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 org.elasticsearch.action.ActionListener;
+
+import java.util.Map;
+
+public interface HeapUsageSupplier {
+
+    /**
+     * This will be used when there are no heap usage suppliers available
+     */
+    HeapUsageSupplier EMPTY = listener -> listener.onResponse(Map.of());
+
+    /**
+     * Get the heap usage for every node in the cluster
+     *
+     * @param listener The listener which will receive the results
+     */
+    void getClusterHeapUsage(ActionListener<Map<String, HeapUsage>> listener);
+}

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

@@ -89,10 +89,12 @@ public class InternalClusterInfoService implements ClusterInfoService, ClusterSt
     private volatile Map<String, DiskUsage> leastAvailableSpaceUsages;
     private volatile Map<String, DiskUsage> mostAvailableSpaceUsages;
     private volatile IndicesStatsSummary indicesStatsSummary;
+    private volatile Map<String, HeapUsage> nodesHeapUsage;
 
     private final ThreadPool threadPool;
     private final Client client;
     private final List<Consumer<ClusterInfo>> listeners = new CopyOnWriteArrayList<>();
+    private final HeapUsageSupplier heapUsageSupplier;
 
     private final Object mutex = new Object();
     private final List<ActionListener<ClusterInfo>> nextRefreshListeners = new ArrayList<>();
@@ -101,12 +103,20 @@ public class InternalClusterInfoService implements ClusterInfoService, ClusterSt
     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,
+        HeapUsageSupplier heapUsageSupplier
+    ) {
         this.leastAvailableSpaceUsages = Map.of();
         this.mostAvailableSpaceUsages = Map.of();
+        this.nodesHeapUsage = Map.of();
         this.indicesStatsSummary = IndicesStatsSummary.EMPTY;
         this.threadPool = threadPool;
         this.client = client;
+        this.heapUsageSupplier = heapUsageSupplier;
         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);
@@ -173,6 +183,7 @@ void execute() {
                 logger.trace("skipping collecting info from cluster, notifying listeners with empty cluster info");
                 leastAvailableSpaceUsages = Map.of();
                 mostAvailableSpaceUsages = Map.of();
+                nodesHeapUsage = Map.of();
                 indicesStatsSummary = IndicesStatsSummary.EMPTY;
                 callListeners();
                 return;
@@ -187,9 +198,26 @@ void execute() {
                 try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
                     fetchIndicesStats();
                 }
+                try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
+                    fetchNodesHeapUsage();
+                }
             }
         }
 
+        private void fetchNodesHeapUsage() {
+            heapUsageSupplier.getClusterHeapUsage(ActionListener.releaseAfter(new ActionListener<>() {
+                @Override
+                public void onResponse(Map<String, HeapUsage> stringHeapUsageMap) {
+                    nodesHeapUsage = stringHeapUsageMap;
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    logger.warn("failed to fetch heap usage for nodes", e);
+                }
+            }, fetchRefs.acquire()));
+        }
+
         private void fetchIndicesStats() {
             final IndicesStatsRequest indicesStatsRequest = new IndicesStatsRequest();
             indicesStatsRequest.clear();
@@ -413,7 +441,8 @@ public ClusterInfo getClusterInfo() {
             indicesStatsSummary.shardSizes,
             indicesStatsSummary.shardDataSetSizes,
             indicesStatsSummary.dataPath,
-            indicesStatsSummary.reservedSpace
+            indicesStatsSummary.reservedSpace,
+            nodesHeapUsage
         );
     }
 

server/src/main/java/org/elasticsearch/node/NodeServiceProvider.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.action.search.OnlinePrewarmingService;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.ClusterInfoService;
+import org.elasticsearch.cluster.HeapUsageSupplier;
 import org.elasticsearch.cluster.InternalClusterInfoService;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.service.ClusterService;
@@ -28,6 +29,7 @@
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.indices.recovery.RecoverySettings;
+import org.elasticsearch.plugins.ClusterPlugin;
 import org.elasticsearch.plugins.PluginsLoader;
 import org.elasticsearch.plugins.PluginsService;
 import org.elasticsearch.readiness.ReadinessService;
@@ -44,6 +46,7 @@
 import org.elasticsearch.transport.TransportService;
 
 import java.util.Map;
+import java.util.Objects;
 import java.util.function.Function;
 import java.util.function.LongSupplier;
 
@@ -74,7 +77,14 @@ ClusterInfoService newClusterInfoService(
         ThreadPool threadPool,
         NodeClient client
     ) {
-        final InternalClusterInfoService service = new InternalClusterInfoService(settings, clusterService, threadPool, client);
+        final HeapUsageSupplier heapUsageSupplier = getHeapUsageSupplier(pluginsService);
+        final InternalClusterInfoService service = new InternalClusterInfoService(
+            settings,
+            clusterService,
+            threadPool,
+            client,
+            heapUsageSupplier
+        );
         if (DiscoveryNode.isMasterNode(settings)) {
             // listen for state changes (this node starts/stops being the elected master, or new nodes are added)
             clusterService.addListener(service);
@@ -146,4 +156,16 @@ void processRecoverySettings(PluginsService pluginsService, ClusterSettings clus
     ReadinessService newReadinessService(PluginsService pluginsService, ClusterService clusterService, Environment environment) {
         return new ReadinessService(clusterService, environment);
     }
+
+    private static HeapUsageSupplier getHeapUsageSupplier(PluginsService pluginsService) {
+        final var heapUsageSuppliers = pluginsService.filterPlugins(ClusterPlugin.class)
+            .map(ClusterPlugin::getHeapUsageSupplier)
+            .filter(Objects::nonNull)
+            .toList();
+        return switch (heapUsageSuppliers.size()) {
+            case 0 -> HeapUsageSupplier.EMPTY;
+            case 1 -> heapUsageSuppliers.getFirst();
+            default -> throw new IllegalArgumentException("multiple plugins define heap usage suppliers, which is not permitted");
+        };
+    }
 }

server/src/main/java/org/elasticsearch/plugins/ClusterPlugin.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.plugins;
 
+import org.elasticsearch.cluster.HeapUsageSupplier;
 import org.elasticsearch.cluster.routing.ShardRoutingRoleStrategy;
 import org.elasticsearch.cluster.routing.allocation.ExistingShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.WriteLoadForecaster;
@@ -41,6 +42,16 @@ default Collection<AllocationDecider> createAllocationDeciders(Settings settings
         return Collections.emptyList();
     }
 
+    /**
+     * Create a {@link HeapUsageSupplier} that will be used to determine the approximate heap usage for all
+     * cluster nodes
+     * <p>
+     * Note: Only a single {@link ClusterPlugin} can define a heap usage supplier.
+     */
+    default HeapUsageSupplier getHeapUsageSupplier() {
+        return null;
+    }
+
     /**
      * Return {@link ShardsAllocator} implementations added by this plugin.
      *

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

@@ -41,10 +41,28 @@ public static ClusterInfo randomClusterInfo() {
             randomShardSizes(),
             randomDataSetSizes(),
             randomRoutingToDataPath(),
-            randomReservedSpace()
+            randomReservedSpace(),
+            randomNodeHeapUsage()
         );
     }
 
+    private static Map<String, HeapUsage> randomNodeHeapUsage() {
+        int numEntries = randomIntBetween(0, 128);
+        Map<String, HeapUsage> nodeHeapUsage = new HashMap<>(numEntries);
+        for (int i = 0; i < numEntries; i++) {
+            String key = randomAlphaOfLength(32);
+            final int totalBytes = randomIntBetween(0, Integer.MAX_VALUE);
+            final HeapUsage diskUsage = new HeapUsage(
+                randomAlphaOfLength(4),
+                randomAlphaOfLength(4),
+                totalBytes,
+                randomIntBetween(0, totalBytes)
+            );
+            nodeHeapUsage.put(key, diskUsage);
+        }
+        return nodeHeapUsage;
+    }
+
     private static Map<String, DiskUsage> randomDiskUsage() {
         int numEntries = randomIntBetween(0, 128);
         Map<String, DiskUsage> builder = new HashMap<>(numEntries);