Take relocating shard into consideration during awareness allocation

Previous fix #12551 counted twice for relocating shard (source and target).
Fix it to consider only target node.

core/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/AwarenessAllocationDecider.java

@@ -184,11 +184,9 @@ public class AwarenessAllocationDecider extends AllocationDecider {
             // build the count of shards per attribute value
             ObjectIntHashMap<String> shardPerAttribute = new ObjectIntHashMap<>();
             for (ShardRouting assignedShard : allocation.routingNodes().assignedShards(shardRouting)) {
-                // if the shard is relocating, then make sure we count it as part of the node it is relocating to
+                if (assignedShard.started() || assignedShard.initializing()) {
-                if (assignedShard.relocating()) {
+                    // Note: this also counts relocation targets as that will be the new location of the shard.
-                    RoutingNode relocationNode = allocation.routingNodes().node(assignedShard.relocatingNodeId());
+                    // Relocation sources should not be counted as the shard is moving away
-                    shardPerAttribute.addTo(relocationNode.node().attributes().get(awarenessAttribute), 1);
-                } else if (assignedShard.started() || assignedShard.initializing()) {
                     RoutingNode routingNode = allocation.routingNodes().node(assignedShard.currentNodeId());
                     shardPerAttribute.addTo(routingNode.node().attributes().get(awarenessAttribute), 1);
                 }

core/src/test/java/org/elasticsearch/cluster/routing/allocation/AwarenessAllocationTests.java

@@ -28,9 +28,13 @@ import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.cluster.routing.RoutingTable;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.ShardRoutingState;
+import org.elasticsearch.cluster.routing.allocation.command.AllocationCommands;
+import org.elasticsearch.cluster.routing.allocation.command.CancelAllocationCommand;
+import org.elasticsearch.cluster.routing.allocation.command.MoveAllocationCommand;
 import org.elasticsearch.cluster.routing.allocation.decider.ClusterRebalanceAllocationDecider;
 import org.elasticsearch.common.logging.ESLogger;
 import org.elasticsearch.common.logging.Loggers;
+import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.test.ESAllocationTestCase;
 import org.junit.Test;
 
@@ -853,6 +857,7 @@ public class AwarenessAllocationTests extends ESAllocationTestCase {
                         .put(newNode("A-1", ImmutableMap.of("zone", "a")))
                         .put(newNode("A-2", ImmutableMap.of("zone", "a")))
                         .put(newNode("A-3", ImmutableMap.of("zone", "a")))
+                        .put(newNode("A-4", ImmutableMap.of("zone", "a")))
                         .put(newNode("B-0", ImmutableMap.of("zone", "b")))
         ).build();
         routingTable = strategy.reroute(clusterState).routingTable();
@@ -866,5 +871,25 @@ public class AwarenessAllocationTests extends ESAllocationTestCase {
         assertThat(clusterState.getRoutingNodes().shardsWithState(STARTED).size(), equalTo(1));
         assertThat(clusterState.getRoutingNodes().shardsWithState(INITIALIZING).size(), equalTo(3));
         assertThat(clusterState.getRoutingNodes().shardsWithState(UNASSIGNED).size(), equalTo(1)); // Unassigned shard is expected.
+
+        // Cancel all initializing shards and move started primary to another node.
+        AllocationCommands commands = new AllocationCommands();
+        String primaryNode = null;
+        for (ShardRouting routing : routingTable.allShards()) {
+            if (routing.primary()) {
+                primaryNode = routing.currentNodeId();
+            } else if (routing.initializing()) {
+                commands.add(new CancelAllocationCommand(routing.shardId(), routing.currentNodeId(), false));
+            }
+        }
+        commands.add(new MoveAllocationCommand(new ShardId("test", 0), primaryNode, "A-4"));
+
+        routingTable = strategy.reroute(clusterState, commands).routingTable();
+        clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
+
+        assertThat(clusterState.getRoutingNodes().shardsWithState(STARTED).size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().shardsWithState(RELOCATING).size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().shardsWithState(INITIALIZING).size(), equalTo(4)); // +1 for relocating shard.
+        assertThat(clusterState.getRoutingNodes().shardsWithState(UNASSIGNED).size(), equalTo(1)); // Still 1 unassigned.
     }
 }