Reapply "Add an option to return early from an allocate call (elastic#134786)" (elastic#135476)

This reverts commit 31f1810.

Author: ywangd

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java

@@ -140,25 +140,83 @@ public BalancedShardsAllocator(
 
     @Override
     public void allocate(RoutingAllocation allocation) {
+        assert allocation.isSimulating() == false || balancerSettings.completeEarlyOnShardAssignmentChange()
+            : "inconsistent states: isSimulating ["
+                + allocation.isSimulating()
+                + "] vs completeEarlyOnShardAssignmentChange ["
+                + balancerSettings.completeEarlyOnShardAssignmentChange()
+                + "]";
         if (allocation.metadata().hasAnyIndices()) {
             // must not use licensed features when just starting up
             writeLoadForecaster.refreshLicense();
         }
 
         assert allocation.ignoreDisable() == false;
+        assert allocation.isSimulating() == false || allocation.routingNodes().hasInactiveShards() == false
+            : "expect no initializing shard, but got " + allocation.routingNodes();
+        // TODO: ES-12943 cannot assert the following because shards moved by commands are not simulated promptly in DesiredBalanceComputer
+        // assert allocation.isSimulating() == false || allocation.routingNodes().getRelocatingShardCount() == 0
+        // : "expect no relocating shard, but got " + allocation.routingNodes();
 
         if (allocation.routingNodes().size() == 0) {
             failAllocationOfNewPrimaries(allocation);
             return;
         }
         final BalancingWeights balancingWeights = balancingWeightsFactory.create();
-        final Balancer balancer = new Balancer(writeLoadForecaster, allocation, balancerSettings.getThreshold(), balancingWeights);
-        balancer.allocateUnassigned();
-        balancer.moveShards();
-        balancer.balance();
+        final Balancer balancer = new Balancer(
+            writeLoadForecaster,
+            allocation,
+            balancerSettings.getThreshold(),
+            balancingWeights,
+            balancerSettings.completeEarlyOnShardAssignmentChange()
+        );
+
+        boolean shardAssigned = false, shardMoved = false, shardBalanced = false;
+        try {
+            shardAssigned = balancer.allocateUnassigned();
+            if (shardAssigned && balancerSettings.completeEarlyOnShardAssignmentChange()) {
+                return;
+            }
+
+            shardMoved = balancer.moveShards();
+            if (shardMoved && balancerSettings.completeEarlyOnShardAssignmentChange()) {
+                return;
+            }
+
+            shardBalanced = balancer.balance();
+        } finally {
+            if (logger.isDebugEnabled()) {
+                logger.debug(
+                    "shards assigned: {}, shards moved: {}, shards balanced: {}, "
+                        + "routingNodes hasInactiveShards [{}], relocation count [{}]",
+                    shardAssigned,
+                    shardMoved,
+                    shardBalanced,
+                    allocation.routingNodes().hasInactiveShards(),
+                    allocation.routingNodes().getRelocatingShardCount()
+                );
+            }
+            assert assertShardAssignmentChanges(allocation, shardAssigned, shardMoved, shardBalanced);
+            // Node weights are calculated after each internal balancing round and saved to the RoutingNodes copy.
+            collectAndRecordNodeWeightStats(balancer, balancingWeights, allocation);
+        }
+    }
 
-        // Node weights are calculated after each internal balancing round and saved to the RoutingNodes copy.
-        collectAndRecordNodeWeightStats(balancer, balancingWeights, allocation);
+    private boolean assertShardAssignmentChanges(
+        RoutingAllocation allocation,
+        boolean shardAssigned,
+        boolean shardMoved,
+        boolean shardBalanced
+    ) {
+        if (allocation.isSimulating() == false) {
+            return true;
+        }
+        assert shardAssigned == false || allocation.routingNodes().hasInactiveShards()
+            : "expect initializing shard, but got " + allocation.routingNodes();
+
+        assert (shardMoved == false && shardBalanced == false) || allocation.routingNodes().getRelocatingShardCount() > 0
+            : "expect relocating shard, but got " + allocation.routingNodes();
+        return true;
     }
 
     private void collectAndRecordNodeWeightStats(Balancer balancer, BalancingWeights balancingWeights, RoutingAllocation allocation) {
@@ -188,7 +246,8 @@ public ShardAllocationDecision decideShardAllocation(final ShardRouting shard, f
             writeLoadForecaster,
             allocation,
             balancerSettings.getThreshold(),
-            balancingWeightsFactory.create()
+            balancingWeightsFactory.create(),
+            balancerSettings.completeEarlyOnShardAssignmentChange()
         );
         AllocateUnassignedDecision allocateUnassignedDecision = AllocateUnassignedDecision.NOT_TAKEN;
         MoveDecision moveDecision = MoveDecision.NOT_TAKEN;
@@ -248,12 +307,14 @@ public static class Balancer {
         private final Map<String, ModelNode> nodes;
         private final BalancingWeights balancingWeights;
         private final NodeSorters nodeSorters;
+        private final boolean completeEarlyOnShardAssignmentChange;
 
         private Balancer(
             WriteLoadForecaster writeLoadForecaster,
             RoutingAllocation allocation,
             float threshold,
-            BalancingWeights balancingWeights
+            BalancingWeights balancingWeights,
+            boolean completeEarlyOnShardAssignmentChange
         ) {
             this.writeLoadForecaster = writeLoadForecaster;
             this.allocation = allocation;
@@ -266,6 +327,7 @@ private Balancer(
             nodes = Collections.unmodifiableMap(buildModelFromAssigned());
             this.nodeSorters = balancingWeights.createNodeSorters(nodesArray(), this);
             this.balancingWeights = balancingWeights;
+            this.completeEarlyOnShardAssignmentChange = completeEarlyOnShardAssignmentChange;
         }
 
         private static long getShardDiskUsageInBytes(ShardRouting shardRouting, IndexMetadata indexMetadata, ClusterInfo clusterInfo) {
@@ -358,7 +420,7 @@ private IndexMetadata indexMetadata(ProjectIndex index) {
          * Balances the nodes on the cluster model according to the weight function.
          * The actual balancing is delegated to {@link #balanceByWeights(NodeSorter)}
          */
-        private void balance() {
+        private boolean balance() {
             if (logger.isTraceEnabled()) {
                 logger.trace("Start balancing cluster");
             }
@@ -371,21 +433,27 @@ private void balance() {
                  * Therefore we only do a rebalance if we have fetched all information.
                  */
                 logger.debug("skipping rebalance due to in-flight shard/store fetches");
-                return;
+                return false;
             }
             if (allocation.deciders().canRebalance(allocation).type() != Type.YES) {
                 logger.trace("skipping rebalance as it is disabled");
-                return;
+                return false;
             }
 
+            boolean shardBalanced = false;
             // Balance each partition
             for (NodeSorter nodeSorter : nodeSorters) {
                 if (nodeSorter.modelNodes.length < 2) { /* skip if we only have one node */
                     logger.trace("skipping rebalance as the partition has single node only");
                     continue;
                 }
-                balanceByWeights(nodeSorter);
+                shardBalanced |= balanceByWeights(nodeSorter);
+                // TODO: We could choose to account shardBalanced separately for each partition since they do not overlap.
+                if (shardBalanced && completeEarlyOnShardAssignmentChange) {
+                    return true;
+                }
             }
+            return shardBalanced;
         }
 
         /**
@@ -531,7 +599,8 @@ private MoveDecision decideRebalance(final ProjectIndex index, final ShardRoutin
          * only, or in other words relocations that move the weight delta closer
          * to {@code 0.0}
          */
-        private void balanceByWeights(NodeSorter sorter) {
+        private boolean balanceByWeights(NodeSorter sorter) {
+            boolean shardBalanced = false;
             final AllocationDeciders deciders = allocation.deciders();
             final ModelNode[] modelNodes = sorter.modelNodes;
             final float[] weights = sorter.weights;
@@ -630,6 +699,15 @@ private void balanceByWeights(NodeSorter sorter) {
                             sorter.sort(0, relevantNodes);
                             lowIdx = 0;
                             highIdx = relevantNodes - 1;
+
+                            shardBalanced = true;
+                            if (completeEarlyOnShardAssignmentChange && routingNodes.getRelocatingShardCount() > 0) {
+                                // ES-12955: Check routingNodes.getRelocatingShardCount() > 0 in case the first relocation is a THROTTLE.
+                                // It should not happen in production, i.e, throttling should not happen unless there is a prior shard
+                                // that is already relocating. But in tests, we have decider like RandomAllocationDecider that can
+                                // randomly return THROTTLE when there is no existing relocation.
+                                return true;
+                            }
                             continue;
                         }
                     }
@@ -651,6 +729,7 @@ private void balanceByWeights(NodeSorter sorter) {
                     }
                 }
             }
+            return shardBalanced;
         }
 
         /**
@@ -721,7 +800,8 @@ protected int comparePivot(int j) {
          * shard is created with an incremented version in the state
          * {@link ShardRoutingState#INITIALIZING}.
          */
-        public void moveShards() {
+        public boolean moveShards() {
+            boolean shardMoved = false;
             // Iterate over the started shards interleaving between nodes, and check if they can remain. In the presence of throttling
             // shard movements, the goal of this iteration order is to achieve a fairer movement of shards from the nodes that are
             // offloading the shards.
@@ -745,10 +825,15 @@ public void moveShards() {
                     if (logger.isTraceEnabled()) {
                         logger.trace("Moved shard [{}] to node [{}]", shardRouting, targetNode.getRoutingNode());
                     }
+                    shardMoved = true;
+                    if (completeEarlyOnShardAssignmentChange) {
+                        return true;
+                    }
                 } else if (moveDecision.isDecisionTaken() && moveDecision.canRemain() == false) {
                     logger.trace("[{}][{}] can't move", shardRouting.index(), shardRouting.id());
                 }
             }
+            return shardMoved;
         }
 
         /**
@@ -888,14 +973,14 @@ private Map<String, ModelNode> buildModelFromAssigned() {
          * Allocates all given shards on the minimal eligible node for the shards index
          * with respect to the weight function. All given shards must be unassigned.
          */
-        private void allocateUnassigned() {
+        private boolean allocateUnassigned() {
             RoutingNodes.UnassignedShards unassigned = routingNodes.unassigned();
             assert nodes.isEmpty() == false;
             if (logger.isTraceEnabled()) {
                 logger.trace("Start allocating unassigned shards");
             }
             if (unassigned.isEmpty()) {
-                return;
+                return false;
             }
 
             /*
@@ -932,6 +1017,7 @@ private void allocateUnassigned() {
             int secondaryLength = 0;
             int primaryLength = primary.length;
             ArrayUtil.timSort(primary, comparator);
+            boolean shardAssignmentChanged = false;
             do {
                 for (int i = 0; i < primaryLength; i++) {
                     ShardRouting shard = primary[i];
@@ -949,6 +1035,7 @@ private void allocateUnassigned() {
 
                         final long shardSize = getExpectedShardSize(shard, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE, allocation);
                         shard = routingNodes.initializeShard(shard, minNode.getNodeId(), null, shardSize, allocation.changes());
+                        shardAssignmentChanged = true;
                         minNode.addShard(index, shard);
                         if (shard.primary() == false) {
                             // copy over the same replica shards to the secondary array so they will get allocated
@@ -972,6 +1059,9 @@ private void allocateUnassigned() {
                             assert allocationDecision.getAllocationStatus() == AllocationStatus.DECIDERS_THROTTLED;
                             final long shardSize = getExpectedShardSize(shard, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE, allocation);
                             minNode.addShard(projectIndex(shard), shard.initialize(minNode.getNodeId(), null, shardSize));
+                            // If we see a throttle decision in simulation, there must be other shards that got assigned before it.
+                            assert allocation.isSimulating() == false || shardAssignmentChanged
+                                : "shard " + shard + " was throttled but no other shards were assigned";
                         } else {
                             if (logger.isTraceEnabled()) {
                                 logger.trace("No Node found to assign shard [{}]", shard);
@@ -994,6 +1084,7 @@ private void allocateUnassigned() {
                 secondaryLength = 0;
             } while (primaryLength > 0);
             // clear everything we have either added it or moved to ignoreUnassigned
+            return shardAssignmentChanged;
         }
 
         private ProjectIndex projectIndex(ShardRouting shardRouting) {

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/BalancerSettings.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.cluster.routing.allocation.allocator;
 
+import org.elasticsearch.cluster.ClusterModule;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Settings;
 
@@ -26,6 +27,7 @@ public class BalancerSettings {
     private volatile float writeLoadBalanceFactor;
     private volatile float diskUsageBalanceFactor;
     private volatile float threshold;
+    private final boolean completeEarlyOnShardAssignmentChange;
 
     public BalancerSettings(Settings settings) {
         this(ClusterSettings.createBuiltInClusterSettings(settings));
@@ -37,6 +39,9 @@ public BalancerSettings(ClusterSettings clusterSettings) {
         clusterSettings.initializeAndWatch(WRITE_LOAD_BALANCE_FACTOR_SETTING, value -> this.writeLoadBalanceFactor = value);
         clusterSettings.initializeAndWatch(DISK_USAGE_BALANCE_FACTOR_SETTING, value -> this.diskUsageBalanceFactor = value);
         clusterSettings.initializeAndWatch(THRESHOLD_SETTING, value -> this.threshold = value);
+        this.completeEarlyOnShardAssignmentChange = ClusterModule.DESIRED_BALANCE_ALLOCATOR.equals(
+            clusterSettings.get(ClusterModule.SHARDS_ALLOCATOR_TYPE_SETTING)
+        );
     }
 
     /**
@@ -67,4 +72,8 @@ public float getDiskUsageBalanceFactor() {
     public float getThreshold() {
         return threshold;
     }
+
+    public boolean completeEarlyOnShardAssignmentChange() {
+        return completeEarlyOnShardAssignmentChange;
+    }
 }