Rework MockInternalClusterInfoService (#90092)

This changes MockInternalClusterInfoService to fake entire data structure rather
 than a getter call. This is required if the structure is copied by some
 code (such as desired balance allocator).

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

@@ -34,19 +34,9 @@ public class ShardStats implements Writeable, ToXContentFragment {
     private final CommitStats commitStats;
     @Nullable
     private final SeqNoStats seqNoStats;
-
     @Nullable
     private final RetentionLeaseStats retentionLeaseStats;
 
-    /**
-     * Gets the current retention lease stats.
-     *
-     * @return the current retention lease stats
-     */
-    public RetentionLeaseStats getRetentionLeaseStats() {
-        return retentionLeaseStats;
-    }
-
     private final String dataPath;
     private final String statePath;
     private final boolean isCustomDataPath;
@@ -67,21 +57,43 @@ public class ShardStats implements Writeable, ToXContentFragment {
     }
 
     public ShardStats(
-        final ShardRouting routing,
+        final ShardRouting shardRouting,
         final ShardPath shardPath,
         final CommonStats commonStats,
         final CommitStats commitStats,
         final SeqNoStats seqNoStats,
         final RetentionLeaseStats retentionLeaseStats
     ) {
-        this.shardRouting = routing;
-        this.dataPath = shardPath.getRootDataPath().toString();
-        this.statePath = shardPath.getRootStatePath().toString();
-        this.isCustomDataPath = shardPath.isCustomDataPath();
-        this.commitStats = commitStats;
+        this(
+            shardRouting,
+            commonStats,
+            commitStats,
+            seqNoStats,
+            retentionLeaseStats,
+            shardPath.getRootDataPath().toString(),
+            shardPath.getRootStatePath().toString(),
+            shardPath.isCustomDataPath()
+        );
+    }
+
+    public ShardStats(
+        ShardRouting shardRouting,
+        CommonStats commonStats,
+        CommitStats commitStats,
+        SeqNoStats seqNoStats,
+        RetentionLeaseStats retentionLeaseStats,
+        String dataPath,
+        String statePath,
+        boolean isCustomDataPath
+    ) {
+        this.shardRouting = shardRouting;
         this.commonStats = commonStats;
+        this.commitStats = commitStats;
         this.seqNoStats = seqNoStats;
         this.retentionLeaseStats = retentionLeaseStats;
+        this.dataPath = dataPath;
+        this.statePath = statePath;
+        this.isCustomDataPath = isCustomDataPath;
     }
 
     @Override
@@ -125,6 +137,15 @@ public class ShardStats implements Writeable, ToXContentFragment {
         return this.seqNoStats;
     }
 
+    /**
+     * Gets the current retention lease stats.
+     *
+     * @return the current retention lease stats
+     */
+    public RetentionLeaseStats getRetentionLeaseStats() {
+        return retentionLeaseStats;
+    }
+
     public String getDataPath() {
         return dataPath;
     }

server/src/main/java/org/elasticsearch/cluster/ClusterInfo.java

@@ -10,6 +10,7 @@ package org.elasticsearch.cluster;
 
 import org.elasticsearch.Version;
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
@@ -227,6 +228,11 @@ public class ClusterInfo implements ToXContentFragment, Writeable {
         return shardRouting.shardId().toString() + "[" + (shardRouting.primary() ? "p" : "r") + "]";
     }
 
+    @Override
+    public String toString() {
+        return Strings.toString(this, true, false);
+    }
+
     /**
      * Represents a data path on a node
      */

server/src/main/java/org/elasticsearch/cluster/InternalClusterInfoService.java

@@ -250,7 +250,7 @@ public class InternalClusterInfoService implements ClusterInfoService, ClusterSt
                                     new HashMap<>();
                                 buildShardLevelInfo(
                                     clusterService.state().routingTable(),
-                                    stats,
+                                    adjustShardStats(stats),
                                     shardSizeByIdentifierBuilder,
                                     shardDataSetSizeBuilder,
                                     dataPathByShardRoutingBuilder,
@@ -430,6 +430,10 @@ public class InternalClusterInfoService implements ClusterInfoService, ClusterSt
         return nodeStats;
     }
 
+    ShardStats[] adjustShardStats(ShardStats[] shardStats) {
+        return shardStats;
+    }
+
     void refreshAsync(ActionListener<ClusterInfo> future) {
         final Runnable newRefresh;
         synchronized (mutex) {

server/src/main/java/org/elasticsearch/index/engine/CommitStats.java

@@ -12,18 +12,14 @@ import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.lucene.Lucene;
-import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.xcontent.ToXContentFragment;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Base64;
 import java.util.Map;
 import java.util.Objects;
 
-import static java.util.Map.entry;
-
 /** a class the returns dynamic information with respect to the last commit point of this shard */
 public final class CommitStats implements Writeable, ToXContentFragment {
 
@@ -42,12 +38,7 @@ public final class CommitStats implements Writeable, ToXContentFragment {
     }
 
     CommitStats(StreamInput in) throws IOException {
-        final int length = in.readVInt();
-        final var entries = new ArrayList<Map.Entry<String, String>>(length);
-        for (int i = length; i > 0; i--) {
-            entries.add(entry(in.readString(), in.readString()));
-        }
-        userData = Maps.ofEntries(entries);
+        userData = in.readImmutableMap(StreamInput::readString, StreamInput::readString);
         generation = in.readLong();
         id = in.readOptionalString();
         numDocs = in.readInt();

test/framework/src/main/java/org/elasticsearch/cluster/MockInternalClusterInfoService.java

@@ -8,6 +8,9 @@
 package org.elasticsearch.cluster;
 
 import org.elasticsearch.action.admin.cluster.node.stats.NodeStats;
+import org.elasticsearch.action.admin.indices.stats.CommonStats;
+import org.elasticsearch.action.admin.indices.stats.CommonStatsFlags;
+import org.elasticsearch.action.admin.indices.stats.ShardStats;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.ShardRouting;
@@ -15,10 +18,12 @@ import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.index.store.StoreStats;
 import org.elasticsearch.monitor.fs.FsInfo;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.threadpool.ThreadPool;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.function.BiFunction;
 import java.util.function.Function;
@@ -50,15 +55,9 @@ public class MockInternalClusterInfoService extends InternalClusterInfoService {
         ClusterInfoServiceUtils.refresh(this);
     }
 
-    @Override
-    public ClusterInfo getClusterInfo() {
-        final ClusterInfo clusterInfo = super.getClusterInfo();
-        return new SizeFakingClusterInfo(clusterInfo);
-    }
-
     @Override
     List<NodeStats> adjustNodesStats(List<NodeStats> nodesStats) {
-        final BiFunction<DiscoveryNode, FsInfo.Path, FsInfo.Path> diskUsageFunctionCopy = this.diskUsageFunction;
+        var diskUsageFunctionCopy = this.diskUsageFunction;
         if (diskUsageFunctionCopy == null) {
             return nodesStats;
         }
@@ -94,27 +93,35 @@ public class MockInternalClusterInfoService extends InternalClusterInfoService {
         }).collect(Collectors.toList());
     }
 
-    class SizeFakingClusterInfo extends ClusterInfo {
-        SizeFakingClusterInfo(ClusterInfo delegate) {
-            super(
-                delegate.getNodeLeastAvailableDiskUsages(),
-                delegate.getNodeMostAvailableDiskUsages(),
-                delegate.shardSizes,
-                delegate.shardDataSetSizes,
-                delegate.routingToDataPath,
-                delegate.reservedSpace
-            );
+    @Override
+    ShardStats[] adjustShardStats(ShardStats[] shardsStats) {
+        var shardSizeFunctionCopy = this.shardSizeFunction;
+        if (shardSizeFunctionCopy == null) {
+            return shardsStats;
         }
 
-        @Override
-        public Long getShardSize(ShardRouting shardRouting) {
-            final Function<ShardRouting, Long> shardSizeFunctionCopy = MockInternalClusterInfoService.this.shardSizeFunction;
-            if (shardSizeFunctionCopy == null) {
-                return super.getShardSize(shardRouting);
-            }
+        return Arrays.stream(shardsStats).map(shardStats -> {
 
-            return shardSizeFunctionCopy.apply(shardRouting);
-        }
+            var shardRouting = shardStats.getShardRouting();
+            var storeStats = new StoreStats(
+                shardSizeFunctionCopy.apply(shardRouting),
+                shardSizeFunctionCopy.apply(shardRouting),
+                shardStats.getStats().store.getReservedSize().getBytes()
+            );
+            var commonStats = new CommonStats(new CommonStatsFlags(CommonStatsFlags.Flag.Store));
+            commonStats.store = storeStats;
+
+            return new ShardStats(
+                shardRouting,
+                commonStats,
+                shardStats.getCommitStats(),
+                shardStats.getSeqNoStats(),
+                shardStats.getRetentionLeaseStats(),
+                shardStats.getDataPath(),
+                shardStats.getStatePath(),
+                shardStats.isCustomDataPath()
+            );
+        }).toArray(ShardStats[]::new);
     }
 
     @Override