Invalidate cached query results if query timed out (#22807)

Today we cache query results even if the query timed out. This is obviously
problematic since results are not complete. Yet, the decision if a query timed
out or not happens too late to simply not cache the result since if we'd just throw
an exception all currently waiting requests with the same request / cache key would
fail with the same exception without the option to access the result or to re-execute.
Instead, this change will allow the request to enter the cache but invalidates it immediately.
Concurrent request might not get executed and return the timed out result which is not absolutely
correct but very likely since identical requests will likely timeout as well. As a side-effect
we won't hammer the node with concurrent slow searches but rather only execute one of them
and return shortly cached result.

Closes #22789

Author: s1monw

core/src/main/java/org/elasticsearch/indices/IndicesRequestCache.java

@@ -130,6 +130,16 @@ BytesReference getOrCompute(CacheEntity cacheEntity, Supplier<BytesReference> lo
         return value;
     }
 
+    /**
+     * Invalidates the given the cache entry for the given key and it's context
+     * @param cacheEntity the cache entity to invalidate for
+     * @param reader the reader to invalidate the cache entry for
+     * @param cacheKey the cache key to invalidate
+     */
+    void invalidate(CacheEntity cacheEntity, DirectoryReader reader, BytesReference cacheKey) {
+        cache.invalidate(new Key(cacheEntity, reader.getVersion(), cacheKey));
+    }
+
     private static class Loader implements CacheLoader<Key, BytesReference> {
 
         private final CacheEntity entity;

core/src/main/java/org/elasticsearch/indices/IndicesService.java

@@ -1134,17 +1134,28 @@ public void loadIntoContext(ShardSearchRequest request, SearchContext context, Q
             queryPhase.execute(context);
             try {
                 context.queryResult().writeToNoId(out);
+
             } catch (IOException e) {
                 throw new AssertionError("Could not serialize response", e);
             }
             loadedFromCache[0] = false;
         });
+
         if (loadedFromCache[0]) {
             // restore the cached query result into the context
             final QuerySearchResult result = context.queryResult();
             StreamInput in = new NamedWriteableAwareStreamInput(bytesReference.streamInput(), namedWriteableRegistry);
             result.readFromWithId(context.id(), in);
             result.shardTarget(context.shardTarget());
+        } else if (context.queryResult().searchTimedOut()) {
+            // we have to invalidate the cache entry if we cached a query result form a request that timed out.
+            // we can't really throw exceptions in the loading part to signal a timed out search to the outside world since if there are
+            // multiple requests that wait for the cache entry to be calculated they'd fail all with the same exception.
+            // instead we all caching such a result for the time being, return the timed out result for all other searches with that cache
+            // key invalidate the result in the thread that caused the timeout. This will end up to be simpler and eventually correct since
+            // running a search that times out concurrently will likely timeout again if it's run while we have this `stale` result in the
+            // cache. One other option is to not cache requests with a timeout at all...
+            indicesRequestCache.invalidate(new IndexShardCacheEntity(context.indexShard()), directoryReader, request.cacheKey());
         }
     }
 

core/src/test/java/org/elasticsearch/indices/IndicesRequestCacheTests.java

@@ -169,7 +169,6 @@ public void testCacheDifferentReaders() throws Exception {
         assertEquals(2, requestCacheStats.stats().getMissCount());
         assertEquals(0, requestCacheStats.stats().getEvictions());
         assertTrue(loader.loadedFromCache);
-        assertTrue(loader.loadedFromCache);
         assertEquals(1, cache.count());
         assertEquals(cacheSize, requestCacheStats.stats().getMemorySize().bytesAsInt());
         assertEquals(1, cache.numRegisteredCloseListeners());
@@ -186,7 +185,6 @@ public void testCacheDifferentReaders() throws Exception {
         assertEquals(2, requestCacheStats.stats().getMissCount());
         assertEquals(0, requestCacheStats.stats().getEvictions());
         assertTrue(loader.loadedFromCache);
-        assertTrue(loader.loadedFromCache);
         assertEquals(0, cache.count());
         assertEquals(0, requestCacheStats.stats().getMemorySize().bytesAsInt());
 
@@ -215,7 +213,7 @@ public void testEviction() throws Exception {
                 new ShardId("foo", "bar", 1));
             TestEntity secondEntity = new TestEntity(requestCacheStats, indexShard);
             Loader secondLoader = new Loader(secondReader, 0);
-            
+
             BytesReference value1 = cache.getOrCompute(entity, loader, reader, termQuery.buildAsBytes());
             assertEquals("foo", value1.streamInput().readString());
             BytesReference value2 = cache.getOrCompute(secondEntity, secondLoader, secondReader, termQuery.buildAsBytes());
@@ -347,6 +345,74 @@ public BytesReference get() {
 
     }
 
+    public void testInvalidate() throws Exception {
+        ShardRequestCache requestCacheStats = new ShardRequestCache();
+        IndicesRequestCache cache = new IndicesRequestCache(Settings.EMPTY);
+        Directory dir = newDirectory();
+        IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig());
+
+        writer.addDocument(newDoc(0, "foo"));
+        DirectoryReader reader = ElasticsearchDirectoryReader.wrap(DirectoryReader.open(writer),
+            new ShardId("foo", "bar", 1));
+        TermQueryBuilder termQuery = new TermQueryBuilder("id", "0");
+        AtomicBoolean indexShard = new AtomicBoolean(true);
+
+        // initial cache
+        TestEntity entity = new TestEntity(requestCacheStats, indexShard);
+        Loader loader = new Loader(reader, 0);
+        BytesReference value = cache.getOrCompute(entity, loader, reader, termQuery.buildAsBytes());
+        assertEquals("foo", value.streamInput().readString());
+        assertEquals(0, requestCacheStats.stats().getHitCount());
+        assertEquals(1, requestCacheStats.stats().getMissCount());
+        assertEquals(0, requestCacheStats.stats().getEvictions());
+        assertFalse(loader.loadedFromCache);
+        assertEquals(1, cache.count());
+
+        // cache hit
+        entity = new TestEntity(requestCacheStats, indexShard);
+        loader = new Loader(reader, 0);
+        value = cache.getOrCompute(entity, loader, reader, termQuery.buildAsBytes());
+        assertEquals("foo", value.streamInput().readString());
+        assertEquals(1, requestCacheStats.stats().getHitCount());
+        assertEquals(1, requestCacheStats.stats().getMissCount());
+        assertEquals(0, requestCacheStats.stats().getEvictions());
+        assertTrue(loader.loadedFromCache);
+        assertEquals(1, cache.count());
+        assertTrue(requestCacheStats.stats().getMemorySize().bytesAsInt() > value.length());
+        assertEquals(1, cache.numRegisteredCloseListeners());
+
+        // load again after invalidate
+        entity = new TestEntity(requestCacheStats, indexShard);
+        loader = new Loader(reader, 0);
+        cache.invalidate(entity, reader,  termQuery.buildAsBytes());
+        value = cache.getOrCompute(entity, loader, reader, termQuery.buildAsBytes());
+        assertEquals("foo", value.streamInput().readString());
+        assertEquals(1, requestCacheStats.stats().getHitCount());
+        assertEquals(2, requestCacheStats.stats().getMissCount());
+        assertEquals(0, requestCacheStats.stats().getEvictions());
+        assertFalse(loader.loadedFromCache);
+        assertEquals(1, cache.count());
+        assertTrue(requestCacheStats.stats().getMemorySize().bytesAsInt() > value.length());
+        assertEquals(1, cache.numRegisteredCloseListeners());
+
+        // release
+        if (randomBoolean()) {
+            reader.close();
+        } else {
+            indexShard.set(false); // closed shard but reader is still open
+            cache.clear(entity);
+        }
+        cache.cleanCache();
+        assertEquals(1, requestCacheStats.stats().getHitCount());
+        assertEquals(2, requestCacheStats.stats().getMissCount());
+        assertEquals(0, requestCacheStats.stats().getEvictions());
+        assertEquals(0, cache.count());
+        assertEquals(0, requestCacheStats.stats().getMemorySize().bytesAsInt());
+
+        IOUtils.close(reader, writer, dir, cache);
+        assertEquals(0, cache.numRegisteredCloseListeners());
+    }
+
     private class TestEntity extends AbstractIndexShardCacheEntity {
         private final AtomicBoolean standInForIndexShard;
         private final ShardRequestCache shardRequestCache;