Collect node thread pool usage for shard balancing

Adds a new transport action to collect usage stats from the
data nodes. ClusterInfoService uses the action to pull thread
pool usage information from the data nodes to the master node
periodically.

Also removes NodeUsageStatsForThreadPoolsCollector as
an interface/plugin and replaces it with a single
class implementation.

Closes ES-12316

Author: DiannaHohensee

server/src/internalClusterTest/java/org/elasticsearch/cluster/ClusterInfoServiceIT.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.admin.cluster.node.stats.TransportNodesStatsAction;
+import org.elasticsearch.action.admin.cluster.node.usage.TransportNodeUsageStatsForThreadPoolsAction;
 import org.elasticsearch.action.admin.indices.stats.IndicesStatsAction;
 import org.elasticsearch.action.support.ActionFilter;
 import org.elasticsearch.action.support.ActionFilters;
@@ -21,6 +22,7 @@
 import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
 import org.elasticsearch.cluster.routing.RoutingTable;
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.allocation.WriteLoadConstraintSettings;
 import org.elasticsearch.cluster.routing.allocation.decider.EnableAllocationDecider;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.Strings;
@@ -39,21 +41,25 @@
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.test.InternalTestCluster;
 import org.elasticsearch.test.transport.MockTransportService;
+import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 import org.hamcrest.Matchers;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import static java.util.Collections.emptySet;
 import static java.util.Collections.singletonList;
 import static java.util.Collections.unmodifiableSet;
+import static org.elasticsearch.cluster.InternalClusterInfoService.INTERNAL_CLUSTER_INFO_UPDATE_INTERVAL_SETTING;
 import static org.elasticsearch.common.util.set.Sets.newHashSet;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.hamcrest.CoreMatchers.equalTo;
@@ -202,7 +208,7 @@ public void testClusterInfoServiceInformationClearOnError() {
         internalCluster().startNodes(
             2,
             // manually control publishing
-            Settings.builder().put(InternalClusterInfoService.INTERNAL_CLUSTER_INFO_UPDATE_INTERVAL_SETTING.getKey(), "60m").build()
+            Settings.builder().put(INTERNAL_CLUSTER_INFO_UPDATE_INTERVAL_SETTING.getKey(), "60m").build()
         );
         prepareCreate("test").setSettings(Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)).get();
         ensureGreen("test");
@@ -334,4 +340,65 @@ public void testClusterInfoServiceInformationClearOnError() {
             );
         }
     }
+
+    public void testClusterInfoIncludesNodeUsageStatsForThreadPools() {
+        var settings = Settings.builder()
+            .put(
+                WriteLoadConstraintSettings.WRITE_LOAD_DECIDER_ENABLED_SETTING.getKey(),
+                WriteLoadConstraintSettings.WriteLoadDeciderStatus.ENABLED
+            )
+            .build();
+        var masterName = internalCluster().startMasterOnlyNode(settings);
+        var dataNodeName = internalCluster().startDataOnlyNode(settings);
+        ensureStableCluster(2);
+        assertEquals(internalCluster().getMasterName(), masterName);
+        assertNotEquals(internalCluster().getMasterName(), dataNodeName);
+        logger.info("---> master node: " + masterName + ", data node: " + dataNodeName);
+
+        // Track when the data node receives a poll from the master for the write thread pool's stats.
+        final MockTransportService dataNodeMockTransportService = MockTransportService.getInstance(dataNodeName);
+        final CountDownLatch nodeThreadPoolStatsPolledByMaster = new CountDownLatch(1);
+        dataNodeMockTransportService.addRequestHandlingBehavior(
+            TransportNodeUsageStatsForThreadPoolsAction.NAME + "[n]",
+            (handler, request, channel, task) -> {
+                handler.messageReceived(request, channel, task);
+
+                if (nodeThreadPoolStatsPolledByMaster.getCount() > 0) {
+                    logger.info("---> Data node received a request for thread pool stats");
+                }
+                nodeThreadPoolStatsPolledByMaster.countDown();
+            }
+        );
+
+        // Do some writes to create some write thread pool activity.
+        final String indexName = randomIdentifier();
+        for (int i = 0; i < randomIntBetween(1, 1000); i++) {
+            index(indexName, Integer.toString(i), Collections.singletonMap("foo", "bar"));
+        }
+
+        // Force a refresh of the ClusterInfo state to collect fresh info from the data nodes.
+        final InternalClusterInfoService masterClusterInfoService = asInstanceOf(
+            InternalClusterInfoService.class,
+            internalCluster().getCurrentMasterNodeInstance(ClusterInfoService.class)
+        );
+        final ClusterInfo clusterInfo = ClusterInfoServiceUtils.refresh(masterClusterInfoService);
+
+        // Verify that the data node received a request for thread pool stats.
+        safeAwait(nodeThreadPoolStatsPolledByMaster);
+
+        final Map<String, NodeUsageStatsForThreadPools> usageStatsForThreadPools = clusterInfo.getNodeUsageStatsForThreadPools();
+        logger.info("---> Thread pool usage stats reported by data nodes to the master: " + usageStatsForThreadPools);
+        assertThat(usageStatsForThreadPools.size(), equalTo(2)); // master and data node
+        var dataNodeId = getNodeId(dataNodeName);
+        var nodeUsageStatsForThreadPool = usageStatsForThreadPools.get(dataNodeId);
+        assertNotNull(nodeUsageStatsForThreadPool);
+        logger.info("---> Data node's thread pool stats: " + nodeUsageStatsForThreadPool);
+
+        assertEquals(dataNodeId, nodeUsageStatsForThreadPool.nodeId());
+        var writeThreadPoolStats = nodeUsageStatsForThreadPool.threadPoolUsageStatsMap().get(ThreadPool.Names.WRITE);
+        assertNotNull("Expected to find stats for the WRITE thread pool", writeThreadPoolStats);
+        assertThat(writeThreadPoolStats.totalThreadPoolThreads(), greaterThan(0));
+        assertThat(writeThreadPoolStats.averageThreadPoolUtilization(), greaterThan(0f));
+        assertThat(writeThreadPoolStats.maxThreadPoolQueueLatencyMillis(), greaterThanOrEqualTo(0L));
+    }
 }

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

@@ -25,7 +25,6 @@
 import org.elasticsearch.cluster.EstimatedHeapUsageCollector;
 import org.elasticsearch.cluster.InternalClusterInfoService;
 import org.elasticsearch.cluster.NodeUsageStatsForThreadPools;
-import org.elasticsearch.cluster.NodeUsageStatsForThreadPoolsCollector;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeUtils;
@@ -91,7 +90,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -133,11 +131,7 @@ public class IndexShardIT extends ESSingleNodeTestCase {
 
     @Override
     protected Collection<Class<? extends Plugin>> getPlugins() {
-        return pluginList(
-            InternalSettingsPlugin.class,
-            BogusEstimatedHeapUsagePlugin.class,
-            BogusNodeUsageStatsForThreadPoolsCollectorPlugin.class
-        );
+        return pluginList(InternalSettingsPlugin.class, BogusEstimatedHeapUsagePlugin.class);
     }
 
     public void testLockTryingToDelete() throws Exception {
@@ -332,8 +326,7 @@ public void testNodeWriteLoadsArePresent() {
             ClusterInfoServiceUtils.refresh(clusterInfoService);
             nodeThreadPoolStats = clusterInfoService.getClusterInfo().getNodeUsageStatsForThreadPools();
 
-            /** Verify that each node has usage stats reported. The test {@link BogusNodeUsageStatsForThreadPoolsCollector} implementation
-             * generates random usage values */
+            /** Verify that each node has usage stats reported. */
             ClusterState state = getInstanceFromNode(ClusterService.class).state();
             assertEquals(state.nodes().size(), nodeThreadPoolStats.size());
             for (DiscoveryNode node : state.nodes()) {
@@ -346,7 +339,7 @@ public void testNodeWriteLoadsArePresent() {
                 assertNotNull(writeThreadPoolStats);
                 assertThat(writeThreadPoolStats.totalThreadPoolThreads(), greaterThanOrEqualTo(0));
                 assertThat(writeThreadPoolStats.averageThreadPoolUtilization(), greaterThanOrEqualTo(0.0f));
-                assertThat(writeThreadPoolStats.averageThreadPoolQueueLatencyMillis(), greaterThanOrEqualTo(0L));
+                assertThat(writeThreadPoolStats.maxThreadPoolQueueLatencyMillis(), greaterThanOrEqualTo(0L));
             }
         } finally {
             updateClusterSettings(
@@ -935,61 +928,4 @@ public ClusterService getClusterService() {
             return clusterService.get();
         }
     }
-
-    /**
-     * A simple {@link NodeUsageStatsForThreadPoolsCollector} implementation that creates and returns random
-     * {@link NodeUsageStatsForThreadPools} for each node in the cluster.
-     * <p>
-     * Note: there's an 'org.elasticsearch.cluster.NodeUsageStatsForThreadPoolsCollector' 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 BogusNodeUsageStatsForThreadPoolsCollector implements NodeUsageStatsForThreadPoolsCollector {
-
-        private final BogusNodeUsageStatsForThreadPoolsCollectorPlugin plugin;
-
-        public BogusNodeUsageStatsForThreadPoolsCollector(BogusNodeUsageStatsForThreadPoolsCollectorPlugin plugin) {
-            this.plugin = plugin;
-        }
-
-        @Override
-        public void collectUsageStats(ActionListener<Map<String, NodeUsageStatsForThreadPools>> listener) {
-            ActionListener.completeWith(
-                listener,
-                () -> plugin.getClusterService()
-                    .state()
-                    .nodes()
-                    .stream()
-                    .collect(Collectors.toUnmodifiableMap(DiscoveryNode::getId, node -> makeRandomNodeUsageStats(node.getId())))
-            );
-        }
-
-        private NodeUsageStatsForThreadPools makeRandomNodeUsageStats(String nodeId) {
-            NodeUsageStatsForThreadPools.ThreadPoolUsageStats writeThreadPoolStats = new NodeUsageStatsForThreadPools.ThreadPoolUsageStats(
-                randomNonNegativeInt(),
-                randomFloat(),
-                randomNonNegativeLong()
-            );
-            Map<String, NodeUsageStatsForThreadPools.ThreadPoolUsageStats> statsForThreadPools = new HashMap<>();
-            statsForThreadPools.put(ThreadPool.Names.WRITE, writeThreadPoolStats);
-            return new NodeUsageStatsForThreadPools(nodeId, statsForThreadPools);
-        }
-    }
-
-    /**
-     * Make a plugin to gain access to the {@link ClusterService} instance.
-     */
-    public static class BogusNodeUsageStatsForThreadPoolsCollectorPlugin 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.NodeUsageStatsForThreadPoolsCollector

@@ -342,6 +342,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion NODE_USAGE_STATS_FOR_THREAD_POOLS_IN_CLUSTER_INFO = def(9_121_0_00);
     public static final TransportVersion ESQL_CATEGORIZE_OPTIONS = def(9_122_0_00);
     public static final TransportVersion ML_INFERENCE_AZURE_AI_STUDIO_RERANK_ADDED = def(9_123_0_00);
+    public static final TransportVersion TRANSPORT_NODE_USAGE_STATS_FOR_THREAD_POOLS_ACTION = def(9_124_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -38,6 +38,7 @@
 import org.elasticsearch.action.admin.cluster.node.tasks.cancel.TransportCancelTasksAction;
 import org.elasticsearch.action.admin.cluster.node.tasks.get.TransportGetTaskAction;
 import org.elasticsearch.action.admin.cluster.node.tasks.list.TransportListTasksAction;
+import org.elasticsearch.action.admin.cluster.node.usage.TransportNodeUsageStatsForThreadPoolsAction;
 import org.elasticsearch.action.admin.cluster.node.usage.TransportNodesUsageAction;
 import org.elasticsearch.action.admin.cluster.remote.RemoteClusterNodesAction;
 import org.elasticsearch.action.admin.cluster.remote.TransportRemoteInfoAction;
@@ -629,6 +630,7 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
         ActionRegistry actions = new ActionRegistry();
 
         actions.register(TransportNodesInfoAction.TYPE, TransportNodesInfoAction.class);
+        actions.register(TransportNodeUsageStatsForThreadPoolsAction.TYPE, TransportNodeUsageStatsForThreadPoolsAction.class);
         actions.register(TransportRemoteInfoAction.TYPE, TransportRemoteInfoAction.class);
         actions.register(TransportNodesCapabilitiesAction.TYPE, TransportNodesCapabilitiesAction.class);
         actions.register(TransportNodesFeaturesAction.TYPE, TransportNodesFeaturesAction.class);