2626import org .elasticsearch .common .unit .ByteSizeValue ;
2727import org .elasticsearch .core .Nullable ;
2828import org .elasticsearch .core .Predicates ;
29+ import org .elasticsearch .index .IndexService ;
2930import org .elasticsearch .index .cache .query .QueryCacheStats ;
31+ import org .elasticsearch .index .shard .IndexShard ;
3032import org .elasticsearch .index .shard .ShardId ;
3133
3234import java .io .Closeable ;
3335import java .io .IOException ;
3436import java .util .Collections ;
37+ import java .util .HashMap ;
3538import java .util .IdentityHashMap ;
3639import java .util .Map ;
3740import java .util .Set ;
3841import java .util .concurrent .ConcurrentHashMap ;
3942import java .util .function .Predicate ;
43+ import java .util .function .Supplier ;
4044
4145public class IndicesQueryCache implements QueryCache , Closeable {
4246
@@ -67,6 +71,40 @@ public class IndicesQueryCache implements QueryCache, Closeable {
6771 private final Map <ShardId , Stats > shardStats = new ConcurrentHashMap <>();
6872 private volatile long sharedRamBytesUsed ;
6973
74+ /**
75+ * Calculates a map of {@link ShardId} to {@link Long} which contains the calculated share of the {@link IndicesQueryCache} shared ram
76+ * size for a given shard (that is, the sum of all the longs is the size of the indices query cache). Since many shards will not
77+ * participate in the cache, shards whose calculated share is zero will not be contained in the map at all. As a consequence, the
78+ * correct pattern for using the returned map will be via {@link Map#getOrDefault(Object, Object)} with a {@code defaultValue} of
79+ * {@code 0L}.
80+ * @return an unmodifiable map from {@link ShardId} to the calculated share of the query cache's shared RAM size for each shard,
81+ * omitting shards with a zero share
82+ */
83+ public static Map <ShardId , Long > getSharedRamSizeForAllShards (IndicesService indicesService ) {
84+ Map <ShardId , Long > shardIdToSharedRam = new HashMap <>();
85+ IndicesQueryCache .CacheTotals cacheTotals = IndicesQueryCache .getCacheTotalsForAllShards (indicesService );
86+ for (IndexService indexService : indicesService ) {
87+ for (IndexShard indexShard : indexService ) {
88+ final var queryCache = indicesService .getIndicesQueryCache ();
89+ long sharedRam = (queryCache == null ) ? 0L : queryCache .getSharedRamSizeForShard (indexShard .shardId (), cacheTotals );
90+ // as a size optimization, only store non-zero values in the map
91+ if (sharedRam > 0L ) {
92+ shardIdToSharedRam .put (indexShard .shardId (), sharedRam );
93+ }
94+ }
95+ }
96+ return Collections .unmodifiableMap (shardIdToSharedRam );
97+ }
98+
99+ public long getCacheSizeForShard (ShardId shardId ) {
100+ Stats stats = shardStats .get (shardId );
101+ return stats != null ? stats .cacheSize : 0L ;
102+ }
103+
104+ public long getSharedRamBytesUsed () {
105+ return sharedRamBytesUsed ;
106+ }
107+
70108 // This is a hack for the fact that the close listener for the
71109 // ShardCoreKeyMap will be called before onDocIdSetEviction
72110 // See onDocIdSetEviction for more info
@@ -89,40 +127,52 @@ private static QueryCacheStats toQueryCacheStatsSafe(@Nullable Stats stats) {
89127 return stats == null ? new QueryCacheStats () : stats .toQueryCacheStats ();
90128 }
91129
92- private long getShareOfAdditionalRamBytesUsed (long itemsInCacheForShard ) {
93- if (sharedRamBytesUsed == 0L ) {
94- return 0L ;
95- }
96-
97- /*
98- * We have some shared ram usage that we try to distribute proportionally to the number of segment-requests in the cache for each
99- * shard.
100- */
101- // TODO avoid looping over all local shards here - see https://github.com/elastic/elasticsearch/issues/97222
130+ /**
131+ * Computes the total cache size in bytes, and the total shard count in the cache for all shards.
132+ * @param indicesService the IndicesService instance to retrieve cache information from
133+ * @return A CacheTotals object containing the computed total number of items in the cache and the number of shards seen in the cache
134+ */
135+ public static CacheTotals getCacheTotalsForAllShards (IndicesService indicesService ) {
136+ IndicesQueryCache queryCache = indicesService .getIndicesQueryCache ();
137+ boolean hasQueryCache = queryCache != null ;
102138 long totalItemsInCache = 0L ;
103139 int shardCount = 0 ;
104- if (itemsInCacheForShard == 0L ) {
105- for (final var stats : shardStats .values ()) {
106- shardCount += 1 ;
107- if (stats .cacheSize > 0L ) {
108- // some shard has nonzero cache footprint, so we apportion the shared size by cache footprint, and this shard has none
109- return 0L ;
110- }
111- }
112- } else {
113- // branchless loop for the common case
114- for (final var stats : shardStats .values ()) {
115- shardCount += 1 ;
116- totalItemsInCache += stats .cacheSize ;
140+ for (final IndexService indexService : indicesService ) {
141+ for (final IndexShard indexShard : indexService ) {
142+ final var shardId = indexShard .shardId ();
143+ long cacheSize = hasQueryCache ? queryCache .getCacheSizeForShard (shardId ) : 0L ;
144+ shardCount ++;
145+ assert cacheSize >= 0 : "Unexpected cache size of " + cacheSize + " for shard " + shardId ;
146+ totalItemsInCache += cacheSize ;
117147 }
118148 }
149+ return new CacheTotals (totalItemsInCache , shardCount );
150+ }
151+
152+ /**
153+ * This method computes the shared RAM size in bytes for the given indexShard.
154+ * @param shardId The shard to compute the shared RAM size for.
155+ * @param cacheTotals Shard totals computed in {@link #getCacheTotalsForAllShards(IndicesService)}.
156+ * @return the shared RAM size in bytes allocated to the given shard, or 0 if unavailable
157+ */
158+ public long getSharedRamSizeForShard (ShardId shardId , CacheTotals cacheTotals ) {
159+ long sharedRamBytesUsed = getSharedRamBytesUsed ();
160+ if (sharedRamBytesUsed == 0L ) {
161+ return 0L ;
162+ }
119163
164+ int shardCount = cacheTotals .shardCount ();
120165 if (shardCount == 0 ) {
121166 // Sometimes it's not possible to do this when there are no shard entries at all, which can happen as the shared ram usage can
122167 // extend beyond the closing of all shards.
123168 return 0L ;
124169 }
125-
170+ /*
171+ * We have some shared ram usage that we try to distribute proportionally to the number of segment-requests in the cache for each
172+ * shard.
173+ */
174+ long totalItemsInCache = cacheTotals .totalItemsInCache ();
175+ long itemsInCacheForShard = getCacheSizeForShard (shardId );
126176 final long additionalRamBytesUsed ;
127177 if (totalItemsInCache == 0 ) {
128178 // all shards have zero cache footprint, so we apportion the size of the shared bytes equally across all shards
@@ -143,10 +193,12 @@ private long getShareOfAdditionalRamBytesUsed(long itemsInCacheForShard) {
143193 return additionalRamBytesUsed ;
144194 }
145195
196+ public record CacheTotals (long totalItemsInCache , int shardCount ) {}
197+
146198 /** Get usage statistics for the given shard. */
147- public QueryCacheStats getStats (ShardId shard ) {
199+ public QueryCacheStats getStats (ShardId shard , Supplier < Long > precomputedSharedRamBytesUsed ) {
148200 final QueryCacheStats queryCacheStats = toQueryCacheStatsSafe (shardStats .get (shard ));
149- queryCacheStats .addRamBytesUsed (getShareOfAdditionalRamBytesUsed ( queryCacheStats . getCacheSize () ));
201+ queryCacheStats .addRamBytesUsed (precomputedSharedRamBytesUsed . get ( ));
150202 return queryCacheStats ;
151203 }
152204
@@ -243,7 +295,7 @@ QueryCacheStats toQueryCacheStats() {
243295 public String toString () {
244296 return "{shardId="
245297 + shardId
246- + ", ramBytedUsed ="
298+ + ", ramBytesUsed ="
247299 + ramBytesUsed
248300 + ", hitCount="
249301 + hitCount
@@ -340,11 +392,7 @@ protected void onDocIdSetCache(Object readerCoreKey, long ramBytesUsed) {
340392 shardStats .cacheCount += 1 ;
341393 shardStats .ramBytesUsed += ramBytesUsed ;
342394
343- StatsAndCount statsAndCount = stats2 .get (readerCoreKey );
344- if (statsAndCount == null ) {
345- statsAndCount = new StatsAndCount (shardStats );
346- stats2 .put (readerCoreKey , statsAndCount );
347- }
395+ StatsAndCount statsAndCount = stats2 .computeIfAbsent (readerCoreKey , ignored -> new StatsAndCount (shardStats ));
348396 statsAndCount .count += 1 ;
349397 }
350398
@@ -357,7 +405,7 @@ protected void onDocIdSetEviction(Object readerCoreKey, int numEntries, long sum
357405 if (numEntries > 0 ) {
358406 // We can't use ShardCoreKeyMap here because its core closed
359407 // listener is called before the listener of the cache which
360- // triggers this eviction. So instead we use use stats2 that
408+ // triggers this eviction. So instead we use stats2 that
361409 // we only evict when nothing is cached anymore on the segment
362410 // instead of relying on close listeners
363411 final StatsAndCount statsAndCount = stats2 .get (readerCoreKey );
0 commit comments