ソースを参照

Improve RoutingNodes API

Currently the RoutingNodes API allows modification of it's internal state outside of the class.
This commit improves the APIs of `RoutingNode` and `RoutingNode` to change internal state
only within the classes itself.

Closes #4458
Simon Willnauer 12 年 前
コミット
30c6f2fa23
39 ファイル変更845 行追加710 行削除
  1. 3 6
      src/main/java/org/elasticsearch/action/mlt/TransportMoreLikeThisAction.java
  2. 20 21
      src/main/java/org/elasticsearch/cluster/routing/RoutingNode.java
  3. 393 210
      src/main/java/org/elasticsearch/cluster/routing/RoutingNodes.java
  4. 46 56
      src/main/java/org/elasticsearch/cluster/routing/allocation/AllocationService.java
  5. 24 59
      src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java
  6. 13 13
      src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/EvenShardsCountAllocator.java
  7. 1 0
      src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/ShardsAllocators.java
  8. 1 1
      src/main/java/org/elasticsearch/cluster/routing/allocation/command/AllocateAllocationCommand.java
  9. 11 15
      src/main/java/org/elasticsearch/cluster/routing/allocation/command/CancelAllocationCommand.java
  10. 3 3
      src/main/java/org/elasticsearch/cluster/routing/allocation/command/MoveAllocationCommand.java
  11. 1 2
      src/main/java/org/elasticsearch/cluster/routing/allocation/decider/AwarenessAllocationDecider.java
  12. 1 8
      src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ConcurrentRebalanceAllocationDecider.java
  13. 2 8
      src/main/java/org/elasticsearch/cluster/routing/allocation/decider/RebalanceOnlyWhenActiveAllocationDecider.java
  14. 2 2
      src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ReplicaAfterPrimaryActiveAllocationDecider.java
  15. 4 8
      src/main/java/org/elasticsearch/cluster/routing/allocation/decider/SameShardAllocationDecider.java
  16. 2 8
      src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ShardsLimitAllocationDecider.java
  17. 2 8
      src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ThrottlingAllocationDecider.java
  18. 4 4
      src/main/java/org/elasticsearch/gateway/blobstore/BlobReuseExistingGatewayAllocator.java
  19. 6 6
      src/main/java/org/elasticsearch/gateway/local/LocalGatewayAllocator.java
  20. 8 9
      src/main/java/org/elasticsearch/indices/cluster/IndicesClusterStateService.java
  21. 6 10
      src/test/java/org/elasticsearch/cluster/ack/AckTests.java
  22. 10 10
      src/test/java/org/elasticsearch/cluster/allocation/ClusterRerouteTests.java
  23. 9 15
      src/test/java/org/elasticsearch/cluster/allocation/SimpleAllocationTests.java
  24. 11 11
      src/test/java/org/elasticsearch/cluster/routing/allocation/AddIncrementallyTests.java
  25. 34 34
      src/test/java/org/elasticsearch/cluster/routing/allocation/AllocationCommandsTests.java
  26. 10 10
      src/test/java/org/elasticsearch/cluster/routing/allocation/AwarenessAllocationTests.java
  27. 11 13
      src/test/java/org/elasticsearch/cluster/routing/allocation/BalanceConfigurationTests.java
  28. 17 17
      src/test/java/org/elasticsearch/cluster/routing/allocation/ClusterRebalanceRoutingTests.java
  29. 32 32
      src/test/java/org/elasticsearch/cluster/routing/allocation/DeadNodesAllocationTests.java
  30. 24 23
      src/test/java/org/elasticsearch/cluster/routing/allocation/FailedShardsRoutingTests.java
  31. 1 1
      src/test/java/org/elasticsearch/cluster/routing/allocation/RandomAllocationDeciderTests.java
  32. 1 1
      src/test/java/org/elasticsearch/cluster/routing/allocation/RebalanceAfterActiveTests.java
  33. 6 6
      src/test/java/org/elasticsearch/cluster/routing/allocation/RoutingNodesIntegrityTests.java
  34. 36 0
      src/test/java/org/elasticsearch/cluster/routing/allocation/RoutingNodesUtils.java
  35. 8 5
      src/test/java/org/elasticsearch/cluster/routing/allocation/SameShardRoutingTests.java
  36. 3 2
      src/test/java/org/elasticsearch/cluster/routing/allocation/ShardsLimitAllocationTests.java
  37. 7 9
      src/test/java/org/elasticsearch/cluster/routing/allocation/SingleShardNoReplicasRoutingTests.java
  38. 28 28
      src/test/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDeciderTests.java
  39. 44 36
      src/test/java/org/elasticsearch/indexlifecycle/IndexLifecycleActionTests.java

+ 3 - 6
src/main/java/org/elasticsearch/action/mlt/TransportMoreLikeThisAction.java

@@ -35,10 +35,7 @@ import org.elasticsearch.action.support.TransportAction;
 import org.elasticsearch.cluster.ClusterService;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.node.DiscoveryNode;
-import org.elasticsearch.cluster.routing.MutableShardRouting;
-import org.elasticsearch.cluster.routing.RoutingNode;
-import org.elasticsearch.cluster.routing.ShardIterator;
-import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.*;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.engine.DocumentMissingException;
@@ -96,13 +93,13 @@ public class TransportMoreLikeThisAction extends TransportAction<MoreLikeThisReq
         // update to the concrete index
         final String concreteIndex = clusterState.metaData().concreteIndex(request.index());
 
-        RoutingNode routingNode = clusterState.getRoutingNodes().nodesToShards().get(clusterService.localNode().getId());
+        Iterable<MutableShardRouting> routingNode = clusterState.getRoutingNodes().routingNodeIter(clusterService.localNode().getId());
         if (routingNode == null) {
             redirect(request, concreteIndex, listener, clusterState);
             return;
         }
         boolean hasIndexLocally = false;
-        for (MutableShardRouting shardRouting : routingNode.shards()) {
+        for (MutableShardRouting shardRouting : routingNode) {
             if (concreteIndex.equals(shardRouting.index())) {
                 hasIndexLocally = true;
                 break;

+ 20 - 21
src/main/java/org/elasticsearch/cluster/routing/RoutingNode.java

@@ -19,10 +19,12 @@
 
 package org.elasticsearch.cluster.routing;
 
+import com.google.common.collect.Iterators;
 import org.elasticsearch.ElasticSearchIllegalStateException;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
 
@@ -52,6 +54,10 @@ public class RoutingNode implements Iterable<MutableShardRouting> {
 
     @Override
     public Iterator<MutableShardRouting> iterator() {
+        return Iterators.unmodifiableIterator(shards.iterator());
+    }
+
+    Iterator<MutableShardRouting> mutableIterator() {
         return shards.iterator();
     }
 
@@ -72,12 +78,8 @@ public class RoutingNode implements Iterable<MutableShardRouting> {
         return this.nodeId;
     }
 
-    /**
-     * Get a list of shards hosted on this node  
-     * @return list of shards
-     */
-    public List<MutableShardRouting> shards() {
-        return this.shards;
+    public int size() {
+        return shards.size();
     }
 
     /**
@@ -150,21 +152,6 @@ public class RoutingNode implements Iterable<MutableShardRouting> {
         return shards;
     }
 
-    /**
-     * Get the number of shard that not match the given states
-     * @param state set states to exclude
-     * @return number of shards which state is listed
-     */
-    public int numberOfShardsNotWithState(ShardRoutingState state) {
-        int count = 0;
-        for (MutableShardRouting shardEntry : this) {
-            if (shardEntry.state() != state) {
-                count++;
-            }
-        }
-        return count;
-    }
-
     /**
      * The number of shards on this node that will not be eventually relocated.
      */
@@ -187,4 +174,16 @@ public class RoutingNode implements Iterable<MutableShardRouting> {
         }
         return sb.toString();
     }
+
+    public MutableShardRouting get(int i) {
+        return shards.get(i) ;
+    }
+
+    public Collection<MutableShardRouting> copyShards() {
+        return new ArrayList<MutableShardRouting>(shards);
+    }
+
+    public boolean isEmpty() {
+        return shards.isEmpty();
+    }
 }

+ 393 - 210
src/main/java/org/elasticsearch/cluster/routing/RoutingNodes.java

@@ -23,12 +23,14 @@ import com.carrotsearch.hppc.ObjectIntOpenHashMap;
 import com.carrotsearch.hppc.cursors.ObjectCursor;
 import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterators;
 import com.google.common.collect.Sets;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.block.ClusterBlocks;
 import org.elasticsearch.cluster.metadata.IndexMetaData;
 import org.elasticsearch.cluster.metadata.MetaData;
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.common.collect.IdentityHashSet;
 import org.elasticsearch.index.shard.ShardId;
 
 import java.util.*;
@@ -51,20 +53,18 @@ public class RoutingNodes implements Iterable<RoutingNode> {
 
     private final Map<String, RoutingNode> nodesToShards = newHashMap();
 
-    private final List<MutableShardRouting> unassigned = newArrayList();
+    private final UnassignedShards unassignedShards = new UnassignedShards();
 
-    private final List<MutableShardRouting> ignoredUnassigned = newArrayList();
+    private final List<MutableShardRouting> ignoredUnassignedShards = newArrayList();
 
-    private final Map<ShardId, List<MutableShardRouting>> replicaSets = newHashMap();
-
-    private int unassignedPrimaryCount = 0;
+    private final Map<ShardId, Set<MutableShardRouting>> assignedShards = newHashMap();
 
     private int inactivePrimaryCount = 0;
 
     private int inactiveShardCount = 0;
 
-    Set<ShardId> relocatingReplicaSets = new HashSet<ShardId>();
-    
+    private int relocatingShards = 0;
+
     private Set<ShardId> clearPostAllocationFlag;
 
     private final Map<String, ObjectIntOpenHashMap<String>> nodesPerAttributeNames = new HashMap<String, ObjectIntOpenHashMap<String>>();
@@ -97,10 +97,10 @@ public class RoutingNodes implements Iterable<RoutingNode> {
                         }
                         MutableShardRouting sr = new MutableShardRouting(shard);
                         entries.add(sr);
-                        addToReplicaSet(sr);
+                        activeShardsAdd(sr);
                         if (shard.relocating()) {
                             entries = nodesToShards.get(shard.relocatingNodeId());
-                            relocatingReplicaSets.add(shard.shardId());
+                            relocatingShards++;
                             if (entries == null) {
                                 entries = newArrayList();
                                 nodesToShards.put(shard.relocatingNodeId(), entries);
@@ -110,7 +110,7 @@ public class RoutingNodes implements Iterable<RoutingNode> {
                             sr = new MutableShardRouting(shard.index(), shard.id(), shard.relocatingNodeId(),
                                     shard.currentNodeId(), shard.primary(), ShardRoutingState.INITIALIZING, shard.version());
                             entries.add(sr);
-                            addToReplicaSet(sr);
+                            activeShardsAdd(sr);
                         } else if (!shard.active()) { // shards that are initializing without being relocated
                             if (shard.primary()) {
                                 inactivePrimaryCount++;
@@ -119,12 +119,8 @@ public class RoutingNodes implements Iterable<RoutingNode> {
                         }
                     } else {
                         MutableShardRouting sr = new MutableShardRouting(shard);
-                        addToReplicaSet(sr);
-                        unassigned.add(sr);
-                        if (shard.primary()) {
-                            unassignedPrimaryCount++;
-                        }
-
+                        activeShardsAdd(sr);
+                        unassignedShards.add(sr);
                     }
                 }
             }
@@ -137,7 +133,7 @@ public class RoutingNodes implements Iterable<RoutingNode> {
 
     @Override
     public Iterator<RoutingNode> iterator() {
-        return nodesToShards.values().iterator();
+        return Iterators.unmodifiableIterator(nodesToShards.values().iterator());
     }
 
     public RoutingTable routingTable() {
@@ -177,27 +173,19 @@ public class RoutingNodes implements Iterable<RoutingNode> {
     }
 
     public boolean hasUnassigned() {
-        return !unassigned.isEmpty();
+        return !unassignedShards.isEmpty();
     }
 
     public List<MutableShardRouting> ignoredUnassigned() {
-        return this.ignoredUnassigned;
+        return this.ignoredUnassignedShards;
     }
 
-    public List<MutableShardRouting> unassigned() {
-        return this.unassigned;
+    public UnassignedShards unassigned() {
+        return this.unassignedShards;
     }
 
-    public List<MutableShardRouting> getUnassigned() {
-        return unassigned();
-    }
-
-    public Map<String, RoutingNode> nodesToShards() {
-        return nodesToShards;
-    }
-
-    public Map<String, RoutingNode> getNodesToShards() {
-        return nodesToShards();
+    public RoutingNodesIterator nodes() {
+        return new RoutingNodesIterator(nodesToShards.values().iterator());
     }
 
     /**
@@ -238,11 +226,11 @@ public class RoutingNodes implements Iterable<RoutingNode> {
     }
 
     public boolean hasUnassignedPrimaries() {
-        return unassignedPrimaryCount > 0;
+        return unassignedShards.numPrimaries() > 0;
     }
 
     public boolean hasUnassignedShards() {
-        return !unassigned.isEmpty();
+        return !unassignedShards.isEmpty();
     }
 
     public boolean hasInactivePrimaries() {
@@ -254,48 +242,46 @@ public class RoutingNodes implements Iterable<RoutingNode> {
     }
 
     public int getRelocatingShardCount() {
-        return relocatingReplicaSets.size();
+        return relocatingShards;
     }
 
-    public MutableShardRouting findPrimaryForReplica(ShardRouting shard) {
+    /**
+     * Returns the active primary shard for the given ShardRouting or <code>null</code> if
+     * no primary is found or the primary is not active.
+     */
+    public MutableShardRouting activePrimary(ShardRouting shard) {
         assert !shard.primary();
-        MutableShardRouting primary = null;
-        for (MutableShardRouting shardRouting : shardsRoutingFor(shard)) {
+        for (MutableShardRouting shardRouting : activeShards(shard.shardId())) {
             if (shardRouting.primary()) {
-                primary = shardRouting;
+                if (shardRouting.active()) {
+                    return shardRouting;
+                }
                 break;
             }
         }
-        assert primary != null;
-        return primary;
-    }
-
-    public List<MutableShardRouting> shardsRoutingFor(ShardRouting shardRouting) {
-        return shardsRoutingFor(shardRouting.index(), shardRouting.id());
-    }
-
-    public List<MutableShardRouting> shardsRoutingFor(String index, int shardId) {
-        ShardId sid = new ShardId(index, shardId);
-        List<MutableShardRouting> shards = replicaSetFor(sid);
-        assert shards != null;
-        // no need to check unassigned array, since the ShardRoutings are in the replica set.
-        return shards;
+        return null;
     }
 
-    public int numberOfShardsOfType(ShardRoutingState state) {
-        int count = 0;
-        for (RoutingNode routingNode : this) {
-            count += routingNode.numberOfShardsWithState(state);
+    /**
+     * Returns <code>true</code> iff all replicas are active for the given shard routing. Otherwise <code>false</code>
+     */
+    public boolean allReplicasActive(ShardRouting shardRouting) {
+        final Set<MutableShardRouting> shards = activeShards(shardRouting.shardId());
+        if (shards.isEmpty() || shards.size() < this.routingTable.index(shardRouting.index()).shard(shardRouting.id()).size()) {
+            return false; // if we are empty nothing is active if we have less than total at least one is unassigned
+        }
+        for (MutableShardRouting shard : shards) {
+            if (!shard.active()) {
+                return false;
+            }
         }
-        return count;
+        return true;
     }
 
     public List<MutableShardRouting> shards(Predicate<MutableShardRouting> predicate) {
         List<MutableShardRouting> shards = newArrayList();
         for (RoutingNode routingNode : this) {
-            List<MutableShardRouting> nodeShards = routingNode.shards();
-            for (int i = 0; i < nodeShards.size(); i++) {
-                MutableShardRouting shardRouting = nodeShards.get(i);
+            for (MutableShardRouting shardRouting : routingNode) {
                 if (predicate.apply(shardRouting)) {
                     shards.add(shardRouting);
                 }
@@ -305,6 +291,7 @@ public class RoutingNodes implements Iterable<RoutingNode> {
     }
 
     public List<MutableShardRouting> shardsWithState(ShardRoutingState... state) {
+        // TODO these are used on tests only - move into utils class
         List<MutableShardRouting> shards = newArrayList();
         for (RoutingNode routingNode : this) {
             shards.addAll(routingNode.shardsWithState(state));
@@ -313,6 +300,7 @@ public class RoutingNodes implements Iterable<RoutingNode> {
     }
 
     public List<MutableShardRouting> shardsWithState(String index, ShardRoutingState... state) {
+        // TODO these are used on tests only - move into utils class
         List<MutableShardRouting> shards = newArrayList();
         for (RoutingNode routingNode : this) {
             shards.addAll(routingNode.shardsWithState(index, state));
@@ -326,63 +314,12 @@ public class RoutingNodes implements Iterable<RoutingNode> {
             sb.append(routingNode.prettyPrint());
         }
         sb.append("---- unassigned\n");
-        for (MutableShardRouting shardEntry : unassigned) {
+        for (MutableShardRouting shardEntry : unassignedShards) {
             sb.append("--------").append(shardEntry.shortSummary()).append('\n');
         }
         return sb.toString();
     }
 
-    /**
-     * calculates RoutingNodes statistics by iterating over all {@link MutableShardRouting}s
-     * in the cluster to ensure the {@link RoutingManager} book-keeping is correct.
-     * For performance reasons, this should only be called from test cases.
-     *
-     * @return true if all counts are the same, false if either of the book-keeping numbers is off.
-     */
-    public boolean assertShardStats() {
-       int unassignedPrimaryCount = 0;
-       int inactivePrimaryCount = 0;
-       int inactiveShardCount = 0;
-       int totalShards = 0;
-
-       Set<ShardId> seenShards = newHashSet();
-
-       for (RoutingNode node : this) {
-           for (MutableShardRouting shard : node) {
-               if (!shard.active()) {
-                   if (!shard.relocating()) {
-                       inactiveShardCount++;
-                       if (shard.primary()){
-                           inactivePrimaryCount++;
-                       }
-                   }
-               }
-               totalShards++;
-               seenShards.add(shard.shardId());
-           }
-       }
-       for (MutableShardRouting shard : unassigned) {
-           if (shard.primary()) {
-               unassignedPrimaryCount++;
-           }
-           totalShards++;
-           seenShards.add(shard.shardId());
-       }
-
-       for (ShardId shardId : seenShards) {
-           assert replicaSetFor(shardId) != null;
-       }
-
-       assert unassignedPrimaryCount == 0 || hasUnassignedPrimaries();
-       assert inactivePrimaryCount == 0 || hasInactivePrimaries();
-       assert inactiveShardCount == 0 || hasInactiveShards();
-       assert hasUnassignedPrimaries() || unassignedPrimaryCount == 0;
-       assert hasInactivePrimaries() || inactivePrimaryCount == 0;
-       assert hasInactiveShards() || inactiveShardCount == 0;
-
-       return true;
-    }
-
     /**
      * Assign a shard to a node. This will increment the inactiveShardCount counter
      * and the inactivePrimaryCount counter if the shard is the primary.
@@ -400,158 +337,404 @@ public class RoutingNodes implements Iterable<RoutingNode> {
      * @param shard the shard to be assigned
      * @param nodeId the nodeId this shard should initialize on or relocate from
      */
-    public void assignShardToNode(MutableShardRouting shard, String nodeId) {
-
+    public void assign(MutableShardRouting shard, String nodeId) {
         // state will not change if the shard is already initializing.
         ShardRoutingState oldState = shard.state();
-        
         shard.assignToNode(nodeId);
         node(nodeId).add(shard);
-
         if (oldState == ShardRoutingState.UNASSIGNED) {
             inactiveShardCount++;
             if (shard.primary()) {
-                unassignedPrimaryCount--;
                 inactivePrimaryCount++;
             }
         }
+
         if (shard.state() == ShardRoutingState.RELOCATING) {
-            // this a HashSet. double add no worry.
-            relocatingReplicaSets.add(shard.shardId()); 
+            relocatingShards++;
         }
-        // possibly double/triple adding it to a replica set doesn't matter
-        // but make sure we know about the shard.
-        addToReplicaSet(shard);
+        activeShardsAdd(shard);
     }
 
     /**
      * Relocate a shard to another node.
-     *
-     * STARTED => RELOCATING
-     *
-     * @param shard the shard to relocate
-     * @param nodeId the node to relocate to
      */
-    public void relocateShard(MutableShardRouting shard, String nodeId) {
-        relocatingReplicaSets.add(shard.shardId());
+    public void relocate(MutableShardRouting shard, String nodeId) {
+        relocatingShards++;
         shard.relocate(nodeId);
     }
 
     /**
-     * Cancels the relocation of a shard.
-     *
-     * RELOCATING => STARTED
-     *
-     * @param shard the shard that was relocating previously and now should be started again.
+     * Mark a shard as started and adjusts internal statistics.
+     */
+    public void started(MutableShardRouting shard) {
+        if (!shard.active() && shard.relocatingNodeId() == null) {
+            inactiveShardCount--;
+            if (shard.primary()) {
+                inactivePrimaryCount--;
+            }
+        } else if (shard.relocating()) {
+            relocatingShards--;
+        }
+        assert !shard.started();
+        shard.moveToStarted();
+    }
+
+    /**
+     * Cancels a relocation of a shard that shard must relocating.
      */
-    public void cancelRelocationForShard(MutableShardRouting shard) {
-        relocatingReplicaSets.remove(shard.shardId());
+    public void cancelRelocation(MutableShardRouting shard) {
+        relocatingShards--;
         shard.cancelRelocation();
     }
 
     /**
-     * Unassigns shard from a node.
-     * Both relocating and started shards that are deallocated need a new 
-     * primary elected.
-     *
-     * RELOCATING   => null
-     * STARTED      => null
-     * INITIALIZING => null
+     * swaps the status of a shard, making replicas primary and vice versa.
      *
-     * @param shard the shard to be unassigned.
+     * @param shards the shard to have its primary status swapped.
      */
-    public void deassignShard(MutableShardRouting shard) {
-        if (shard.state() == ShardRoutingState.RELOCATING) {
-            cancelRelocationForShard(shard);
+    public void swapPrimaryFlag(MutableShardRouting... shards) {
+        for (MutableShardRouting shard : shards) {
+            if (shard.primary()) {
+                shard.moveFromPrimary();
+                if (shard.unassigned()) {
+                    unassignedShards.primaries--;
+                }
+            } else {
+                shard.moveToPrimary();
+                if (shard.unassigned()) {
+                    unassignedShards.primaries++;
+                }
+            }
         }
-        if (shard.primary())
-            unassignedPrimaryCount++;
-        shard.deassignNode();
+    }
+
+    private static final Set<MutableShardRouting> EMPTY = Collections.emptySet();
+
+    private Set<MutableShardRouting> activeShards(ShardId shardId) {
+        final Set<MutableShardRouting> replicaSet = assignedShards.get(shardId);
+        return replicaSet == null ? EMPTY : Collections.unmodifiableSet(replicaSet);
     }
 
     /**
-     * Mark a shard as started.
-     * Decreases the counters and marks a replication complete or failed,
-     * which is the same for accounting in this class.
-     *
-     * INITIALIZING => STARTED
-     * RELOCATIng   => STARTED
-     *
-     * @param shard the shard to be marked as started
+     * Cancels the give shard from the Routing nodes internal statistics and cancels
+     * the relocation if the shard is relocating.
+     * @param shard
      */
-    public void markShardStarted(MutableShardRouting shard) {
-        if (!relocatingReplicaSets.contains(shard.shardId()) && shard.state() == ShardRoutingState.INITIALIZING) {
+    private void remove(MutableShardRouting shard) {
+        if (!shard.active() && shard.relocatingNodeId() == null) {
             inactiveShardCount--;
+            assert inactiveShardCount >= 0;
             if (shard.primary()) {
                 inactivePrimaryCount--;
             }
+        } else if (shard.relocating()) {
+            cancelRelocation(shard);
         }
-        if (shard.state() == ShardRoutingState.INITIALIZING 
-             && shard.relocatingNodeId() != null) {
-            relocatingReplicaSets.remove(shard.shardId());
+        activeShardsRemove(shard);
+    }
+
+    private void activeShardsAdd(MutableShardRouting shard) {
+        if (shard.unassigned()) {
+            // no unassigned
+            return;
         }
-        shard.moveToStarted();
+        Set<MutableShardRouting> replicaSet = assignedShards.get(shard.shardId());
+        if (replicaSet == null) {
+            replicaSet = new IdentityHashSet<MutableShardRouting>();
+            assignedShards.put(shard.shardId(), replicaSet);
+        }
+        replicaSet.add(shard);
     }
 
-    /**
-     * Return a list of shards belonging to a replica set
-     * 
-     * @param shard the shard for which to retrieve the replica set
-     * @return an unmodifiable List of the replica set
-     */
-    public List<MutableShardRouting> replicaSetFor(MutableShardRouting shard) {
-        return replicaSetFor(shard.shardId());
+    private void activeShardsRemove(MutableShardRouting shard) {
+        Set<MutableShardRouting> replicaSet = assignedShards.get(shard.shardId());
+        if (replicaSet != null) {
+            if (replicaSet.contains(shard)) {
+                replicaSet.remove(shard);
+            } else {
+                assert false : "Illegal state";
+                Iterator<MutableShardRouting> iterator = replicaSet.iterator();
+                while(iterator.hasNext()) {
+                    if (shard.equals(iterator.next())) {
+                        iterator.remove();
+                    }
+                }
+            }
+        }
     }
 
-    /**
-     * Return a list of shards belonging to a replica set
-     * 
-     * @param shardId the {@link ShardId} for which to retrieve the replica set
-     * @return an unmodifiable List of the replica set
-     */
-    public List<MutableShardRouting> replicaSetFor(ShardId shardId) {
-        List<MutableShardRouting> replicaSet = replicaSets.get(shardId);
-        assert replicaSet != null;
-        return Collections.unmodifiableList(replicaSet);
+    public boolean isKnown(DiscoveryNode node) {
+        return nodesToShards.containsKey(node.getId());
     }
 
-    /**
-     * Let this class know about a shard, which it then sorts into 
-     * its replica set. Package private as only {@link RoutingNodes} 
-     * should notify this class of shards during initialization.
-     *
-     * @param shard the shard to be sorted into its replica set
-     */
-    private void addToReplicaSet(MutableShardRouting shard) {
-        List<MutableShardRouting> replicaSet = replicaSets.get(shard.shardId());
-        if (replicaSet == null) {
-            replicaSet = new ArrayList<MutableShardRouting>();
-            replicaSets.put(shard.shardId(), replicaSet);
+    public void addNode(DiscoveryNode node) {
+        RoutingNode routingNode = new RoutingNode(node.id(), node);
+        nodesToShards.put(routingNode.nodeId(), routingNode);
+    }
+
+    public RoutingNodeIterator routingNodeIter(String nodeId) {
+        final RoutingNode routingNode = nodesToShards.get(nodeId);
+        if (routingNode == null) {
+            return null;
         }
-        replicaSet.add(shard);
+        assert assertShardStats(this);
+        return new RoutingNodeIterator(routingNode);
     }
 
-    /**
-     * marks a replica set as relocating. 
-     *
-     * @param shard a member of the relocating replica set
-     */
-    private void markRelocating(MutableShardRouting shard) {
-        relocatingReplicaSets.add(shard.shardId());
+    public RoutingNode[] toArray() {
+        return nodesToShards.values().toArray(new RoutingNode[nodesToShards.size()]);
+    }
+
+    public final static class UnassignedShards implements Iterable<MutableShardRouting>  {
+
+        private final List<MutableShardRouting> unassigned;
+        private int primaries = 0;
+        private long transactionId = 0;
+        private final UnassignedShards source;
+        private final long sourceTransactionId;
+
+        public UnassignedShards(UnassignedShards other) {
+            source = other;
+            sourceTransactionId = other.transactionId;
+            unassigned = new ArrayList<MutableShardRouting>(other.unassigned);
+            primaries = other.primaries;
+        }
+
+        public UnassignedShards() {
+            unassigned = new ArrayList<MutableShardRouting>();
+            source = null;
+            sourceTransactionId = -1;
+        }
+
+        public void add(MutableShardRouting mutableShardRouting) {
+            if(mutableShardRouting.primary()) {
+                primaries++;
+            }
+            unassigned.add(mutableShardRouting);
+            transactionId++;
+        }
+
+        public void addAll(Collection<MutableShardRouting> mutableShardRoutings) {
+            for (MutableShardRouting r : mutableShardRoutings) {
+                add(r);
+            }
+        }
+
+        public int size() {
+            return unassigned.size();
+        }
+
+        public int numPrimaries() {
+            return primaries;
+        }
+
+        @Override
+        public Iterator<MutableShardRouting> iterator() {
+            final Iterator<MutableShardRouting> iterator = unassigned.iterator();
+            return new Iterator<MutableShardRouting>() {
+                private  MutableShardRouting current;
+                @Override
+                public boolean hasNext() {
+                    return iterator.hasNext();
+                }
+
+                @Override
+                public MutableShardRouting next() {
+                    return current = iterator.next();
+                }
+
+                @Override
+                public void remove() {
+                    iterator.remove();
+                    if (current.primary()) {
+                        primaries--;
+                    }
+                    transactionId++;
+                }
+            };
+        }
+
+        public boolean isEmpty() {
+            return unassigned.isEmpty();
+        }
+
+        public void shuffle() {
+            Collections.shuffle(unassigned);
+        }
+
+        public void clear() {
+            transactionId++;
+            unassigned.clear();
+            primaries = 0;
+        }
+
+        public void transactionEnd(UnassignedShards shards) {
+           assert shards.source == this && shards.sourceTransactionId == transactionId :
+                   "Expected ID: " + shards.sourceTransactionId + " actual: " + transactionId + " Expected Source: " + shards.source + " actual: " + this;
+           transactionId++;
+           this.unassigned.clear();
+           this.unassigned.addAll(shards.unassigned);
+           this.primaries = shards.primaries;
+        }
+
+        public UnassignedShards transactionBegin() {
+            return new UnassignedShards(this);
+        }
+
+        public void copyAll(Collection<MutableShardRouting> others) {
+            others.addAll(unassigned);
+        }
     }
 
+
     /**
-     * swaps the status of a shard, making replicas primary and vice versa.
-     * 
-     * @param shard the shard to have its primary status swapped.
+     * Calculates RoutingNodes statistics by iterating over all {@link MutableShardRouting}s
+     * in the cluster to ensure the book-keeping is correct.
+     * For performance reasons, this should only be called from asserts
+     *
+     * @return this method always returns <code>true</code> or throws an assertion error. If assertion are not enabled
+     *         this method does nothing.
      */
-    public void changePrimaryStatusForShard(MutableShardRouting... shards) {
-        for (MutableShardRouting shard : shards) {
+    public static boolean assertShardStats(RoutingNodes routingNodes) {
+        boolean run = false;
+        assert (run = true); // only run if assertions are enabled!
+        if (!run) {
+            return true;
+        }
+        int unassignedPrimaryCount = 0;
+        int inactivePrimaryCount = 0;
+        int inactiveShardCount = 0;
+        int relocating = 0;
+        final Set<ShardId> seenShards = newHashSet();
+        Map<String, Integer> indicesAndShards = new HashMap<String, Integer>();
+        for (RoutingNode node : routingNodes) {
+            for (MutableShardRouting shard : node) {
+                if (!shard.active() && shard.relocatingNodeId() == null) {
+                    if (!shard.relocating()) {
+                        inactiveShardCount++;
+                        if (shard.primary()) {
+                            inactivePrimaryCount++;
+                        }
+                    }
+                }
+                if (shard.relocating()) {
+                    relocating++;
+                }
+                seenShards.add(shard.shardId());
+                Integer i = indicesAndShards.get(shard.index());
+                if (i == null) {
+                    i = shard.id();
+                }
+                indicesAndShards.put(shard.index(), Math.max(i, shard.id()));
+            }
+        }
+        // Assert that the active shard routing are identical.
+        Set<Map.Entry<String, Integer>> entries = indicesAndShards.entrySet();
+        Set<MutableShardRouting> shards = newHashSet();
+        for (Map.Entry<String, Integer> e : entries) {
+            String index = e.getKey();
+            for (int i = 0; i < e.getValue(); i++) {
+                for (RoutingNode routingNode : routingNodes) {
+                    for (MutableShardRouting shardRouting : routingNode) {
+                        if (shardRouting.index().equals(index) && shardRouting.id() == i) {
+                            shards.add(shardRouting);
+                        }
+                    }
+                }
+                Set<MutableShardRouting> mutableShardRoutings = routingNodes.activeShards(new ShardId(index, i));
+                for (MutableShardRouting r : mutableShardRoutings) {
+                    assert shards.contains(r);
+                    shards.remove(r);
+                }
+                assert shards.isEmpty();
+            }
+        }
+
+        for (MutableShardRouting shard : routingNodes.unassigned()) {
             if (shard.primary()) {
-                shard.moveFromPrimary();
-            } else {
-                shard.moveToPrimary();
+                unassignedPrimaryCount++;
             }
+            seenShards.add(shard.shardId());
+        }
+
+        assert unassignedPrimaryCount == routingNodes.unassignedShards.numPrimaries() :
+                "Unassigned primaries is [" + unassignedPrimaryCount + "] but RoutingNodes returned unassigned primaries [" + routingNodes.unassigned().numPrimaries() + "]";
+        assert inactivePrimaryCount == routingNodes.inactivePrimaryCount :
+                "Inactive Primary count [" + inactivePrimaryCount + "] but RoutingNodes returned inactive primaries [" + routingNodes.inactivePrimaryCount + "]";
+        assert inactiveShardCount == routingNodes.inactiveShardCount :
+                "Inactive Shard count [" + inactiveShardCount + "] but RoutingNodes returned inactive shards [" + routingNodes.inactiveShardCount + "]";
+        assert routingNodes.getRelocatingShardCount() == relocating : "Relocating shards mismatch [" + routingNodes.getRelocatingShardCount() + "] but expected [" + relocating + "]";
+        return true;
+    }
+
+
+    public class RoutingNodesIterator implements Iterator<RoutingNode>, Iterable<MutableShardRouting> {
+        private RoutingNode current;
+        private final Iterator<RoutingNode> delegate;
+
+        public RoutingNodesIterator(Iterator<RoutingNode> iterator) {
+            delegate = iterator;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return delegate.hasNext();
+        }
+
+        @Override
+        public RoutingNode next() {
+            return current = delegate.next();
+        }
+
+        public RoutingNodeIterator nodeShards() {
+            return new RoutingNodeIterator(current);
+        }
+
+        @Override
+        public void remove() {
+           delegate.remove();
+        }
+
+        @Override
+        public Iterator<MutableShardRouting> iterator() {
+            return nodeShards();
+        }
+    }
+
+    public final class RoutingNodeIterator implements Iterator<MutableShardRouting>, Iterable<MutableShardRouting> {
+        private final RoutingNode iterable;
+        private MutableShardRouting shard;
+        private final Iterator<MutableShardRouting> delegate;
+
+        public RoutingNodeIterator(RoutingNode iterable) {
+            this.delegate = iterable.mutableIterator();
+            this.iterable = iterable;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return delegate.hasNext();
+        }
+
+        @Override
+        public MutableShardRouting next() {
+            return shard = delegate.next();
+        }
+
+        public void remove() {
+            delegate.remove();
+            RoutingNodes.this.remove(shard);
+        }
+
+        @Override
+        public Iterator<MutableShardRouting> iterator() {
+            return iterable.iterator();
+        }
+
+        public void moveToUnassigned() {
+            iterator().remove();
+            unassigned().add(new MutableShardRouting(shard.index(), shard.id(),
+                    null, shard.primary(), ShardRoutingState.UNASSIGNED, shard.version() + 1));
         }
     }
 }

+ 46 - 56
src/main/java/org/elasticsearch/cluster/routing/allocation/AllocationService.java

@@ -39,7 +39,6 @@ import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.node.settings.NodeSettingsService;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 
@@ -88,7 +87,7 @@ public class AllocationService extends AbstractComponent {
     public RoutingAllocation.Result applyStartedShards(ClusterState clusterState, List<? extends ShardRouting> startedShards, boolean withReroute) {
         RoutingNodes routingNodes = clusterState.routingNodes();
         // shuffle the unassigned nodes, just so we won't have things like poison failed shards
-        Collections.shuffle(routingNodes.unassigned());
+        routingNodes.unassigned().shuffle();
         StartedRerouteAllocation allocation = new StartedRerouteAllocation(allocationDeciders, routingNodes, clusterState.nodes(), startedShards, clusterInfoService.getClusterInfo());
         boolean changed = applyStartedShards(routingNodes, startedShards);
         if (!changed) {
@@ -113,7 +112,7 @@ public class AllocationService extends AbstractComponent {
     public RoutingAllocation.Result applyFailedShards(ClusterState clusterState, List<ShardRouting> failedShards) {
         RoutingNodes routingNodes = clusterState.routingNodes();
         // shuffle the unassigned nodes, just so we won't have things like poison failed shards
-        Collections.shuffle(routingNodes.unassigned());
+        routingNodes.unassigned().shuffle();
         FailedRerouteAllocation allocation = new FailedRerouteAllocation(allocationDeciders, routingNodes, clusterState.nodes(), failedShards, clusterInfoService.getClusterInfo());
         boolean changed = false;
         for (ShardRouting failedShard : failedShards) {
@@ -152,7 +151,7 @@ public class AllocationService extends AbstractComponent {
     public RoutingAllocation.Result reroute(ClusterState clusterState) {
         RoutingNodes routingNodes = clusterState.routingNodes();
         // shuffle the unassigned nodes, just so we won't have things like poison failed shards
-        Collections.shuffle(routingNodes.unassigned());
+        routingNodes.unassigned().shuffle();
         RoutingAllocation allocation = new RoutingAllocation(allocationDeciders, routingNodes, clusterState.nodes(), clusterInfoService.getClusterInfo());
         if (!reroute(allocation)) {
             return new RoutingAllocation.Result(false, clusterState.routingTable(), allocation.explanation());
@@ -168,7 +167,7 @@ public class AllocationService extends AbstractComponent {
     public RoutingAllocation.Result rerouteWithNoReassign(ClusterState clusterState) {
         RoutingNodes routingNodes = clusterState.routingNodes();
         // shuffle the unassigned nodes, just so we won't have things like poison failed shards
-        Collections.shuffle(routingNodes.unassigned());
+        routingNodes.unassigned().shuffle();
         RoutingAllocation allocation = new RoutingAllocation(allocationDeciders, routingNodes, clusterState.nodes(), clusterInfoService.getClusterInfo());
         boolean changed = false;
         // first, clear from the shards any node id they used to belong to that is now dead
@@ -211,6 +210,7 @@ public class AllocationService extends AbstractComponent {
 
         // rebalance
         changed |= shardsAllocators.rebalance(allocation);
+        assert RoutingNodes.assertShardStats(allocation.routingNodes());
         return changed;
     }
 
@@ -221,15 +221,15 @@ public class AllocationService extends AbstractComponent {
         List<MutableShardRouting> shards = new ArrayList<MutableShardRouting>();
         int index = 0;
         boolean found = true;
-        RoutingNodes routingNodes = allocation.routingNodes();
+        final RoutingNodes routingNodes = allocation.routingNodes();
         while (found) {
             found = false;
             for (RoutingNode routingNode : routingNodes) {
-                if (index >= routingNode.shards().size()) {
+                if (index >= routingNode.size()) {
                     continue;
                 }
                 found = true;
-                shards.add(routingNode.shards().get(index));
+                shards.add(routingNode.get(index));
             }
             index++;
         }
@@ -239,7 +239,7 @@ public class AllocationService extends AbstractComponent {
             if (!shardRouting.started()) {
                 continue;
             }
-            RoutingNode routingNode = routingNodes.node(shardRouting.currentNodeId());
+            final RoutingNode routingNode = routingNodes.node(shardRouting.currentNodeId());
             Decision decision = allocation.deciders().canRemain(shardRouting, routingNode, allocation);
             if (decision.type() == Decision.Type.NO) {
                 logger.debug("[{}][{}] allocated on [{}], but can no longer be allocated on it, moving...", shardRouting.index(), shardRouting.id(), routingNode.node());
@@ -247,6 +247,7 @@ public class AllocationService extends AbstractComponent {
                 if (!moved) {
                     logger.debug("[{}][{}] can't move", shardRouting.index(), shardRouting.id());
                 } else {
+                    assert RoutingNodes.assertShardStats(allocation.routingNodes());
                     changed = true;
                 }
             }
@@ -261,15 +262,15 @@ public class AllocationService extends AbstractComponent {
             if (shardEntry.primary() && !shardEntry.assignedToNode()) {
                 boolean elected = false;
                 // primary and not assigned, go over and find a replica that is assigned and active (since it might be relocating)
-                for (RoutingNode routingNode : routingNodes.nodesToShards().values()) {
+                for (RoutingNode routingNode : routingNodes) {
 
-                    for (MutableShardRouting shardEntry2 : routingNode.shards()) {
+                    for (MutableShardRouting shardEntry2 : routingNode) {
                         if (shardEntry.shardId().equals(shardEntry2.shardId()) && shardEntry2.active()) {
                             assert shardEntry2.assignedToNode();
                             assert !shardEntry2.primary();
 
                             changed = true;
-                            routingNodes.changePrimaryStatusForShard( shardEntry, shardEntry2 );
+                            routingNodes.swapPrimaryFlag(shardEntry, shardEntry2);
 
                             if (shardEntry2.relocatingNodeId() != null) {
                                 // its also relocating, make sure to move the other routing to primary
@@ -277,7 +278,7 @@ public class AllocationService extends AbstractComponent {
                                 if (node != null) {
                                     for (MutableShardRouting shardRouting : node) {
                                         if (shardRouting.shardId().equals(shardEntry2.shardId()) && !shardRouting.primary()) {
-                                            routingNodes.changePrimaryStatusForShard( shardRouting );
+                                            routingNodes.swapPrimaryFlag(shardRouting);
                                             break;
                                         }
                                     }
@@ -300,8 +301,8 @@ public class AllocationService extends AbstractComponent {
         List<ShardRouting> shardsToFail = null;
         for (MutableShardRouting shardEntry : routingNodes.unassigned()) {
             if (shardEntry.primary() && !shardEntry.assignedToNode()) {
-                for (RoutingNode routingNode : routingNodes.nodesToShards().values()) {
-                    for (MutableShardRouting shardEntry2 : routingNode.shards()) {
+                for (RoutingNode routingNode : routingNodes) {
+                    for (MutableShardRouting shardEntry2 : routingNode) {
                         if (shardEntry.shardId().equals(shardEntry2.shardId()) && !shardEntry2.active()) {
                             changed = true;
                             if (shardsToFail == null) {
@@ -326,19 +327,18 @@ public class AllocationService extends AbstractComponent {
      * new nodes);
      */
     private void applyNewNodes(RoutingAllocation allocation) {
-        RoutingNodes routingNodes = allocation.routingNodes();
+        final RoutingNodes routingNodes = allocation.routingNodes();
         for (ObjectCursor<DiscoveryNode> cursor : allocation.nodes().dataNodes().values()) {
             DiscoveryNode node = cursor.value;
-            if (!routingNodes.nodesToShards().containsKey(node.id())) {
-                RoutingNode routingNode = new RoutingNode(node.id(), node);
-                routingNodes.nodesToShards().put(node.id(), routingNode);
+            if (!routingNodes.isKnown(node)) {
+                routingNodes.addNode(node);
             }
         }
     }
 
     private boolean deassociateDeadNodes(RoutingAllocation allocation) {
         boolean changed = false;
-        for (Iterator<RoutingNode> it = allocation.routingNodes().nodesToShards().values().iterator(); it.hasNext(); ) {
+        for (RoutingNodes.RoutingNodesIterator it = allocation.routingNodes().nodes(); it.hasNext(); ) {
             RoutingNode node = it.next();
             if (allocation.nodes().dataNodes().containsKey(node.nodeId())) {
                 // its a live node, continue
@@ -346,7 +346,7 @@ public class AllocationService extends AbstractComponent {
             }
             changed = true;
             // now, go over all the shards routing on the node, and fail them
-            for (MutableShardRouting shardRouting : new ArrayList<MutableShardRouting>(node.shards())) {
+            for (MutableShardRouting shardRouting : node.copyShards()) {
                 applyFailedShard(allocation, shardRouting, false);
             }
             // its a dead node, remove it, note, its important to remove it *after* we apply failed shard
@@ -365,14 +365,14 @@ public class AllocationService extends AbstractComponent {
             // retrieve the relocating node id before calling startedShard().
             String relocatingNodeId = null;
 
-            RoutingNode currentRoutingNode = routingNodes.nodesToShards().get(startedShard.currentNodeId());
+            RoutingNodes.RoutingNodeIterator currentRoutingNode = routingNodes.routingNodeIter(startedShard.currentNodeId());
             if (currentRoutingNode != null) {
                 for (MutableShardRouting shard : currentRoutingNode) {
                     if (shard.shardId().equals(startedShard.shardId())) {
                         relocatingNodeId = shard.relocatingNodeId();
                         if (!shard.started()) {
                             dirty = true;
-                            routingNodes.markShardStarted( shard );
+                            routingNodes.started(shard);
                         }
                         break;
                     }
@@ -387,15 +387,14 @@ public class AllocationService extends AbstractComponent {
                 continue;
             }
 
-            RoutingNode sourceRoutingNode = routingNodes.nodesToShards().get(relocatingNodeId);
+            RoutingNodes.RoutingNodeIterator sourceRoutingNode = routingNodes.routingNodeIter(relocatingNodeId);
             if (sourceRoutingNode != null) {
-                Iterator<MutableShardRouting> shardsIter = sourceRoutingNode.iterator();
-                while (shardsIter.hasNext()) {
-                    MutableShardRouting shard = shardsIter.next();
+                while (sourceRoutingNode.hasNext()) {
+                    MutableShardRouting shard = sourceRoutingNode.next();
                     if (shard.shardId().equals(startedShard.shardId())) {
                         if (shard.relocating()) {
                             dirty = true;
-                            shardsIter.remove();
+                            sourceRoutingNode.remove();
                             break;
                         }
                     }
@@ -420,23 +419,19 @@ public class AllocationService extends AbstractComponent {
         }
 
         RoutingNodes routingNodes = allocation.routingNodes();
-
         if (failedShard.relocatingNodeId() != null) {
             // the shard is relocating, either in initializing (recovery from another node) or relocating (moving to another node)
             if (failedShard.state() == INITIALIZING) {
                 // the shard is initializing and recovering from another node
                 boolean dirty = false;
                 // first, we need to cancel the current node that is being initialized
-                RoutingNode initializingNode = routingNodes.node(failedShard.currentNodeId());
+                RoutingNodes.RoutingNodeIterator initializingNode = routingNodes.routingNodeIter(failedShard.currentNodeId());
                 if (initializingNode != null) {
-                    for (Iterator<MutableShardRouting> it = initializingNode.iterator(); it.hasNext(); ) {
-                        MutableShardRouting shardRouting = it.next();
+                    while(initializingNode.hasNext()) {
+                        MutableShardRouting shardRouting = initializingNode.next();
                         if (shardRouting.equals(failedShard)) {
                             dirty = true;
-                            it.remove();
-
-                            routingNodes.deassignShard( shardRouting );
-
+                            initializingNode.remove();
                             if (addToIgnoreList) {
                                 // make sure we ignore this shard on the relevant node
                                 allocation.addIgnoreShardForNode(failedShard.shardId(), failedShard.currentNodeId());
@@ -450,11 +445,10 @@ public class AllocationService extends AbstractComponent {
                     // now, find the node that we are relocating *from*, and cancel its relocation
                     RoutingNode relocatingFromNode = routingNodes.node(failedShard.relocatingNodeId());
                     if (relocatingFromNode != null) {
-                        for (Iterator<MutableShardRouting> it = relocatingFromNode.iterator(); it.hasNext(); ) {
-                            MutableShardRouting shardRouting = it.next();
-                            if (shardRouting.shardId().equals(failedShard.shardId()) && shardRouting.state() == RELOCATING) {
+                        for (MutableShardRouting shardRouting : relocatingFromNode) {
+                            if (shardRouting.shardId().equals(failedShard.shardId()) && shardRouting.relocating()) {
                                 dirty = true;
-                                routingNodes.cancelRelocationForShard( shardRouting );
+                                routingNodes.cancelRelocation(shardRouting);
                                 break;
                             }
                         }
@@ -467,14 +461,13 @@ public class AllocationService extends AbstractComponent {
                 // first, we need to cancel the current relocation from the current node
                 // now, find the node that we are recovering from, cancel the relocation, remove it from the node
                 // and add it to the unassigned shards list...
-                RoutingNode relocatingFromNode = routingNodes.node(failedShard.currentNodeId());
+                RoutingNodes.RoutingNodeIterator relocatingFromNode = routingNodes.routingNodeIter(failedShard.currentNodeId());
                 if (relocatingFromNode != null) {
-                    for (Iterator<MutableShardRouting> it = relocatingFromNode.iterator(); it.hasNext(); ) {
-                        MutableShardRouting shardRouting = it.next();
+                    while(relocatingFromNode.hasNext()) {
+                        MutableShardRouting shardRouting = relocatingFromNode.next();
                         if (shardRouting.equals(failedShard)) {
                             dirty = true;
-                            routingNodes.cancelRelocationForShard( shardRouting );
-                            it.remove();
+                            relocatingFromNode.remove();
                             if (addToIgnoreList) {
                                 // make sure we ignore this shard on the relevant node
                                 allocation.addIgnoreShardForNode(failedShard.shardId(), failedShard.currentNodeId());
@@ -488,14 +481,13 @@ public class AllocationService extends AbstractComponent {
                 }
                 if (dirty) {
                     // next, we need to find the target initializing shard that is recovering from, and remove it...
-                    RoutingNode initializingNode = routingNodes.node(failedShard.relocatingNodeId());
+                    RoutingNodes.RoutingNodeIterator initializingNode = routingNodes.routingNodeIter(failedShard.relocatingNodeId());
                     if (initializingNode != null) {
-                        for (Iterator<MutableShardRouting> it = initializingNode.iterator(); it.hasNext(); ) {
-                            MutableShardRouting shardRouting = it.next();
+                        while (initializingNode.hasNext()) {
+                            MutableShardRouting shardRouting = initializingNode.next();
                             if (shardRouting.shardId().equals(failedShard.shardId()) && shardRouting.state() == INITIALIZING) {
                                 dirty = true;
-                                routingNodes.deassignShard( shardRouting );
-                                it.remove();
+                                initializingNode.remove();
                             }
                         }
                     }
@@ -507,19 +499,17 @@ public class AllocationService extends AbstractComponent {
         } else {
             // the shard is not relocating, its either started, or initializing, just cancel it and move on...
             boolean dirty = false;
-            RoutingNode node = routingNodes.node(failedShard.currentNodeId());
+            RoutingNodes.RoutingNodeIterator node = routingNodes.routingNodeIter(failedShard.currentNodeId());
             if (node != null) {
-                for (Iterator<MutableShardRouting> it = node.iterator(); it.hasNext(); ) {
-                    MutableShardRouting shardRouting = it.next();
+                while(node.hasNext()) {
+                    MutableShardRouting shardRouting = node.next();
                     if (shardRouting.equals(failedShard)) {
                         dirty = true;
                         if (addToIgnoreList) {
                             // make sure we ignore this shard on the relevant node
                             allocation.addIgnoreShardForNode(failedShard.shardId(), failedShard.currentNodeId());
                         }
-
-                        it.remove();
-
+                        node.remove();
                         // move all the shards matching the failed shard to the end of the unassigned list
                         // so we give a chance for other allocations and won't create poison failed allocations
                         // that can keep other shards from being allocated (because of limits applied on how many

+ 24 - 59
src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java

@@ -23,10 +23,7 @@ import com.google.common.base.Predicate;
 import org.apache.lucene.util.IntroSorter;
 import org.elasticsearch.ElasticSearchIllegalArgumentException;
 import org.elasticsearch.cluster.metadata.MetaData;
-import org.elasticsearch.cluster.routing.MutableShardRouting;
-import org.elasticsearch.cluster.routing.RoutingNode;
-import org.elasticsearch.cluster.routing.RoutingNodes;
-import org.elasticsearch.cluster.routing.ShardRoutingState;
+import org.elasticsearch.cluster.routing.*;
 import org.elasticsearch.cluster.routing.allocation.FailedRerouteAllocation;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.cluster.routing.allocation.StartedRerouteAllocation;
@@ -266,6 +263,7 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
         private final Map<String, ModelNode> nodes = new HashMap<String, ModelNode>();
         private final HashSet<String> indices = new HashSet<String>();
         private final RoutingAllocation allocation;
+        private final RoutingNodes routingNodes;
         private final WeightFunction weight;
 
         private final float threshold;
@@ -284,10 +282,11 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
             this.allocation = allocation;
             this.weight = weight;
             this.threshold = threshold;
-            for (RoutingNode node : allocation.routingNodes()) {
+            this.routingNodes = allocation.routingNodes();
+            for (RoutingNode node : routingNodes) {
                 nodes.put(node.nodeId(), new ModelNode(node.nodeId()));
             }
-            metaData = allocation.routingNodes().metaData();
+            metaData = routingNodes.metaData();
         }
 
         /**
@@ -335,7 +334,7 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
             return new NodeSorter(nodesArray(), weight, this);
         }
 
-        private boolean initialize(RoutingNodes routing, List<MutableShardRouting> unassigned) {
+        private boolean initialize(RoutingNodes routing, RoutingNodes.UnassignedShards unassigned) {
             if (logger.isTraceEnabled()) {
                 logger.trace("Start distributing Shards");
             }
@@ -366,8 +365,8 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
             if (logger.isTraceEnabled()) {
                 logger.trace("Start balancing cluster");
             }
-            final TransactionalList<MutableShardRouting> unassigned = new TransactionalList<MutableShardRouting>(allocation.routingNodes().unassigned());
-            boolean changed = initialize(allocation.routingNodes(), unassigned);
+            final RoutingNodes.UnassignedShards unassigned = routingNodes.unassigned().transactionBegin();
+            boolean changed = initialize(routingNodes, unassigned);
             NodeSorter sorter = newNodeSorter();
             if (nodes.size() > 1) { /* skip if we only have one node */
                 for (String index : buildWeightOrderedIndidces(Operation.BALANCE, sorter)) {
@@ -445,7 +444,7 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
                     }
                 }
             }
-            unassigned.commit();
+            routingNodes.unassigned().transactionEnd(unassigned);
             return changed;
         }
 
@@ -520,8 +519,8 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
             if (logger.isTraceEnabled()) {
                 logger.trace("Try moving shard [{}] from [{}]", shard, node);
             }
-            final TransactionalList<MutableShardRouting> unassigned = new TransactionalList<MutableShardRouting>(allocation.routingNodes().unassigned());
-            boolean changed = initialize(allocation.routingNodes(), unassigned);
+            final RoutingNodes.UnassignedShards unassigned = routingNodes.unassigned().transactionBegin();
+            boolean changed = initialize(routingNodes, unassigned);
 
             final ModelNode sourceNode = nodes.get(node.nodeId());
             assert sourceNode != null;
@@ -540,15 +539,15 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
                 if (currentNode.getNodeId().equals(node.nodeId())) {
                     continue;
                 }
-                RoutingNode target = allocation.routingNodes().node(currentNode.getNodeId());
+                RoutingNode target = routingNodes.node(currentNode.getNodeId());
                 Decision decision = allocation.deciders().canAllocate(shard, target, allocation);
                 if (decision.type() == Type.YES) { // TODO maybe we can respect throttling here too?
                     sourceNode.removeShard(shard);
                     final MutableShardRouting initializingShard = new MutableShardRouting(shard.index(), shard.id(), currentNode.getNodeId(),
                             shard.currentNodeId(), shard.restoreSource(), shard.primary(), INITIALIZING, shard.version() + 1);
                     currentNode.addShard(initializingShard, decision);
-                    allocation.routingNodes().assignShardToNode( initializingShard, target.nodeId() );
-                    allocation.routingNodes().relocateShard( shard, target.nodeId() ); // set the node to relocate after we added the initializing shard
+                    routingNodes.assign(initializingShard, target.nodeId());
+                    routingNodes.relocate(shard, target.nodeId()); // set the node to relocate after we added the initializing shard
                     if (logger.isTraceEnabled()) {
                         logger.trace("Moved shard [{}] to node [{}]", shard, currentNode.getNodeId());
                     }
@@ -556,7 +555,7 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
                     break;
                 }
             }
-            unassigned.commit();
+            routingNodes.unassigned().transactionEnd(unassigned);
             return changed;
         }
 
@@ -589,7 +588,7 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
          * Allocates all given shards on the minimal eligable node for the shards index
          * with respect to the weight function. All given shards must be unassigned.
          */
-        private boolean allocateUnassigned(List<MutableShardRouting> unassigned, List<MutableShardRouting> ignoredUnassigned) {
+        private boolean allocateUnassigned(RoutingNodes.UnassignedShards unassigned, List<MutableShardRouting> ignoredUnassigned) {
             assert !nodes.isEmpty();
             if (logger.isTraceEnabled()) {
                 logger.trace("Start allocating unassigned shards");
@@ -603,7 +602,6 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
              * TODO: We could be smarter here and group the shards by index and then
              * use the sorter to save some iterations. 
              */
-            final RoutingNodes routingNodes = allocation.routingNodes();
             final AllocationDeciders deciders = allocation.deciders();
             final Set<MutableShardRouting> currentRound = new TreeSet<MutableShardRouting>(new Comparator<MutableShardRouting>() {
                 @Override
@@ -704,7 +702,7 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
                             if (logger.isTraceEnabled()) {
                                 logger.trace("Assigned shard [{}] to [{}]", shard, minNode.getNodeId());
                             }
-                            routingNodes.assignShardToNode( shard, routingNodes.node(minNode.getNodeId()).nodeId() );
+                            routingNodes.assign(shard, routingNodes.node(minNode.getNodeId()).nodeId());
                             changed = true;
                             continue; // don't add to ignoreUnassigned
                         }
@@ -717,7 +715,7 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
                     ignoredUnassigned.add(shard);
                 }
                 if (!iterationChanged && !unassigned.isEmpty()) {
-                    ignoredUnassigned.addAll(unassigned);
+                    unassigned.copyAll(ignoredUnassigned);
                     unassigned.clear();
                     return changed;
                 }
@@ -740,7 +738,7 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
                     logger.trace("Try relocating shard for index index [{}] from node [{}] to node [{}]", idx, maxNode.getNodeId(),
                             minNode.getNodeId());
                 }
-                final RoutingNode node = allocation.routingNodes().node(minNode.getNodeId());
+                final RoutingNode node = routingNodes.node(minNode.getNodeId());
                 MutableShardRouting candidate = null;
                 final AllocationDeciders deciders = allocation.deciders();
                 /* make a copy since we modify this list in the loop */
@@ -781,14 +779,14 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
                         }
                         /* now allocate on the cluster - if we are started we need to relocate the shard */
                         if (candidate.started()) {
-                            RoutingNode lowRoutingNode = allocation.routingNodes().node(minNode.getNodeId());
-                            allocation.routingNodes().assignShardToNode(new MutableShardRouting(candidate.index(), candidate.id(), lowRoutingNode.nodeId(), candidate
-                                    .currentNodeId(), candidate.restoreSource(), candidate.primary(), INITIALIZING, candidate.version() + 1), lowRoutingNode.nodeId() );
-                            allocation.routingNodes().relocateShard( candidate, lowRoutingNode.nodeId());
+                            RoutingNode lowRoutingNode = routingNodes.node(minNode.getNodeId());
+                            routingNodes.assign(new MutableShardRouting(candidate.index(), candidate.id(), lowRoutingNode.nodeId(), candidate
+                                    .currentNodeId(), candidate.restoreSource(), candidate.primary(), INITIALIZING, candidate.version() + 1), lowRoutingNode.nodeId());
+                            routingNodes.relocate(candidate, lowRoutingNode.nodeId());
 
                         } else {
                             assert candidate.unassigned();
-                            allocation.routingNodes().assignShardToNode( candidate, allocation.routingNodes().node(minNode.getNodeId()).nodeId() );
+                            routingNodes.assign(candidate, routingNodes.node(minNode.getNodeId()).nodeId());
                         }
                         return true;
 
@@ -1043,37 +1041,4 @@ public class BalancedShardsAllocator extends AbstractComponent implements Shards
             return weights[weights.length - 1] - weights[0];
         }
     }
-    
-    /**
-     * A list that makes a full copy of the original list and applies all
-     * modification to the copied list once {@link TransactionalList#commit()}
-     * is called.
-     * 
-     */
-    @SuppressWarnings("serial")
-    private static final class TransactionalList<T> extends ArrayList<T> {
-        
-        private final List<T> originalList;
-        private List<T> assertingList; // only with assert
-
-        TransactionalList(List<T> originalList) {
-            super(originalList);
-            assert copyAsseringList(originalList);
-            this.originalList = originalList;
-        }
-        
-        private boolean copyAsseringList(List<T> orig) {
-            this.assertingList = new ArrayList<T>(orig);
-            return true;
-        }
-        
-        public void commit() {
-            /* Ensure that the actual source list is not modified while
-             * the transaction is running */
-            assert assertingList.equals(originalList) : "The list was modified outside of the scope";
-            originalList.clear();
-            originalList.addAll(this);    
-            
-        }
-    }
 }

+ 13 - 13
src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/EvenShardsCountAllocator.java

@@ -102,13 +102,13 @@ public class EvenShardsCountAllocator extends AbstractComponent implements Shard
 
                 Decision decision = allocation.deciders().canAllocate(shard, node, allocation);
                 if (decision.type() == Decision.Type.YES) {
-                    int numberOfShardsToAllocate = routingNodes.requiredAverageNumberOfShardsPerNode() - node.shards().size();
+                    int numberOfShardsToAllocate = routingNodes.requiredAverageNumberOfShardsPerNode() - node.size();
                     if (numberOfShardsToAllocate <= 0) {
                         continue;
                     }
 
                     changed = true;
-                    allocation.routingNodes().assignShardToNode( shard, node.nodeId() );
+                    allocation.routingNodes().assign(shard, node.nodeId());
                     unassignedIterator.remove();
                     break;
                 }
@@ -123,7 +123,7 @@ public class EvenShardsCountAllocator extends AbstractComponent implements Shard
                 Decision decision = allocation.deciders().canAllocate(shard, routingNode, allocation);
                 if (decision.type() == Decision.Type.YES) {
                     changed = true;
-                    allocation.routingNodes().assignShardToNode( shard, routingNode.nodeId() );
+                    allocation.routingNodes().assign(shard, routingNode.nodeId());
                     it.remove();
                     break;
                 }
@@ -156,7 +156,7 @@ public class EvenShardsCountAllocator extends AbstractComponent implements Shard
                     continue;
                 }
 
-                if (lowRoutingNode.shards().size() >= averageNumOfShards) {
+                if (lowRoutingNode.size() >= averageNumOfShards) {
                     lowIndex++;
                     continue;
                 }
@@ -173,11 +173,11 @@ public class EvenShardsCountAllocator extends AbstractComponent implements Shard
                     Decision allocateDecision = allocation.deciders().canAllocate(startedShard, lowRoutingNode, allocation);
                     if (allocateDecision.type() == Decision.Type.YES) {
                         changed = true;
-                        allocation.routingNodes().assignShardToNode(new MutableShardRouting(startedShard.index(), startedShard.id(),
+                        allocation.routingNodes().assign(new MutableShardRouting(startedShard.index(), startedShard.id(),
                                 lowRoutingNode.nodeId(), startedShard.currentNodeId(), startedShard.restoreSource(),
-                                startedShard.primary(), INITIALIZING, startedShard.version() + 1), lowRoutingNode.nodeId() );
+                                startedShard.primary(), INITIALIZING, startedShard.version() + 1), lowRoutingNode.nodeId());
 
-                        allocation.routingNodes().relocateShard( startedShard, lowRoutingNode.nodeId() );
+                        allocation.routingNodes().relocate(startedShard, lowRoutingNode.nodeId());
                         relocated = true;
                         relocationPerformed = true;
                         break;
@@ -210,11 +210,11 @@ public class EvenShardsCountAllocator extends AbstractComponent implements Shard
             }
             Decision decision = allocation.deciders().canAllocate(shardRouting, nodeToCheck, allocation);
             if (decision.type() == Decision.Type.YES) {
-                allocation.routingNodes().assignShardToNode(new MutableShardRouting(shardRouting.index(), shardRouting.id(),
+                allocation.routingNodes().assign(new MutableShardRouting(shardRouting.index(), shardRouting.id(),
                         nodeToCheck.nodeId(), shardRouting.currentNodeId(), shardRouting.restoreSource(),
-                        shardRouting.primary(), INITIALIZING, shardRouting.version() + 1), nodeToCheck.nodeId() );
+                        shardRouting.primary(), INITIALIZING, shardRouting.version() + 1), nodeToCheck.nodeId());
 
-                allocation.routingNodes().relocateShard( shardRouting, nodeToCheck.nodeId() );
+                allocation.routingNodes().relocate(shardRouting, nodeToCheck.nodeId());
                 changed = true;
                 break;
             }
@@ -227,13 +227,13 @@ public class EvenShardsCountAllocator extends AbstractComponent implements Shard
         // create count per node id, taking into account relocations
         final ObjectIntOpenHashMap<String> nodeCounts = new ObjectIntOpenHashMap<String>();
         for (RoutingNode node : allocation.routingNodes()) {
-            for (int i = 0; i < node.shards().size(); i++) {
-                ShardRouting shardRouting = node.shards().get(i);
+            for (int i = 0; i < node.size(); i++) {
+                ShardRouting shardRouting = node.get(i);
                 String nodeId = shardRouting.relocating() ? shardRouting.relocatingNodeId() : shardRouting.currentNodeId();
                 nodeCounts.addTo(nodeId, 1);
             }
         }
-        RoutingNode[] nodes = allocation.routingNodes().nodesToShards().values().toArray(new RoutingNode[allocation.routingNodes().nodesToShards().values().size()]);
+        RoutingNode[] nodes = allocation.routingNodes().toArray();
         Arrays.sort(nodes, new Comparator<RoutingNode>() {
             @Override
             public int compare(RoutingNode o1, RoutingNode o2) {

+ 1 - 0
src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/ShardsAllocators.java

@@ -21,6 +21,7 @@ package org.elasticsearch.cluster.routing.allocation.allocator;
 
 import org.elasticsearch.cluster.routing.MutableShardRouting;
 import org.elasticsearch.cluster.routing.RoutingNode;
+import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.cluster.routing.allocation.FailedRerouteAllocation;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.cluster.routing.allocation.StartedRerouteAllocation;

+ 1 - 1
src/main/java/org/elasticsearch/cluster/routing/allocation/command/AllocateAllocationCommand.java

@@ -193,7 +193,7 @@ public class AllocateAllocationCommand implements AllocationCommand {
                 continue;
             }
             it.remove();
-            allocation.routingNodes().assignShardToNode( shardRouting, routingNode.nodeId() );
+            allocation.routingNodes().assign(shardRouting, routingNode.nodeId());
             if (shardRouting.primary()) {
                 // we need to clear the post allocation flag, since its an explicit allocation of the primary shard
                 // and we want to force allocate it (and create a new index for it)

+ 11 - 15
src/main/java/org/elasticsearch/cluster/routing/allocation/command/CancelAllocationCommand.java

@@ -25,6 +25,7 @@ import org.elasticsearch.ElasticSearchParseException;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.MutableShardRouting;
 import org.elasticsearch.cluster.routing.RoutingNode;
+import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -35,7 +36,6 @@ import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.shard.ShardId;
 
 import java.io.IOException;
-import java.util.Iterator;
 
 import static org.elasticsearch.cluster.routing.ShardRoutingState.INITIALIZING;
 import static org.elasticsearch.cluster.routing.ShardRoutingState.RELOCATING;
@@ -161,9 +161,8 @@ public class CancelAllocationCommand implements AllocationCommand {
     @Override
     public void execute(RoutingAllocation allocation) throws ElasticSearchException {
         DiscoveryNode discoNode = allocation.nodes().resolveNode(node);
-
         boolean found = false;
-        for (Iterator<MutableShardRouting> it = allocation.routingNodes().node(discoNode.id()).iterator(); it.hasNext(); ) {
+        for (RoutingNodes.RoutingNodeIterator it = allocation.routingNodes().routingNodeIter(discoNode.id()); it.hasNext(); ) {
             MutableShardRouting shardRouting = it.next();
             if (!shardRouting.shardId().equals(shardId)) {
                 continue;
@@ -173,35 +172,33 @@ public class CancelAllocationCommand implements AllocationCommand {
                 if (shardRouting.initializing()) {
                     // the shard is initializing and recovering from another node, simply cancel the recovery
                     it.remove();
-                    allocation.routingNodes().deassignShard( shardRouting );
                     // and cancel the relocating state from the shard its being relocated from
                     RoutingNode relocatingFromNode = allocation.routingNodes().node(shardRouting.relocatingNodeId());
                     if (relocatingFromNode != null) {
                         for (MutableShardRouting fromShardRouting : relocatingFromNode) {
+                            assert shardRouting.state() != RELOCATING;
                             if (fromShardRouting.shardId().equals(shardRouting.shardId()) && shardRouting.state() == RELOCATING) {
-                                allocation.routingNodes().cancelRelocationForShard( fromShardRouting );
+                                // NOCOMMIT @Shay the check on the shardRouting.state() == RELOCATING can never be true right this should be fromShardRouting?
+                                allocation.routingNodes().cancelRelocation(fromShardRouting);
                                 break;
                             }
                         }
                     }
                 } else if (shardRouting.relocating()) {
+
                     // the shard is relocating to another node, cancel the recovery on the other node, and deallocate this one
                     if (!allowPrimary && shardRouting.primary()) {
                         // can't cancel a primary shard being initialized
                         throw new ElasticSearchIllegalArgumentException("[cancel_allocation] can't cancel " + shardId + " on node " + discoNode + ", shard is primary and initializing its state");
                     }
-                    it.remove();
-                    allocation.routingNodes().unassigned().add(new MutableShardRouting(shardRouting.index(), shardRouting.id(),
-                            null, shardRouting.primary(), ShardRoutingState.UNASSIGNED, shardRouting.version() + 1));
-
+                    it.moveToUnassigned();
                     // now, go and find the shard that is initializing on the target node, and cancel it as well...
-                    RoutingNode initializingNode = allocation.routingNodes().node(shardRouting.relocatingNodeId());
+                    RoutingNodes.RoutingNodeIterator initializingNode = allocation.routingNodes().routingNodeIter(shardRouting.relocatingNodeId());
                     if (initializingNode != null) {
-                        for (Iterator<MutableShardRouting> itX = initializingNode.iterator(); itX.hasNext(); ) {
-                            MutableShardRouting initializingShardRouting = itX.next();
+                        while (initializingNode.hasNext()) {
+                            MutableShardRouting initializingShardRouting = initializingNode.next();
                             if (initializingShardRouting.shardId().equals(shardRouting.shardId()) && initializingShardRouting.state() == INITIALIZING) {
-                                allocation.routingNodes().deassignShard( shardRouting );
-                                itX.remove();
+                                initializingNode.remove();
                             }
                         }
                     }
@@ -217,7 +214,6 @@ public class CancelAllocationCommand implements AllocationCommand {
                         null, shardRouting.primary(), ShardRoutingState.UNASSIGNED, shardRouting.version() + 1));
             }
         }
-
         if (!found) {
             throw new ElasticSearchIllegalArgumentException("[cancel_allocation] can't cancel " + shardId + ", failed to find it on node " + discoNode);
         }

+ 3 - 3
src/main/java/org/elasticsearch/cluster/routing/allocation/command/MoveAllocationCommand.java

@@ -167,11 +167,11 @@ public class MoveAllocationCommand implements AllocationCommand {
                 // its being throttled, maybe have a flag to take it into account and fail? for now, just do it since the "user" wants it...
             }
 
-            allocation.routingNodes().assignShardToNode(new MutableShardRouting(shardRouting.index(), shardRouting.id(),
+            allocation.routingNodes().assign(new MutableShardRouting(shardRouting.index(), shardRouting.id(),
                     toRoutingNode.nodeId(), shardRouting.currentNodeId(), shardRouting.restoreSource(),
-                    shardRouting.primary(), ShardRoutingState.INITIALIZING, shardRouting.version() + 1), toRoutingNode.nodeId() );
+                    shardRouting.primary(), ShardRoutingState.INITIALIZING, shardRouting.version() + 1), toRoutingNode.nodeId());
 
-            allocation.routingNodes().relocateShard( shardRouting, toRoutingNode.nodeId() );
+            allocation.routingNodes().relocate(shardRouting, toRoutingNode.nodeId());
         }
 
         if (!found) {

+ 1 - 2
src/main/java/org/elasticsearch/cluster/routing/allocation/decider/AwarenessAllocationDecider.java

@@ -185,8 +185,7 @@ public class AwarenessAllocationDecider extends AllocationDecider {
             // build the count of shards per attribute value
             ObjectIntOpenHashMap<String> shardPerAttribute = new ObjectIntOpenHashMap<String>();
             for (RoutingNode routingNode : allocation.routingNodes()) {
-                for (int i = 0; i < routingNode.shards().size(); i++) {
-                    MutableShardRouting nodeShardRouting = routingNode.shards().get(i);
+                for (MutableShardRouting nodeShardRouting : routingNode) {
                     if (nodeShardRouting.shardId().equals(shardRouting.shardId())) {
                         // if the shard is relocating, then make sure we count it as part of the node it is relocating to
                         if (nodeShardRouting.relocating()) {

+ 1 - 8
src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ConcurrentRebalanceAllocationDecider.java

@@ -19,18 +19,12 @@
 
 package org.elasticsearch.cluster.routing.allocation.decider;
 
-import org.elasticsearch.cluster.routing.MutableShardRouting;
-import org.elasticsearch.cluster.routing.RoutingNode;
-import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.cluster.routing.ShardRouting;
-import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.node.settings.NodeSettingsService;
 
-import java.util.List;
-
 /**
  * Similar to the {@link ClusterRebalanceAllocationDecider} this
  * {@link AllocationDecider} controls the number of currently in-progress
@@ -73,8 +67,7 @@ public class ConcurrentRebalanceAllocationDecider extends AllocationDecider {
         if (clusterConcurrentRebalance == -1) {
             return Decision.YES;
         }
-        int rebalance = allocation.routingNodes().getRelocatingShardCount();
-        if (rebalance >= clusterConcurrentRebalance) {
+        if (allocation.routingNodes().getRelocatingShardCount() >= clusterConcurrentRebalance) {
             return Decision.NO;
         }
         return Decision.YES;

+ 2 - 8
src/main/java/org/elasticsearch/cluster/routing/allocation/decider/RebalanceOnlyWhenActiveAllocationDecider.java

@@ -19,14 +19,11 @@
 
 package org.elasticsearch.cluster.routing.allocation.decider;
 
-import org.elasticsearch.cluster.routing.MutableShardRouting;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
 
-import java.util.List;
-
 /**
  * Only allow rebalancing when all shards are active within the shard replication group.
  */
@@ -39,13 +36,10 @@ public class RebalanceOnlyWhenActiveAllocationDecider extends AllocationDecider
 
     @Override
     public Decision canRebalance(ShardRouting shardRouting, RoutingAllocation allocation) {
-        List<MutableShardRouting> shards = allocation.routingNodes().shardsRoutingFor(shardRouting);
         // its ok to check for active here, since in relocation, a shard is split into two in routing
         // nodes, once relocating, and one initializing
-        for (int i = 0; i < shards.size(); i++) {
-            if (!shards.get(i).active()) {
-                return Decision.NO;
-            }
+        if (!allocation.routingNodes().allReplicasActive(shardRouting)) {
+            return Decision.NO;
         }
         return Decision.YES;
     }

+ 2 - 2
src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ReplicaAfterPrimaryActiveAllocationDecider.java

@@ -41,8 +41,8 @@ public class ReplicaAfterPrimaryActiveAllocationDecider extends AllocationDecide
         if (shardRouting.primary()) {
             return Decision.YES;
         }
-        MutableShardRouting primary = allocation.routingNodes().findPrimaryForReplica(shardRouting);
-        if (primary == null || !primary.active()) {
+        MutableShardRouting primary = allocation.routingNodes().activePrimary(shardRouting);
+        if (primary == null) {
             return Decision.NO;
         }
         return Decision.YES;

+ 4 - 8
src/main/java/org/elasticsearch/cluster/routing/allocation/decider/SameShardAllocationDecider.java

@@ -26,8 +26,6 @@ import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
 
-import java.util.List;
-
 /**
  * An allocation decider that prevents the more than prevents multiple instances
  * of the same shard to be allocated on a single <tt>host</tt>. The cluster setting can
@@ -54,10 +52,9 @@ public class SameShardAllocationDecider extends AllocationDecider {
 
     @Override
     public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, RoutingAllocation allocation) {
-        List<MutableShardRouting> shards = node.shards();
-        for (int i = 0; i < shards.size(); i++) {
+        for (MutableShardRouting nodeShard : node) {
             // we do not allow for two shards of the same shard id to exists on the same node
-            if (shards.get(i).shardId().equals(shardRouting.shardId())) {
+            if (nodeShard.shardId().equals(shardRouting.shardId())) {
                 return Decision.NO;
             }
         }
@@ -71,9 +68,8 @@ public class SameShardAllocationDecider extends AllocationDecider {
                     if (!checkNode.node().address().sameHost(node.node().address())) {
                         continue;
                     }
-                    shards = checkNode.shards();
-                    for (int i = 0; i < shards.size(); i++) {
-                        if (shards.get(i).shardId().equals(shardRouting.shardId())) {
+                    for (MutableShardRouting nodeShard : checkNode) {
+                        if (nodeShard.shardId().equals(shardRouting.shardId())) {
                             return Decision.NO;
                         }
                     }

+ 2 - 8
src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ShardsLimitAllocationDecider.java

@@ -28,8 +28,6 @@ import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
 
-import java.util.List;
-
 /**
  * This {@link AllocationDecider} limits the number of shards per node on a per
  * index basis. The allocator prevents a single node to hold more than
@@ -71,9 +69,7 @@ public class ShardsLimitAllocationDecider extends AllocationDecider {
         }
 
         int nodeCount = 0;
-        List<MutableShardRouting> shards = node.shards();
-        for (int i = 0; i < shards.size(); i++) {
-            MutableShardRouting nodeShard = shards.get(i);
+        for (MutableShardRouting nodeShard : node) {
             if (!nodeShard.index().equals(shardRouting.index())) {
                 continue;
             }
@@ -98,9 +94,7 @@ public class ShardsLimitAllocationDecider extends AllocationDecider {
         }
 
         int nodeCount = 0;
-        List<MutableShardRouting> shards = node.shards();
-        for (int i = 0; i < shards.size(); i++) {
-            MutableShardRouting nodeShard = shards.get(i);
+        for (MutableShardRouting nodeShard : node) {;
             if (!nodeShard.index().equals(shardRouting.index())) {
                 continue;
             }

+ 2 - 8
src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ThrottlingAllocationDecider.java

@@ -28,8 +28,6 @@ import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.node.settings.NodeSettingsService;
 
-import java.util.List;
-
 /**
  * {@link ThrottlingAllocationDecider} controls the recovery process per node in
  * the cluster. It exposes two settings via the cluster update API that allow
@@ -79,9 +77,7 @@ public class ThrottlingAllocationDecider extends AllocationDecider {
                 // primary is unassigned, means we are going to do recovery from gateway
                 // count *just the primary* currently doing recovery on the node and check against concurrent_recoveries
                 int primariesInRecovery = 0;
-                List<MutableShardRouting> shards = node.shards();
-                for (int i = 0; i < shards.size(); i++) {
-                    MutableShardRouting shard = shards.get(i);
+                for (MutableShardRouting shard : node) {;
                     // when a primary shard is INITIALIZING, it can be because of *initial recovery* or *relocation from another node*
                     // we only count initial recoveries here, so we need to make sure that relocating node is null
                     if (shard.state() == ShardRoutingState.INITIALIZING && shard.primary() && shard.relocatingNodeId() == null) {
@@ -100,9 +96,7 @@ public class ThrottlingAllocationDecider extends AllocationDecider {
 
         // count the number of recoveries on the node, its for both target (INITIALIZING) and source (RELOCATING)
         int currentRecoveries = 0;
-        List<MutableShardRouting> shards = node.shards();
-        for (int i = 0; i < shards.size(); i++) {
-            MutableShardRouting shard = shards.get(i);
+        for (MutableShardRouting shard : node) {
             if (shard.state() == ShardRoutingState.INITIALIZING || shard.state() == ShardRoutingState.RELOCATING) {
                 currentRecoveries++;
             }

+ 4 - 4
src/main/java/org/elasticsearch/gateway/blobstore/BlobReuseExistingGatewayAllocator.java

@@ -214,8 +214,9 @@ public class BlobReuseExistingGatewayAllocator extends AbstractComponent impleme
                 } else {
                     // if its backup, see if there is a primary that *is* allocated, and try and assign a location that is closest to it
                     // note, since we replicate operations, this might not be the same (different flush intervals)
-                    MutableShardRouting primaryShard = routingNodes.findPrimaryForReplica(shard);
-                    if (primaryShard != null && primaryShard.active()) {
+                    MutableShardRouting primaryShard = routingNodes.activePrimary(shard);
+                    if (primaryShard != null) {
+                        assert primaryShard.active();
                         DiscoveryNode primaryNode = nodes.get(primaryShard.currentNodeId());
                         if (primaryNode != null) {
                             TransportNodesListShardStoreMetaData.StoreFilesMetaData primaryNodeStore = shardStores.get(primaryNode);
@@ -252,12 +253,11 @@ public class BlobReuseExistingGatewayAllocator extends AbstractComponent impleme
                     }
                     // we found a match
                     changed = true;
-                    allocation.routingNodes().assignShardToNode( shard, lastNodeMatched.nodeId() );
+                    allocation.routingNodes().assign(shard, lastNodeMatched.nodeId());
                     unassignedIterator.remove();
                 }
             }
         }
-
         return changed;
     }
 

+ 6 - 6
src/main/java/org/elasticsearch/gateway/local/LocalGatewayAllocator.java

@@ -214,7 +214,7 @@ public class LocalGatewayAllocator extends AbstractComponent implements GatewayA
                     // we found a match
                     changed = true;
                     // make sure we create one with the version from the recovered state
-                    allocation.routingNodes().assignShardToNode(new MutableShardRouting(shard, highestVersion), node.nodeId());
+                    allocation.routingNodes().assign(new MutableShardRouting(shard, highestVersion), node.nodeId());
                     unassignedIterator.remove();
 
                     // found a node, so no throttling, no "no", and break out of the loop
@@ -234,7 +234,7 @@ public class LocalGatewayAllocator extends AbstractComponent implements GatewayA
                     // we found a match
                     changed = true;
                     // make sure we create one with the version from the recovered state
-                    allocation.routingNodes().assignShardToNode(new MutableShardRouting(shard, highestVersion), node.nodeId());
+                    allocation.routingNodes().assign(new MutableShardRouting(shard, highestVersion), node.nodeId());
                     unassignedIterator.remove();
                 }
             } else {
@@ -311,8 +311,9 @@ public class LocalGatewayAllocator extends AbstractComponent implements GatewayA
                 }
 
                 if (!shard.primary()) {
-                    MutableShardRouting primaryShard = routingNodes.findPrimaryForReplica(shard);
-                    if (primaryShard != null && primaryShard.active()) {
+                    MutableShardRouting primaryShard = routingNodes.activePrimary(shard);
+                    if (primaryShard != null) {
+                        assert primaryShard.active();
                         DiscoveryNode primaryNode = nodes.get(primaryShard.currentNodeId());
                         if (primaryNode != null) {
                             TransportNodesListShardStoreMetaData.StoreFilesMetaData primaryNodeStore = shardStores.get(primaryNode);
@@ -351,12 +352,11 @@ public class LocalGatewayAllocator extends AbstractComponent implements GatewayA
                     }
                     // we found a match
                     changed = true;
-                    allocation.routingNodes().assignShardToNode( shard, lastNodeMatched.nodeId() );
+                    allocation.routingNodes().assign(shard, lastNodeMatched.nodeId());
                     unassignedIterator.remove();
                 }
             }
         }
-
         return changed;
     }
 

+ 8 - 9
src/main/java/org/elasticsearch/indices/cluster/IndicesClusterStateService.java

@@ -263,7 +263,7 @@ public class IndicesClusterStateService extends AbstractLifecycleComponent<Indic
     }
 
     private void applyDeletedShards(final ClusterChangedEvent event) {
-        RoutingNode routingNode = event.state().readOnlyRoutingNodes().nodesToShards().get(event.state().nodes().localNodeId());
+        RoutingNodes.RoutingNodeIterator routingNode = event.state().readOnlyRoutingNodes().routingNodeIter(event.state().nodes().localNodeId());
         if (routingNode == null) {
             return;
         }
@@ -276,8 +276,7 @@ public class IndicesClusterStateService extends AbstractLifecycleComponent<Indic
             }
             // now, go over and delete shards that needs to get deleted
             newShardIds.clear();
-            List<MutableShardRouting> shards = routingNode.shards();
-            for (MutableShardRouting shard : shards) {
+            for (MutableShardRouting shard : routingNode) {
                 if (shard.index().equals(index)) {
                     newShardIds.add(shard.id());
                 }
@@ -304,7 +303,7 @@ public class IndicesClusterStateService extends AbstractLifecycleComponent<Indic
 
     private void applyNewIndices(final ClusterChangedEvent event) {
         // we only create indices for shards that are allocated
-        RoutingNode routingNode = event.state().readOnlyRoutingNodes().nodesToShards().get(event.state().nodes().localNodeId());
+        RoutingNodes.RoutingNodeIterator routingNode = event.state().readOnlyRoutingNodes().routingNodeIter(event.state().nodes().localNodeId());
         if (routingNode == null) {
             return;
         }
@@ -502,14 +501,14 @@ public class IndicesClusterStateService extends AbstractLifecycleComponent<Indic
         }
 
         RoutingTable routingTable = event.state().routingTable();
-        RoutingNode routingNodes = event.state().readOnlyRoutingNodes().nodesToShards().get(event.state().nodes().localNodeId());
-        if (routingNodes == null) {
+        RoutingNodes.RoutingNodeIterator routingNode = event.state().readOnlyRoutingNodes().routingNodeIter(event.state().nodes().localNodeId());;
+        if (routingNode == null) {
             failedShards.clear();
             return;
         }
         DiscoveryNodes nodes = event.state().nodes();
 
-        for (final ShardRouting shardRouting : routingNodes) {
+        for (final ShardRouting shardRouting : routingNode) {
             final IndexService indexService = indicesService.indexService(shardRouting.index());
             if (indexService == null) {
                 // got deleted on us, ignore
@@ -564,8 +563,8 @@ public class IndicesClusterStateService extends AbstractLifecycleComponent<Indic
 
     private void cleanFailedShards(final ClusterChangedEvent event) {
         RoutingTable routingTable = event.state().routingTable();
-        RoutingNode routingNodes = event.state().readOnlyRoutingNodes().nodesToShards().get(event.state().nodes().localNodeId());
-        if (routingNodes == null) {
+        RoutingNodes.RoutingNodeIterator routingNode = event.state().readOnlyRoutingNodes().routingNodeIter(event.state().nodes().localNodeId());;
+        if (routingNode == null) {
             failedShards.clear();
             return;
         }

+ 6 - 10
src/test/java/org/elasticsearch/cluster/ack/AckTests.java

@@ -188,8 +188,7 @@ public class AckTests extends ElasticsearchIntegrationTest {
 
         for (Client client : clients()) {
             ClusterState clusterState = getLocalClusterState(client);
-            RoutingNode routingNode = clusterState.routingNodes().nodesToShards().get(moveAllocationCommand.fromNode());
-            for (MutableShardRouting mutableShardRouting : routingNode) {
+            for (MutableShardRouting mutableShardRouting : clusterState.routingNodes().routingNodeIter(moveAllocationCommand.fromNode())) {
                 //if the shard that we wanted to move is still on the same node, it must be relocating
                 if (mutableShardRouting.shardId().equals(moveAllocationCommand.shardId())) {
                     assertThat(mutableShardRouting.relocating(), equalTo(true));
@@ -197,9 +196,8 @@ public class AckTests extends ElasticsearchIntegrationTest {
 
             }
 
-            routingNode = clusterState.routingNodes().nodesToShards().get(moveAllocationCommand.toNode());
             boolean found = false;
-            for (MutableShardRouting mutableShardRouting : routingNode) {
+            for (MutableShardRouting mutableShardRouting :  clusterState.routingNodes().routingNodeIter(moveAllocationCommand.toNode())) {
                 if (mutableShardRouting.shardId().equals(moveAllocationCommand.shardId())) {
                     assertThat(mutableShardRouting.state(), anyOf(equalTo(ShardRoutingState.INITIALIZING), equalTo(ShardRoutingState.STARTED)));
                     found = true;
@@ -241,9 +239,8 @@ public class AckTests extends ElasticsearchIntegrationTest {
         //testing only on master with the latest cluster state as we didn't make any change thus we cannot guarantee that
         //all nodes hold the same cluster state version. We only know there was no need to change anything, thus no need for ack on this update.
         ClusterStateResponse clusterStateResponse = client().admin().cluster().prepareState().get();
-        RoutingNode routingNode = clusterStateResponse.getState().routingNodes().nodesToShards().get(moveAllocationCommand.fromNode());
         boolean found = false;
-        for (MutableShardRouting mutableShardRouting : routingNode) {
+        for (MutableShardRouting mutableShardRouting :  clusterStateResponse.getState().routingNodes().routingNodeIter(moveAllocationCommand.fromNode())) {
             //the shard that we wanted to move is still on the same node, as we had dryRun flag
             if (mutableShardRouting.shardId().equals(moveAllocationCommand.shardId())) {
                 assertThat(mutableShardRouting.started(), equalTo(true));
@@ -253,8 +250,7 @@ public class AckTests extends ElasticsearchIntegrationTest {
         }
         assertThat(found, equalTo(true));
 
-        routingNode = clusterStateResponse.getState().routingNodes().nodesToShards().get(moveAllocationCommand.toNode());
-        for (MutableShardRouting mutableShardRouting : routingNode) {
+        for (MutableShardRouting mutableShardRouting : clusterStateResponse.getState().routingNodes().routingNodeIter(moveAllocationCommand.toNode())) {
             if (mutableShardRouting.shardId().equals(moveAllocationCommand.shardId())) {
                 fail("shard [" + mutableShardRouting + "] shouldn't be on node [" + moveAllocationCommand.toString() + "]");
             }
@@ -281,11 +277,11 @@ public class AckTests extends ElasticsearchIntegrationTest {
         String toNodeId = null;
         MutableShardRouting shardToBeMoved = null;
         ClusterStateResponse clusterStateResponse = client().admin().cluster().prepareState().get();
-        for (RoutingNode routingNode : clusterStateResponse.getState().routingNodes().nodesToShards().values()) {
+        for (RoutingNode routingNode : clusterStateResponse.getState().routingNodes()) {
             if (routingNode.node().isDataNode()) {
                 if (fromNodeId == null && routingNode.numberOfOwningShards() > 0) {
                     fromNodeId = routingNode.nodeId();
-                    shardToBeMoved = routingNode.shards().get(randomInt(routingNode.shards().size()-1));
+                    shardToBeMoved = routingNode.get(randomInt(routingNode.size()-1));
                 } else {
                     toNodeId = routingNode.nodeId();
                 }

+ 10 - 10
src/test/java/org/elasticsearch/cluster/allocation/ClusterRerouteTests.java

@@ -76,7 +76,7 @@ public class ClusterRerouteTests extends ElasticsearchIntegrationTest {
                 .setDryRun(true)
                 .execute().actionGet().getState();
         assertThat(state.routingNodes().unassigned().size(), equalTo(1));
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).shards().get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
 
         logger.info("--> get the state, verify nothing changed because of the dry run");
         state = client().admin().cluster().prepareState().execute().actionGet().getState();
@@ -87,7 +87,7 @@ public class ClusterRerouteTests extends ElasticsearchIntegrationTest {
                 .add(new AllocateAllocationCommand(new ShardId("test", 0), node_1, true))
                 .execute().actionGet().getState();
         assertThat(state.routingNodes().unassigned().size(), equalTo(1));
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).shards().get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
 
         ClusterHealthResponse healthResponse = client().admin().cluster().prepareHealth().setWaitForEvents(Priority.LANGUID).setWaitForYellowStatus().execute().actionGet();
         assertThat(healthResponse.isTimedOut(), equalTo(false));
@@ -95,15 +95,15 @@ public class ClusterRerouteTests extends ElasticsearchIntegrationTest {
         logger.info("--> get the state, verify shard 1 primary allocated");
         state = client().admin().cluster().prepareState().execute().actionGet().getState();
         assertThat(state.routingNodes().unassigned().size(), equalTo(1));
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).shards().get(0).state(), equalTo(ShardRoutingState.STARTED));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).get(0).state(), equalTo(ShardRoutingState.STARTED));
 
         logger.info("--> move shard 1 primary from node1 to node2");
         state = client().admin().cluster().prepareReroute()
                 .add(new MoveAllocationCommand(new ShardId("test", 0), node_1, node_2))
                 .execute().actionGet().getState();
 
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).shards().get(0).state(), equalTo(ShardRoutingState.RELOCATING));
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_2).id()).shards().get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).get(0).state(), equalTo(ShardRoutingState.RELOCATING));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_2).id()).get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
 
 
         healthResponse = client().admin().cluster().prepareHealth().setWaitForEvents(Priority.LANGUID).setWaitForYellowStatus().setWaitForRelocatingShards(0).execute().actionGet();
@@ -112,7 +112,7 @@ public class ClusterRerouteTests extends ElasticsearchIntegrationTest {
         logger.info("--> get the state, verify shard 1 primary moved from node1 to node2");
         state = client().admin().cluster().prepareState().execute().actionGet().getState();
         assertThat(state.routingNodes().unassigned().size(), equalTo(1));
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_2).id()).shards().get(0).state(), equalTo(ShardRoutingState.STARTED));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_2).id()).get(0).state(), equalTo(ShardRoutingState.STARTED));
     }
 
     @Test
@@ -143,7 +143,7 @@ public class ClusterRerouteTests extends ElasticsearchIntegrationTest {
                 .add(new AllocateAllocationCommand(new ShardId("test", 0), node_1, true))
                 .execute().actionGet().getState();
         assertThat(state.routingNodes().unassigned().size(), equalTo(1));
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).shards().get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
 
         healthResponse = client().admin().cluster().prepareHealth().setWaitForEvents(Priority.LANGUID).setWaitForYellowStatus().execute().actionGet();
         assertThat(healthResponse.isTimedOut(), equalTo(false));
@@ -151,7 +151,7 @@ public class ClusterRerouteTests extends ElasticsearchIntegrationTest {
         logger.info("--> get the state, verify shard 1 primary allocated");
         state = client().admin().cluster().prepareState().execute().actionGet().getState();
         assertThat(state.routingNodes().unassigned().size(), equalTo(1));
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).shards().get(0).state(), equalTo(ShardRoutingState.STARTED));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).get(0).state(), equalTo(ShardRoutingState.STARTED));
 
         client().prepareIndex("test", "type", "1").setSource("field", "value").setRefresh(true).execute().actionGet();
 
@@ -176,7 +176,7 @@ public class ClusterRerouteTests extends ElasticsearchIntegrationTest {
                 .add(new AllocateAllocationCommand(new ShardId("test", 0), node_1, true))
                 .execute().actionGet().getState();
         assertThat(state.routingNodes().unassigned().size(), equalTo(1));
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).shards().get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
 
         healthResponse = client().admin().cluster().prepareHealth().setWaitForEvents(Priority.LANGUID).setWaitForYellowStatus().execute().actionGet();
         assertThat(healthResponse.isTimedOut(), equalTo(false));
@@ -184,7 +184,7 @@ public class ClusterRerouteTests extends ElasticsearchIntegrationTest {
         logger.info("--> get the state, verify shard 1 primary allocated");
         state = client().admin().cluster().prepareState().execute().actionGet().getState();
         assertThat(state.routingNodes().unassigned().size(), equalTo(1));
-        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).shards().get(0).state(), equalTo(ShardRoutingState.STARTED));
+        assertThat(state.routingNodes().node(state.nodes().resolveNode(node_1).id()).get(0).state(), equalTo(ShardRoutingState.STARTED));
 
     }
 

+ 9 - 15
src/test/java/org/elasticsearch/cluster/allocation/SimpleAllocationTests.java

@@ -24,9 +24,6 @@ import org.elasticsearch.common.Priority;
 import org.elasticsearch.test.ElasticsearchIntegrationTest;
 import org.junit.Test;
 
-import java.util.Map;
-import java.util.Map.Entry;
-
 import static org.elasticsearch.common.settings.ImmutableSettings.settingsBuilder;
 import static org.hamcrest.Matchers.equalTo;
 
@@ -45,10 +42,9 @@ public class SimpleAllocationTests extends ElasticsearchIntegrationTest {
         ensureGreen();            
         ClusterState state = client().admin().cluster().prepareState().execute().actionGet().getState();
         assertThat(state.routingNodes().unassigned().size(), equalTo(0));
-        Map<String, RoutingNode> nodesToShards = state.routingNodes().getNodesToShards();
-        for (Entry<String, RoutingNode> entry : nodesToShards.entrySet()) {
-            if (!entry.getValue().shards().isEmpty()) { 
-                assertThat(entry.getValue().shards().size(), equalTo(2));
+        for (RoutingNode node : state.routingNodes()) {
+            if (!node.isEmpty()) {
+                assertThat(node.size(), equalTo(2));
             }
         }
         client().admin().indices().prepareUpdateSettings("test").setSettings(settingsBuilder().put("index.number_of_replicas", 0)).execute().actionGet();
@@ -56,10 +52,9 @@ public class SimpleAllocationTests extends ElasticsearchIntegrationTest {
         state = client().admin().cluster().prepareState().execute().actionGet().getState();
 
         assertThat(state.routingNodes().unassigned().size(), equalTo(0));
-        nodesToShards = state.routingNodes().getNodesToShards();
-        for (Entry<String, RoutingNode> entry : nodesToShards.entrySet()) {
-            if (!entry.getValue().shards().isEmpty()) {
-                assertThat(entry.getValue().shards().size(), equalTo(1));
+        for (RoutingNode node : state.routingNodes()) {
+            if (!node.isEmpty()) {
+                assertThat(node.size(), equalTo(1));
             }
         }
         
@@ -76,10 +71,9 @@ public class SimpleAllocationTests extends ElasticsearchIntegrationTest {
         state = client().admin().cluster().prepareState().execute().actionGet().getState();
 
         assertThat(state.routingNodes().unassigned().size(), equalTo(0));
-        nodesToShards = state.routingNodes().getNodesToShards();
-        for (Entry<String, RoutingNode> entry : nodesToShards.entrySet()) {
-            if (!entry.getValue().shards().isEmpty()) {
-                assertThat(entry.getValue().shards().size(), equalTo(4));
+        for (RoutingNode node : state.routingNodes()) {
+            if (!node.isEmpty()) {
+                assertThat(node.size(), equalTo(4));
             }
         }
     }

+ 11 - 11
src/test/java/org/elasticsearch/cluster/routing/allocation/AddIncrementallyTests.java

@@ -34,12 +34,12 @@ public class AddIncrementallyTests extends ElasticsearchTestCase {
 
         ClusterState clusterState = initCluster(service, 1, 3, 3, 1);
         assertThat(clusterState.routingNodes().node("node0").shardsWithState(STARTED).size(), Matchers.equalTo(9));
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(9));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(9));
         int nodeOffset = 1;
         clusterState = addNodes(clusterState, service, 1, nodeOffset++);
         assertThat(clusterState.routingNodes().node("node0").shardsWithState(STARTED).size(), Matchers.equalTo(9));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), Matchers.equalTo(9));
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(0));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(0));
         assertNumIndexShardsPerNode(clusterState, Matchers.equalTo(3));
         clusterState = addNodes(clusterState, service, 1, nodeOffset++);
         assertNumIndexShardsPerNode(clusterState, Matchers.equalTo(2));
@@ -50,23 +50,23 @@ public class AddIncrementallyTests extends ElasticsearchTestCase {
         assertNumIndexShardsPerNode(clusterState, Matchers.equalTo(2));
 
         clusterState = addIndex(clusterState, service, 3, 2, 3);
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(2));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(2));
         assertNumIndexShardsPerNode(clusterState, "test3", Matchers.equalTo(2));
         assertNumIndexShardsPerNode(clusterState, Matchers.lessThanOrEqualTo(2));
 
         clusterState = addIndex(clusterState, service, 4, 2, 3);
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(4));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(4));
         assertNumIndexShardsPerNode(clusterState, "test4", Matchers.equalTo(2));
         assertNumIndexShardsPerNode(clusterState, Matchers.lessThanOrEqualTo(2));
         clusterState = addNodes(clusterState, service, 1, nodeOffset++);
         assertNumIndexShardsPerNode(clusterState, Matchers.lessThanOrEqualTo(2));
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(0));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(0));
         clusterState = removeNodes(clusterState, service, 1);
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(4));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(4));
         assertNumIndexShardsPerNode(clusterState, Matchers.lessThanOrEqualTo(2));
         clusterState = addNodes(clusterState, service, 1, nodeOffset++);
         assertNumIndexShardsPerNode(clusterState, Matchers.lessThanOrEqualTo(2));
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(0));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(0));
         logger.debug("ClusterState: {}", clusterState.getRoutingNodes().prettyPrint());
     }
 
@@ -79,12 +79,12 @@ public class AddIncrementallyTests extends ElasticsearchTestCase {
 
         ClusterState clusterState = initCluster(service, 1, 3, 3, 1);
         assertThat(clusterState.routingNodes().node("node0").shardsWithState(STARTED).size(), Matchers.equalTo(9));
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(9));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(9));
         int nodeOffset = 1;
         clusterState = addNodes(clusterState, service, 1, nodeOffset++);
         assertThat(clusterState.routingNodes().node("node0").shardsWithState(STARTED).size(), Matchers.equalTo(9));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), Matchers.equalTo(9));
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(0));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(0));
         assertNumIndexShardsPerNode(clusterState, Matchers.equalTo(3));
 
         logger.info("now, start one more node, check that rebalancing will happen because we set it to always");
@@ -150,12 +150,12 @@ public class AddIncrementallyTests extends ElasticsearchTestCase {
 
         ClusterState clusterState = initCluster(service, 1, 3, 3, 1);
         assertThat(clusterState.routingNodes().node("node0").shardsWithState(STARTED).size(), Matchers.equalTo(9));
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(9));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(9));
         int nodeOffset = 1;
         clusterState = addNodes(clusterState, service, 1, nodeOffset++);
         assertThat(clusterState.routingNodes().node("node0").shardsWithState(STARTED).size(), Matchers.equalTo(9));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), Matchers.equalTo(9));
-        assertThat(clusterState.routingNodes().getUnassigned().size(), Matchers.equalTo(0));
+        assertThat(clusterState.routingNodes().unassigned().size(), Matchers.equalTo(0));
         assertNumIndexShardsPerNode(clusterState, Matchers.equalTo(3));
 
         logger.info("now, start one more node, check that rebalancing will happen because we set it to always");

+ 34 - 34
src/test/java/org/elasticsearch/cluster/routing/allocation/AllocationCommandsTests.java

@@ -87,15 +87,15 @@ public class AllocationCommandsTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(new MoveAllocationCommand(new ShardId("test", 0), existingNodeId, toNodeId)));
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node(existingNodeId).shards().get(0).state(), equalTo(ShardRoutingState.RELOCATING));
-        assertThat(clusterState.routingNodes().node(toNodeId).shards().get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
+        assertThat(clusterState.routingNodes().node(existingNodeId).get(0).state(), equalTo(ShardRoutingState.RELOCATING));
+        assertThat(clusterState.routingNodes().node(toNodeId).get(0).state(), equalTo(ShardRoutingState.INITIALIZING));
 
         logger.info("finish moving the shard");
         rerouteResult = allocation.applyStartedShards(clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
-        assertThat(clusterState.routingNodes().node(existingNodeId).shards().isEmpty(), equalTo(true));
-        assertThat(clusterState.routingNodes().node(toNodeId).shards().get(0).state(), equalTo(ShardRoutingState.STARTED));
+        assertThat(clusterState.routingNodes().node(existingNodeId).isEmpty(), equalTo(true));
+        assertThat(clusterState.routingNodes().node(toNodeId).get(0).state(), equalTo(ShardRoutingState.STARTED));
     }
 
     @Test
@@ -135,16 +135,16 @@ public class AllocationCommandsTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(new AllocateAllocationCommand(new ShardId("test", 0), "node1", true)));
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(INITIALIZING).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(0));
 
         logger.info("--> start the primary shard");
         rerouteResult = allocation.applyStartedShards(clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(0));
 
         logger.info("--> allocate the replica shard on the primary shard node, should fail");
         try {
@@ -157,18 +157,18 @@ public class AllocationCommandsTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(new AllocateAllocationCommand(new ShardId("test", 0), "node2", false)));
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node2").shardsWithState(INITIALIZING).size(), equalTo(1));
 
 
         logger.info("--> start the replica shard");
         rerouteResult = allocation.applyStartedShards(clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node2").shardsWithState(STARTED).size(), equalTo(1));
 
         logger.info("--> verify that we fail when there are no unassigned shards");
@@ -209,9 +209,9 @@ public class AllocationCommandsTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(new AllocateAllocationCommand(new ShardId("test", 0), "node1", true)));
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(INITIALIZING).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(0));
 
         logger.info("--> cancel primary allocation, make sure it fails...");
         try {
@@ -223,9 +223,9 @@ public class AllocationCommandsTests extends ElasticsearchTestCase {
         logger.info("--> start the primary shard");
         rerouteResult = allocation.applyStartedShards(clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(0));
 
         logger.info("--> cancel primary allocation, make sure it fails...");
         try {
@@ -238,27 +238,27 @@ public class AllocationCommandsTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(new AllocateAllocationCommand(new ShardId("test", 0), "node2", false)));
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node2").shardsWithState(INITIALIZING).size(), equalTo(1));
 
         logger.info("--> cancel the relocation allocation");
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(new CancelAllocationCommand(new ShardId("test", 0), "node2", false)));
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));
-        assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node3").size(), equalTo(0));
 
         logger.info("--> allocate the replica shard on on the second node");
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(new AllocateAllocationCommand(new ShardId("test", 0), "node2", false)));
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node2").shardsWithState(INITIALIZING).size(), equalTo(1));
 
         logger.info("--> cancel the primary being replicated, make sure it fails");
@@ -271,34 +271,34 @@ public class AllocationCommandsTests extends ElasticsearchTestCase {
         logger.info("--> start the replica shard");
         rerouteResult = allocation.applyStartedShards(clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node2").shardsWithState(STARTED).size(), equalTo(1));
 
         logger.info("--> cancel allocation of the replica shard");
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(new CancelAllocationCommand(new ShardId("test", 0), "node2", false)));
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));
-        assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node3").size(), equalTo(0));
 
         logger.info("--> allocate the replica shard on on the second node");
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(new AllocateAllocationCommand(new ShardId("test", 0), "node2", false)));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
         assertThat(rerouteResult.changed(), equalTo(true));
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node2").shardsWithState(INITIALIZING).size(), equalTo(1));
         logger.info("--> start the replica shard");
         rerouteResult = allocation.applyStartedShards(clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
         assertThat(clusterState.routingNodes().node("node2").shardsWithState(STARTED).size(), equalTo(1));
 
         logger.info("--> cancel the primary allocation (with allow_primary set to true)");
@@ -306,8 +306,8 @@ public class AllocationCommandsTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
         assertThat(rerouteResult.changed(), equalTo(true));
         assertThat(clusterState.routingNodes().node("node2").shardsWithState(STARTED).get(0).primary(), equalTo(true));
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(0));
-        assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node3").size(), equalTo(0));
     }
 
     @Test

+ 10 - 10
src/test/java/org/elasticsearch/cluster/routing/allocation/AwarenessAllocationTests.java

@@ -348,9 +348,9 @@ public class AwarenessAllocationTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
         assertThat(clusterState.routingNodes().shardsWithState(ShardRoutingState.STARTED).size(), equalTo(20));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(10));
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(5));
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(5));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(10));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(5));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(5));
 
         logger.info("--> do another reroute, make sure nothing moves");
         assertThat(strategy.reroute(clusterState).routingTable(), sameInstance(clusterState.routingTable()));
@@ -370,10 +370,10 @@ public class AwarenessAllocationTests extends ElasticsearchTestCase {
             clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
         }
         assertThat(clusterState.routingNodes().shardsWithState(ShardRoutingState.STARTED).size(), equalTo(20));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(5));
-        assertThat(clusterState.getRoutingNodes().node("node4").shards().size(), equalTo(5));
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(5));
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(5));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(5));
+        assertThat(clusterState.getRoutingNodes().node("node4").size(), equalTo(5));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(5));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(5));
 
         logger.info("--> do another reroute, make sure nothing moves");
         assertThat(strategy.reroute(clusterState).routingTable(), sameInstance(clusterState.routingTable()));
@@ -821,8 +821,8 @@ public class AwarenessAllocationTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
         assertThat(clusterState.routingNodes().shardsWithState(ShardRoutingState.STARTED).size(), equalTo(10));
-        assertThat(clusterState.getRoutingNodes().node("A-1").shards().size(), equalTo(2));
-        assertThat(clusterState.getRoutingNodes().node("A-0").shards().size(), equalTo(3));
-        assertThat(clusterState.getRoutingNodes().node("B-0").shards().size(), equalTo(5));
+        assertThat(clusterState.getRoutingNodes().node("A-1").size(), equalTo(2));
+        assertThat(clusterState.getRoutingNodes().node("A-0").size(), equalTo(3));
+        assertThat(clusterState.getRoutingNodes().node("B-0").size(), equalTo(5));
     }
 }

+ 11 - 13
src/test/java/org/elasticsearch/cluster/routing/allocation/BalanceConfigurationTests.java

@@ -41,8 +41,6 @@ import org.elasticsearch.test.ElasticsearchTestCase;
 import org.hamcrest.Matchers;
 import org.junit.Test;
 
-import java.util.List;
-
 import static org.elasticsearch.cluster.routing.ShardRoutingState.INITIALIZING;
 import static org.elasticsearch.cluster.routing.ShardRoutingState.STARTED;
 import static org.elasticsearch.cluster.routing.allocation.RoutingAllocationTests.newNode;
@@ -403,43 +401,43 @@ public class BalanceConfigurationTests extends ElasticsearchTestCase {
              */
             @Override
             public boolean allocateUnassigned(RoutingAllocation allocation) {
-                List<MutableShardRouting> unassigned = allocation.routingNodes().unassigned();
+                RoutingNodes.UnassignedShards unassigned = allocation.routingNodes().unassigned();
                 boolean changed = !unassigned.isEmpty();
                 for (MutableShardRouting sr : unassigned) {
                     switch (sr.id()) {
                         case 0:
                             if (sr.primary()) {
-                                allocation.routingNodes().assignShardToNode( sr, "node1" );
+                                allocation.routingNodes().assign(sr, "node1");
                             } else {
-                                allocation.routingNodes().assignShardToNode( sr, "node0" );
+                                allocation.routingNodes().assign(sr, "node0");
                             }
                             break;
                         case 1:
                             if (sr.primary()) {
-                                allocation.routingNodes().assignShardToNode( sr, "node1" );
+                                allocation.routingNodes().assign(sr, "node1");
                             } else {
-                                allocation.routingNodes().assignShardToNode( sr, "node2" );
+                                allocation.routingNodes().assign(sr, "node2");
                             }
                             break;
                         case 2:
                             if (sr.primary()) {
-                                allocation.routingNodes().assignShardToNode( sr, "node3" );
+                                allocation.routingNodes().assign(sr, "node3");
                             } else {
-                                allocation.routingNodes().assignShardToNode( sr, "node2" );
+                                allocation.routingNodes().assign(sr, "node2");
                             }
                             break;
                         case 3:
                             if (sr.primary()) {
-                                allocation.routingNodes().assignShardToNode( sr, "node3" );
+                                allocation.routingNodes().assign(sr, "node3");
                             } else {
-                                allocation.routingNodes().assignShardToNode( sr, "node1" );
+                                allocation.routingNodes().assign(sr, "node1");
                             }
                             break;
                         case 4:
                             if (sr.primary()) {
-                                allocation.routingNodes().assignShardToNode( sr, "node2" );
+                                allocation.routingNodes().assign(sr, "node2");
                             } else {
-                                allocation.routingNodes().assignShardToNode( sr, "node0" );
+                                allocation.routingNodes().assign(sr, "node0");
                             }
                             break;
                     }

+ 17 - 17
src/test/java/org/elasticsearch/cluster/routing/allocation/ClusterRebalanceRoutingTests.java

@@ -121,8 +121,8 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
         routingNodes = clusterState.routingNodes();
 
-        assertThat(routingNodes.node("node3").shards().size(), equalTo(1));
-        assertThat(routingNodes.node("node3").shards().get(0).shardId().index().name(), equalTo("test1"));
+        assertThat(routingNodes.node("node3").size(), equalTo(1));
+        assertThat(routingNodes.node("node3").get(0).shardId().index().name(), equalTo("test1"));
     }
 
 
@@ -149,13 +149,13 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
         for (int i = 0; i < routingTable.index("test1").shards().size(); i++) {
-            assertThat(routingTable.index("test1").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test1").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test1").shard(i).primaryShard().state(), equalTo(INITIALIZING));
             assertThat(routingTable.index("test1").shard(i).replicaShards().get(0).state(), equalTo(UNASSIGNED));
         }
 
         for (int i = 0; i < routingTable.index("test2").shards().size(); i++) {
-            assertThat(routingTable.index("test2").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test2").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test2").shard(i).primaryShard().state(), equalTo(INITIALIZING));
             assertThat(routingTable.index("test2").shard(i).replicaShards().get(0).state(), equalTo(UNASSIGNED));
         }
@@ -168,13 +168,13 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         routingNodes = clusterState.routingNodes();
 
         for (int i = 0; i < routingTable.index("test1").shards().size(); i++) {
-            assertThat(routingTable.index("test1").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test1").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test1").shard(i).primaryShard().state(), equalTo(STARTED));
             assertThat(routingTable.index("test1").shard(i).replicaShards().get(0).state(), equalTo(INITIALIZING));
         }
 
         for (int i = 0; i < routingTable.index("test2").shards().size(); i++) {
-            assertThat(routingTable.index("test2").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test2").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test2").shard(i).primaryShard().state(), equalTo(INITIALIZING));
             assertThat(routingTable.index("test2").shard(i).replicaShards().get(0).state(), equalTo(UNASSIGNED));
         }
@@ -187,13 +187,13 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         routingNodes = clusterState.routingNodes();
 
         for (int i = 0; i < routingTable.index("test1").shards().size(); i++) {
-            assertThat(routingTable.index("test1").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test1").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test1").shard(i).primaryShard().state(), equalTo(STARTED));
             assertThat(routingTable.index("test1").shard(i).replicaShards().get(0).state(), equalTo(STARTED));
         }
 
         for (int i = 0; i < routingTable.index("test2").shards().size(); i++) {
-            assertThat(routingTable.index("test2").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test2").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test2").shard(i).primaryShard().state(), equalTo(INITIALIZING));
             assertThat(routingTable.index("test2").shard(i).replicaShards().get(0).state(), equalTo(UNASSIGNED));
         }
@@ -206,13 +206,13 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         routingNodes = clusterState.routingNodes();
 
         for (int i = 0; i < routingTable.index("test1").shards().size(); i++) {
-            assertThat(routingTable.index("test1").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test1").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test1").shard(i).primaryShard().state(), equalTo(STARTED));
             assertThat(routingTable.index("test1").shard(i).replicaShards().get(0).state(), equalTo(STARTED));
         }
 
         for (int i = 0; i < routingTable.index("test2").shards().size(); i++) {
-            assertThat(routingTable.index("test2").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test2").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test2").shard(i).primaryShard().state(), equalTo(STARTED));
             assertThat(routingTable.index("test2").shard(i).replicaShards().get(0).state(), equalTo(INITIALIZING));
         }
@@ -226,8 +226,8 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
         routingNodes = clusterState.routingNodes();
 
-        assertThat(routingNodes.node("node3").shards().size(), equalTo(1));
-        assertThat(routingNodes.node("node3").shards().get(0).shardId().index().name(), equalTo("test1"));
+        assertThat(routingNodes.node("node3").size(), equalTo(1));
+        assertThat(routingNodes.node("node3").get(0).shardId().index().name(), equalTo("test1"));
     }
 
     @Test
@@ -311,7 +311,7 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
         routingNodes = clusterState.routingNodes();
 
-        assertThat(routingNodes.node("node3").shards().isEmpty(), equalTo(true));
+        assertThat(routingNodes.node("node3").isEmpty(), equalTo(true));
     }
 
     @Test
@@ -433,8 +433,8 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
         routingNodes = clusterState.routingNodes();
 
-        assertThat(routingNodes.node("node3").shards().size(), equalTo(1));
-        assertThat(routingNodes.node("node3").shards().get(0).shardId().index().name(), equalTo("test1"));
+        assertThat(routingNodes.node("node3").size(), equalTo(1));
+        assertThat(routingNodes.node("node3").get(0).shardId().index().name(), equalTo("test1"));
     }
 
     @Test
@@ -518,7 +518,7 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
         routingNodes = clusterState.routingNodes();
 
-        assertThat(routingNodes.node("node3").shards().isEmpty(), equalTo(true));
+        assertThat(routingNodes.node("node3").isEmpty(), equalTo(true));
     }
 
     @Test
@@ -621,6 +621,6 @@ public class ClusterRebalanceRoutingTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
         routingNodes = clusterState.routingNodes();
 
-        assertThat(routingNodes.node("node3").shards().isEmpty(), equalTo(true));
+        assertThat(routingNodes.node("node3").isEmpty(), equalTo(true));
     }
 }

+ 32 - 32
src/test/java/org/elasticsearch/cluster/routing/allocation/DeadNodesAllocationTests.java

@@ -75,10 +75,10 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
         logger.info("--> verifying all is allocated");
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node1").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").get(0).state(), equalTo(STARTED));
 
         logger.info("--> fail node with primary");
         String nodeIdToFail = clusterState.routingTable().index("test").shard(0).primaryShard().currentNodeId();
@@ -90,8 +90,8 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState);
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
-        assertThat(clusterState.routingNodes().node(nodeIdRemaining).shards().get(0).primary(), equalTo(true));
-        assertThat(clusterState.routingNodes().node(nodeIdRemaining).shards().get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node(nodeIdRemaining).get(0).primary(), equalTo(true));
+        assertThat(clusterState.routingNodes().node(nodeIdRemaining).get(0).state(), equalTo(STARTED));
     }
 
     @Test
@@ -127,10 +127,10 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
         logger.info("--> verifying all is allocated");
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node1").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").get(0).state(), equalTo(STARTED));
 
         logger.info("--> adding additional node");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder(clusterState.nodes())
@@ -139,11 +139,11 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState);
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node1").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node3").size(), equalTo(0));
 
         String origPrimaryNodeId = clusterState.routingTable().index("test").shard(0).primaryShard().currentNodeId();
         String origReplicaNodeId = clusterState.routingTable().index("test").shard(0).replicaShards().get(0).currentNodeId();
@@ -154,8 +154,8 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         );
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).shards().get(0).state(), equalTo(RELOCATING));
-        assertThat(clusterState.routingNodes().node("node3").shards().get(0).state(), equalTo(INITIALIZING));
+        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).get(0).state(), equalTo(RELOCATING));
+        assertThat(clusterState.routingNodes().node("node3").get(0).state(), equalTo(INITIALIZING));
 
         logger.info("--> fail primary shard recovering instance on node3 being initialized by killing node3");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder()
@@ -165,8 +165,8 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState);
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
-        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node(origReplicaNodeId).shards().get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node(origReplicaNodeId).get(0).state(), equalTo(STARTED));
     }
 
     @Test
@@ -202,10 +202,10 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
         logger.info("--> verifying all is allocated");
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node1").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").get(0).state(), equalTo(STARTED));
 
         logger.info("--> adding additional node");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder(clusterState.nodes())
@@ -214,11 +214,11 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState);
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node1").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node3").size(), equalTo(0));
 
         String origPrimaryNodeId = clusterState.routingTable().index("test").shard(0).primaryShard().currentNodeId();
         String origReplicaNodeId = clusterState.routingTable().index("test").shard(0).replicaShards().get(0).currentNodeId();
@@ -229,8 +229,8 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         );
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).shards().get(0).state(), equalTo(RELOCATING));
-        assertThat(clusterState.routingNodes().node("node3").shards().get(0).state(), equalTo(INITIALIZING));
+        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).get(0).state(), equalTo(RELOCATING));
+        assertThat(clusterState.routingNodes().node("node3").get(0).state(), equalTo(INITIALIZING));
 
         logger.info("--> fail primary shard recovering instance on 'origPrimaryNodeId' being relocated");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder()
@@ -240,7 +240,7 @@ public class DeadNodesAllocationTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState);
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
-        assertThat(clusterState.routingNodes().node(origReplicaNodeId).shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node3").shards().get(0).state(), equalTo(INITIALIZING));
+        assertThat(clusterState.routingNodes().node(origReplicaNodeId).get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node3").get(0).state(), equalTo(INITIALIZING));
     }
 }

+ 24 - 23
src/test/java/org/elasticsearch/cluster/routing/allocation/FailedShardsRoutingTests.java

@@ -79,10 +79,10 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
         logger.info("--> verifying all is allocated");
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node1").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").get(0).state(), equalTo(STARTED));
 
         logger.info("--> adding additional node");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder(clusterState.nodes())
@@ -91,11 +91,11 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
         rerouteResult = allocation.reroute(clusterState);
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
-        assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node1").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.routingNodes().node("node2").shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node("node1").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node1").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.routingNodes().node("node2").get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node3").size(), equalTo(0));
 
         String origPrimaryNodeId = clusterState.routingTable().index("test").shard(0).primaryShard().currentNodeId();
         String origReplicaNodeId = clusterState.routingTable().index("test").shard(0).replicaShards().get(0).currentNodeId();
@@ -106,15 +106,15 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
         );
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).shards().get(0).state(), equalTo(RELOCATING));
-        assertThat(clusterState.routingNodes().node("node3").shards().get(0).state(), equalTo(INITIALIZING));
+        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).get(0).state(), equalTo(RELOCATING));
+        assertThat(clusterState.routingNodes().node("node3").get(0).state(), equalTo(INITIALIZING));
 
         logger.info("--> fail primary shard recovering instance on node3 being initialized");
-        rerouteResult = allocation.applyFailedShard(clusterState, new ImmutableShardRouting(clusterState.routingNodes().node("node3").shards().get(0)));
+        rerouteResult = allocation.applyFailedShard(clusterState, new ImmutableShardRouting(clusterState.routingNodes().node("node3").get(0)));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
-        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).shards().get(0).state(), equalTo(STARTED));
-        assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));
+        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node("node3").size(), equalTo(0));
 
         logger.info("--> moving primary shard to node3");
         rerouteResult = allocation.reroute(clusterState, new AllocationCommands(
@@ -122,15 +122,15 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
         );
         assertThat(rerouteResult.changed(), equalTo(true));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
-        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).shards().get(0).state(), equalTo(RELOCATING));
-        assertThat(clusterState.routingNodes().node("node3").shards().get(0).state(), equalTo(INITIALIZING));
+        assertThat(clusterState.routingNodes().node(origPrimaryNodeId).get(0).state(), equalTo(RELOCATING));
+        assertThat(clusterState.routingNodes().node("node3").get(0).state(), equalTo(INITIALIZING));
 
         logger.info("--> fail primary shard recovering instance on node1 being relocated");
-        rerouteResult = allocation.applyFailedShard(clusterState, new ImmutableShardRouting(clusterState.routingNodes().node(origPrimaryNodeId).shards().get(0)));
+        rerouteResult = allocation.applyFailedShard(clusterState, new ImmutableShardRouting(clusterState.routingNodes().node(origPrimaryNodeId).get(0)));
         clusterState = ClusterState.builder(clusterState).routingTable(rerouteResult.routingTable()).build();
 
         // check promotion of replica to primary
-        assertThat(clusterState.routingNodes().node(origReplicaNodeId).shards().get(0).state(), equalTo(STARTED));
+        assertThat(clusterState.routingNodes().node(origReplicaNodeId).get(0).state(), equalTo(STARTED));
         assertThat(clusterState.routingTable().index("test").shard(0).primaryShard().currentNodeId(), equalTo(origReplicaNodeId));
         assertThat(clusterState.routingTable().index("test").shard(0).replicaShards().get(0).currentNodeId(), anyOf(equalTo(origPrimaryNodeId), equalTo("node3")));
     }
@@ -162,6 +162,7 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
 
         logger.info("Start the shards (primaries)");
         RoutingNodes routingNodes = clusterState.routingNodes();
+
         prevRoutingTable = routingTable;
         routingTable = strategy.applyStartedShards(clusterState, routingNodes.shardsWithState(INITIALIZING)).routingTable();
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
@@ -170,7 +171,7 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
         assertThat(routingTable.index("test").shards().size(), equalTo(1));
         for (int i = 0; i < routingTable.index("test").shards().size(); i++) {
             assertThat(routingTable.index("test").shard(i).size(), equalTo(2));
-            assertThat(routingTable.index("test").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test").shard(i).primaryShard().state(), equalTo(STARTED));
             assertThat(routingTable.index("test").shard(i).primaryShard().currentNodeId(), anyOf(equalTo("node1"), equalTo("node2")));
             assertThat(routingTable.index("test").shard(i).replicaShards().size(), equalTo(1));
@@ -188,7 +189,7 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
         assertThat(routingTable.index("test").shards().size(), equalTo(1));
         for (int i = 0; i < routingTable.index("test").shards().size(); i++) {
             assertThat(routingTable.index("test").shard(i).size(), equalTo(2));
-            assertThat(routingTable.index("test").shard(i).shards().size(), equalTo(2));
+            assertThat(routingTable.index("test").shard(i).size(), equalTo(2));
             assertThat(routingTable.index("test").shard(i).primaryShard().state(), equalTo(STARTED));
             assertThat(routingTable.index("test").shard(i).primaryShard().currentNodeId(), anyOf(equalTo("node1"), equalTo("node2")));
             assertThat(routingTable.index("test").shard(i).replicaShards().size(), equalTo(1));
@@ -205,7 +206,7 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
 
         assertThat(prevRoutingTable != routingTable, equalTo(true));
         assertThat(routingTable.index("test").shard(0).size(), equalTo(2));
-        assertThat(routingTable.index("test").shard(0).shards().size(), equalTo(2));
+        assertThat(routingTable.index("test").shard(0).size(), equalTo(2));
         assertThat(routingTable.index("test").shard(0).primaryShard().currentNodeId(), not(equalTo(shardToFail.currentNodeId())));
         assertThat(routingTable.index("test").shard(0).primaryShard().state(), equalTo(STARTED));
         assertThat(routingTable.index("test").shard(0).primaryShard().currentNodeId(), anyOf(equalTo("node1"), equalTo("node2")));
@@ -408,7 +409,7 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
 
 
         logger.info("Fail the shards on node 3");
-        ShardRouting shardToFail = routingNodes.node("node3").shards().get(0);
+        ShardRouting shardToFail = routingNodes.node("node3").get(0);
         routingNodes = clusterState.routingNodes();
         prevRoutingTable = routingTable;
         routingTable = strategy.applyFailedShard(clusterState, new ImmutableShardRouting(shardToFail)).routingTable();
@@ -423,6 +424,6 @@ public class FailedShardsRoutingTests extends ElasticsearchTestCase {
         assertThat(routingNodes.node("node2").numberOfShardsWithState(STARTED), lessThan(3));
         assertThat(routingNodes.node("node3").numberOfShardsWithState(INITIALIZING), equalTo(1));
         // make sure the failedShard is not INITIALIZING again on node3
-        assertThat(routingNodes.node("node3").shards().get(0).shardId(), not(equalTo(shardToFail.shardId())));
+        assertThat(routingNodes.node("node3").get(0).shardId(), not(equalTo(shardToFail.shardId())));
     }
 }

+ 1 - 1
src/test/java/org/elasticsearch/cluster/routing/allocation/RandomAllocationDeciderTests.java

@@ -156,7 +156,7 @@ public class RandomAllocationDeciderTests extends ElasticsearchTestCase {
             if (clusterState.getRoutingNodes().node("NODE_" + i) == null) {
                 continue;
             }
-            assertThat(clusterState.getRoutingNodes().node("NODE_" + i).shards().size(), Matchers.anyOf(
+            assertThat(clusterState.getRoutingNodes().node("NODE_" + i).size(), Matchers.anyOf(
                     Matchers.anyOf(equalTo((shards / numNodes) + 1), equalTo((shards / numNodes) - 1), equalTo((shards / numNodes))),
                     Matchers.allOf(Matchers.greaterThanOrEqualTo(lowerBound), Matchers.lessThanOrEqualTo(upperBound))));
         }

+ 1 - 1
src/test/java/org/elasticsearch/cluster/routing/allocation/RebalanceAfterActiveTests.java

@@ -145,7 +145,7 @@ public class RebalanceAfterActiveTests extends ElasticsearchTestCase {
         assertThat(routingTable.shardsWithState(STARTED).size(), equalTo(10));
         // make sure we have an even relocation
         for (RoutingNode routingNode : routingNodes) {
-            assertThat(routingNode.shards().size(), equalTo(1));
+            assertThat(routingNode.size(), equalTo(1));
         }
     }
 }

+ 6 - 6
src/test/java/org/elasticsearch/cluster/routing/allocation/RoutingNodesIntegrityTests.java

@@ -32,11 +32,11 @@ import org.elasticsearch.common.settings.ImmutableSettings;
 import org.elasticsearch.test.ElasticsearchTestCase;
 import org.junit.Test;
 
-import static org.elasticsearch.cluster.routing.ShardRoutingState.*;
+import static org.elasticsearch.cluster.routing.ShardRoutingState.INITIALIZING;
+import static org.elasticsearch.cluster.routing.ShardRoutingState.STARTED;
 import static org.elasticsearch.cluster.routing.allocation.RoutingAllocationTests.newNode;
 import static org.elasticsearch.common.settings.ImmutableSettings.settingsBuilder;
 import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.nullValue;
 
 /**
  *
@@ -95,7 +95,7 @@ public class RoutingNodesIntegrityTests extends ElasticsearchTestCase {
         routingTable = strategy.applyStartedShards(clusterState, routingNodes.shardsWithState(INITIALIZING)).routingTable();
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
         routingNodes = clusterState.routingNodes();
-        
+
         logger.info("Reroute, nothing should change");
         prevRoutingTable = routingTable;
         routingTable = strategy.reroute(clusterState).routingTable();
@@ -230,7 +230,7 @@ public class RoutingNodesIntegrityTests extends ElasticsearchTestCase {
         logger.info("Adding three node and performing rerouting");
         clusterState = ClusterState.builder(clusterState)
                 .nodes(DiscoveryNodes.builder().put(newNode("node1")).put(newNode("node2")).put(newNode("node3"))).build();
-        
+
         RoutingNodes routingNodes = clusterState.routingNodes();
         assertThat(assertShardStats(routingNodes), equalTo(true));
         assertThat(routingNodes.hasInactiveShards(), equalTo(false));
@@ -370,7 +370,7 @@ public class RoutingNodesIntegrityTests extends ElasticsearchTestCase {
         assertThat(routingNodes.node("node1").shardsWithState("test1", STARTED).size(), equalTo(2));
         assertThat(routingNodes.node("node2").shardsWithState("test1", STARTED).size(), equalTo(2));
         assertThat(routingNodes.node("node3").shardsWithState("test1", STARTED).size(), equalTo(2));
-        
+
         logger.info("kill one node");
         IndexShardRoutingTable indexShardRoutingTable = routingTable.index("test").shard(0);
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder(clusterState.nodes()).remove(indexShardRoutingTable.primaryShard().currentNodeId())).build();
@@ -411,6 +411,6 @@ public class RoutingNodesIntegrityTests extends ElasticsearchTestCase {
     }
 
     private boolean assertShardStats(RoutingNodes routingNodes) {
-        return routingNodes.assertShardStats();
+        return RoutingNodes.assertShardStats(routingNodes);
     }
 }

+ 36 - 0
src/test/java/org/elasticsearch/cluster/routing/allocation/RoutingNodesUtils.java

@@ -0,0 +1,36 @@
+/*
+ * Licensed to ElasticSearch and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. ElasticSearch licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.cluster.routing.allocation;
+
+import org.elasticsearch.cluster.routing.RoutingNode;
+import org.elasticsearch.cluster.routing.RoutingNodes;
+import org.elasticsearch.cluster.routing.ShardRoutingState;
+
+/**
+ */
+public class RoutingNodesUtils {
+
+    public static int numberOfShardsOfType(RoutingNodes nodes, ShardRoutingState state) {
+        int count = 0;
+        for (RoutingNode routingNode : nodes) {
+            count += routingNode.numberOfShardsWithState(state);
+        }
+        return count;
+    }
+}

+ 8 - 5
src/test/java/org/elasticsearch/cluster/routing/allocation/SameShardRoutingTests.java

@@ -12,10 +12,12 @@ import org.elasticsearch.common.logging.ESLogger;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.common.transport.InetSocketTransportAddress;
 import org.elasticsearch.test.ElasticsearchTestCase;
+import org.elasticsearch.test.junit.annotations.TestLogging;
 import org.junit.Test;
 
 import static org.elasticsearch.cluster.routing.ShardRoutingState.INITIALIZING;
 import static org.elasticsearch.cluster.routing.allocation.RoutingAllocationTests.newNode;
+import static org.elasticsearch.cluster.routing.allocation.RoutingNodesUtils.numberOfShardsOfType;
 import static org.elasticsearch.common.settings.ImmutableSettings.settingsBuilder;
 import static org.hamcrest.Matchers.equalTo;
 
@@ -26,6 +28,7 @@ public class SameShardRoutingTests extends ElasticsearchTestCase {
     private final ESLogger logger = Loggers.getLogger(SameShardRoutingTests.class);
 
     @Test
+    @TestLogging("cluster.routing.allocation:TRACE")
     public void sameHost() {
         AllocationService strategy = new AllocationService(settingsBuilder().put(SameShardAllocationDecider.SAME_HOST_SETTING, true).build());
 
@@ -45,14 +48,14 @@ public class SameShardRoutingTests extends ElasticsearchTestCase {
         routingTable = strategy.reroute(clusterState).routingTable();
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
-        assertThat(clusterState.readOnlyRoutingNodes().numberOfShardsOfType(ShardRoutingState.INITIALIZING), equalTo(2));
+        assertThat(numberOfShardsOfType(clusterState.readOnlyRoutingNodes(), ShardRoutingState.INITIALIZING), equalTo(2));
 
         logger.info("--> start all primary shards, no replica will be started since its on the same host");
         routingTable = strategy.applyStartedShards(clusterState, clusterState.readOnlyRoutingNodes().shardsWithState(INITIALIZING)).routingTable();
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
-        assertThat(clusterState.readOnlyRoutingNodes().numberOfShardsOfType(ShardRoutingState.STARTED), equalTo(2));
-        assertThat(clusterState.readOnlyRoutingNodes().numberOfShardsOfType(ShardRoutingState.INITIALIZING), equalTo(0));
+        assertThat(numberOfShardsOfType(clusterState.readOnlyRoutingNodes(), ShardRoutingState.STARTED), equalTo(2));
+        assertThat(numberOfShardsOfType(clusterState.readOnlyRoutingNodes(), ShardRoutingState.INITIALIZING), equalTo(0));
 
         logger.info("--> add another node, with a different host, replicas will be allocating");
         clusterState = ClusterState.builder(clusterState).nodes(DiscoveryNodes.builder(clusterState.nodes())
@@ -60,8 +63,8 @@ public class SameShardRoutingTests extends ElasticsearchTestCase {
         routingTable = strategy.reroute(clusterState).routingTable();
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
-        assertThat(clusterState.readOnlyRoutingNodes().numberOfShardsOfType(ShardRoutingState.STARTED), equalTo(2));
-        assertThat(clusterState.readOnlyRoutingNodes().numberOfShardsOfType(ShardRoutingState.INITIALIZING), equalTo(2));
+        assertThat(numberOfShardsOfType(clusterState.readOnlyRoutingNodes(), ShardRoutingState.STARTED), equalTo(2));
+        assertThat(numberOfShardsOfType(clusterState.readOnlyRoutingNodes(), ShardRoutingState.INITIALIZING), equalTo(2));
         for (MutableShardRouting shardRouting : clusterState.readOnlyRoutingNodes().shardsWithState(INITIALIZING)) {
             assertThat(shardRouting.currentNodeId(), equalTo("node3"));
         }

+ 3 - 2
src/test/java/org/elasticsearch/cluster/routing/allocation/ShardsLimitAllocationTests.java

@@ -36,6 +36,7 @@ import org.junit.Test;
 
 import static org.elasticsearch.cluster.routing.ShardRoutingState.*;
 import static org.elasticsearch.cluster.routing.allocation.RoutingAllocationTests.newNode;
+import static org.elasticsearch.cluster.routing.allocation.RoutingNodesUtils.numberOfShardsOfType;
 import static org.elasticsearch.common.settings.ImmutableSettings.settingsBuilder;
 import static org.hamcrest.Matchers.equalTo;
 
@@ -123,7 +124,7 @@ public class ShardsLimitAllocationTests extends ElasticsearchTestCase {
         routingTable = strategy.applyStartedShards(clusterState, routingNodes.shardsWithState(INITIALIZING)).routingTable();
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
-        assertThat(clusterState.readOnlyRoutingNodes().numberOfShardsOfType(STARTED), equalTo(5));
+        assertThat(numberOfShardsOfType(clusterState.readOnlyRoutingNodes(), STARTED), equalTo(5));
 
         logger.info("add another index with 5 shards");
         metaData = MetaData.builder(metaData)
@@ -147,7 +148,7 @@ public class ShardsLimitAllocationTests extends ElasticsearchTestCase {
         routingTable = strategy.applyStartedShards(clusterState, routingNodes.shardsWithState(INITIALIZING)).routingTable();
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
-        assertThat(clusterState.readOnlyRoutingNodes().numberOfShardsOfType(STARTED), equalTo(10));
+        assertThat(numberOfShardsOfType(clusterState.readOnlyRoutingNodes(), STARTED), equalTo(10));
 
         for (MutableShardRouting shardRouting : clusterState.readOnlyRoutingNodes().node("node1")) {
             assertThat(shardRouting.index(), equalTo("test"));

+ 7 - 9
src/test/java/org/elasticsearch/cluster/routing/allocation/SingleShardNoReplicasRoutingTests.java

@@ -24,10 +24,7 @@ import org.elasticsearch.cluster.metadata.IndexMetaData;
 import org.elasticsearch.cluster.metadata.MetaData;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
-import org.elasticsearch.cluster.routing.MutableShardRouting;
-import org.elasticsearch.cluster.routing.RoutingNode;
-import org.elasticsearch.cluster.routing.RoutingNodes;
-import org.elasticsearch.cluster.routing.RoutingTable;
+import org.elasticsearch.cluster.routing.*;
 import org.elasticsearch.common.logging.ESLogger;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.test.ElasticsearchTestCase;
@@ -40,6 +37,7 @@ import static com.google.common.collect.Lists.newArrayList;
 import static com.google.common.collect.Sets.newHashSet;
 import static org.elasticsearch.cluster.routing.ShardRoutingState.*;
 import static org.elasticsearch.cluster.routing.allocation.RoutingAllocationTests.newNode;
+import static org.elasticsearch.cluster.routing.allocation.RoutingNodesUtils.numberOfShardsOfType;
 import static org.elasticsearch.common.settings.ImmutableSettings.settingsBuilder;
 import static org.hamcrest.Matchers.*;
 
@@ -263,7 +261,7 @@ public class SingleShardNoReplicasRoutingTests extends ElasticsearchTestCase {
         Set<String> encounteredIndices = newHashSet();
         for (RoutingNode routingNode : routingNodes) {
             assertThat(routingNode.numberOfShardsWithState(STARTED), equalTo(0));
-            assertThat(routingNode.shards().size(), equalTo(2));
+            assertThat(routingNode.size(), equalTo(2));
             // make sure we still have 2 shards initializing per node on the only 25 nodes
             int nodeIndex = Integer.parseInt(routingNode.nodeId().substring("node".length()));
             assertThat(nodeIndex, lessThan(25));
@@ -358,7 +356,7 @@ public class SingleShardNoReplicasRoutingTests extends ElasticsearchTestCase {
             assertThat(routingTable.index("test" + i).shard(0).shards().get(0).state(), equalTo(INITIALIZING));
         }
         RoutingNodes routingNodes = clusterState.routingNodes();
-        assertThat(routingNodes.numberOfShardsOfType(INITIALIZING), equalTo(numberOfIndices));
+        assertThat(numberOfShardsOfType(routingNodes, INITIALIZING), equalTo(numberOfIndices));
         assertThat(routingNodes.node("node1").numberOfShardsWithState(INITIALIZING), anyOf(equalTo(3), equalTo(4)));
         assertThat(routingNodes.node("node2").numberOfShardsWithState(INITIALIZING), anyOf(equalTo(3), equalTo(4)));
         assertThat(routingNodes.node("node2").numberOfShardsWithState(INITIALIZING), anyOf(equalTo(3), equalTo(4)));
@@ -387,8 +385,8 @@ public class SingleShardNoReplicasRoutingTests extends ElasticsearchTestCase {
             assertThat(routingTable.index("test" + i).shard(0).shards().get(0).state(), anyOf(equalTo(RELOCATING), equalTo(STARTED)));
         }
         routingNodes = clusterState.routingNodes();
-        assertThat("4 source shard routing are relocating", routingNodes.numberOfShardsOfType(RELOCATING), equalTo(4));
-        assertThat("4 target shard routing are initializing", routingNodes.numberOfShardsOfType(INITIALIZING), equalTo(4));
+        assertThat("4 source shard routing are relocating", numberOfShardsOfType(routingNodes, RELOCATING), equalTo(4));
+        assertThat("4 target shard routing are initializing", numberOfShardsOfType(routingNodes, INITIALIZING), equalTo(4));
 
         logger.info("Now, mark the relocated as started");
         prevRoutingTable = routingTable;
@@ -404,7 +402,7 @@ public class SingleShardNoReplicasRoutingTests extends ElasticsearchTestCase {
             assertThat(routingTable.index("test" + i).shard(0).shards().get(0).state(), anyOf(equalTo(RELOCATING), equalTo(STARTED)));
         }
         routingNodes = clusterState.routingNodes();
-        assertThat(routingNodes.numberOfShardsOfType(STARTED), equalTo(numberOfIndices));
+        assertThat(numberOfShardsOfType(routingNodes, STARTED), equalTo(numberOfIndices));
         for (RoutingNode routingNode : routingNodes) {
             assertThat(routingNode.numberOfShardsWithState(STARTED), equalTo(2));
         }

+ 28 - 28
src/test/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDeciderTests.java

@@ -115,7 +115,7 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
         // Assert that we're able to start the primary
         assertThat(clusterState.routingNodes().shardsWithState(ShardRoutingState.STARTED).size(), equalTo(1));
         // Assert that node1 didn't get any shards because its disk usage is too high
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
 
         logger.info("--> start the shards (replicas)");
         routingTable = strategy.applyStartedShards(clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING)).routingTable();
@@ -145,9 +145,9 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
         logShardStates(clusterState);
         // Assert that the replica couldn't be started since node1 doesn't have enough space
         assertThat(clusterState.routingNodes().shardsWithState(ShardRoutingState.STARTED).size(), equalTo(2));
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(1));
 
         logger.info("--> changing decider settings");
 
@@ -176,9 +176,9 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
 
         // Shards remain started
         assertThat(clusterState.routingNodes().shardsWithState(STARTED).size(), equalTo(2));
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(1));
 
         logger.info("--> changing settings again");
 
@@ -207,10 +207,10 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
         logShardStates(clusterState);
         // Shards remain started
         assertThat(clusterState.routingNodes().shardsWithState(STARTED).size(), equalTo(2));
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
         // Shard hasn't been moved off of node2 yet because there's nowhere for it to go
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(1));
 
         logger.info("--> adding node4");
 
@@ -230,11 +230,11 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
         logShardStates(clusterState);
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
         // Node4 is available now, so the shard is moved off of node2
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(0));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(1));
-        assertThat(clusterState.getRoutingNodes().node("node4").shards().size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node4").size(), equalTo(1));
     }
 
     @Test
@@ -304,7 +304,7 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
         // Assert that we're able to start the primary
         assertThat(clusterState.routingNodes().shardsWithState(ShardRoutingState.STARTED).size(), equalTo(1));
         // Assert that node1 didn't get any shards because its disk usage is too high
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
 
         logger.info("--> start the shards (replicas)");
         routingTable = strategy.applyStartedShards(clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING)).routingTable();
@@ -334,9 +334,9 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
         logShardStates(clusterState);
         // Assert that the replica couldn't be started since node1 doesn't have enough space
         assertThat(clusterState.routingNodes().shardsWithState(ShardRoutingState.STARTED).size(), equalTo(2));
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(1));
 
         logger.info("--> changing decider settings");
 
@@ -365,9 +365,9 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
 
         // Shards remain started
         assertThat(clusterState.routingNodes().shardsWithState(STARTED).size(), equalTo(2));
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(1));
 
         logger.info("--> changing settings again");
 
@@ -396,10 +396,10 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
         logShardStates(clusterState);
         // Shards remain started
         assertThat(clusterState.routingNodes().shardsWithState(STARTED).size(), equalTo(2));
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
         // Shard hasn't been moved off of node2 yet because there's nowhere for it to go
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(1));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(1));
 
         logger.info("--> adding node4");
 
@@ -419,11 +419,11 @@ public class DiskThresholdDeciderTests extends ElasticsearchTestCase {
         clusterState = ClusterState.builder(clusterState).routingTable(routingTable).build();
 
         logShardStates(clusterState);
-        assertThat(clusterState.getRoutingNodes().node("node1").shards().size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node1").size(), equalTo(0));
         // Node4 is available now, so the shard is moved off of node2
-        assertThat(clusterState.getRoutingNodes().node("node2").shards().size(), equalTo(0));
-        assertThat(clusterState.getRoutingNodes().node("node3").shards().size(), equalTo(1));
-        assertThat(clusterState.getRoutingNodes().node("node4").shards().size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node2").size(), equalTo(0));
+        assertThat(clusterState.getRoutingNodes().node("node3").size(), equalTo(1));
+        assertThat(clusterState.getRoutingNodes().node("node4").size(), equalTo(1));
     }
 
     @Test

+ 44 - 36
src/test/java/org/elasticsearch/indexlifecycle/IndexLifecycleActionTests.java

@@ -20,6 +20,9 @@
 package org.elasticsearch.indexlifecycle;
 
 import com.carrotsearch.randomizedtesting.annotations.Nightly;
+import com.google.common.base.Function;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Sets;
 import org.apache.lucene.util.LuceneTestCase.Slow;
 import org.elasticsearch.action.admin.cluster.health.ClusterHealthResponse;
 import org.elasticsearch.action.admin.cluster.health.ClusterHealthStatus;
@@ -27,6 +30,7 @@ import org.elasticsearch.action.admin.indices.create.CreateIndexResponse;
 import org.elasticsearch.action.admin.indices.delete.DeleteIndexResponse;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.routing.RoutingNode;
+import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.discovery.Discovery;
@@ -36,7 +40,6 @@ import org.elasticsearch.test.ElasticsearchIntegrationTest.Scope;
 import org.elasticsearch.test.TestCluster;
 import org.junit.Test;
 
-import java.util.Map;
 import java.util.Set;
 
 import static org.elasticsearch.client.Requests.*;
@@ -89,7 +92,7 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
         assertThat(clusterHealth.getStatus(), equalTo(ClusterHealthStatus.YELLOW));
 
         ClusterState clusterState = client().admin().cluster().prepareState().get().getState();
-        RoutingNode routingNodeEntry1 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node1);
+        RoutingNode routingNodeEntry1 = clusterState.readOnlyRoutingNodes().node(node1);
         assertThat(routingNodeEntry1.numberOfShardsWithState(STARTED), equalTo(11));
 
         logger.info("Starting server2");
@@ -121,11 +124,11 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
 
 
         clusterState = client().admin().cluster().prepareState().get().getState();
-        assertNodesPresent(clusterState.readOnlyRoutingNodes().nodesToShards(), node1, node2);
-        routingNodeEntry1 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node1);
+        assertNodesPresent(clusterState.readOnlyRoutingNodes(), node1, node2);
+        routingNodeEntry1 = clusterState.readOnlyRoutingNodes().node(node1);
         assertThat(routingNodeEntry1.numberOfShardsWithState(RELOCATING), equalTo(0));
         assertThat(routingNodeEntry1.numberOfShardsWithState(STARTED), equalTo(11));
-        RoutingNode routingNodeEntry2 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node2);
+        RoutingNode routingNodeEntry2 = clusterState.readOnlyRoutingNodes().node(node2);
         assertThat(routingNodeEntry2.numberOfShardsWithState(INITIALIZING), equalTo(0));
         assertThat(routingNodeEntry2.numberOfShardsWithState(STARTED), equalTo(11));
 
@@ -157,11 +160,11 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
 
 
         clusterState = client().admin().cluster().prepareState().get().getState();
-        assertNodesPresent(clusterState.readOnlyRoutingNodes().nodesToShards(), node1, node2, node3);
+        assertNodesPresent(clusterState.readOnlyRoutingNodes(), node1, node2, node3);
 
-        routingNodeEntry1 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node1);
-        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node2);
-        RoutingNode routingNodeEntry3 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node3);
+        routingNodeEntry1 = clusterState.readOnlyRoutingNodes().node(node1);
+        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().node(node2);
+        RoutingNode routingNodeEntry3 = clusterState.readOnlyRoutingNodes().node(node3);
 
         assertThat(routingNodeEntry1.numberOfShardsWithState(STARTED) + routingNodeEntry2.numberOfShardsWithState(STARTED) + routingNodeEntry3.numberOfShardsWithState(STARTED), equalTo(22));
 
@@ -197,9 +200,9 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
         assertThat(clusterHealth.getActivePrimaryShards(), equalTo(11));
 
         clusterState = client().admin().cluster().prepareState().get().getState();
-        assertNodesPresent(clusterState.readOnlyRoutingNodes().nodesToShards(), node3, node2);
-        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node2);
-        routingNodeEntry3 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node3);
+        assertNodesPresent(clusterState.readOnlyRoutingNodes(), node3, node2);
+        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().node(node2);
+        routingNodeEntry3 = clusterState.readOnlyRoutingNodes().node(node3);
 
         assertThat(routingNodeEntry2.numberOfShardsWithState(STARTED) + routingNodeEntry3.numberOfShardsWithState(STARTED), equalTo(22));
 
@@ -216,12 +219,12 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
         assertThat(deleteIndexResponse.isAcknowledged(), equalTo(true));
 
         clusterState = client().admin().cluster().prepareState().get().getState();
-        assertNodesPresent(clusterState.readOnlyRoutingNodes().nodesToShards(), node3, node2);
-        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node2);
-        assertThat(routingNodeEntry2.shards().isEmpty(), equalTo(true));
+        assertNodesPresent(clusterState.readOnlyRoutingNodes(), node3, node2);
+        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().node(node2);
+        assertThat(routingNodeEntry2.isEmpty(), equalTo(true));
 
-        routingNodeEntry3 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node3);
-        assertThat(routingNodeEntry3.shards().isEmpty(), equalTo(true));
+        routingNodeEntry3 = clusterState.readOnlyRoutingNodes().node(node3);
+        assertThat(routingNodeEntry3.isEmpty(), equalTo(true));
     }
 
     private String getLocalNodeId(String name) {
@@ -262,8 +265,8 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
         assertThat(clusterHealth.getActivePrimaryShards(), equalTo(11));
 
         ClusterState clusterState = client().admin().cluster().prepareState().get().getState();
-        assertNodesPresent(clusterState.readOnlyRoutingNodes().nodesToShards(), node1);
-        RoutingNode routingNodeEntry1 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node1);
+        assertNodesPresent(clusterState.readOnlyRoutingNodes(), node1);
+        RoutingNode routingNodeEntry1 = clusterState.readOnlyRoutingNodes().node(node1);
         assertThat(routingNodeEntry1.numberOfShardsWithState(STARTED), equalTo(11));
 
         // start another server
@@ -292,11 +295,11 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
 
 
         clusterState = client().admin().cluster().prepareState().get().getState();
-        assertNodesPresent(clusterState.readOnlyRoutingNodes().nodesToShards(), node1, node2);
-        routingNodeEntry1 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node1);
+        assertNodesPresent(clusterState.readOnlyRoutingNodes(), node1, node2);
+        routingNodeEntry1 = clusterState.readOnlyRoutingNodes().node(node1);
         assertThat(routingNodeEntry1.numberOfShardsWithState(RELOCATING), equalTo(0));
         assertThat(routingNodeEntry1.numberOfShardsWithState(STARTED), anyOf(equalTo(6), equalTo(5)));
-        RoutingNode routingNodeEntry2 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node2);
+        RoutingNode routingNodeEntry2 = clusterState.readOnlyRoutingNodes().node(node2);
         assertThat(routingNodeEntry2.numberOfShardsWithState(INITIALIZING), equalTo(0));
         assertThat(routingNodeEntry2.numberOfShardsWithState(STARTED), anyOf(equalTo(5), equalTo(6)));
 
@@ -325,10 +328,10 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
 
 
         clusterState = client().admin().cluster().prepareState().get().getState();
-        assertNodesPresent(clusterState.readOnlyRoutingNodes().nodesToShards(), node1, node2, node3);
-        routingNodeEntry1 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node1);
-        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node2);
-        RoutingNode routingNodeEntry3 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node3);
+        assertNodesPresent(clusterState.readOnlyRoutingNodes(), node1, node2, node3);
+        routingNodeEntry1 = clusterState.readOnlyRoutingNodes().node(node1);
+        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().node(node2);
+        RoutingNode routingNodeEntry3 = clusterState.readOnlyRoutingNodes().node(node3);
 
         assertThat(routingNodeEntry1.numberOfShardsWithState(STARTED) + routingNodeEntry2.numberOfShardsWithState(STARTED) + routingNodeEntry3.numberOfShardsWithState(STARTED), equalTo(11));
 
@@ -366,10 +369,10 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
         assertThat(clusterHealth.getActivePrimaryShards(), equalTo(11));
 
         clusterState = client().admin().cluster().prepareState().get().getState();
-        assertNodesPresent(clusterState.readOnlyRoutingNodes().nodesToShards(), node3, node2);
+        assertNodesPresent(clusterState.readOnlyRoutingNodes(), node3, node2);
 
-        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node2);
-        routingNodeEntry3 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node3);
+        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().node(node2);
+        routingNodeEntry3 = clusterState.readOnlyRoutingNodes().node(node3);
 
         assertThat(routingNodeEntry2.numberOfShardsWithState(STARTED) + routingNodeEntry3.numberOfShardsWithState(STARTED), equalTo(11));
 
@@ -387,17 +390,22 @@ public class IndexLifecycleActionTests extends ElasticsearchIntegrationTest {
         assertThat(deleteIndexResponse.isAcknowledged(), equalTo(true));
 
         clusterState = client().admin().cluster().prepareState().get().getState();
-        assertNodesPresent(clusterState.readOnlyRoutingNodes().nodesToShards(), node3, node2);
+        assertNodesPresent(clusterState.readOnlyRoutingNodes(), node3, node2);
 
-        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node2);
-        assertThat(routingNodeEntry2.shards().isEmpty(), equalTo(true));
+        routingNodeEntry2 = clusterState.readOnlyRoutingNodes().node(node2);
+        assertThat(routingNodeEntry2.isEmpty(), equalTo(true));
 
-        routingNodeEntry3 = clusterState.readOnlyRoutingNodes().nodesToShards().get(node3);
-        assertThat(routingNodeEntry3.shards().isEmpty(), equalTo(true));
+        routingNodeEntry3 = clusterState.readOnlyRoutingNodes().node(node3);
+        assertThat(routingNodeEntry3.isEmpty(), equalTo(true));
     }
 
-    private void assertNodesPresent(Map<String, ?> nodesToShards, String...nodes) {
-        Set<String> keySet = nodesToShards.keySet();
+    private void assertNodesPresent(RoutingNodes routingNodes, String...nodes) {
+        final Set<String> keySet = Sets.newHashSet(Iterables.transform(routingNodes, new Function<RoutingNode, String>() {
+            @Override
+            public String apply(RoutingNode input) {
+                return input.nodeId();
+            }
+        }));
         assertThat(keySet, containsInAnyOrder(nodes));
     }
 }