From 2d9b3a69e8f71ab6545a09dcd9ad8a025c1ced3e Mon Sep 17 00:00:00 2001
From: David Turner <david.turner@elastic.co>
Date: Mon, 17 Jun 2019 17:13:04 +0100
Subject: [PATCH] Relocation targets are assigned shards too (#43276)

Adds relocation targets to the output of
`IndexShardRoutingTable#assignedShards`.
---
 .../cluster/routing/IndexShardRoutingTable.java       | 11 +++++------
 .../routing/allocation/IndexMetaDataUpdater.java      |  6 +++++-
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/IndexShardRoutingTable.java b/server/src/main/java/org/elasticsearch/cluster/routing/IndexShardRoutingTable.java
index c35b7d810ca..97cf855b8b6 100644
--- a/server/src/main/java/org/elasticsearch/cluster/routing/IndexShardRoutingTable.java
+++ b/server/src/main/java/org/elasticsearch/cluster/routing/IndexShardRoutingTable.java
@@ -105,6 +105,10 @@ public class IndexShardRoutingTable implements Iterable<ShardRouting> {
                 // create the target initializing shard routing on the node the shard is relocating to
                 allInitializingShards.add(shard.getTargetRelocatingShard());
                 allAllocationIds.add(shard.getTargetRelocatingShard().allocationId().getId());
+
+                assert shard.assignedToNode() : "relocating from unassigned " + shard;
+                assert shard.getTargetRelocatingShard().assignedToNode() : "relocating to unassigned " + shard.getTargetRelocatingShard();
+                assignedShards.add(shard.getTargetRelocatingShard());
             }
             if (shard.assignedToNode()) {
                 assignedShards.add(shard);
@@ -211,7 +215,7 @@ public class IndexShardRoutingTable implements Iterable<ShardRouting> {
     }
 
     /**
-     * Returns a {@link List} of assigned shards
+     * Returns a {@link List} of assigned shards, including relocation targets
      *
      * @return a {@link List} of shards
      */
@@ -518,11 +522,6 @@ public class IndexShardRoutingTable implements Iterable<ShardRouting> {
             if (shardRouting.allocationId().getId().equals(allocationId)) {
                 return shardRouting;
             }
-            if (shardRouting.relocating()) {
-                if (shardRouting.getTargetRelocatingShard().allocationId().getId().equals(allocationId)) {
-                    return shardRouting.getTargetRelocatingShard();
-                }
-            }
         }
         return null;
     }
diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/allocation/IndexMetaDataUpdater.java b/server/src/main/java/org/elasticsearch/cluster/routing/allocation/IndexMetaDataUpdater.java
index 54625a15e8d..18d3b45fd9c 100644
--- a/server/src/main/java/org/elasticsearch/cluster/routing/allocation/IndexMetaDataUpdater.java
+++ b/server/src/main/java/org/elasticsearch/cluster/routing/allocation/IndexMetaDataUpdater.java
@@ -194,9 +194,13 @@ public class IndexMetaDataUpdater extends RoutingChangesObserver.AbstractRouting
             // of replicas was decreased while shards were unassigned.
             int maxActiveShards = oldIndexMetaData.getNumberOfReplicas() + 1; // +1 for the primary
             IndexShardRoutingTable newShardRoutingTable = newRoutingTable.shardRoutingTable(shardId);
+            assert newShardRoutingTable.assignedShards().stream()
+                .filter(ShardRouting::isRelocationTarget).map(s -> s.allocationId().getId()).noneMatch(inSyncAllocationIds::contains)
+                : newShardRoutingTable.assignedShards() + " vs " + inSyncAllocationIds;
             if (inSyncAllocationIds.size() > oldInSyncAllocationIds.size() && inSyncAllocationIds.size() > maxActiveShards) {
                 // trim entries that have no corresponding shard routing in the cluster state (i.e. trim unavailable copies)
-                List<ShardRouting> assignedShards = newShardRoutingTable.assignedShards();
+                List<ShardRouting> assignedShards = newShardRoutingTable.assignedShards()
+                    .stream().filter(s -> s.isRelocationTarget() == false).collect(Collectors.toList());
                 assert assignedShards.size() <= maxActiveShards :
                     "cannot have more assigned shards " + assignedShards + " than maximum possible active shards " + maxActiveShards;
                 Set<String> assignedAllocations = assignedShards.stream().map(s -> s.allocationId().getId()).collect(Collectors.toSet());