First attempt at improving statsByShard performance

masseyke · masseyke · commit 168c6c53793f · 2025-07-08T16:09:31.000-05:00
diff --git a/server/src/main/java/org/elasticsearch/action/admin/indices/stats/CommonStats.java b/server/src/main/java/org/elasticsearch/action/admin/indices/stats/CommonStats.java
@@ -155,6 +155,18 @@ public CommonStats(CommonStatsFlags flags) {
      * Filters the given flags for {@link CommonStatsFlags#SHARD_LEVEL} flags and calculates the corresponding statistics.
      */
     public static CommonStats getShardLevelStats(IndicesQueryCache indicesQueryCache, IndexShard indexShard, CommonStatsFlags flags) {
+        return getShardLevelStats(indicesQueryCache, indexShard, flags, null);
+    }
+
+    /**
+     * Overload that takes a precomputed shared RAM map for O(N) stats.
+     */
+    public static CommonStats getShardLevelStats(
+        IndicesQueryCache indicesQueryCache,
+        IndexShard indexShard,
+        CommonStatsFlags flags,
+        java.util.Map<org.elasticsearch.index.shard.ShardId, Long> precomputedSharedRam
+    ) {
         // Filter shard level flags
         CommonStatsFlags filteredFlags = flags.clone();
         for (CommonStatsFlags.Flag flag : filteredFlags.getFlags()) {
@@ -174,7 +186,16 @@ public static CommonStats getShardLevelStats(IndicesQueryCache indicesQueryCache
                     case Refresh -> stats.refresh = indexShard.refreshStats();
                     case Flush -> stats.flush = indexShard.flushStats();
                     case Warmer -> stats.warmer = indexShard.warmerStats();
-                    case QueryCache -> stats.queryCache = indicesQueryCache.getStats(indexShard.shardId());
+                    case QueryCache -> {
+                        if (precomputedSharedRam != null && precomputedSharedRam.containsKey(indexShard.shardId())) {
+                            stats.queryCache = indicesQueryCache.getStats(
+                                indexShard.shardId(),
+                                precomputedSharedRam.get(indexShard.shardId())
+                            );
+                        } else {
+                            stats.queryCache = indicesQueryCache.getStats(indexShard.shardId());
+                        }
+                    }
                     case FieldData -> stats.fieldData = indexShard.fieldDataStats(flags.fieldDataFields());
                     case Completion -> stats.completion = indexShard.completionStats(flags.completionDataFields());
                     case Segments -> stats.segments = indexShard.segmentStats(
diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesQueryCache.java b/server/src/main/java/org/elasticsearch/indices/IndicesQueryCache.java
@@ -32,6 +32,7 @@
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.Set;
@@ -67,6 +68,17 @@ public class IndicesQueryCache implements QueryCache, Closeable {
     private final Map<ShardId, Stats> shardStats = new ConcurrentHashMap<>();
     private volatile long sharedRamBytesUsed;
 
+    // Package-private for IndicesService efficient stats collection
+    long getCacheSizeForShard(ShardId shardId) {
+        Stats stats = shardStats.get(shardId);
+        return stats != null ? stats.cacheSize : 0L;
+    }
+
+    // Package-private for IndicesService efficient stats collection
+    long getSharedRamBytesUsed() {
+        return sharedRamBytesUsed;
+    }
+
     // This is a hack for the fact that the close listener for the
     // ShardCoreKeyMap will be called before onDocIdSetEviction
     // See onDocIdSetEviction for more info
@@ -139,6 +151,57 @@ public QueryCacheStats getStats(ShardId shard) {
         return queryCacheStats;
     }
 
+    /**
+     * Overload to allow passing in a precomputed shared RAM split for this shard.
+     */
+    public QueryCacheStats getStats(ShardId shard, long precomputedSharedRamBytesUsed) {
+        final QueryCacheStats queryCacheStats = toQueryCacheStatsSafe(shardStats.get(shard));
+        queryCacheStats.addRamBytesUsed(precomputedSharedRamBytesUsed);
+        return queryCacheStats;
+    }
+
+    /**
+     * Precompute the shared RAM split for all shards, returning a map of ShardId to the additional shared RAM bytes used.
+     * This avoids O(N^2) when collecting stats for all shards.
+     */
+    public Map<ShardId, Long> computeAllShardSharedRamBytesUsed() {
+        Map<ShardId, Long> result = new HashMap<>();
+        if (sharedRamBytesUsed == 0L) {
+            for (ShardId shardId : shardStats.keySet()) {
+                result.put(shardId, 0L);
+            }
+            return result;
+        }
+        long totalSize = 0L;
+        int shardCount = 0;
+        boolean anyNonZero = false;
+        for (Stats stats : shardStats.values()) {
+            shardCount += 1;
+            if (stats.cacheSize > 0L) {
+                anyNonZero = true;
+                totalSize += stats.cacheSize;
+            }
+        }
+        if (shardCount == 0) {
+            return result;
+        }
+        if (anyNonZero == false) {
+            // All shards have zero cache footprint, apportion equally
+            long perShard = Math.round((double) sharedRamBytesUsed / shardCount);
+            for (ShardId shardId : shardStats.keySet()) {
+                result.put(shardId, perShard);
+            }
+        } else {
+            // Apportion proportionally to cache footprint
+            for (Map.Entry<ShardId, Stats> entry : shardStats.entrySet()) {
+                long cacheSize = entry.getValue().cacheSize;
+                long ram = (totalSize == 0) ? 0L : Math.round((double) sharedRamBytesUsed * cacheSize / totalSize);
+                result.put(entry.getKey(), ram);
+            }
+        }
+        return result;
+    }
+
     @Override
     public Weight doCache(Weight weight, QueryCachingPolicy policy) {
         while (weight instanceof CachingWeightWrapper) {
diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesService.java b/server/src/main/java/org/elasticsearch/indices/IndicesService.java
@@ -167,6 +167,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.EnumMap;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -519,33 +520,81 @@ static Map<Index, CommonStats> statsByIndex(final IndicesService indicesService,
     }
 
     static Map<Index, List<IndexShardStats>> statsByShard(final IndicesService indicesService, final CommonStatsFlags flags) {
-        final Map<Index, List<IndexShardStats>> statsByShard = new HashMap<>();
-
+        IndicesQueryCache queryCache = indicesService.getIndicesQueryCache();
+        // First pass: gather all shards, cache sizes, and compute totals
+        class ShardCacheInfo {
+            final IndexService indexService;
+            final IndexShard indexShard;
+            final org.elasticsearch.index.shard.ShardId shardId;
+            final long cacheSize;
+
+            ShardCacheInfo(IndexService is, IndexShard shard, org.elasticsearch.index.shard.ShardId id, long size) {
+                this.indexService = is;
+                this.indexShard = shard;
+                this.shardId = id;
+                this.cacheSize = size;
+            }
+        }
+        List<ShardCacheInfo> shardInfos = new ArrayList<>();
+        long totalSize = 0L;
+        int shardCount = 0;
+        boolean anyNonZero = false;
         for (final IndexService indexService : indicesService) {
             for (final IndexShard indexShard : indexService) {
-                try {
-                    final IndexShardStats indexShardStats = indicesService.indexShardStats(indicesService, indexShard, flags);
-
-                    if (indexShardStats == null) {
-                        continue;
-                    }
-
-                    if (statsByShard.containsKey(indexService.index()) == false) {
-                        statsByShard.put(indexService.index(), arrayAsArrayList(indexShardStats));
-                    } else {
-                        statsByShard.get(indexService.index()).add(indexShardStats);
-                    }
-                } catch (IllegalIndexShardStateException | AlreadyClosedException e) {
-                    // we can safely ignore illegal state on ones that are closing for example
-                    logger.trace(() -> format("%s ignoring shard stats", indexShard.shardId()), e);
+                org.elasticsearch.index.shard.ShardId shardId = indexShard.shardId();
+                long cacheSize = queryCache.getCacheSizeForShard(shardId);
+                shardInfos.add(new ShardCacheInfo(indexService, indexShard, shardId, cacheSize));
+                shardCount++;
+                if (cacheSize > 0L) {
+                    anyNonZero = true;
+                    totalSize += cacheSize;
                 }
             }
         }
-
+        long sharedRamBytesUsed = queryCache.getSharedRamBytesUsed();
+        final Map<Index, List<IndexShardStats>> statsByShard = new HashMap<>();
+        // Second pass: build stats, compute shared RAM on the fly
+        for (ShardCacheInfo info : shardInfos) {
+            long sharedRam = 0L;
+            if (sharedRamBytesUsed != 0L) {
+                if (anyNonZero == false) {
+                    sharedRam = Math.round((double) sharedRamBytesUsed / shardCount);
+                } else if (totalSize != 0) {
+                    sharedRam = Math.round((double) sharedRamBytesUsed * info.cacheSize / totalSize);
+                }
+            }
+            try {
+                final IndexShardStats indexShardStats = indicesService.indexShardStats(
+                    indicesService,
+                    info.indexShard,
+                    flags,
+                    Collections.singletonMap(info.shardId, sharedRam)
+                );
+                if (indexShardStats == null) {
+                    continue;
+                }
+                if (statsByShard.containsKey(info.indexService.index()) == false) {
+                    statsByShard.put(info.indexService.index(), arrayAsArrayList(indexShardStats));
+                } else {
+                    statsByShard.get(info.indexService.index()).add(indexShardStats);
+                }
+            } catch (IllegalIndexShardStateException | AlreadyClosedException e) {
+                logger.trace(() -> format("%s ignoring shard stats", info.shardId), e);
+            }
+        }
         return statsByShard;
     }
 
     IndexShardStats indexShardStats(final IndicesService indicesService, final IndexShard indexShard, final CommonStatsFlags flags) {
+        return indexShardStats(indicesService, indexShard, flags, null);
+    }
+
+    IndexShardStats indexShardStats(
+        final IndicesService indicesService,
+        final IndexShard indexShard,
+        final CommonStatsFlags flags,
+        Map<org.elasticsearch.index.shard.ShardId, Long> precomputedSharedRam
+    ) {
         if (indexShard.routingEntry() == null) {
             return null;
         }
@@ -570,7 +619,7 @@ IndexShardStats indexShardStats(final IndicesService indicesService, final Index
                 new ShardStats(
                     indexShard.routingEntry(),
                     indexShard.shardPath(),
-                    CommonStats.getShardLevelStats(indicesService.getIndicesQueryCache(), indexShard, flags),
+                    CommonStats.getShardLevelStats(indicesService.getIndicesQueryCache(), indexShard, flags, precomputedSharedRam),
                     commitStats,
                     seqNoStats,
                     retentionLeaseStats,
