Remove nested loop in IndicesStatsResponse (#40988) (#41138)

This commit removes nested loop in `getIndices`.

server/src/main/java/org/elasticsearch/action/admin/indices/stats/IndexStats.java

@@ -108,4 +108,24 @@ public class IndexStats implements Iterable<IndexShardStats> {
         primary = stats;
         return stats;
     }
+
+    public static class IndexStatsBuilder {
+        private final String indexName;
+        private final String uuid;
+        private final List<ShardStats> shards = new ArrayList<>();
+
+        public IndexStatsBuilder(String indexName, String uuid) {
+            this.indexName = indexName;
+            this.uuid = uuid;
+        }
+
+        public IndexStatsBuilder add(ShardStats shardStats) {
+            shards.add(shardStats);
+            return this;
+        }
+
+        public IndexStats build() {
+            return new IndexStats(indexName, uuid, shards.toArray(new ShardStats[shards.size()]));
+        }
+    }
 }

server/src/main/java/org/elasticsearch/action/admin/indices/stats/IndicesStatsResponse.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.action.admin.indices.stats;
 
+import org.elasticsearch.action.admin.indices.stats.IndexStats.IndexStatsBuilder;
 import org.elasticsearch.action.support.DefaultShardOperationFailedException;
 import org.elasticsearch.action.support.broadcast.BroadcastResponse;
 import org.elasticsearch.cluster.routing.ShardRouting;
@@ -29,12 +30,10 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.index.Index;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
+import java.util.stream.Collectors;
 
 import static java.util.Collections.unmodifiableMap;
 
@@ -83,26 +82,17 @@ public class IndicesStatsResponse extends BroadcastResponse {
         if (indicesStats != null) {
             return indicesStats;
         }
-        Map<String, IndexStats> indicesStats = new HashMap<>();
 
-        Set<Index> indices = new HashSet<>();
+        final Map<String, IndexStatsBuilder> indexToIndexStatsBuilder = new HashMap<>();
         for (ShardStats shard : shards) {
-            indices.add(shard.getShardRouting().index());
+            Index index = shard.getShardRouting().index();
+            IndexStatsBuilder indexStatsBuilder = indexToIndexStatsBuilder.computeIfAbsent(index.getName(),
+                    k -> new IndexStatsBuilder(k, index.getUUID()));
+            indexStatsBuilder.add(shard);
         }
 
-        for (Index index : indices) {
+        indicesStats = indexToIndexStatsBuilder.entrySet().stream()
-            List<ShardStats> shards = new ArrayList<>();
+                .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().build()));
-            String indexName = index.getName();
-            for (ShardStats shard : this.shards) {
-                if (shard.getShardRouting().getIndexName().equals(indexName)) {
-                    shards.add(shard);
-                }
-            }
-            indicesStats.put(
-                indexName, new IndexStats(indexName, index.getUUID(), shards.toArray(new ShardStats[shards.size()]))
-            );
-        }
-        this.indicesStats = indicesStats;
         return indicesStats;
     }
 

server/src/test/java/org/elasticsearch/action/admin/indices/stats/IndicesStatsResponseTests.java

@@ -19,16 +19,30 @@
 
 package org.elasticsearch.action.admin.indices.stats;
 
+import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.ShardRoutingState;
+import org.elasticsearch.cluster.routing.TestShardRouting;
+import org.elasticsearch.common.UUIDs;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.json.JsonXContent;
+import org.elasticsearch.index.Index;
+import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.index.shard.ShardPath;
 import org.elasticsearch.test.ESTestCase;
 
+import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
 
 import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.is;
 import static org.hamcrest.object.HasToString.hasToString;
 
-
 public class IndicesStatsResponseTests extends ESTestCase {
 
     public void testInvalidLevel() {
@@ -42,4 +56,59 @@ public class IndicesStatsResponseTests extends ESTestCase {
             hasToString(containsString("level parameter must be one of [cluster] or [indices] or [shards] but was [" + level + "]")));
     }
 
+    public void testGetIndices() {
+        List<ShardStats> shards = new ArrayList<>();
+        int noOfIndexes = randomIntBetween(2, 5);
+        List<String> expectedIndexes = new ArrayList<>();
+        Map<String, AtomicLong> expectedIndexToPrimaryShardsCount = new HashMap<>();
+        Map<String, AtomicLong> expectedIndexToTotalShardsCount = new HashMap<>();
+
+        for (int indCnt = 0; indCnt < noOfIndexes; indCnt++) {
+            Index index = createIndex(randomAlphaOfLength(9));
+            expectedIndexes.add(index.getName());
+            int numShards = randomIntBetween(1, 5);
+            for (int shardId = 0; shardId < numShards; shardId++) {
+                ShardId shId = new ShardId(index, shardId);
+                Path path = createTempDir().resolve("indices").resolve(index.getUUID()).resolve(String.valueOf(shardId));
+                ShardPath shardPath = new ShardPath(false, path, path, shId);
+                ShardRouting routing = createShardRouting(index, shId, (shardId == 0));
+                shards.add(new ShardStats(routing, shardPath, null, null, null, null));
+                AtomicLong primaryShardsCounter = expectedIndexToPrimaryShardsCount.computeIfAbsent(index.getName(),
+                        k -> new AtomicLong(0L));
+                if (routing.primary()) {
+                    primaryShardsCounter.incrementAndGet();
+                }
+                AtomicLong shardsCounter = expectedIndexToTotalShardsCount.computeIfAbsent(index.getName(), k -> new AtomicLong(0L));
+                shardsCounter.incrementAndGet();
+            }
+        }
+        final IndicesStatsResponse indicesStatsResponse = new IndicesStatsResponse(shards.toArray(new ShardStats[shards.size()]), 0, 0, 0,
+                null);
+        Map<String, IndexStats> indexStats = indicesStatsResponse.getIndices();
+
+        assertThat(indexStats.size(), is(noOfIndexes));
+        assertThat(indexStats.keySet(), containsInAnyOrder(expectedIndexes.toArray(new String[0])));
+
+        for (String index : indexStats.keySet()) {
+            IndexStats stat = indexStats.get(index);
+            ShardStats[] shardStats = stat.getShards();
+            long primaryCount = 0L;
+            long totalCount = shardStats.length;
+            for (ShardStats shardStat : shardStats) {
+                if (shardStat.getShardRouting().primary()) {
+                    primaryCount++;
+                }
+            }
+            assertThat(primaryCount, is(expectedIndexToPrimaryShardsCount.get(index).get()));
+            assertThat(totalCount, is(expectedIndexToTotalShardsCount.get(index).get()));
+        }
+    }
+
+    private ShardRouting createShardRouting(Index index, ShardId shardId, boolean isPrimary) {
+        return TestShardRouting.newShardRouting(shardId, randomAlphaOfLength(4), isPrimary, ShardRoutingState.STARTED);
+    }
+
+    private Index createIndex(String indexName) {
+        return new Index(indexName, UUIDs.base64UUID());
+    }
 }