Tie break search shard iterator comparisons on cluster alias (#38853)

`SearchShardIterator` inherits its `compareTo` implementation from `PlainShardIterator`. That is good in most of the cases, as such comparisons are based on the shard id which is unique, even when searching against indices with same names across multiple clusters (thanks to the index uuid being different). In case though the same cluster is registered multiple times with different aliases, the shard id is exactly the same, hence remote results will be returned before local ones with same shard id objects. That is because remote iterators are added before local ones, and we use a stable sorting method in `GroupShardIterators` constructor.

This PR enhances `compareTo` for `SearchShardIterator` to tie break on cluster alias and introduces consistent `equals` and `hashcode` methods. This allows to remove a TODO in `SearchResponseMerger` which otherwise has to handle this special case specifically. Also, while at it I added missing tests around equals/hashcode and compareTo and expanded existing ones.

docs/reference/modules/cross-cluster-search.asciidoc

@@ -150,10 +150,10 @@ will be prefixed with their remote cluster name:
     "max_score": 1,
     "hits": [
       {
-        "_index": "cluster_one:twitter",
+        "_index": "twitter",
         "_type": "_doc",
         "_id": "0",
-        "_score": 1,
+        "_score": 2,
         "_source": {
           "user": "kimchy",
           "date": "2009-11-15T14:12:12",
@@ -162,10 +162,10 @@ will be prefixed with their remote cluster name:
         }
       },
       {
-        "_index": "twitter",
+        "_index": "cluster_one:twitter",
         "_type": "_doc",
         "_id": "0",
-        "_score": 2,
+        "_score": 1,
         "_source": {
           "user": "kimchy",
           "date": "2009-11-15T14:12:12",
@@ -243,10 +243,10 @@ GET /cluster_one:twitter,cluster_two:twitter,twitter/_search <1>
     "max_score": 1,
     "hits": [
       {
-        "_index": "cluster_one:twitter",
+        "_index": "twitter",
         "_type": "_doc",
         "_id": "0",
-        "_score": 1,
+        "_score": 2,
         "_source": {
           "user": "kimchy",
           "date": "2009-11-15T14:12:12",
@@ -255,10 +255,10 @@ GET /cluster_one:twitter,cluster_two:twitter,twitter/_search <1>
         }
       },
       {
-        "_index": "twitter",
+        "_index": "cluster_one:twitter",
         "_type": "_doc",
         "_id": "0",
-        "_score": 2,
+        "_score": 1,
         "_source": {
           "user": "kimchy",
           "date": "2009-11-15T14:12:12",

server/src/main/java/org/elasticsearch/action/search/SearchResponseMerger.java

@@ -39,7 +39,6 @@ import org.elasticsearch.search.internal.InternalSearchResponse;
 import org.elasticsearch.search.profile.ProfileShardResult;
 import org.elasticsearch.search.profile.SearchProfileShardResults;
 import org.elasticsearch.search.suggest.Suggest;
-import org.elasticsearch.transport.RemoteClusterAware;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -368,17 +367,7 @@ final class SearchResponseMerger {
             if (shardIdCompareTo != 0) {
                 return shardIdCompareTo;
             }
-            int clusterAliasCompareTo = clusterAlias.compareTo(o.clusterAlias);
-            if (clusterAliasCompareTo != 0) {
-                //TODO we may want to fix this, CCS returns remote results before local ones (TransportSearchAction#mergeShardsIterators)
-                if (clusterAlias.equals(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY)) {
-                    return 1;
-                }
-                if (o.clusterAlias.equals(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY)) {
-                    return -1;
-                }
-            }
-            return clusterAliasCompareTo;
+            return clusterAlias.compareTo(o.clusterAlias);
         }
     }
 }

server/src/main/java/org/elasticsearch/action/search/SearchShardIterator.java

@@ -21,12 +21,14 @@ package org.elasticsearch.action.search;
 
 import org.elasticsearch.action.OriginalIndices;
 import org.elasticsearch.cluster.routing.PlainShardIterator;
+import org.elasticsearch.cluster.routing.ShardIterator;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.search.SearchShardTarget;
 
 import java.util.List;
+import java.util.Objects;
 
 /**
  * Extension of {@link PlainShardIterator} used in the search api, which also holds the {@link OriginalIndices}
@@ -93,4 +95,43 @@ public final class SearchShardIterator extends PlainShardIterator {
     boolean skip() {
         return skip;
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        if (super.equals(o) == false) {
+            return false;
+        }
+        SearchShardIterator that = (SearchShardIterator) o;
+        return Objects.equals(clusterAlias, that.clusterAlias);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), clusterAlias);
+    }
+
+    @Override
+    public int compareTo(ShardIterator o) {
+        int superCompareTo = super.compareTo(o);
+        if (superCompareTo != 0 || (o instanceof SearchShardIterator == false)) {
+            return superCompareTo;
+        }
+        SearchShardIterator searchShardIterator = (SearchShardIterator)o;
+        if (clusterAlias == null && searchShardIterator.getClusterAlias() == null) {
+            return 0;
+        }
+        if (clusterAlias == null) {
+            return -1;
+        }
+        if (searchShardIterator.getClusterAlias() == null) {
+            return 1;
+        }
+        return clusterAlias.compareTo(searchShardIterator.getClusterAlias());
+    }
 }

server/src/test/java/org/elasticsearch/action/OriginalIndicesTests.java

@@ -53,7 +53,7 @@ public class OriginalIndicesTests extends ESTestCase {
         }
     }
 
-    private static OriginalIndices randomOriginalIndices() {
+    public static OriginalIndices randomOriginalIndices() {
         int numIndices = randomInt(10);
         String[] indices = new String[numIndices];
         for (int j = 0; j < indices.length; j++) {

server/src/test/java/org/elasticsearch/action/search/SearchResponseMergerTests.java

@@ -632,12 +632,6 @@ public class SearchResponseMergerTests extends ESTestCase {
             }
             int clusterAliasCompareTo = aShard.getClusterAlias().compareTo(bShard.getClusterAlias());
             if (clusterAliasCompareTo != 0) {
-                if (aShard.getClusterAlias().equals(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY)) {
-                    return 1;
-                }
-                if (bShard.getClusterAlias().equals(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY)) {
-                    return -1;
-                }
                 return clusterAliasCompareTo;
             }
             return Integer.compare(a.docId(), b.docId());

server/src/test/java/org/elasticsearch/action/search/SearchShardIteratorTests.java

@@ -20,12 +20,19 @@
 package org.elasticsearch.action.search;
 
 import org.elasticsearch.action.OriginalIndices;
+import org.elasticsearch.action.OriginalIndicesTests;
 import org.elasticsearch.action.support.IndicesOptions;
+import org.elasticsearch.cluster.routing.GroupShardsIteratorTests;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.EqualsHashCodeTestUtils;
+import org.hamcrest.Matchers;
 
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.List;
 
 public class SearchShardIteratorTests extends ESTestCase {
 
@@ -64,4 +71,79 @@ public class SearchShardIteratorTests extends ESTestCase {
         assertEquals(nodeId, searchShardTarget.getNodeId());
         assertSame(originalIndices, searchShardTarget.getOriginalIndices());
     }
+
+    public void testEqualsAndHashcode() {
+        EqualsHashCodeTestUtils.checkEqualsAndHashCode(randomSearchShardIterator(), s -> new SearchShardIterator(s.getClusterAlias(),
+            s.shardId(), s.getShardRoutings(), s.getOriginalIndices()), s -> {
+            if (randomBoolean()) {
+                String clusterAlias;
+                if (s.getClusterAlias() == null) {
+                    clusterAlias = randomAlphaOfLengthBetween(5, 10);
+                } else {
+                    clusterAlias = randomBoolean() ? null : s.getClusterAlias() + randomAlphaOfLength(3);
+                }
+                return new SearchShardIterator(clusterAlias, s.shardId(), s.getShardRoutings(), s.getOriginalIndices());
+            } else {
+                ShardId shardId = new ShardId(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLength(10),
+                    randomIntBetween(0, Integer.MAX_VALUE));
+                return new SearchShardIterator(s.getClusterAlias(), shardId, s.getShardRoutings(), s.getOriginalIndices());
+            }
+        });
+    }
+
+    public void testCompareTo() {
+        String[] clusters = generateRandomStringArray(2, 10, false, false);
+        Arrays.sort(clusters);
+        String[] indices = generateRandomStringArray(3, 10, false, false);
+        Arrays.sort(indices);
+        String[] uuids = generateRandomStringArray(3, 10, false, false);
+        Arrays.sort(uuids);
+        List<SearchShardIterator> shardIterators = new ArrayList<>();
+        int numShards = randomIntBetween(1, 5);
+        for (int i = 0; i < numShards; i++) {
+            for (String index : indices) {
+                for (String uuid : uuids) {
+                    ShardId shardId = new ShardId(index, uuid, i);
+                    shardIterators.add(new SearchShardIterator(null, shardId, GroupShardsIteratorTests.randomShardRoutings(shardId),
+                        OriginalIndicesTests.randomOriginalIndices()));
+                    for (String cluster : clusters) {
+                        shardIterators.add(new SearchShardIterator(cluster, shardId, GroupShardsIteratorTests.randomShardRoutings(shardId),
+                            OriginalIndicesTests.randomOriginalIndices()));
+                    }
+
+                }
+            }
+        }
+        for (int i = 0; i < shardIterators.size(); i++) {
+            SearchShardIterator currentIterator = shardIterators.get(i);
+            for (int j = i + 1; j < shardIterators.size(); j++) {
+                SearchShardIterator greaterIterator = shardIterators.get(j);
+                assertThat(currentIterator, Matchers.lessThan(greaterIterator));
+                assertThat(greaterIterator, Matchers.greaterThan(currentIterator));
+                assertNotEquals(currentIterator, greaterIterator);
+            }
+            for (int j = i - 1; j >= 0; j--) {
+                SearchShardIterator smallerIterator = shardIterators.get(j);
+                assertThat(smallerIterator, Matchers.lessThan(currentIterator));
+                assertThat(currentIterator, Matchers.greaterThan(smallerIterator));
+                assertNotEquals(currentIterator, smallerIterator);
+            }
+        }
+    }
+
+    public void testCompareToEqualItems() {
+        SearchShardIterator shardIterator1 = randomSearchShardIterator();
+        SearchShardIterator shardIterator2 = new SearchShardIterator(shardIterator1.getClusterAlias(), shardIterator1.shardId(),
+            shardIterator1.getShardRoutings(), shardIterator1.getOriginalIndices());
+        assertEquals(shardIterator1, shardIterator2);
+        assertEquals(0, shardIterator1.compareTo(shardIterator2));
+        assertEquals(0, shardIterator2.compareTo(shardIterator1));
+    }
+
+    private static SearchShardIterator randomSearchShardIterator() {
+        String clusterAlias = randomBoolean() ? null : randomAlphaOfLengthBetween(5, 10);
+        ShardId shardId = new ShardId(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLength(10), randomIntBetween(0, Integer.MAX_VALUE));
+        return new SearchShardIterator(clusterAlias, shardId, GroupShardsIteratorTests.randomShardRoutings(shardId),
+            OriginalIndicesTests.randomOriginalIndices());
+    }
 }

server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java

@@ -25,11 +25,13 @@ import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.LatchedActionListener;
 import org.elasticsearch.action.OriginalIndices;
+import org.elasticsearch.action.OriginalIndicesTests;
 import org.elasticsearch.action.admin.cluster.shards.ClusterSearchShardsGroup;
 import org.elasticsearch.action.admin.cluster.shards.ClusterSearchShardsResponse;
 import org.elasticsearch.action.support.IndicesOptions;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.GroupShardsIterator;
+import org.elasticsearch.cluster.routing.GroupShardsIteratorTests;
 import org.elasticsearch.cluster.routing.PlainShardIterator;
 import org.elasticsearch.cluster.routing.ShardIterator;
 import org.elasticsearch.cluster.routing.ShardRouting;
@@ -39,6 +41,7 @@ import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.transport.TransportAddress;
+import org.elasticsearch.index.Index;
 import org.elasticsearch.index.query.InnerHitBuilder;
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.TermsQueryBuilder;
@@ -71,6 +74,7 @@ import org.elasticsearch.transport.TransportRequestOptions;
 import org.elasticsearch.transport.TransportService;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -85,7 +89,6 @@ import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
-import static org.elasticsearch.cluster.routing.ShardRoutingState.STARTED;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.awaitLatch;
 import static org.hamcrest.CoreMatchers.containsString;
 import static org.hamcrest.CoreMatchers.instanceOf;
@@ -101,92 +104,96 @@ public class TransportSearchActionTests extends ESTestCase {
         ThreadPool.terminate(threadPool, 10, TimeUnit.SECONDS);
     }
 
+    private static SearchShardIterator createSearchShardIterator(int id, Index index,
+                                                                 OriginalIndices originalIndices, String clusterAlias) {
+        ShardId shardId = new ShardId(index, id);
+        List<ShardRouting> shardRoutings = GroupShardsIteratorTests.randomShardRoutings(shardId);
+        return new SearchShardIterator(clusterAlias, shardId, shardRoutings, originalIndices);
+    }
+
     public void testMergeShardsIterators() {
-        List<ShardIterator> localShardIterators = new ArrayList<>();
-        {
-            ShardId shardId = new ShardId("local_index", "local_index_uuid", 0);
-            ShardRouting shardRouting = TestShardRouting.newShardRouting(shardId, "local_node", true, STARTED);
-            ShardIterator shardIterator = new PlainShardIterator(shardId, Collections.singletonList(shardRouting));
-            localShardIterators.add(shardIterator);
+        Index[] indices = new Index[randomIntBetween(1, 10)];
+        for (int i = 0; i < indices.length; i++) {
+            if (randomBoolean() && i > 0) {
+                Index existingIndex = indices[randomIntBetween(0, i - 1)];
+                indices[i] = new Index(existingIndex.getName(), randomAlphaOfLength(10));
+            } else {
+                indices[i] = new Index(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLength(10));
+            }
         }
-        {
-            ShardId shardId2 = new ShardId("local_index_2", "local_index_2_uuid", 1);
-            ShardRouting shardRouting2 = TestShardRouting.newShardRouting(shardId2, "local_node", true, STARTED);
-            ShardIterator shardIterator2 = new PlainShardIterator(shardId2, Collections.singletonList(shardRouting2));
-            localShardIterators.add(shardIterator2);
+        Arrays.sort(indices, (o1, o2) -> {
+            int nameCompareTo = o1.getName().compareTo(o2.getName());
+            if (nameCompareTo == 0) {
+                return o1.getUUID().compareTo(o2.getUUID());
+            }
+            return nameCompareTo;
+        });
+        String[] remoteClusters = new String[randomIntBetween(1, 3)];
+        for (int i = 0; i < remoteClusters.length; i++) {
+            remoteClusters[i] = randomAlphaOfLengthBetween(5, 10);
         }
-        GroupShardsIterator<ShardIterator> localShardsIterator = new GroupShardsIterator<>(localShardIterators);
-
-        OriginalIndices localIndices = new OriginalIndices(new String[]{"local_alias", "local_index_2"},
-                SearchRequest.DEFAULT_INDICES_OPTIONS);
+        Arrays.sort(remoteClusters);
 
-        OriginalIndices remoteIndices = new OriginalIndices(new String[]{"remote_alias", "remote_index_2"},
-                IndicesOptions.strictExpandOpen());
+        List<SearchShardIterator> expected = new ArrayList<>();
+        String localClusterAlias = randomAlphaOfLengthBetween(5, 10);
+        OriginalIndices localIndices = OriginalIndicesTests.randomOriginalIndices();
+        List<ShardIterator> localShardIterators = new ArrayList<>();
         List<SearchShardIterator> remoteShardIterators = new ArrayList<>();
-        {
-            ShardId remoteShardId = new ShardId("remote_index", "remote_index_uuid", 2);
-            ShardRouting remoteShardRouting = TestShardRouting.newShardRouting(remoteShardId, "remote_node", true, STARTED);
-            SearchShardIterator remoteShardIterator = new SearchShardIterator("remote", remoteShardId,
-                    Collections.singletonList(remoteShardRouting), remoteIndices);
-            remoteShardIterators.add(remoteShardIterator);
-        }
-        {
-            ShardId remoteShardId2 = new ShardId("remote_index_2", "remote_index_2_uuid", 3);
-            ShardRouting remoteShardRouting2 = TestShardRouting.newShardRouting(remoteShardId2, "remote_node", true, STARTED);
-            SearchShardIterator remoteShardIterator2 = new SearchShardIterator("remote", remoteShardId2,
-                    Collections.singletonList(remoteShardRouting2), remoteIndices);
-            remoteShardIterators.add(remoteShardIterator2);
+        int numShards = randomIntBetween(0, 10);
+        for (int i = 0; i < numShards; i++) {
+            int numIndices = randomIntBetween(0, indices.length);
+            for (int j = 0; j < numIndices; j++) {
+                Index index = indices[j];
+                boolean localIndex = randomBoolean();
+                if (localIndex) {
+                    SearchShardIterator localIterator = createSearchShardIterator(i, index, localIndices, localClusterAlias);
+                    localShardIterators.add(new PlainShardIterator(localIterator.shardId(), localIterator.getShardRoutings()));
+                    if (rarely()) {
+                        String remoteClusterAlias = randomFrom(remoteClusters);
+                        //simulate scenario where the local cluster is also registered as a remote one
+                        SearchShardIterator remoteIterator = createSearchShardIterator(i, index,
+                            OriginalIndicesTests.randomOriginalIndices(), remoteClusterAlias);
+                        remoteShardIterators.add(remoteIterator);
+                        assert remoteClusterAlias.equals(localClusterAlias) == false;
+                        if (remoteClusterAlias.compareTo(localClusterAlias) < 0) {
+                            expected.add(remoteIterator);
+                            expected.add(localIterator);
+                        } else {
+                            expected.add(localIterator);
+                            expected.add(remoteIterator);
+                        }
+                    } else {
+                        expected.add(localIterator);
+                    }
+                } else if (rarely()) {
+                    int numClusters = randomIntBetween(1, remoteClusters.length);
+                    for (int k = 0; k < numClusters; k++) {
+                        //simulate scenario where the same cluster is registered multiple times with different aliases
+                        String clusterAlias = remoteClusters[k];
+                        SearchShardIterator iterator = createSearchShardIterator(i, index, OriginalIndicesTests.randomOriginalIndices(),
+                            clusterAlias);
+                        expected.add(iterator);
+                        remoteShardIterators.add(iterator);
+                    }
+                } else {
+                    SearchShardIterator iterator = createSearchShardIterator(i, index, OriginalIndicesTests.randomOriginalIndices(),
+                        randomFrom(remoteClusters));
+                    expected.add(iterator);
+                    remoteShardIterators.add(iterator);
+                }
+            }
         }
-        OriginalIndices remoteIndices2 = new OriginalIndices(new String[]{"remote_index_3"}, IndicesOptions.strictExpand());
 
-        {
-            ShardId remoteShardId3 = new ShardId("remote_index_3", "remote_index_3_uuid", 4);
-            ShardRouting remoteShardRouting3 = TestShardRouting.newShardRouting(remoteShardId3, "remote_node", true, STARTED);
-            SearchShardIterator remoteShardIterator3 = new SearchShardIterator("remote", remoteShardId3,
-                    Collections.singletonList(remoteShardRouting3), remoteIndices2);
-            remoteShardIterators.add(remoteShardIterator3);
-        }
+        Collections.shuffle(localShardIterators, random());
+        Collections.shuffle(remoteShardIterators, random());
 
-        String localClusterAlias = randomBoolean() ? null : "local";
-        GroupShardsIterator<SearchShardIterator> searchShardIterators = TransportSearchAction.mergeShardsIterators(localShardsIterator,
-                localIndices, localClusterAlias, remoteShardIterators);
-
-        assertEquals(searchShardIterators.size(), 5);
-        int i = 0;
-        for (SearchShardIterator searchShardIterator : searchShardIterators) {
-            switch(i++) {
-                case 0:
-                    assertEquals("local_index", searchShardIterator.shardId().getIndexName());
-                    assertEquals(0, searchShardIterator.shardId().getId());
-                    assertSame(localIndices, searchShardIterator.getOriginalIndices());
-                    assertEquals(localClusterAlias, searchShardIterator.getClusterAlias());
-                    break;
-                case 1:
-                    assertEquals("local_index_2", searchShardIterator.shardId().getIndexName());
-                    assertEquals(1, searchShardIterator.shardId().getId());
-                    assertSame(localIndices, searchShardIterator.getOriginalIndices());
-                    assertEquals(localClusterAlias, searchShardIterator.getClusterAlias());
-                    break;
-                case 2:
-                    assertEquals("remote_index", searchShardIterator.shardId().getIndexName());
-                    assertEquals(2, searchShardIterator.shardId().getId());
-                    assertSame(remoteIndices, searchShardIterator.getOriginalIndices());
-                    assertEquals("remote", searchShardIterator.getClusterAlias());
-                    break;
-                case 3:
-                    assertEquals("remote_index_2", searchShardIterator.shardId().getIndexName());
-                    assertEquals(3, searchShardIterator.shardId().getId());
-                    assertSame(remoteIndices, searchShardIterator.getOriginalIndices());
-                    assertEquals("remote", searchShardIterator.getClusterAlias());
-                    break;
-                case 4:
-                    assertEquals("remote_index_3", searchShardIterator.shardId().getIndexName());
-                    assertEquals(4, searchShardIterator.shardId().getId());
-                    assertSame(remoteIndices2, searchShardIterator.getOriginalIndices());
-                    assertEquals("remote", searchShardIterator.getClusterAlias());
-                    break;
-            }
+        GroupShardsIterator<SearchShardIterator> groupShardsIterator = TransportSearchAction.mergeShardsIterators(
+            new GroupShardsIterator<>(localShardIterators), localIndices, localClusterAlias, remoteShardIterators);
+        List<SearchShardIterator> result = new ArrayList<>();
+        for (SearchShardIterator searchShardIterator : groupShardsIterator) {
+            result.add(searchShardIterator);
         }
+        assertEquals(expected, result);
     }
 
     public void testProcessRemoteShards() {

server/src/test/java/org/elasticsearch/cluster/routing/GroupShardsIteratorTests.java

@@ -20,6 +20,8 @@
 package org.elasticsearch.cluster.routing;
 
 import org.apache.lucene.util.CollectionUtil;
+import org.elasticsearch.action.OriginalIndicesTests;
+import org.elasticsearch.action.search.SearchShardIterator;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.test.ESTestCase;
@@ -29,20 +31,44 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 
+import static org.elasticsearch.cluster.routing.ShardRoutingState.STARTED;
+
 public class GroupShardsIteratorTests extends ESTestCase {
 
+    public static List<ShardRouting> randomShardRoutings(ShardId shardId) {
+        return randomShardRoutings(shardId, randomIntBetween(0, 2));
+    }
+
+    private static List<ShardRouting> randomShardRoutings(ShardId shardId, int numReplicas) {
+        List<ShardRouting> shardRoutings = new ArrayList<>();
+        shardRoutings.add(TestShardRouting.newShardRouting(shardId, randomAlphaOfLengthBetween(5, 10), true, STARTED));
+        for (int j = 0; j < numReplicas; j++) {
+            shardRoutings.add(TestShardRouting.newShardRouting(shardId, randomAlphaOfLengthBetween(5, 10), false, STARTED));
+        }
+        return shardRoutings;
+    }
+
     public void testSize() {
         List<ShardIterator> list = new ArrayList<>();
         Index index = new Index("foo", "na");
-
-        list.add(new PlainShardIterator(new ShardId(index, 0), Arrays.asList(newRouting(index, 0, true), newRouting(index, 0, true),
-            newRouting(index, 0, true))));
+        {
+            ShardId shardId = new ShardId(index, 0);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId, 2)));
+        }
         list.add(new PlainShardIterator(new ShardId(index, 1), Collections.emptyList()));
-        list.add(new PlainShardIterator(new ShardId(index, 2), Arrays.asList(newRouting(index, 2, true))));
+        {
+            ShardId shardId = new ShardId(index, 2);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId, 0)));
+        }
         index = new Index("foo_1", "na");
-
-        list.add(new PlainShardIterator(new ShardId(index, 0), Arrays.asList(newRouting(index, 0, true))));
-        list.add(new PlainShardIterator(new ShardId(index, 1), Arrays.asList(newRouting(index, 1, true))));
+        {
+            ShardId shardId = new ShardId(index, 0);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId, 0)));
+        }
+        {
+            ShardId shardId = new ShardId(index, 1);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId, 0)));
+        }
         GroupShardsIterator iter = new GroupShardsIterator<>(list);
         assertEquals(7, iter.totalSizeWith1ForEmpty());
         assertEquals(5, iter.size());
@@ -52,21 +78,35 @@ public class GroupShardsIteratorTests extends ESTestCase {
     public void testIterate() {
         List<ShardIterator> list = new ArrayList<>();
         Index index = new Index("foo", "na");
-
-        list.add(new PlainShardIterator(new ShardId(index, 0), Arrays.asList(newRouting(index, 0, true), newRouting(index, 0, true),
-            newRouting(index, 0, true))));
+        {
+            ShardId shardId = new ShardId(index, 0);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId)));
+        }
         list.add(new PlainShardIterator(new ShardId(index, 1), Collections.emptyList()));
-        list.add(new PlainShardIterator(new ShardId(index, 2), Arrays.asList(newRouting(index, 2, true))));
-
-        list.add(new PlainShardIterator(new ShardId(index, 0), Arrays.asList(newRouting(index, 0, true))));
-        list.add(new PlainShardIterator(new ShardId(index, 1), Arrays.asList(newRouting(index, 1, true))));
-
+        {
+            ShardId shardId = new ShardId(index, 2);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId)));
+        }
+        {
+            ShardId shardId = new ShardId(index, 0);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId)));
+        }
+        {
+            ShardId shardId = new ShardId(index, 1);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId)));
+        }
         index = new Index("foo_2", "na");
-        list.add(new PlainShardIterator(new ShardId(index, 0), Arrays.asList(newRouting(index, 0, true))));
-        list.add(new PlainShardIterator(new ShardId(index, 1), Arrays.asList(newRouting(index, 1, true))));
+        {
+            ShardId shardId = new ShardId(index, 0);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId)));
+        }
+        {
+            ShardId shardId = new ShardId(index, 1);
+            list.add(new PlainShardIterator(shardId, randomShardRoutings(shardId)));
+        }
 
         Collections.shuffle(list, random());
-        ArrayList<ShardIterator> actualIterators = new ArrayList<>();
+        List<ShardIterator> actualIterators = new ArrayList<>();
         GroupShardsIterator<ShardIterator> iter = new GroupShardsIterator<>(list);
         for (ShardIterator shardsIterator : iter) {
             actualIterators.add(shardsIterator);
@@ -75,13 +115,39 @@ public class GroupShardsIteratorTests extends ESTestCase {
         assertEquals(actualIterators, list);
     }
 
-    public ShardRouting newRouting(Index index, int id, boolean started) {
-        ShardRouting shardRouting = ShardRouting.newUnassigned(new ShardId(index, id), true,
-            RecoverySource.EmptyStoreRecoverySource.INSTANCE, new UnassignedInfo(UnassignedInfo.Reason.INDEX_CREATED, "foo"));
-        shardRouting = ShardRoutingHelper.initialize(shardRouting, "some node");
-        if (started) {
-            shardRouting = ShardRoutingHelper.moveToStarted(shardRouting);
+    public void testOrderingWithSearchShardIterators() {
+        String[] indices = generateRandomStringArray(10, 10, false, false);
+        Arrays.sort(indices);
+        String[] uuids = generateRandomStringArray(5, 10, false, false);
+        Arrays.sort(uuids);
+        String[] clusters = generateRandomStringArray(5, 10, false, false);
+        Arrays.sort(clusters);
+
+        List<SearchShardIterator> expected = new ArrayList<>();
+        int numShards = randomIntBetween(1, 10);
+        for (int i = 0; i < numShards; i++) {
+            for (String index : indices) {
+                for (String uuid : uuids) {
+                    ShardId shardId = new ShardId(index, uuid, i);
+                    SearchShardIterator shardIterator = new SearchShardIterator(null, shardId,
+                        GroupShardsIteratorTests.randomShardRoutings(shardId), OriginalIndicesTests.randomOriginalIndices());
+                    expected.add(shardIterator);
+                    for (String cluster : clusters) {
+                        SearchShardIterator remoteIterator = new SearchShardIterator(cluster, shardId,
+                            GroupShardsIteratorTests.randomShardRoutings(shardId), OriginalIndicesTests.randomOriginalIndices());
+                        expected.add(remoteIterator);
+                    }
+                }
+            }
+        }
+
+        List<SearchShardIterator> shuffled = new ArrayList<>(expected);
+        Collections.shuffle(shuffled, random());
+        List<ShardIterator> actualIterators = new ArrayList<>();
+        GroupShardsIterator<SearchShardIterator> iter = new GroupShardsIterator<>(shuffled);
+        for (SearchShardIterator searchShardIterator : iter) {
+            actualIterators.add(searchShardIterator);
         }
-        return shardRouting;
+        assertEquals(expected, actualIterators);
     }
 }

server/src/test/java/org/elasticsearch/cluster/routing/PlainShardIteratorTests.java

@@ -19,26 +19,83 @@
 
 package org.elasticsearch.cluster.routing;
 
-import org.elasticsearch.index.Index;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.EqualsHashCodeTestUtils;
+import org.hamcrest.Matchers;
 
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 
 public class PlainShardIteratorTests extends ESTestCase {
 
-    public void testEquals() {
-        Index index = new Index("a", "b");
-        ShardId shardId = new ShardId(index, 1);
-        ShardId shardId2 = new ShardId(index, 2);
-        PlainShardIterator iterator1 = new PlainShardIterator(shardId, new ArrayList<>());
-        PlainShardIterator iterator2 = new PlainShardIterator(shardId, new ArrayList<>());
-        PlainShardIterator iterator3 = new PlainShardIterator(shardId2, new ArrayList<>());
-        String s = "Some other random object";
-        assertEquals(iterator1, iterator1);
-        assertEquals(iterator1, iterator2);
-        assertNotEquals(iterator1, null);
-        assertNotEquals(iterator1, s);
-        assertNotEquals(iterator1, iterator3);
+    public void testEqualsAndHashCode() {
+        EqualsHashCodeTestUtils.checkEqualsAndHashCode(randomPlainShardIterator(),
+            i -> new PlainShardIterator(i.shardId(), i.getShardRoutings()),
+            i -> {
+                ShardId shardId;
+                switch(randomIntBetween(0, 2)) {
+                    case 0:
+                        shardId = new ShardId(i.shardId().getIndex(), i.shardId().getId() + randomIntBetween(1, 1000));
+                        break;
+                    case 1:
+                        shardId = new ShardId(i.shardId().getIndexName(),
+                            i.shardId().getIndex().getUUID() + randomAlphaOfLengthBetween(1, 3), i.shardId().getId());
+                        break;
+                    case 2:
+                        shardId = new ShardId(i.shardId().getIndexName() + randomAlphaOfLengthBetween(1, 3),
+                            i.shardId().getIndex().getUUID(), i.shardId().getId());
+                        break;
+                    default:
+                        throw new UnsupportedOperationException();
+                }
+                return new PlainShardIterator(shardId, i.getShardRoutings());
+            });
+    }
+
+    public void testCompareTo() {
+        String[] indices = generateRandomStringArray(3, 10, false, false);
+        Arrays.sort(indices);
+        String[] uuids = generateRandomStringArray(3, 10, false, false);
+        Arrays.sort(uuids);
+        List<PlainShardIterator> shardIterators = new ArrayList<>();
+        int numShards = randomIntBetween(1, 5);
+        for (int i = 0; i < numShards; i++) {
+            for (String index : indices) {
+                for (String uuid : uuids) {
+                    ShardId shardId = new ShardId(index, uuid, i);
+                    shardIterators.add(new PlainShardIterator(shardId, GroupShardsIteratorTests.randomShardRoutings(shardId)));
+                }
+            }
+        }
+        for (int i = 0; i < shardIterators.size(); i++) {
+            PlainShardIterator currentIterator = shardIterators.get(i);
+            for (int j = i + 1; j < shardIterators.size(); j++) {
+                PlainShardIterator greaterIterator = shardIterators.get(j);
+                assertThat(currentIterator, Matchers.lessThan(greaterIterator));
+                assertThat(greaterIterator, Matchers.greaterThan(currentIterator));
+                assertNotEquals(currentIterator, greaterIterator);
+            }
+            for (int j = i - 1; j >= 0; j--) {
+                PlainShardIterator smallerIterator = shardIterators.get(j);
+                assertThat(smallerIterator, Matchers.lessThan(currentIterator));
+                assertThat(currentIterator, Matchers.greaterThan(smallerIterator));
+                assertNotEquals(currentIterator, smallerIterator);
+            }
+        }
+    }
+
+    public void testCompareToEqualItems() {
+        PlainShardIterator shardIterator1 = randomPlainShardIterator();
+        PlainShardIterator shardIterator2 = new PlainShardIterator(shardIterator1.shardId(), shardIterator1.getShardRoutings());
+        assertEquals(shardIterator1, shardIterator2);
+        assertEquals(0, shardIterator1.compareTo(shardIterator2));
+        assertEquals(0, shardIterator2.compareTo(shardIterator1));
+    }
+
+    private static PlainShardIterator randomPlainShardIterator() {
+        ShardId shardId = new ShardId(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLength(10), randomIntBetween(1, Integer.MAX_VALUE));
+        return new PlainShardIterator(shardId, GroupShardsIteratorTests.randomShardRoutings(shardId));
     }
 }