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 ;
@@ -67,6 +70,38 @@ public class IndicesQueryCache implements QueryCache, Closeable {
6770 private final Map <ShardId , Stats > shardStats = new ConcurrentHashMap <>();
6871 private volatile long sharedRamBytesUsed ;
6972
73+ /**
74+ * Calculates a map of {@link ShardId} to {@link Long} which contains the calculated share of the {@link IndicesQueryCache} shared ram
75+ * 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
76+ * participate in the cache, shards whose calculated share is zero will not be contained in the map at all. As a consequence, the
77+ * correct pattern for using the returned map will be via {@link Map#getOrDefault(Object, Object)} with a {@code defaultValue} of
78+ * {@code 0L}.
79+ */
80+ public static Map <ShardId , Long > getSharedRamSizeForAllShards (IndicesService indicesService ) {
81+ Map <ShardId , Long > shardIdToSharedRam = new HashMap <>();
82+ IndicesQueryCache .CacheTotals cacheTotals = IndicesQueryCache .getCacheTotalsForAllShards (indicesService );
83+ for (IndexService indexService : indicesService ) {
84+ for (IndexShard indexShard : indexService ) {
85+ final var queryCache = indicesService .getIndicesQueryCache ();
86+ long sharedRam = (queryCache == null ) ? 0L : queryCache .getSharedRamSizeForShard (indexShard .shardId (), cacheTotals );
87+ // as a size optimization, only store non-zero values in the map
88+ if (sharedRam > 0L ) {
89+ shardIdToSharedRam .put (indexShard .shardId (), sharedRam );
90+ }
91+ }
92+ }
93+ return Collections .unmodifiableMap (shardIdToSharedRam );
94+ }
95+
96+ public long getCacheSizeForShard (ShardId shardId ) {
97+ Stats stats = shardStats .get (shardId );
98+ return stats != null ? stats .cacheSize : 0L ;
99+ }
100+
101+ public long getSharedRamBytesUsed () {
102+ return sharedRamBytesUsed ;
103+ }
104+
70105 // This is a hack for the fact that the close listener for the
71106 // ShardCoreKeyMap will be called before onDocIdSetEviction
72107 // See onDocIdSetEviction for more info
@@ -89,40 +124,58 @@ private static QueryCacheStats toQueryCacheStatsSafe(@Nullable Stats stats) {
89124 return stats == null ? new QueryCacheStats () : stats .toQueryCacheStats ();
90125 }
91126
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
127+ /**
128+ * This computes the total cache size in bytes, and the total shard count in the cache for all shards.
129+ * @param indicesService
130+ * @return A CacheTotals object containing the computed total number of items in the cache and the number of shards seen in the cache
131+ */
132+ private static CacheTotals getCacheTotalsForAllShards (IndicesService indicesService ) {
133+ IndicesQueryCache queryCache = indicesService .getIndicesQueryCache ();
134+ boolean hasQueryCache = queryCache != null ;
102135 long totalItemsInCache = 0L ;
103136 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 ;
137+ for (final IndexService indexService : indicesService ) {
138+ for (final IndexShard indexShard : indexService ) {
139+ final var shardId = indexShard .shardId ();
140+ long cacheSize = hasQueryCache ? queryCache .getCacheSizeForShard (shardId ) : 0L ;
141+ shardCount ++;
142+ assert cacheSize >= 0 : "Unexpected cache size of " + cacheSize + " for shard " + shardId ;
143+ totalItemsInCache += cacheSize ;
117144 }
118145 }
146+ return new CacheTotals (totalItemsInCache , shardCount );
147+ }
119148
149+ public static long getSharedRamSizeForShard (IndicesService indicesService , ShardId shardId ) {
150+ IndicesQueryCache .CacheTotals cacheTotals = IndicesQueryCache .getCacheTotalsForAllShards (indicesService );
151+ final var queryCache = indicesService .getIndicesQueryCache ();
152+ return (queryCache == null ) ? 0L : queryCache .getSharedRamSizeForShard (shardId , cacheTotals );
153+ }
154+
155+ /**
156+ * This method computes the shared RAM size in bytes for the given indexShard.
157+ * @param shardId The shard to compute the shared RAM size for
158+ * @param cacheTotals Shard totals computed in getCacheTotalsForAllShards()
159+ * @return the shared RAM size in bytes allocated to the given shard, or 0 if unavailable
160+ */
161+ private long getSharedRamSizeForShard (ShardId shardId , CacheTotals cacheTotals ) {
162+ long sharedRamBytesUsed = getSharedRamBytesUsed ();
163+ if (sharedRamBytesUsed == 0L ) {
164+ return 0L ;
165+ }
166+
167+ int shardCount = cacheTotals .shardCount ();
120168 if (shardCount == 0 ) {
121169 // 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
122170 // extend beyond the closing of all shards.
123171 return 0L ;
124172 }
125-
173+ /*
174+ * We have some shared ram usage that we try to distribute proportionally to the number of segment-requests in the cache for each
175+ * shard.
176+ */
177+ long totalItemsInCache = cacheTotals .totalItemsInCache ();
178+ long itemsInCacheForShard = getCacheSizeForShard (shardId );
126179 final long additionalRamBytesUsed ;
127180 if (totalItemsInCache == 0 ) {
128181 // all shards have zero cache footprint, so we apportion the size of the shared bytes equally across all shards
@@ -143,10 +196,12 @@ private long getShareOfAdditionalRamBytesUsed(long itemsInCacheForShard) {
143196 return additionalRamBytesUsed ;
144197 }
145198
199+ private record CacheTotals (long totalItemsInCache , int shardCount ) {}
200+
146201 /** Get usage statistics for the given shard. */
147- public QueryCacheStats getStats (ShardId shard ) {
202+ public QueryCacheStats getStats (ShardId shard , long precomputedSharedRamBytesUsed ) {
148203 final QueryCacheStats queryCacheStats = toQueryCacheStatsSafe (shardStats .get (shard ));
149- queryCacheStats .addRamBytesUsed (getShareOfAdditionalRamBytesUsed ( queryCacheStats . getCacheSize ()) );
204+ queryCacheStats .addRamBytesUsed (precomputedSharedRamBytesUsed );
150205 return queryCacheStats ;
151206 }
152207
@@ -243,7 +298,7 @@ QueryCacheStats toQueryCacheStats() {
243298 public String toString () {
244299 return "{shardId="
245300 + shardId
246- + ", ramBytedUsed ="
301+ + ", ramBytesUsed ="
247302 + ramBytesUsed
248303 + ", hitCount="
249304 + hitCount
@@ -340,11 +395,7 @@ protected void onDocIdSetCache(Object readerCoreKey, long ramBytesUsed) {
340395 shardStats .cacheCount += 1 ;
341396 shardStats .ramBytesUsed += ramBytesUsed ;
342397
343- StatsAndCount statsAndCount = stats2 .get (readerCoreKey );
344- if (statsAndCount == null ) {
345- statsAndCount = new StatsAndCount (shardStats );
346- stats2 .put (readerCoreKey , statsAndCount );
347- }
398+ StatsAndCount statsAndCount = stats2 .computeIfAbsent (readerCoreKey , ignored -> new StatsAndCount (shardStats ));
348399 statsAndCount .count += 1 ;
349400 }
350401
@@ -357,7 +408,7 @@ protected void onDocIdSetEviction(Object readerCoreKey, int numEntries, long sum
357408 if (numEntries > 0 ) {
358409 // We can't use ShardCoreKeyMap here because its core closed
359410 // listener is called before the listener of the cache which
360- // triggers this eviction. So instead we use use stats2 that
411+ // triggers this eviction. So instead we use stats2 that
361412 // we only evict when nothing is cached anymore on the segment
362413 // instead of relying on close listeners
363414 final StatsAndCount statsAndCount = stats2 .get (readerCoreKey );
0 commit comments