Reapply "Track shardStarted events for simulation in DesiredBalanceComputer" (#135597)

This PR reapplies both #133630 and #135052 with a performance bug fix.
The original PR #133630 had a severe impact on throughput for index
creation. It was reverted with #135369. The flamegraph suggests the
system spent a lot time to compute shard assignments on ClusterInfo
instantiation time. But that is unnecessary since they do not change
within a single polling interval. This PR fixes it by reuse the last
value and avoid recomputation.

Copying the original commit message here

If a shard starts on the target node before the next ClusterInfo
polling, today we don't include it for the simulation. With this PR, we
track shards that can potentially start within one ClusterInfo polling
cycle so that they are always included in simulation. The tracking is
reset when a new ClusterInfo arrives.

Resolves: ES-12723

Author: ywangd

docs/changelog/135597.yaml

@@ -0,0 +1,5 @@
+pr: 135597
+summary: Track `shardStarted` events for simulation in `DesiredBalanceComputer`
+area: Allocation
+type: enhancement
+issues: []

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

@@ -21,12 +21,15 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Map;
@@ -69,6 +72,7 @@ public class ClusterInfo implements ChunkedToXContent, Writeable, ExpectedShardS
     final Map<ShardId, Double> shardWriteLoads;
     // max heap size per node ID
     final Map<String, ByteSizeValue> maxHeapSizePerNode;
+    private final Map<ShardId, Set<String>> shardToNodeIds;
 
     protected ClusterInfo() {
         this(Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of(), Map.of());
@@ -99,6 +103,34 @@ public ClusterInfo(
         Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools,
         Map<ShardId, Double> shardWriteLoads,
         Map<String, ByteSizeValue> maxHeapSizePerNode
+    ) {
+        this(
+            leastAvailableSpaceUsage,
+            mostAvailableSpaceUsage,
+            shardSizes,
+            shardDataSetSizes,
+            dataPath,
+            reservedSpace,
+            estimatedHeapUsages,
+            nodeUsageStatsForThreadPools,
+            shardWriteLoads,
+            maxHeapSizePerNode,
+            computeShardToNodeIds(dataPath)
+        );
+    }
+
+    private ClusterInfo(
+        Map<String, DiskUsage> leastAvailableSpaceUsage,
+        Map<String, DiskUsage> mostAvailableSpaceUsage,
+        Map<String, Long> shardSizes,
+        Map<ShardId, Long> shardDataSetSizes,
+        Map<NodeAndShard, String> dataPath,
+        Map<NodeAndPath, ReservedSpace> reservedSpace,
+        Map<String, EstimatedHeapUsage> estimatedHeapUsages,
+        Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools,
+        Map<ShardId, Double> shardWriteLoads,
+        Map<String, ByteSizeValue> maxHeapSizePerNode,
+        Map<ShardId, Set<String>> shardToNodeIds
     ) {
         this.leastAvailableSpaceUsage = Map.copyOf(leastAvailableSpaceUsage);
         this.mostAvailableSpaceUsage = Map.copyOf(mostAvailableSpaceUsage);
@@ -110,6 +142,7 @@ public ClusterInfo(
         this.nodeUsageStatsForThreadPools = Map.copyOf(nodeUsageStatsForThreadPools);
         this.shardWriteLoads = Map.copyOf(shardWriteLoads);
         this.maxHeapSizePerNode = Map.copyOf(maxHeapSizePerNode);
+        this.shardToNodeIds = shardToNodeIds;
     }
 
     public ClusterInfo(StreamInput in) throws IOException {
@@ -139,6 +172,46 @@ public ClusterInfo(StreamInput in) throws IOException {
         } else {
             this.maxHeapSizePerNode = Map.of();
         }
+        this.shardToNodeIds = computeShardToNodeIds(dataPath);
+    }
+
+    ClusterInfo updateWith(
+        Map<String, DiskUsage> leastAvailableSpaceUsage,
+        Map<String, DiskUsage> mostAvailableSpaceUsage,
+        Map<String, Long> shardSizes,
+        Map<NodeAndPath, ReservedSpace> reservedSpace,
+        Map<String, EstimatedHeapUsage> estimatedHeapUsages,
+        Map<String, NodeUsageStatsForThreadPools> nodeUsageStatsForThreadPools
+    ) {
+        return new ClusterInfo(
+            leastAvailableSpaceUsage,
+            mostAvailableSpaceUsage,
+            shardSizes,
+            shardDataSetSizes,
+            dataPath,
+            reservedSpace,
+            estimatedHeapUsages,
+            nodeUsageStatsForThreadPools,
+            shardWriteLoads,
+            maxHeapSizePerNode,
+            shardToNodeIds
+        );
+    }
+
+    private static Map<ShardId, Set<String>> computeShardToNodeIds(Map<NodeAndShard, String> dataPath) {
+        if (dataPath.isEmpty()) {
+            return Map.of();
+        }
+        final var shardToNodeIds = new HashMap<ShardId, Set<String>>();
+        for (NodeAndShard nodeAndShard : dataPath.keySet()) {
+            shardToNodeIds.computeIfAbsent(nodeAndShard.shardId, ignore -> new HashSet<>()).add(nodeAndShard.nodeId);
+        }
+        return Collections.unmodifiableMap(Maps.transformValues(shardToNodeIds, Collections::unmodifiableSet));
+    }
+
+    public Set<String> getNodeIdsForShard(ShardId shardId) {
+        assert shardToNodeIds != null : "shardToNodeIds not computed for simulations, make sure this ClusterInfo is from polling";
+        return shardToNodeIds.getOrDefault(shardId, Set.of());
     }
 
     @Override
@@ -325,6 +398,16 @@ public Map<String, ByteSizeValue> getMaxHeapSizePerNode() {
         return this.maxHeapSizePerNode;
     }
 
+    /**
+     * Return true if the shard has moved since the time ClusterInfo was created.
+     */
+    public boolean hasShardMoved(ShardRouting shardRouting) {
+        // We use dataPath to find out whether a shard is allocated on a node.
+        // TODO: DataPath is sent with disk usages but thread pool usage is sent separately so that local shard allocation
+        // may change between the two calls.
+        return getDataPath(shardRouting) == null;
+    }
+
     /**
      * Method that incorporates the ShardId for the shard into a string that
      * includes a 'p' or 'r' depending on whether the shard is a primary.

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

@@ -9,12 +9,14 @@
 
 package org.elasticsearch.cluster;
 
-import org.elasticsearch.cluster.ClusterInfo.NodeAndShard;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.elasticsearch.cluster.routing.ShardMovementWriteLoadSimulator;
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.UnassignedInfo;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
-import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.CopyOnFirstWriteMap;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.shard.ShardId;
 
 import java.util.HashMap;
@@ -25,29 +27,26 @@
 import static org.elasticsearch.cluster.routing.ExpectedShardSizeEstimator.getExpectedShardSize;
 import static org.elasticsearch.cluster.routing.ExpectedShardSizeEstimator.shouldReserveSpaceForInitializingShard;
 import static org.elasticsearch.cluster.routing.ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE;
+import static org.elasticsearch.cluster.routing.UnassignedInfo.Reason.REINITIALIZED;
 
 public class ClusterInfoSimulator {
 
+    private static final Logger logger = LogManager.getLogger(ClusterInfoSimulator.class);
+
     private final RoutingAllocation allocation;
 
     private final Map<String, DiskUsage> leastAvailableSpaceUsage;
     private final Map<String, DiskUsage> mostAvailableSpaceUsage;
     private final CopyOnFirstWriteMap<String, Long> shardSizes;
-    private final Map<ShardId, Long> shardDataSetSizes;
-    private final Map<NodeAndShard, String> dataPath;
     private final Map<String, EstimatedHeapUsage> estimatedHeapUsages;
-    private final Map<String, ByteSizeValue> maxHeapSizePerNode;
     private final ShardMovementWriteLoadSimulator shardMovementWriteLoadSimulator;
 
     public ClusterInfoSimulator(RoutingAllocation allocation) {
         this.allocation = allocation;
         this.leastAvailableSpaceUsage = getAdjustedDiskSpace(allocation, allocation.clusterInfo().getNodeLeastAvailableDiskUsages());
         this.mostAvailableSpaceUsage = getAdjustedDiskSpace(allocation, allocation.clusterInfo().getNodeMostAvailableDiskUsages());
         this.shardSizes = new CopyOnFirstWriteMap<>(allocation.clusterInfo().shardSizes);
-        this.shardDataSetSizes = Map.copyOf(allocation.clusterInfo().shardDataSetSizes);
-        this.dataPath = Map.copyOf(allocation.clusterInfo().dataPath);
         this.estimatedHeapUsages = allocation.clusterInfo().getEstimatedHeapUsages();
-        this.maxHeapSizePerNode = Map.copyOf(allocation.clusterInfo().maxHeapSizePerNode);
         this.shardMovementWriteLoadSimulator = new ShardMovementWriteLoadSimulator(allocation);
     }
 
@@ -95,7 +94,7 @@ private static Map<String, DiskUsage> getAdjustedDiskSpace(RoutingAllocation all
      * Balance is later recalculated with a refreshed cluster info containing actual shards placement.
      */
     public void simulateShardStarted(ShardRouting shard) {
-        assert shard.initializing();
+        assert shard.initializing() : "expected an initializing shard, but got: " + shard;
 
         var project = allocation.metadata().projectFor(shard.index());
         var size = getExpectedShardSize(
@@ -122,6 +121,36 @@ public void simulateShardStarted(ShardRouting shard) {
         shardMovementWriteLoadSimulator.simulateShardStarted(shard);
     }
 
+    /**
+     * This method simulates starting an already started shard with an optional {@code sourceNodeId} in case of a relocation.
+     * @param startedShard The shard to simulate. Must be started already.
+     * @param sourceNodeId The source node ID if the shard started as a result of relocation. {@code null} otherwise.
+     */
+    public void simulateAlreadyStartedShard(ShardRouting startedShard, @Nullable String sourceNodeId) {
+        assert startedShard.started() : "expected an already started shard, but got: " + startedShard;
+        if (logger.isDebugEnabled()) {
+            logger.debug(
+                "simulated started shard {} on node [{}] as a {}",
+                startedShard.shardId(),
+                startedShard.currentNodeId(),
+                sourceNodeId != null ? "relocating shard from node [" + sourceNodeId + "]" : "new shard"
+            );
+        }
+        final long expectedShardSize = startedShard.getExpectedShardSize();
+        if (sourceNodeId != null) {
+            final var relocatingShard = startedShard.moveToUnassigned(new UnassignedInfo(REINITIALIZED, "simulation"))
+                .initialize(sourceNodeId, null, expectedShardSize)
+                .moveToStarted(expectedShardSize)
+                .relocate(startedShard.currentNodeId(), expectedShardSize)
+                .getTargetRelocatingShard();
+            simulateShardStarted(relocatingShard);
+        } else {
+            final var initializingShard = startedShard.moveToUnassigned(new UnassignedInfo(REINITIALIZED, "simulation"))
+                .initialize(startedShard.currentNodeId(), null, expectedShardSize);
+            simulateShardStarted(initializingShard);
+        }
+    }
+
     private void modifyDiskUsage(String nodeId, long freeDelta) {
         if (freeDelta == 0) {
             return;
@@ -156,17 +185,14 @@ private static long withinRange(long min, long max, long value) {
     }
 
     public ClusterInfo getClusterInfo() {
-        return new ClusterInfo(
-            leastAvailableSpaceUsage,
-            mostAvailableSpaceUsage,
-            shardSizes.toImmutableMap(),
-            shardDataSetSizes,
-            dataPath,
-            Map.of(),
-            estimatedHeapUsages,
-            shardMovementWriteLoadSimulator.simulatedNodeUsageStatsForThreadPools(),
-            allocation.clusterInfo().getShardWriteLoads(),
-            maxHeapSizePerNode
-        );
+        return allocation.clusterInfo()
+            .updateWith(
+                leastAvailableSpaceUsage,
+                mostAvailableSpaceUsage,
+                shardSizes.toImmutableMap(),
+                Map.of(),
+                estimatedHeapUsages,
+                shardMovementWriteLoadSimulator.simulatedNodeUsageStatsForThreadPools()
+            );
     }
 }

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

@@ -100,13 +100,6 @@ public class InternalClusterInfoService implements ClusterInfoService, ClusterSt
     private volatile TimeValue updateFrequency;
     private volatile TimeValue fetchTimeout;
 
-    private volatile Map<String, DiskUsage> leastAvailableSpaceUsages;
-    private volatile Map<String, DiskUsage> mostAvailableSpaceUsages;
-    private volatile Map<String, ByteSizeValue> maxHeapPerNode;
-    private volatile Map<String, Long> estimatedHeapUsagePerNode;
-    private volatile Map<String, NodeUsageStatsForThreadPools> nodeThreadPoolUsageStatsPerNode;
-    private volatile IndicesStatsSummary indicesStatsSummary;
-
     private final ThreadPool threadPool;
     private final Client client;
     private final Supplier<ClusterState> clusterStateSupplier;
@@ -120,6 +113,8 @@ public class InternalClusterInfoService implements ClusterInfoService, ClusterSt
     private AsyncRefresh currentRefresh;
     private RefreshScheduler refreshScheduler;
 
+    private volatile ClusterInfo currentClusterInfo = ClusterInfo.EMPTY;
+
     @SuppressWarnings("this-escape")
     public InternalClusterInfoService(
         Settings settings,
@@ -129,12 +124,6 @@ public InternalClusterInfoService(
         EstimatedHeapUsageCollector estimatedHeapUsageCollector,
         NodeUsageStatsForThreadPoolsCollector nodeUsageStatsForThreadPoolsCollector
     ) {
-        this.leastAvailableSpaceUsages = Map.of();
-        this.mostAvailableSpaceUsages = Map.of();
-        this.maxHeapPerNode = Map.of();
-        this.estimatedHeapUsagePerNode = Map.of();
-        this.nodeThreadPoolUsageStatsPerNode = Map.of();
-        this.indicesStatsSummary = IndicesStatsSummary.EMPTY;
         this.threadPool = threadPool;
         this.client = client;
         this.estimatedHeapUsageCollector = estimatedHeapUsageCollector;
@@ -208,6 +197,13 @@ public void clusterChanged(ClusterChangedEvent event) {
 
     private class AsyncRefresh {
 
+        private volatile Map<String, DiskUsage> leastAvailableSpaceUsages;
+        private volatile Map<String, DiskUsage> mostAvailableSpaceUsages;
+        private volatile Map<String, ByteSizeValue> maxHeapPerNode;
+        private volatile Map<String, Long> estimatedHeapUsagePerNode;
+        private volatile Map<String, NodeUsageStatsForThreadPools> nodeThreadPoolUsageStatsPerNode;
+        private volatile IndicesStatsSummary indicesStatsSummary;
+
         private final List<ActionListener<ClusterInfo>> thisRefreshListeners;
         private final RefCountingRunnable fetchRefs = new RefCountingRunnable(this::callListeners);
 
@@ -453,7 +449,7 @@ public void onFailure(Exception e) {
         private void callListeners() {
             try {
                 logger.trace("stats all received, computing cluster info and notifying listeners");
-                final ClusterInfo clusterInfo = getClusterInfo();
+                final ClusterInfo clusterInfo = updateAndGetCurrentClusterInfo();
                 boolean anyListeners = false;
                 for (final Consumer<ClusterInfo> listener : listeners) {
                     anyListeners = true;
@@ -473,6 +469,30 @@ private void callListeners() {
                 onRefreshComplete(this);
             }
         }
+
+        private ClusterInfo updateAndGetCurrentClusterInfo() {
+            final Map<String, EstimatedHeapUsage> estimatedHeapUsages = new HashMap<>();
+            maxHeapPerNode.forEach((nodeId, maxHeapSize) -> {
+                final Long estimatedHeapUsage = estimatedHeapUsagePerNode.get(nodeId);
+                if (estimatedHeapUsage != null) {
+                    estimatedHeapUsages.put(nodeId, new EstimatedHeapUsage(nodeId, maxHeapSize.getBytes(), estimatedHeapUsage));
+                }
+            });
+            final var newClusterInfo = new ClusterInfo(
+                leastAvailableSpaceUsages,
+                mostAvailableSpaceUsages,
+                indicesStatsSummary.shardSizes,
+                indicesStatsSummary.shardDataSetSizes,
+                indicesStatsSummary.dataPath,
+                indicesStatsSummary.reservedSpace,
+                estimatedHeapUsages,
+                nodeThreadPoolUsageStatsPerNode,
+                indicesStatsSummary.shardWriteLoads(),
+                maxHeapPerNode
+            );
+            currentClusterInfo = newClusterInfo;
+            return newClusterInfo;
+        }
     }
 
     private void onRefreshComplete(AsyncRefresh completedRefresh) {
@@ -537,27 +557,7 @@ private boolean shouldRefresh() {
 
     @Override
     public ClusterInfo getClusterInfo() {
-        final IndicesStatsSummary indicesStatsSummary = this.indicesStatsSummary; // single volatile read
-        final Map<String, EstimatedHeapUsage> estimatedHeapUsages = new HashMap<>();
-        final var currentMaxHeapPerNode = this.maxHeapPerNode; // Make sure we use a consistent view
-        currentMaxHeapPerNode.forEach((nodeId, maxHeapSize) -> {
-            final Long estimatedHeapUsage = estimatedHeapUsagePerNode.get(nodeId);
-            if (estimatedHeapUsage != null) {
-                estimatedHeapUsages.put(nodeId, new EstimatedHeapUsage(nodeId, maxHeapSize.getBytes(), estimatedHeapUsage));
-            }
-        });
-        return new ClusterInfo(
-            leastAvailableSpaceUsages,
-            mostAvailableSpaceUsages,
-            indicesStatsSummary.shardSizes,
-            indicesStatsSummary.shardDataSetSizes,
-            indicesStatsSummary.dataPath,
-            indicesStatsSummary.reservedSpace,
-            estimatedHeapUsages,
-            nodeThreadPoolUsageStatsPerNode,
-            indicesStatsSummary.shardWriteLoads(),
-            currentMaxHeapPerNode
-        );
+        return currentClusterInfo;
     }
 
     // allow tests to adjust the node stats on receipt