Simulate shards moved by explicit commands

In DesiredBalancerComputer, shard movements should be simulated by
starting any shards that are initializatng. Previously the explicitly
moved shards are not covered. This PR fixes it.

Resolves: ES-12943

Author: ywangd

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

@@ -156,9 +156,8 @@ public void allocate(RoutingAllocation allocation) {
         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();
+        assert allocation.isSimulating() == false || allocation.routingNodes().getRelocatingShardCount() == 0
+            : "expect no relocating shard, but got " + allocation.routingNodes();
 
         if (allocation.routingNodes().size() == 0) {
             failAllocationOfNewPrimaries(allocation);

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

@@ -36,6 +36,7 @@
 import org.elasticsearch.index.shard.ShardId;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
@@ -291,7 +292,16 @@ public DesiredBalance compute(
         while ((commands = pendingDesiredBalanceMoves.poll()) != null) {
             for (MoveAllocationCommand command : commands) {
                 try {
-                    command.execute(routingAllocation, false);
+                    final var rerouteExplanation = command.execute(routingAllocation, false);
+                    if (rerouteExplanation.decisions().type() != Decision.Type.NO) {
+                        final ShardRouting[] initializingShards = routingNodes.node(command.toNode()).initializing();
+                        assert initializingShards.length == 1 && command.fromNode().equals(initializingShards[0].relocatingNodeId())
+                            : "expect one relocating shard, but got : " + List.of(initializingShards);
+                        Arrays.stream(initializingShards).forEach(shard -> {
+                            clusterInfoSimulator.simulateShardStarted(shard);
+                            routingNodes.startShard(shard, changes, 0L);
+                        });
+                    }
                 } catch (RuntimeException e) {
                     logger.debug(
                         () -> "move shard ["

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputerTests.java

@@ -54,6 +54,7 @@
 import org.elasticsearch.common.time.TimeProviderUtils;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.Maps;
+import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.core.Tuple;
@@ -566,7 +567,21 @@ public void testNoDataNodes() {
     }
 
     public void testAppliesMoveCommands() {
-        var desiredBalanceComputer = createDesiredBalanceComputer();
+        var desiredBalanceComputer = createDesiredBalanceComputer(new ShardsAllocator() {
+            @Override
+            public void allocate(RoutingAllocation allocation) {
+                assertThat(
+                    "unexpected relocating shards: " + allocation.routingNodes(),
+                    allocation.routingNodes().getRelocatingShardCount(),
+                    equalTo(0)
+                );
+            }
+
+            @Override
+            public ShardAllocationDecision decideShardAllocation(ShardRouting shard, RoutingAllocation allocation) {
+                throw new AssertionError("only used for allocation explain");
+            }
+        });
         var clusterState = createInitialClusterState(3);
         var index = clusterState.metadata().getProject().index(TEST_INDEX).getIndex();
 
@@ -578,25 +593,29 @@ public void testAppliesMoveCommands() {
         }
         clusterState = rebuildRoutingTable(clusterState, routingNodes);
 
-        var desiredBalance = desiredBalanceComputer.compute(
-            DesiredBalance.BECOME_MASTER_INITIAL,
-            createInput(clusterState),
-            queue(
-                new MoveAllocationCommand(index.getName(), 0, "node-1", "node-2"),
-                new MoveAllocationCommand(index.getName(), 1, "node-1", "node-2")
-            ),
-            input -> true
-        );
+        final var dataNodeIds = clusterState.nodes().getDataNodes().keySet();
+        for (var nodeId : List.of("node-0", "node-1")) {
+            var desiredBalance = desiredBalanceComputer.compute(
+                DesiredBalance.BECOME_MASTER_INITIAL,
+                createInput(clusterState),
+                queue(
+                    new MoveAllocationCommand(index.getName(), 0, nodeId, "node-2"),
+                    new MoveAllocationCommand(index.getName(), 1, nodeId, "node-2")
+                ),
+                input -> true
+            );
 
-        assertDesiredAssignments(
-            desiredBalance,
-            Map.of(
-                new ShardId(index, 0),
-                new ShardAssignment(Set.of("node-0", "node-2"), 2, 0, 0),
-                new ShardId(index, 1),
-                new ShardAssignment(Set.of("node-0", "node-2"), 2, 0, 0)
-            )
-        );
+            final Set<String> expectedNodeIds = Sets.difference(dataNodeIds, Set.of(nodeId));
+            assertDesiredAssignments(
+                desiredBalance,
+                Map.of(
+                    new ShardId(index, 0),
+                    new ShardAssignment(expectedNodeIds, 2, 0, 0),
+                    new ShardId(index, 1),
+                    new ShardAssignment(expectedNodeIds, 2, 0, 0)
+                )
+            );
+        }
     }
 
     public void testDesiredBalanceShouldConvergeInABigCluster() {