Add allocation write load stats to write thread pool (#130373)

Instrument the WRITE thread pool to collect:
1. Thread pool utilization for periodic collection
2. Thread pool max queue latency for periodic collection
This supports collection of node-level write load
stats, which will be collected by the balancer in
a future patch.

Relates ES-12233

Author: DiannaHohensee

server/src/internalClusterTest/java/org/elasticsearch/threadpool/SimpleThreadPoolIT.java

@@ -234,19 +234,29 @@ public void assertValid(TestTelemetryPlugin testTelemetryPlugin, String metricSu
         }
     }
 
-    public void testWriteThreadpoolEwmaAlphaSetting() {
+    public void testWriteThreadpoolsEwmaAlphaSetting() {
         Settings settings = Settings.EMPTY;
-        var ewmaAlpha = DEFAULT_INDEX_AUTOSCALING_EWMA_ALPHA;
+        var executionEwmaAlpha = DEFAULT_INDEX_AUTOSCALING_EWMA_ALPHA;
         if (randomBoolean()) {
-            ewmaAlpha = randomDoubleBetween(0.0, 1.0, true);
-            settings = Settings.builder().put(WRITE_THREAD_POOLS_EWMA_ALPHA_SETTING.getKey(), ewmaAlpha).build();
+            executionEwmaAlpha = randomDoubleBetween(0.0, 1.0, true);
+            settings = Settings.builder().put(WRITE_THREAD_POOLS_EWMA_ALPHA_SETTING.getKey(), executionEwmaAlpha).build();
         }
         var nodeName = internalCluster().startNode(settings);
         var threadPool = internalCluster().getInstance(ThreadPool.class, nodeName);
+
+        // Verify that the write thread pools all use the tracking executor.
         for (var name : List.of(ThreadPool.Names.WRITE, ThreadPool.Names.SYSTEM_WRITE, ThreadPool.Names.SYSTEM_CRITICAL_WRITE)) {
             assertThat(threadPool.executor(name), instanceOf(TaskExecutionTimeTrackingEsThreadPoolExecutor.class));
             final var executor = (TaskExecutionTimeTrackingEsThreadPoolExecutor) threadPool.executor(name);
-            assertThat(Double.compare(executor.getEwmaAlpha(), ewmaAlpha), CoreMatchers.equalTo(0));
+            assertThat(Double.compare(executor.getExecutionEwmaAlpha(), executionEwmaAlpha), CoreMatchers.equalTo(0));
+
+            // Only the WRITE thread pool should enable further tracking.
+            if (name.equals(ThreadPool.Names.WRITE) == false) {
+                assertFalse(executor.trackingMaxQueueLatency());
+            } else {
+                // Verify that the WRITE thread pool has extra tracking enabled.
+                assertTrue(executor.trackingMaxQueueLatency());
+            }
         }
     }
 }

server/src/main/java/org/elasticsearch/common/util/concurrent/EsExecutors.java

@@ -577,24 +577,53 @@ public void rejectedExecution(Runnable task, ThreadPoolExecutor executor) {
     }
 
     public static class TaskTrackingConfig {
-        // This is a random starting point alpha. TODO: revisit this with actual testing and/or make it configurable
-        public static final double DEFAULT_EWMA_ALPHA = 0.3;
+        public static final double DEFAULT_EXECUTION_TIME_EWMA_ALPHA_FOR_TEST = 0.3;
 
         private final boolean trackExecutionTime;
         private final boolean trackOngoingTasks;
-        private final double ewmaAlpha;
+        private final boolean trackMaxQueueLatency;
+        private final double executionTimeEwmaAlpha;
+
+        public static final TaskTrackingConfig DO_NOT_TRACK = new TaskTrackingConfig(
+            false,
+            false,
+            false,
+            DEFAULT_EXECUTION_TIME_EWMA_ALPHA_FOR_TEST
+        );
+        public static final TaskTrackingConfig DEFAULT = new TaskTrackingConfig(
+            true,
+            false,
+            false,
+            DEFAULT_EXECUTION_TIME_EWMA_ALPHA_FOR_TEST
+        );
 
-        public static final TaskTrackingConfig DO_NOT_TRACK = new TaskTrackingConfig(false, false, DEFAULT_EWMA_ALPHA);
-        public static final TaskTrackingConfig DEFAULT = new TaskTrackingConfig(true, false, DEFAULT_EWMA_ALPHA);
+        public TaskTrackingConfig(boolean trackOngoingTasks, double executionTimeEWMAAlpha) {
+            this(true, trackOngoingTasks, false, executionTimeEWMAAlpha);
+        }
 
-        public TaskTrackingConfig(boolean trackOngoingTasks, double ewmaAlpha) {
-            this(true, trackOngoingTasks, ewmaAlpha);
+        /**
+         * Execution tracking enabled constructor, with extra options to enable further specialized tracking.
+         */
+        public TaskTrackingConfig(boolean trackOngoingTasks, boolean trackMaxQueueLatency, double executionTimeEwmaAlpha) {
+            this(true, trackOngoingTasks, trackMaxQueueLatency, executionTimeEwmaAlpha);
         }
 
-        private TaskTrackingConfig(boolean trackExecutionTime, boolean trackOngoingTasks, double EWMAAlpha) {
+        /**
+         * @param trackExecutionTime Whether to track execution stats
+         * @param trackOngoingTasks Whether to track ongoing task execution time, not just finished tasks
+         * @param trackMaxQueueLatency Whether to track max queue latency.
+         * @param executionTimeEWMAAlpha The alpha seed for execution time EWMA (ExponentiallyWeightedMovingAverage).
+         */
+        private TaskTrackingConfig(
+            boolean trackExecutionTime,
+            boolean trackOngoingTasks,
+            boolean trackMaxQueueLatency,
+            double executionTimeEWMAAlpha
+        ) {
             this.trackExecutionTime = trackExecutionTime;
             this.trackOngoingTasks = trackOngoingTasks;
-            this.ewmaAlpha = EWMAAlpha;
+            this.trackMaxQueueLatency = trackMaxQueueLatency;
+            this.executionTimeEwmaAlpha = executionTimeEWMAAlpha;
         }
 
         public boolean trackExecutionTime() {
@@ -605,8 +634,12 @@ public boolean trackOngoingTasks() {
             return trackOngoingTasks;
         }
 
-        public double getEwmaAlpha() {
-            return ewmaAlpha;
+        public boolean trackMaxQueueLatency() {
+            return trackMaxQueueLatency;
+        }
+
+        public double getExecutionTimeEwmaAlpha() {
+            return executionTimeEwmaAlpha;
         }
     }
 

server/src/main/java/org/elasticsearch/common/util/concurrent/TaskExecutionTimeTrackingEsThreadPoolExecutor.java

@@ -13,6 +13,8 @@
 import org.elasticsearch.common.metrics.ExponentialBucketHistogram;
 import org.elasticsearch.common.util.concurrent.EsExecutors.TaskTrackingConfig;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
 import org.elasticsearch.telemetry.metric.DoubleWithAttributes;
 import org.elasticsearch.telemetry.metric.Instrument;
 import org.elasticsearch.telemetry.metric.LongWithAttributes;
@@ -27,6 +29,7 @@
 import java.util.concurrent.RejectedExecutionHandler;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.LongAccumulator;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Function;
 
@@ -37,6 +40,7 @@
  * An extension to thread pool executor, which tracks statistics for the task execution time.
  */
 public final class TaskExecutionTimeTrackingEsThreadPoolExecutor extends EsThreadPoolExecutor {
+    private static final Logger logger = LogManager.getLogger(TaskExecutionTimeTrackingEsThreadPoolExecutor.class);
 
     public static final int QUEUE_LATENCY_HISTOGRAM_BUCKETS = 18;
     private static final int[] LATENCY_PERCENTILES_TO_REPORT = { 50, 90, 99 };
@@ -47,9 +51,17 @@ public final class TaskExecutionTimeTrackingEsThreadPoolExecutor extends EsThrea
     private final boolean trackOngoingTasks;
     // The set of currently running tasks and the timestamp of when they started execution in the Executor.
     private final Map<Runnable, Long> ongoingTasks = new ConcurrentHashMap<>();
-    private volatile long lastPollTime = System.nanoTime();
-    private volatile long lastTotalExecutionTime = 0;
     private final ExponentialBucketHistogram queueLatencyMillisHistogram = new ExponentialBucketHistogram(QUEUE_LATENCY_HISTOGRAM_BUCKETS);
+    private final boolean trackMaxQueueLatency;
+    private LongAccumulator maxQueueLatencyMillisSinceLastPoll = new LongAccumulator(Long::max, 0);
+
+    public enum UtilizationTrackingPurpose {
+        APM,
+        ALLOCATION,
+    }
+
+    private volatile UtilizationTracker apmUtilizationTracker = new UtilizationTracker();
+    private volatile UtilizationTracker allocationUtilizationTracker = new UtilizationTracker();
 
     TaskExecutionTimeTrackingEsThreadPoolExecutor(
         String name,
@@ -65,9 +77,11 @@ public final class TaskExecutionTimeTrackingEsThreadPoolExecutor extends EsThrea
         TaskTrackingConfig trackingConfig
     ) {
         super(name, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler, contextHolder);
+
         this.runnableWrapper = runnableWrapper;
-        this.executionEWMA = new ExponentiallyWeightedMovingAverage(trackingConfig.getEwmaAlpha(), 0);
+        this.executionEWMA = new ExponentiallyWeightedMovingAverage(trackingConfig.getExecutionTimeEwmaAlpha(), 0);
         this.trackOngoingTasks = trackingConfig.trackOngoingTasks();
+        this.trackMaxQueueLatency = trackingConfig.trackMaxQueueLatency();
     }
 
     public List<Instrument> setupMetrics(MeterRegistry meterRegistry, String threadPoolName) {
@@ -95,7 +109,7 @@ public List<Instrument> setupMetrics(MeterRegistry meterRegistry, String threadP
                 ThreadPool.THREAD_POOL_METRIC_PREFIX + threadPoolName + THREAD_POOL_METRIC_NAME_UTILIZATION,
                 "fraction of maximum thread time utilized for " + threadPoolName,
                 "fraction",
-                () -> new DoubleWithAttributes(pollUtilization(), Map.of())
+                () -> new DoubleWithAttributes(pollUtilization(UtilizationTrackingPurpose.APM), Map.of())
             )
         );
     }
@@ -136,37 +150,49 @@ public int getCurrentQueueSize() {
         return getQueue().size();
     }
 
+    public long getMaxQueueLatencyMillisSinceLastPollAndReset() {
+        if (trackMaxQueueLatency == false) {
+            return 0;
+        }
+        return maxQueueLatencyMillisSinceLastPoll.getThenReset();
+    }
+
     /**
-     * Returns the fraction of the maximum possible thread time that was actually used since the last time
-     * this method was called.
+     * Returns the fraction of the maximum possible thread time that was actually used since the last time this method was called.
+     * There are two periodic pulling mechanisms that access utilization reporting: {@link UtilizationTrackingPurpose} distinguishes the
+     * caller.
      *
-     * @return the utilization as a fraction, in the range [0, 1]
+     * @return the utilization as a fraction, in the range [0, 1]. This may return >1 if a task completed in the time range but started
+     * earlier, contributing a larger execution time.
      */
-    public double pollUtilization() {
-        final long currentTotalExecutionTimeNanos = totalExecutionTime.sum();
-        final long currentPollTimeNanos = System.nanoTime();
-
-        final long totalExecutionTimeSinceLastPollNanos = currentTotalExecutionTimeNanos - lastTotalExecutionTime;
-        final long timeSinceLastPoll = currentPollTimeNanos - lastPollTime;
-        final long maximumExecutionTimeSinceLastPollNanos = timeSinceLastPoll * getMaximumPoolSize();
-        final double utilizationSinceLastPoll = (double) totalExecutionTimeSinceLastPollNanos / maximumExecutionTimeSinceLastPollNanos;
-
-        lastTotalExecutionTime = currentTotalExecutionTimeNanos;
-        lastPollTime = currentPollTimeNanos;
-        return utilizationSinceLastPoll;
+    public double pollUtilization(UtilizationTrackingPurpose utilizationTrackingPurpose) {
+        switch (utilizationTrackingPurpose) {
+            case APM:
+                return apmUtilizationTracker.pollUtilization();
+            case ALLOCATION:
+                return allocationUtilizationTracker.pollUtilization();
+            default:
+                throw new IllegalStateException("No operation defined for [" + utilizationTrackingPurpose + "]");
+        }
     }
 
     @Override
     protected void beforeExecute(Thread t, Runnable r) {
         if (trackOngoingTasks) {
             ongoingTasks.put(r, System.nanoTime());
         }
+
         assert super.unwrap(r) instanceof TimedRunnable : "expected only TimedRunnables in queue";
         final TimedRunnable timedRunnable = (TimedRunnable) super.unwrap(r);
         timedRunnable.beforeExecute();
         final long taskQueueLatency = timedRunnable.getQueueTimeNanos();
         assert taskQueueLatency >= 0;
-        queueLatencyMillisHistogram.addObservation(TimeUnit.NANOSECONDS.toMillis(taskQueueLatency));
+        var queueLatencyMillis = TimeUnit.NANOSECONDS.toMillis(taskQueueLatency);
+        queueLatencyMillisHistogram.addObservation(queueLatencyMillis);
+
+        if (trackMaxQueueLatency) {
+            maxQueueLatencyMillisSinceLastPoll.accumulate(queueLatencyMillis);
+        }
     }
 
     @Override
@@ -222,7 +248,39 @@ public Map<Runnable, Long> getOngoingTasks() {
     }
 
     // Used for testing
-    public double getEwmaAlpha() {
+    public double getExecutionEwmaAlpha() {
         return executionEWMA.getAlpha();
     }
+
+    // Used for testing
+    public boolean trackingMaxQueueLatency() {
+        return trackMaxQueueLatency;
+    }
+
+    /**
+     * Supports periodic polling for thread pool utilization. Tracks state since the last polling request so that the average utilization
+     * since the last poll can be calculated for the next polling request.
+     *
+     * Uses the difference of {@link #totalExecutionTime} since the last polling request to determine how much activity has occurred.
+     */
+    private class UtilizationTracker {
+        long lastPollTime = System.nanoTime();
+        long lastTotalExecutionTime = 0;
+
+        public synchronized double pollUtilization() {
+            final long currentTotalExecutionTimeNanos = totalExecutionTime.sum();
+            final long currentPollTimeNanos = System.nanoTime();
+
+            final long totalExecutionTimeSinceLastPollNanos = currentTotalExecutionTimeNanos - lastTotalExecutionTime;
+            final long timeSinceLastPoll = currentPollTimeNanos - lastPollTime;
+
+            final long maximumExecutionTimeSinceLastPollNanos = timeSinceLastPoll * getMaximumPoolSize();
+            final double utilizationSinceLastPoll = (double) totalExecutionTimeSinceLastPollNanos / maximumExecutionTimeSinceLastPollNanos;
+
+            lastTotalExecutionTime = currentTotalExecutionTimeNanos;
+            lastPollTime = currentPollTimeNanos;
+
+            return utilizationSinceLastPoll;
+        }
+    }
 }

server/src/main/java/org/elasticsearch/threadpool/DefaultBuiltInExecutorBuilders.java

@@ -56,7 +56,7 @@ public Map<String, ExecutorBuilder> getBuilders(Settings settings, int allocated
                 allocatedProcessors,
                 // 10,000 for all nodes with 8 cores or fewer. Scale up once we have more than 8 cores.
                 Math.max(allocatedProcessors * 750, 10000),
-                new EsExecutors.TaskTrackingConfig(true, indexAutoscalingEWMA)
+                new EsExecutors.TaskTrackingConfig(true, true, indexAutoscalingEWMA)
             )
         );
         int searchOrGetThreadPoolSize = ThreadPool.searchOrGetThreadPoolSize(allocatedProcessors);