Removing the allocationId

Author: drempapis

server/src/main/java/org/elasticsearch/action/search/load/ShardSearchLoadStats.java

@@ -19,16 +19,14 @@
 
 /**
  * ShardSearchLoadStats class represents the statistics of a shard in an index.
- * It contains information such as the index name, shard ID, allocation ID, and EWMA rate.
+ * It contains information such as the index name, shard ID, and search load.
  */
 public class ShardSearchLoadStats implements Writeable {
 
     private final String indexName;
 
     private final Integer shardId;
 
-    private final String allocationId;
-
     private final Double searchLoad;
 
     /**
@@ -41,7 +39,6 @@ public ShardSearchLoadStats(StreamInput in) throws IOException {
         assert Transports.assertNotTransportThread("O(#shards) work must always fork to an appropriate executor");
         this.indexName = in.readString();
         this.shardId = in.readVInt();
-        this.allocationId = in.readString();
         this.searchLoad = in.readDouble();
     }
 
@@ -50,13 +47,11 @@ public ShardSearchLoadStats(StreamInput in) throws IOException {
      *
      * @param indexName   the name of the index
      * @param shardId     the ID of the shard
-     * @param allocationId the allocation ID of the shard
      * @param searchLoad the search load of the shard
      */
-    public ShardSearchLoadStats(String indexName, Integer shardId, String allocationId, Double searchLoad) {
+    public ShardSearchLoadStats(String indexName, Integer shardId, Double searchLoad) {
         this.indexName = indexName;
         this.shardId = shardId;
-        this.allocationId = allocationId;
         this.searchLoad = searchLoad;
     }
 
@@ -67,13 +62,12 @@ public boolean equals(Object o) {
         ShardSearchLoadStats that = (ShardSearchLoadStats) o;
         return Objects.equals(indexName, that.indexName)
             && Objects.equals(shardId, that.shardId)
-            && Objects.equals(allocationId, that.allocationId)
             && Objects.equals(searchLoad, that.searchLoad);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(indexName, shardId, allocationId, searchLoad);
+        return Objects.hash(indexName, shardId, searchLoad);
     }
 
     /**
@@ -94,15 +88,6 @@ public Integer getShardId() {
         return this.shardId;
     }
 
-    /**
-     * Returns the allocation ID of the shard.
-     *
-     * @return the allocation ID
-     */
-    public String getAllocationId() {
-        return this.allocationId;
-    }
-
     /**
      * Returns the search load of the shard.
      *
@@ -116,7 +101,6 @@ public Double getSearchLoad() {
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(indexName);
         out.writeVInt(shardId);
-        out.writeString(allocationId);
         out.writeDouble(searchLoad);
     }
 }

server/src/main/java/org/elasticsearch/action/search/load/ShardSearchLoadStatsResponse.java

@@ -16,9 +16,14 @@
 import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 
 /**
  * Response to a shard stats request.
@@ -39,7 +44,7 @@ public class ShardSearchLoadStatsResponse extends ChunkedBroadcastResponse {
     }
 
     /**
-     * Constructor to create a ShardStatsResponse object with the given parameters.
+     * Constructor to create a ShardSearchLoadStatsResponse object with the given parameters.
      *
      * @param shards          the array of shard stats
      * @param totalShards     the total number of shards
@@ -55,17 +60,18 @@ public class ShardSearchLoadStatsResponse extends ChunkedBroadcastResponse {
         List<DefaultShardOperationFailedException> shardFailures
     ) {
         super(totalShards, successfulShards, failedShards, shardFailures);
-        this.shards = shards;
-        Objects.requireNonNull(shards);
+        this.shards = aggregateSearchLoadByShard(Objects.requireNonNull(shards));
     }
 
+
+
     /**
-     * Returns the array of shard stats.
+     * Returns a copy of the array of shard stats.
      *
      * @return the array of shard stats
      */
     public ShardSearchLoadStats[] getShards() {
-        return shards;
+        return Arrays.copyOf(shards, shards.length);
     }
 
     @Override
@@ -76,6 +82,26 @@ public void writeTo(StreamOutput out) throws IOException {
 
     @Override
     protected Iterator<ToXContent> customXContentChunks(ToXContent.Params params) {
-        return null;
+        return Collections.emptyIterator();
     }
+
+    private ShardSearchLoadStats[] aggregateSearchLoadByShard(ShardSearchLoadStats[] shards) {
+        Map<ShardKey, Double> aggregated = new HashMap<>();
+
+        for (var stat : shards) {
+            var key = new ShardKey(stat.getIndexName(), stat.getShardId());
+            var load = Optional.ofNullable(stat.getSearchLoad()).orElse(0.0);
+            aggregated.merge(key, load, Double::sum);
+        }
+
+        return aggregated.entrySet().stream()
+                .map(e -> new ShardSearchLoadStats(
+                        e.getKey().indexName(),
+                        e.getKey().shardId(),
+                        e.getValue()
+                ))
+                .toArray(ShardSearchLoadStats[]::new);
+    }
+
+    private record ShardKey(String indexName, int shardId) {}
 }

server/src/main/java/org/elasticsearch/action/search/load/TransportShardSearchLoadStatsAction.java

@@ -171,7 +171,6 @@ protected void shardOperation(Request request, ShardRouting shardRouting, Task t
             return new ShardSearchLoadStats(
                 shardId.getIndex().getName(),
                 shardId.getId(),
-                shardRouting.allocationId().getId(),
                 indexShard.getSearchLoadRate()
             );
         });