Iterate directly over contents of RoutingNode

Today `RoutingNode#initializing`, `RoutingNode#started` and
`RoutingNode#relocating` all work by copying the result into a freshly
allocated array. These collections are immutable and all the
(production) callers immediately iterate over the result so the copy is
unnecessary. Moreover the `started` collection in particular might be
huge.

This commit adjusts these methods to return an iterator over the
underlying collection instead, asserting that the resulting iterators
are never mutated.

Closes ES-13364

Author: DaveCTurner

server/src/main/java/org/elasticsearch/cluster/routing/RoutingNode.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.cluster.routing;
 
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.Index;
@@ -213,18 +214,16 @@ void remove(ShardRouting shard) {
         assert invariant();
     }
 
-    private static final ShardRouting[] EMPTY_SHARD_ROUTING_ARRAY = new ShardRouting[0];
-
-    public ShardRouting[] initializing() {
-        return initializingShards.toArray(EMPTY_SHARD_ROUTING_ARRAY);
+    public Iterator<ShardRouting> initializing() {
+        return Iterators.assertReadOnly(initializingShards.iterator());
     }
 
-    public ShardRouting[] relocating() {
-        return relocatingShards.toArray(EMPTY_SHARD_ROUTING_ARRAY);
+    public Iterator<ShardRouting> relocating() {
+        return Iterators.assertReadOnly(relocatingShards.iterator());
     }
 
-    public ShardRouting[] started() {
-        return startedShards.toArray(EMPTY_SHARD_ROUTING_ARRAY);
+    public Iterator<ShardRouting> started() {
+        return Iterators.assertReadOnly(startedShards.iterator());
     }
 
     /**
@@ -313,6 +312,8 @@ public String toString() {
         return sb.toString();
     }
 
+    private static final ShardRouting[] EMPTY_SHARD_ROUTING_ARRAY = new ShardRouting[0];
+
     public ShardRouting[] copyShards() {
         return shards.values().toArray(EMPTY_SHARD_ROUTING_ARRAY);
     }

server/src/main/java/org/elasticsearch/cluster/routing/allocation/RoutingAllocation.java

@@ -175,7 +175,9 @@ private static Map<String, Long> unaccountedSearchableSnapshotSizes(ClusterState
                 DiskUsage usage = clusterInfo.getNodeMostAvailableDiskUsages().get(node.nodeId());
                 ClusterInfo.ReservedSpace reservedSpace = clusterInfo.getReservedSpace(node.nodeId(), usage != null ? usage.path() : "");
                 long totalSize = 0;
-                for (ShardRouting shard : node.started()) {
+                final var startedShardsIterator = node.started();
+                while (startedShardsIterator.hasNext()) {
+                    final var shard = startedShardsIterator.next();
                     if (shard.getExpectedShardSize() > 0
                         && clusterState.metadata().indexMetadata(shard.index()).isSearchableSnapshot()
                         && reservedSpace.containsShardId(shard.shardId()) == false

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

@@ -25,6 +25,7 @@
 import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator.ShardAllocationExplainer;
 import org.elasticsearch.cluster.routing.allocation.command.MoveAllocationCommand;
 import org.elasticsearch.cluster.routing.allocation.decider.Decision;
+import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.metrics.MeanMetric;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Setting;
@@ -38,6 +39,7 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
@@ -294,12 +296,13 @@ public DesiredBalance compute(
                     final var rerouteExplanation = command.execute(routingAllocation, false);
                     assert rerouteExplanation.decisions().type() != Decision.Type.NO : "should have thrown for NO decision";
                     if (rerouteExplanation.decisions().type() != Decision.Type.NO) {
-                        final ShardRouting[] initializingShards = routingNodes.node(
+                        final Iterator<ShardRouting> initializingShardsIterator = routingNodes.node(
                             routingAllocation.nodes().resolveNode(command.toNode()).getId()
                         ).initializing();
-                        assert initializingShards.length == 1
-                            : "expect exactly one relocating shard, but got: " + List.of(initializingShards);
-                        final var initializingShard = initializingShards[0];
+                        assert initializingShardsIterator.hasNext();
+                        final var initializingShard = initializingShardsIterator.next();
+                        assert initializingShardsIterator.hasNext() == false
+                            : "expect exactly one relocating shard, but got: " + Iterators.toList(initializingShardsIterator);
                         assert routingAllocation.nodes()
                             .resolveNode(command.fromNode())
                             .getId()
@@ -535,11 +538,11 @@ static void maybeSimulateAlreadyStartedShards(
         // Find all shards that are started in RoutingNodes but have no data on corresponding node in ClusterInfo
         final var startedShards = new ArrayList<ShardRouting>();
         for (var routingNode : routingNodes) {
-            for (var shardRouting : routingNode.started()) {
+            routingNode.started().forEachRemaining(shardRouting -> {
                 if (clusterInfo.hasShardMoved(shardRouting)) {
                     startedShards.add(shardRouting);
                 }
-            }
+            });
         }
         if (startedShards.isEmpty()) {
             return;

server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDecider.java

@@ -108,7 +108,9 @@ public static long sizeOfUnaccountedShards(
         // no longer initializing because their recovery failed or was cancelled.
 
         // Where reserved space is unavailable (e.g. stats are out-of-sync) compute a conservative estimate for initialising shards
-        for (ShardRouting routing : node.initializing()) {
+        final var initializingShardsIterator = node.initializing();
+        while (initializingShardsIterator.hasNext()) {
+            final var routing = initializingShardsIterator.next();
             // Space needs to be reserved only when initializing shards that are going to use additional space
             // that is not yet accounted for by `reservedSpace` in case of lengthy recoveries
             if (shouldReserveSpaceForInitializingShard(routing, metadata) && reservedSpace.containsShardId(routing.shardId()) == false) {
@@ -135,7 +137,9 @@ public static long sizeOfUnaccountedShards(
         totalSize += sizeOfUnaccountableSearchableSnapshotShards;
 
         if (subtractShardsMovingAway) {
-            for (ShardRouting routing : node.relocating()) {
+            final var relocatingShardsIterator = node.relocating();
+            while (relocatingShardsIterator.hasNext()) {
+                final var routing = relocatingShardsIterator.next();
                 if (dataPath.equals(clusterInfo.getDataPath(routing))) {
                     ProjectMetadata project = metadata.projectFor(routing.index());
                     totalSize -= getExpectedShardSize(

server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ThrottlingAllocationDecider.java

@@ -127,11 +127,17 @@ public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, Routing
 
             // Count the primaries currently doing recovery on the node, to ensure the primariesInitialRecoveries setting is obeyed.
             int primariesInRecovery = 0;
-            for (ShardRouting shard : node.initializing()) {
+            final var initializingShardsIterator = node.initializing();
+            while (initializingShardsIterator.hasNext()) {
+                final var shard = initializingShardsIterator.next();
                 // 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.primary() && shard.relocatingNodeId() == null) {
                     primariesInRecovery++;
+                    if (allocation.debugDecision() && primariesInRecovery >= primariesInitialRecoveries) {
+                        // bail out early if we don't need the final total
+                        return THROTTLE;
+                    }
                 }
             }
             if (primariesInRecovery >= primariesInitialRecoveries) {

server/src/main/java/org/elasticsearch/common/collect/Iterators.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.common.collect;
 
+import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.Nullable;
 
 import java.util.ArrayList;
@@ -507,6 +508,42 @@ public T next() {
         }
     }
 
+    /**
+     * Adds a wrapper around {@code iterator} which asserts that {@link Iterator#remove()} is not called.
+     */
+    public static <T> Iterator<T> assertReadOnly(final Iterator<T> iterator) {
+        return Assertions.ENABLED ? new AssertReadOnlyIterator<>(Objects.requireNonNull(iterator)) : iterator;
+    }
+
+    private static class AssertReadOnlyIterator<T> implements Iterator<T> {
+
+        private final Iterator<T> delegate;
+
+        AssertReadOnlyIterator(Iterator<T> delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return delegate.hasNext();
+        }
+
+        @Override
+        public T next() {
+            return delegate.next();
+        }
+
+        @Override
+        public void forEachRemaining(Consumer<? super T> action) {
+            delegate.forEachRemaining(action);
+        }
+
+        @Override
+        public void remove() {
+            throw new AssertionError();
+        }
+    }
+
     public static <T> boolean equals(Iterator<? extends T> iterator1, Iterator<? extends T> iterator2, BiPredicate<T, T> itemComparer) {
         if (iterator1 == null) {
             return iterator2 == null;

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

@@ -18,9 +18,8 @@
 import org.elasticsearch.test.ESTestCase;
 
 import java.net.InetAddress;
-import java.util.Arrays;
+import java.util.HashSet;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.emptySet;
@@ -161,7 +160,9 @@ public void testReturnStartedShards() {
     }
 
     private static Set<ShardId> startedShardsSet(RoutingNode routingNode) {
-        return Arrays.stream(routingNode.started()).map(ShardRouting::shardId).collect(Collectors.toSet());
+        final var result = new HashSet<ShardId>();
+        routingNode.started().forEachRemaining(shardRouting -> result.add(shardRouting.shardId()));
+        return result;
     }
 
 }

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

@@ -51,6 +51,7 @@
 import org.elasticsearch.cluster.routing.allocation.decider.ThrottlingAllocationDecider;
 import org.elasticsearch.common.Randomness;
 import org.elasticsearch.common.UUIDs;
+import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.time.TimeProvider;
@@ -103,10 +104,10 @@
 import static org.hamcrest.Matchers.aMapWithSize;
 import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.anyOf;
-import static org.hamcrest.Matchers.arrayWithSize;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.everyItem;
 import static org.hamcrest.Matchers.hasEntry;
+import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.mockito.Mockito.mock;
@@ -581,7 +582,7 @@ public void allocate(RoutingAllocation allocation) {
                     allocation.routingNodes().getRelocatingShardCount(),
                     equalTo(0)
                 );
-                assertThat(allocation.routingNodes().node("node-2").started(), arrayWithSize(2));
+                assertThat(Iterators.toList(allocation.routingNodes().node("node-2").started()), hasSize(2));
             }
 
             @Override

server/src/test/java/org/elasticsearch/common/collect/IteratorsTests.java

@@ -328,6 +328,19 @@ public void testCycle() {
         }
     }
 
+    public void testAssertReadOnly() {
+        assumeTrue("assertions enabled", Assertions.ENABLED);
+        final List<Integer> innerList = new ArrayList<>(List.of(1, 2, 3, 4));
+        assertTrue(Iterators.equals(innerList.iterator(), Iterators.assertReadOnly(innerList.iterator()), Objects::equals));
+
+        final var readonly = Iterators.assertReadOnly(innerList.iterator());
+        assertTrue(readonly.hasNext());
+        assertEquals(Integer.valueOf(1), readonly.next());
+        expectThrows(AssertionError.class, readonly::remove);
+
+        assertEquals(List.of(1, 2, 3, 4), innerList);
+    }
+
     public void testEquals() {
         final BiPredicate<Object, Object> notCalled = (a, b) -> { throw new AssertionError("not called"); };