Fix simulator for new cluster info

Author: nicktindall

server/src/main/java/org/elasticsearch/cluster/routing/ShardMovementWriteLoadSimulator.java

@@ -18,7 +18,9 @@
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.threadpool.ThreadPool;
 
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -86,13 +88,17 @@ private static NodeUsageStatsForThreadPools.ThreadPoolUsageStats replaceWritePoo
     ) {
         final NodeUsageStatsForThreadPools.ThreadPoolUsageStats writeThreadPoolStats = value.threadPoolUsageStatsMap()
             .get(ThreadPool.Names.WRITE);
-        return new NodeUsageStatsForThreadPools.ThreadPoolUsageStats(
-            writeThreadPoolStats.totalThreadPoolThreads(),
-            (float) Math.max(
-                (writeThreadPoolStats.averageThreadPoolUtilization() + (writeLoadDelta / writeThreadPoolStats.totalThreadPoolThreads())),
-                0.0
-            ),
-            writeThreadPoolStats.averageThreadPoolQueueLatencyMillis()
+        final float newWritePoolUtilization = (float) Math.max(
+            (writeThreadPoolStats.utilizationSamples().getLast().utilization() + (writeLoadDelta / writeThreadPoolStats.numberOfThreads())),
+            0.0
         );
+        final List<NodeUsageStatsForThreadPools.UtilizationSample> newUtilizationSamples = new ArrayList<>(
+            writeThreadPoolStats.utilizationSamples()
+        );
+        final NodeUsageStatsForThreadPools.UtilizationSample previousUtilization = newUtilizationSamples.removeLast();
+        newUtilizationSamples.add(
+            new NodeUsageStatsForThreadPools.UtilizationSample(previousUtilization.instant(), newWritePoolUtilization)
+        );
+        return new NodeUsageStatsForThreadPools.ThreadPoolUsageStats(writeThreadPoolStats.numberOfThreads(), newUtilizationSamples);
     }
 }

server/src/test/java/org/elasticsearch/cluster/routing/ShardMovementWriteLoadSimulatorTests.java

@@ -21,6 +21,7 @@
 import org.elasticsearch.test.ESTestCase;
 import org.hamcrest.Matchers;
 
+import java.time.Instant;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -87,16 +88,16 @@ public void testMovementOfAShardWillMoveThreadPoolUtilisation() {
         assertThat(calculatedNodeUsageStats, Matchers.aMapWithSize(2));
 
         final var shardWriteLoad = allocation.clusterInfo().getShardWriteLoads().get(randomShard.shardId());
-        final var expectedUtilisationReductionAtSource = shardWriteLoad / originalNode0ThreadPoolStats.totalThreadPoolThreads();
-        final var expectedUtilisationIncreaseAtDestination = shardWriteLoad / originalNode1ThreadPoolStats.totalThreadPoolThreads();
+        final var expectedUtilisationReductionAtSource = shardWriteLoad / originalNode0ThreadPoolStats.numberOfThreads();
+        final var expectedUtilisationIncreaseAtDestination = shardWriteLoad / originalNode1ThreadPoolStats.numberOfThreads();
 
         // Some node_0 utilization should have been moved to node_1
-        if (expectedUtilisationReductionAtSource > originalNode0ThreadPoolStats.averageThreadPoolUtilization()) {
+        if (expectedUtilisationReductionAtSource > originalNode0ThreadPoolStats.utilizationSamples().getLast().utilization()) {
             // We don't return utilization less than zero because that makes no sense
             assertThat(getAverageWritePoolUtilization(shardMovementWriteLoadSimulator, "node_0"), equalTo(0.0f));
         } else {
             assertThat(
-                (double) originalNode0ThreadPoolStats.averageThreadPoolUtilization() - getAverageWritePoolUtilization(
+                (double) originalNode0ThreadPoolStats.utilizationSamples().getLast().utilization() - getAverageWritePoolUtilization(
                     shardMovementWriteLoadSimulator,
                     "node_0"
                 ),
@@ -105,7 +106,9 @@ public void testMovementOfAShardWillMoveThreadPoolUtilisation() {
         }
         assertThat(
             (double) getAverageWritePoolUtilization(shardMovementWriteLoadSimulator, "node_1") - originalNode1ThreadPoolStats
-                .averageThreadPoolUtilization(),
+                .utilizationSamples()
+                .getLast()
+                .utilization(),
             closeTo(expectedUtilisationIncreaseAtDestination, 0.001f)
         );
 
@@ -117,11 +120,11 @@ public void testMovementOfAShardWillMoveThreadPoolUtilisation() {
         // The utilization numbers should return to their original values
         assertThat(
             getAverageWritePoolUtilization(shardMovementWriteLoadSimulator, "node_0"),
-            equalTo(originalNode0ThreadPoolStats.averageThreadPoolUtilization())
+            equalTo(originalNode0ThreadPoolStats.utilizationSamples().getLast().utilization())
         );
         assertThat(
             getAverageWritePoolUtilization(shardMovementWriteLoadSimulator, "node_1"),
-            equalTo(originalNode1ThreadPoolStats.averageThreadPoolUtilization())
+            equalTo(originalNode1ThreadPoolStats.utilizationSamples().getLast().utilization())
         );
     }
 
@@ -150,14 +153,18 @@ public void testMovementBetweenNodesWithNoThreadPoolAndWriteLoadStats() {
     private float getAverageWritePoolUtilization(ShardMovementWriteLoadSimulator shardMovementWriteLoadSimulator, String nodeId) {
         final var generatedNodeUsageStates = shardMovementWriteLoadSimulator.simulatedNodeUsageStatsForThreadPools();
         final var node0WritePoolStats = generatedNodeUsageStates.get(nodeId).threadPoolUsageStatsMap().get("write");
-        return node0WritePoolStats.averageThreadPoolUtilization();
+        return node0WritePoolStats.utilizationSamples().getLast().utilization();
     }
 
     private NodeUsageStatsForThreadPools.ThreadPoolUsageStats randomThreadPoolUsageStats() {
         return new NodeUsageStatsForThreadPools.ThreadPoolUsageStats(
             randomIntBetween(4, 16),
-            randomBoolean() ? 0.0f : randomFloatBetween(0.1f, 1.0f, true),
-            randomLongBetween(0, 60_000)
+            List.of(
+                new NodeUsageStatsForThreadPools.UtilizationSample(
+                    randomInstantBetween(Instant.MIN, Instant.MAX),
+                    randomBoolean() ? 0.0f : randomFloatBetween(0.1f, 1.0f, true)
+                )
+            )
         );
     }