Skip to content

Commit 3873deb

Browse files
masseykemasseyke
committed
adding unit tests
1 parent 21ac4f4 commit 3873deb

File tree

3 files changed

+95
-30
lines changed

3 files changed

+95
-30
lines changed

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,9 +117,7 @@ protected void shardOperation(IndicesStatsRequest request, ShardRouting shardRou
117117
IndicesQueryCache.CacheTotals cacheTotals = IndicesQueryCache.getCacheTotalsForAllShards(indicesService);
118118
IndexService indexService = indicesService.indexServiceSafe(shardRouting.shardId().getIndex());
119119
IndexShard indexShard = indexService.getShard(shardRouting.shardId().id());
120-
ShardId shardId = indexShard.shardId();
121120
long sharedRam = IndicesQueryCache.getSharedRamSizeForShard(indicesService.getIndicesQueryCache(), indexShard, cacheTotals);
122-
123121
CommonStats commonStats = CommonStats.getShardLevelStats(
124122
indicesService.getIndicesQueryCache(),
125123
indexShard,

‎server/src/main/java/org/elasticsearch/indices/IndicesQueryCache.java‎

Lines changed: 36 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -69,13 +69,11 @@ public class IndicesQueryCache implements QueryCache, Closeable {
6969
private final Map<ShardId, Stats> shardStats = new ConcurrentHashMap<>();
7070
private volatile long sharedRamBytesUsed;
7171

72-
// Package-private for IndicesService efficient stats collection
7372
public long getCacheSizeForShard(ShardId shardId) {
7473
Stats stats = shardStats.get(shardId);
7574
return stats != null ? stats.cacheSize : 0L;
7675
}
7776

78-
// Package-private for IndicesService efficient stats collection
7977
public long getSharedRamBytesUsed() {
8078
return sharedRamBytesUsed;
8179
}
@@ -102,24 +100,11 @@ private static QueryCacheStats toQueryCacheStatsSafe(@Nullable Stats stats) {
102100
return stats == null ? new QueryCacheStats() : stats.toQueryCacheStats();
103101
}
104102

105-
/** Get usage statistics for the given shard. */
106-
public QueryCacheStats getStats(ShardId shard, long precomputedSharedRamBytesUsed) {
107-
final QueryCacheStats queryCacheStats = toQueryCacheStatsSafe(shardStats.get(shard));
108-
queryCacheStats.addRamBytesUsed(precomputedSharedRamBytesUsed);
109-
return queryCacheStats;
110-
}
111-
112-
@Override
113-
public Weight doCache(Weight weight, QueryCachingPolicy policy) {
114-
while (weight instanceof CachingWeightWrapper) {
115-
weight = ((CachingWeightWrapper) weight).in;
116-
}
117-
final Weight in = cache.doCache(weight, policy);
118-
// We wrap the weight to track the readers it sees and map them with
119-
// the shards they belong to
120-
return new CachingWeightWrapper(in);
121-
}
122-
103+
/**
104+
* This computes the total cache size in bytes, and the total shard count in the cache for all shards.
105+
* @param indicesService
106+
* @return A CacheTotals object containing the computed total cache size and shard count
107+
*/
123108
public static CacheTotals getCacheTotalsForAllShards(IndicesService indicesService) {
124109
IndicesQueryCache queryCache = indicesService.getIndicesQueryCache();
125110
boolean hasQueryCache = queryCache != null;
@@ -129,17 +114,22 @@ public static CacheTotals getCacheTotalsForAllShards(IndicesService indicesServi
129114
for (final IndexShard indexShard : indexService) {
130115
long cacheSize = hasQueryCache ? queryCache.getCacheSizeForShard(indexShard.shardId()) : 0L;
131116
shardCount++;
132-
if (cacheSize > 0L) {
133-
totalSize += cacheSize;
134-
}
117+
assert cacheSize >= 0 : "Unexpected cache size of " + cacheSize + " for shard " + indexShard.shardId();
118+
totalSize += cacheSize;
135119
}
136120
}
137-
long sharedRamBytesUsed = hasQueryCache ? queryCache.getSharedRamBytesUsed() : 0L;
138-
return new CacheTotals(totalSize, shardCount, sharedRamBytesUsed);
121+
return new CacheTotals(totalSize, shardCount);
139122
}
140123

124+
/**
125+
* This method computes the shared RAM size in bytes for the given indexShard.
126+
* @param queryCache
127+
* @param indexShard The shard to compute the shared RAM size for
128+
* @param cacheTotals Shard totals computed in getCacheTotalsForAllShards()
129+
* @return
130+
*/
141131
public static long getSharedRamSizeForShard(IndicesQueryCache queryCache, IndexShard indexShard, CacheTotals cacheTotals) {
142-
long sharedRamBytesUsed = cacheTotals.sharedRamBytesUsed();
132+
long sharedRamBytesUsed = queryCache != null ? queryCache.getSharedRamBytesUsed() : 0L;
143133
if (sharedRamBytesUsed == 0L) {
144134
return 0L;
145135
}
@@ -152,7 +142,7 @@ public static long getSharedRamSizeForShard(IndicesQueryCache queryCache, IndexS
152142
}
153143

154144
long totalSize = cacheTotals.totalSize();
155-
long cacheSize = queryCache != null ? queryCache.getCacheSizeForShard(indexShard.shardId()) : 0L;
145+
long cacheSize = queryCache.getCacheSizeForShard(indexShard.shardId());
156146
final long additionalRamBytesUsed;
157147
if (totalSize == 0) {
158148
// all shards have zero cache footprint, so we apportion the size of the shared bytes equally across all shards
@@ -165,7 +155,25 @@ public static long getSharedRamSizeForShard(IndicesQueryCache queryCache, IndexS
165155
return additionalRamBytesUsed;
166156
}
167157

168-
public record CacheTotals(long totalSize, int shardCount, long sharedRamBytesUsed) {}
158+
public record CacheTotals(long totalSize, int shardCount) {}
159+
160+
/** Get usage statistics for the given shard. */
161+
public QueryCacheStats getStats(ShardId shard, long precomputedSharedRamBytesUsed) {
162+
final QueryCacheStats queryCacheStats = toQueryCacheStatsSafe(shardStats.get(shard));
163+
queryCacheStats.addRamBytesUsed(precomputedSharedRamBytesUsed);
164+
return queryCacheStats;
165+
}
166+
167+
@Override
168+
public Weight doCache(Weight weight, QueryCachingPolicy policy) {
169+
while (weight instanceof CachingWeightWrapper) {
170+
weight = ((CachingWeightWrapper) weight).in;
171+
}
172+
final Weight in = cache.doCache(weight, policy);
173+
// We wrap the weight to track the readers it sees and map them with
174+
// the shards they belong to
175+
return new CachingWeightWrapper(in);
176+
}
169177

170178
private class CachingWeightWrapper extends Weight {
171179

‎server/src/test/java/org/elasticsearch/indices/IndicesQueryCacheTests.java‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,19 @@
2929
import org.elasticsearch.common.lucene.index.ElasticsearchDirectoryReader;
3030
import org.elasticsearch.common.settings.Settings;
3131
import org.elasticsearch.core.IOUtils;
32+
import org.elasticsearch.index.IndexService;
3233
import org.elasticsearch.index.cache.query.QueryCacheStats;
3334
import org.elasticsearch.index.cache.query.TrivialQueryCachingPolicy;
35+
import org.elasticsearch.index.shard.IndexShard;
3436
import org.elasticsearch.index.shard.ShardId;
3537
import org.elasticsearch.test.ESTestCase;
3638

3739
import java.io.IOException;
40+
import java.util.List;
41+
42+
import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
43+
import static org.mockito.Mockito.mock;
44+
import static org.mockito.Mockito.when;
3845

3946
public class IndicesQueryCacheTests extends ESTestCase {
4047

@@ -404,4 +411,56 @@ public void testDelegatesScorerSupplier() throws Exception {
404411
cache.onClose(shard);
405412
cache.close();
406413
}
414+
public void testGetCacheTotalsForAllShards() throws Exception {
415+
ShardId shardId1 = new ShardId("index", "_na_", 0);
416+
ShardId shardId2 = new ShardId("index", "_na_", 1);
417+
418+
IndexShard shard1 = mock(IndexShard.class);
419+
IndexShard shard2 = mock(IndexShard.class);
420+
when(shard1.shardId()).thenReturn(shardId1);
421+
when(shard2.shardId()).thenReturn(shardId2);
422+
423+
IndexService indexService = mock(IndexService.class, RETURNS_DEEP_STUBS);
424+
when(indexService.iterator()).thenReturn(List.of(shard1, shard2).iterator());
425+
426+
IndicesService indicesService = mock(IndicesService.class, RETURNS_DEEP_STUBS);
427+
when(indicesService.iterator()).thenReturn(List.of(indexService).iterator());
428+
IndicesQueryCache queryCache = mock(IndicesQueryCache.class);
429+
when(indicesService.getIndicesQueryCache()).thenReturn(queryCache);
430+
when(queryCache.getCacheSizeForShard(shardId1)).thenReturn(100L);
431+
when(queryCache.getCacheSizeForShard(shardId2)).thenReturn(200L);
432+
433+
IndicesQueryCache.CacheTotals totals = IndicesQueryCache.getCacheTotalsForAllShards(indicesService);
434+
assertEquals(300L, totals.totalSize());
435+
assertEquals(2, totals.shardCount());
436+
}
437+
438+
public void testGetSharedRamSizeForShard() {
439+
ShardId shardId1 = new ShardId("index", "_na_", 0);
440+
ShardId shardId2 = new ShardId("index", "_na_", 1);
441+
IndexShard shard1 = mock(IndexShard.class);
442+
IndexShard shard2 = mock(IndexShard.class);
443+
when(shard1.shardId()).thenReturn(shardId1);
444+
when(shard2.shardId()).thenReturn(shardId2);
445+
446+
IndicesQueryCache.CacheTotals totals = new IndicesQueryCache.CacheTotals(300L, 2);
447+
IndicesQueryCache queryCache = mock(IndicesQueryCache.class);
448+
// Case 1: sharedRamBytesUsed = 0
449+
when(queryCache.getSharedRamBytesUsed()).thenReturn(0L);
450+
long sharedRam = IndicesQueryCache.getSharedRamSizeForShard(queryCache, shard1, totals);
451+
assertEquals(0L, sharedRam);
452+
// Case 2: sharedRamBytesUsed > 0, totalSize > 0, proportional
453+
when(queryCache.getSharedRamBytesUsed()).thenReturn(600L);
454+
when(queryCache.getCacheSizeForShard(shardId1)).thenReturn(100L);
455+
long sharedRam1 = IndicesQueryCache.getSharedRamSizeForShard(queryCache, shard1, totals);
456+
assertEquals(200L, sharedRam1);
457+
when(queryCache.getCacheSizeForShard(shardId2)).thenReturn(200L);
458+
long sharedRam2 = IndicesQueryCache.getSharedRamSizeForShard(queryCache, shard2, totals);
459+
assertEquals(400L, sharedRam2);
460+
// Case 3: totalSize == 0, shared equally
461+
IndicesQueryCache.CacheTotals zeroTotals = new IndicesQueryCache.CacheTotals(0L, 2);
462+
when(queryCache.getSharedRamBytesUsed()).thenReturn(600L);
463+
long sharedRamEq = IndicesQueryCache.getSharedRamSizeForShard(queryCache, shard1, zeroTotals);
464+
assertEquals(300L, sharedRamEq);
465+
}
407466
}

0 commit comments

Comments
 (0)