Add node write load to the ClusterInfo

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

@@ -8,6 +8,8 @@
  */
 package org.elasticsearch.index.shard;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.store.LockObtainFailedException;
 import org.apache.lucene.util.SetOnce;
@@ -22,13 +24,16 @@
 import org.elasticsearch.cluster.EstimatedHeapUsage;
 import org.elasticsearch.cluster.EstimatedHeapUsageCollector;
 import org.elasticsearch.cluster.InternalClusterInfoService;
+import org.elasticsearch.cluster.NodeWriteLoad;
+import org.elasticsearch.cluster.WriteLoadCollector;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeUtils;
 import org.elasticsearch.cluster.routing.RecoverySource;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.UnassignedInfo;
+import org.elasticsearch.cluster.routing.allocation.WriteLoadConstraintSettings;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.UUIDs;
@@ -117,14 +122,16 @@
 import static org.hamcrest.Matchers.either;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 public class IndexShardIT extends ESSingleNodeTestCase {
+    private static final Logger logger = LogManager.getLogger(IndexShardIT.class);
 
     @Override
     protected Collection<Class<? extends Plugin>> getPlugins() {
-        return pluginList(InternalSettingsPlugin.class, BogusEstimatedHeapUsagePlugin.class);
+        return pluginList(InternalSettingsPlugin.class, BogusEstimatedHeapUsagePlugin.class, BogusWriteLoadCollectorPlugin.class);
     }
 
     public void testLockTryingToDelete() throws Exception {
@@ -295,6 +302,47 @@ public void testHeapUsageEstimateIsPresent() {
         }
     }
 
+    public void testNodeWriteLoadsArePresent() {
+        InternalClusterInfoService clusterInfoService = (InternalClusterInfoService) getInstanceFromNode(ClusterInfoService.class);
+        ClusterInfoServiceUtils.refresh(clusterInfoService);
+        Map<String, NodeWriteLoad> nodeWriteLoads = clusterInfoService.getClusterInfo().getNodeWriteLoads();
+        assertNotNull(nodeWriteLoads);
+        /** Not collecting stats yet because allocation write load stats collection is disabled by default.
+         *  see {@link WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING} */
+        assertTrue(nodeWriteLoads.isEmpty());
+
+        // Enable collection for node write loads.
+        updateClusterSettings(
+            Settings.builder()
+                .put(
+                    WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(),
+                    WriteLoadConstraintSettings.WriteLoadDeciderStatus.ENABLED
+                )
+                .build()
+        );
+        try {
+            // Force a ClusterInfo refresh to run collection of the node write loads.
+            ClusterInfoServiceUtils.refresh(clusterInfoService);
+            nodeWriteLoads = clusterInfoService.getClusterInfo().getNodeWriteLoads();
+
+            /** Verify that each node has a write load reported. The test {@link BogusWriteLoadCollector} generates random load values */
+            ClusterState state = getInstanceFromNode(ClusterService.class).state();
+            assertEquals(state.nodes().size(), nodeWriteLoads.size());
+            for (DiscoveryNode node : state.nodes()) {
+                assertTrue(nodeWriteLoads.containsKey(node.getId()));
+                NodeWriteLoad nodeWriteLoad = nodeWriteLoads.get(node.getId());
+                assertThat(nodeWriteLoad.nodeId(), equalTo(node.getId()));
+                assertThat(nodeWriteLoad.totalWriteThreadPoolThreads(), greaterThanOrEqualTo(0));
+                assertThat(nodeWriteLoad.percentWriteThreadPoolUtilization(), greaterThanOrEqualTo(0));
+                assertThat(nodeWriteLoad.maxTaskTimeInWriteQueueMillis(), greaterThanOrEqualTo(0L));
+            }
+        } finally {
+            updateClusterSettings(
+                Settings.builder().putNull(WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey()).build()
+            );
+        }
+    }
+
     public void testIndexCanChangeCustomDataPath() throws Exception {
         final String index = "test-custom-data-path";
         final Path sharedDataPath = getInstanceFromNode(Environment.class).sharedDataDir().resolve(randomAsciiLettersOfLength(10));
@@ -875,4 +923,55 @@ public ClusterService getClusterService() {
             return clusterService.get();
         }
     }
+
+    /**
+     * A simple {@link WriteLoadCollector} implementation that creates and returns random {@link NodeWriteLoad} for each node in the
+     * cluster.
+     * <p>
+     * Note: there's an 'org.elasticsearch.cluster.WriteLoadCollector' file that declares this implementation so that the plugin system can
+     * pick it up and use it for the test set-up.
+     */
+    public static class BogusWriteLoadCollector implements WriteLoadCollector {
+
+        private final BogusWriteLoadCollectorPlugin plugin;
+
+        public BogusWriteLoadCollector(BogusWriteLoadCollectorPlugin plugin) {
+            this.plugin = plugin;
+        }
+
+        @Override
+        public void collectWriteLoads(ActionListener<Map<String, NodeWriteLoad>> listener) {
+            ActionListener.completeWith(
+                listener,
+                () -> plugin.getClusterService()
+                    .state()
+                    .nodes()
+                    .stream()
+                    .collect(
+                        Collectors.toUnmodifiableMap(
+                            DiscoveryNode::getId,
+                            node -> new NodeWriteLoad(node.getId(), randomNonNegativeInt(), randomNonNegativeInt(), randomNonNegativeLong())
+                        )
+                    )
+            );
+        }
+    }
+
+    /**
+     * Make a plugin to gain access to the {@link ClusterService} instance.
+     */
+    public static class BogusWriteLoadCollectorPlugin 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.WriteLoadCollector

@@ -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$BogusWriteLoadCollector

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

@@ -326,6 +326,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ML_INFERENCE_COHERE_API_VERSION = def(9_110_0_00);
     public static final TransportVersion ESQL_PROFILE_INCLUDE_PLAN = def(9_111_0_00);
     public static final TransportVersion MAPPINGS_IN_DATA_STREAMS = def(9_112_0_00);
+    public static final TransportVersion NODE_WRITE_LOAD_IN_CLUSTER_INFO = def(9_113_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -58,9 +58,10 @@ public class ClusterInfo implements ChunkedToXContent, Writeable {
     final Map<NodeAndShard, String> dataPath;
     final Map<NodeAndPath, ReservedSpace> reservedSpace;
     final Map<String, EstimatedHeapUsage> estimatedHeapUsages;
+    final Map<String, NodeWriteLoad> nodeWriteLoads;
 
     protected ClusterInfo() {
-        this(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());
     }
 
     /**
@@ -73,6 +74,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 nodeWriteLoads estimated node-level write load broken down by node
      * @see #shardIdentifierFromRouting
      */
     public ClusterInfo(
@@ -82,7 +84,8 @@ public ClusterInfo(
         Map<ShardId, Long> shardDataSetSizes,
         Map<NodeAndShard, String> dataPath,
         Map<NodeAndPath, ReservedSpace> reservedSpace,
-        Map<String, EstimatedHeapUsage> estimatedHeapUsages
+        Map<String, EstimatedHeapUsage> estimatedHeapUsages,
+        Map<String, NodeWriteLoad> nodeWriteLoads
     ) {
         this.leastAvailableSpaceUsage = Map.copyOf(leastAvailableSpaceUsage);
         this.mostAvailableSpaceUsage = Map.copyOf(mostAvailableSpaceUsage);
@@ -91,6 +94,7 @@ public ClusterInfo(
         this.dataPath = Map.copyOf(dataPath);
         this.reservedSpace = Map.copyOf(reservedSpace);
         this.estimatedHeapUsages = Map.copyOf(estimatedHeapUsages);
+        this.nodeWriteLoads = Map.copyOf(nodeWriteLoads);
     }
 
     public ClusterInfo(StreamInput in) throws IOException {
@@ -107,6 +111,11 @@ public ClusterInfo(StreamInput in) throws IOException {
         } else {
             this.estimatedHeapUsages = Map.of();
         }
+        if (in.getTransportVersion().onOrAfter(TransportVersions.NODE_WRITE_LOAD_IN_CLUSTER_INFO)) {
+            this.nodeWriteLoads = in.readImmutableMap(NodeWriteLoad::new);
+        } else {
+            this.nodeWriteLoads = Map.of();
+        }
     }
 
     @Override
@@ -124,6 +133,9 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().onOrAfter(TransportVersions.HEAP_USAGE_IN_CLUSTER_INFO)) {
             out.writeMap(this.estimatedHeapUsages, StreamOutput::writeWriteable);
         }
+        if (out.getTransportVersion().onOrAfter(TransportVersions.NODE_WRITE_LOAD_IN_CLUSTER_INFO)) {
+            out.writeMap(this.nodeWriteLoads, StreamOutput::writeWriteable);
+        }
     }
 
     /**
@@ -204,8 +216,8 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
                 return builder.endObject(); // NodeAndPath
             }),
             endArray() // end "reserved_sizes"
-            // NOTE: We don't serialize estimatedHeapUsages at this stage, to avoid
-            // committing to API payloads until the feature is settled
+            // NOTE: We don't serialize estimatedHeapUsages/nodeWriteLoads at this stage, to avoid
+            // committing to API payloads until the features are settled
         );
     }
 
@@ -220,6 +232,13 @@ public Map<String, EstimatedHeapUsage> getEstimatedHeapUsages() {
         return estimatedHeapUsages;
     }
 
+    /**
+     * Returns a map containing the node-level write load estimate for each node by node ID.
+     */
+    public Map<String, NodeWriteLoad> getNodeWriteLoads() {
+        return nodeWriteLoads;
+    }
+
     /**
      * 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.
@@ -311,12 +330,21 @@ public boolean equals(Object o) {
             && shardSizes.equals(that.shardSizes)
             && shardDataSetSizes.equals(that.shardDataSetSizes)
             && dataPath.equals(that.dataPath)
-            && reservedSpace.equals(that.reservedSpace);
+            && reservedSpace.equals(that.reservedSpace)
+            && nodeWriteLoads.equals(that.nodeWriteLoads);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(leastAvailableSpaceUsage, mostAvailableSpaceUsage, shardSizes, shardDataSetSizes, dataPath, reservedSpace);
+        return Objects.hash(
+            leastAvailableSpaceUsage,
+            mostAvailableSpaceUsage,
+            shardSizes,
+            shardDataSetSizes,
+            dataPath,
+            reservedSpace,
+            nodeWriteLoads
+        );
     }
 
     @Override

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

@@ -34,6 +34,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, NodeWriteLoad> nodeWriteLoads;
 
     public ClusterInfoSimulator(RoutingAllocation allocation) {
         this.allocation = allocation;
@@ -43,6 +44,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.nodeWriteLoads = allocation.clusterInfo().getNodeWriteLoads();
     }
 
     /**
@@ -156,7 +158,8 @@ public ClusterInfo getClusterInfo() {
             shardDataSetSizes,
             dataPath,
             Map.of(),
-            estimatedHeapUsages
+            estimatedHeapUsages,
+            nodeWriteLoads
         );
     }
 }