Optimize IndicesSegmentResponse#getIndices (#80064)

Replaces a double-nested loop over all segments and indices with a map
construction. Also tidies up some related code in `IndexSegments`.

Co-authored-by: David Turner <[email protected]>

Author: mushao999

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/indices.segments/10_basic.yml

@@ -40,6 +40,50 @@
   - match:   { indices.index1.shards.0.0.routing.primary: true}
   - match:   { indices.index1.shards.0.0.segments._0.num_docs: 1}
 
+---
+"segments from multiple indices":
+
+  - do:
+      indices.create:
+        index: index1
+        body:
+          settings:
+            number_of_shards: "1"
+            number_of_replicas: "0"
+
+  - do:
+      indices.create:
+        index: index2
+        body:
+          settings:
+            number_of_shards: "1"
+            number_of_replicas: "0"
+
+  - do:
+      index:
+        index: index1
+        body: { foo: bar }
+        refresh: true
+
+  - do:
+      index:
+        index: index2
+        body: { foo: bar }
+        refresh: true
+
+  - do:
+      cluster.health:
+        wait_for_status: green
+
+  - do:
+      indices.segments: {}
+
+  - gte: { _shards.total: 2 }
+  - match:   { indices.index1.shards.0.0.routing.primary: true}
+  - match:   { indices.index1.shards.0.0.segments._0.num_docs: 1}
+  - match:   { indices.index2.shards.0.0.routing.primary: true}
+  - match:   { indices.index2.shards.0.0.segments._0.num_docs: 1}
+
 ---
 "closed segments test":
   - skip:

server/src/main/java/org/elasticsearch/action/admin/indices/segments/IndexSegments.java

@@ -20,20 +20,15 @@ public class IndexSegments implements Iterable<IndexShardSegments> {
 
     private final Map<Integer, IndexShardSegments> indexShards;
 
-    IndexSegments(String index, ShardSegments[] shards) {
+    IndexSegments(String index, List<ShardSegments> shards) {
         this.index = index;
 
-        Map<Integer, List<ShardSegments>> tmpIndexShards = new HashMap<>();
+        final Map<Integer, List<ShardSegments>> segmentsByShardId = new HashMap<>();
         for (ShardSegments shard : shards) {
-            List<ShardSegments> lst = tmpIndexShards.get(shard.getShardRouting().id());
-            if (lst == null) {
-                lst = new ArrayList<>();
-                tmpIndexShards.put(shard.getShardRouting().id(), lst);
-            }
-            lst.add(shard);
+            segmentsByShardId.computeIfAbsent(shard.getShardRouting().id(), k -> new ArrayList<>()).add(shard);
         }
         indexShards = new HashMap<>();
-        for (Map.Entry<Integer, List<ShardSegments>> entry : tmpIndexShards.entrySet()) {
+        for (Map.Entry<Integer, List<ShardSegments>> entry : segmentsByShardId.entrySet()) {
             indexShards.put(
                 entry.getKey(),
                 new IndexShardSegments(

server/src/main/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentResponse.java

@@ -27,17 +27,15 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Set;
 
 public class IndicesSegmentResponse extends BroadcastResponse {
 
     private final ShardSegments[] shards;
 
-    private Map<String, IndexSegments> indicesSegments;
+    private volatile Map<String, IndexSegments> indicesSegments;
 
     IndicesSegmentResponse(StreamInput in) throws IOException {
         super(in);
@@ -61,19 +59,12 @@ public Map<String, IndexSegments> getIndices() {
         }
         Map<String, IndexSegments> indicesSegments = new HashMap<>();
 
-        Set<String> indices = new HashSet<>();
+        final Map<String, List<ShardSegments>> segmentsByIndex = new HashMap<>();
         for (ShardSegments shard : shards) {
-            indices.add(shard.getShardRouting().getIndexName());
+            segmentsByIndex.computeIfAbsent(shard.getShardRouting().getIndexName(), k -> new ArrayList<>()).add(shard);
         }
-
-        for (String indexName : indices) {
-            List<ShardSegments> shards = new ArrayList<>();
-            for (ShardSegments shard : this.shards) {
-                if (shard.getShardRouting().getIndexName().equals(indexName)) {
-                    shards.add(shard);
-                }
-            }
-            indicesSegments.put(indexName, new IndexSegments(indexName, shards.toArray(new ShardSegments[shards.size()])));
+        for (Map.Entry<String, List<ShardSegments>> entry : segmentsByIndex.entrySet()) {
+            indicesSegments.put(entry.getKey(), new IndexSegments(entry.getKey(), entry.getValue()));
         }
         this.indicesSegments = indicesSegments;
         return indicesSegments;