Fix IRC timeout bug

Signed-off-by: Peter Alfonsi <petealft@amazon.com>

Author: Peter Alfonsi

server/src/internalClusterTest/java/org/opensearch/indices/IndicesRequestCacheIT.java

@@ -34,6 +34,12 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
+import org.apache.lucene.index.Term;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.TermQuery;
+import org.apache.lucene.search.Weight;
 import org.opensearch.action.admin.cluster.health.ClusterHealthResponse;
 import org.opensearch.action.admin.cluster.node.stats.NodeStats;
 import org.opensearch.action.admin.cluster.node.stats.NodesStatsResponse;
@@ -56,7 +62,10 @@
 import org.opensearch.env.NodeEnvironment;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.cache.request.RequestCacheStats;
+import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.index.query.QueryShardContext;
+import org.opensearch.index.query.TermQueryBuilder;
 import org.opensearch.search.aggregations.bucket.global.GlobalAggregationBuilder;
 import org.opensearch.search.aggregations.bucket.histogram.DateHistogramInterval;
 import org.opensearch.search.aggregations.bucket.histogram.Histogram;
@@ -65,6 +74,7 @@
 import org.opensearch.test.ParameterizedStaticSettingsOpenSearchIntegTestCase;
 import org.opensearch.test.hamcrest.OpenSearchAssertions;
 
+import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.time.ZoneId;
@@ -768,6 +778,61 @@ public void testDeleteAndCreateSameIndexShardOnSameNode() throws Exception {
         assertTrue(stats.getMemorySizeInBytes() == 0);
     }
 
+    public void testTimedOutQuery() throws Exception {
+        // A timed out query should be cached and then invalidated
+        Client client = client();
+        String index = "index";
+        assertAcked(
+            client.admin()
+                .indices()
+                .prepareCreate(index)
+                .setMapping("k", "type=keyword")
+                .setSettings(
+                    Settings.builder()
+                        .put(IndicesRequestCache.INDEX_CACHE_REQUEST_ENABLED_SETTING.getKey(), true)
+                        .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+                        .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+                        // Disable index refreshing to avoid cache being invalidated mid-test
+                        .put(IndexSettings.INDEX_REFRESH_INTERVAL_SETTING.getKey(), TimeValue.timeValueMillis(-1))
+                )
+                .get()
+        );
+        indexRandom(true, client.prepareIndex(index).setSource("k", "hello"));
+        ensureSearchable(index);
+        // Force merge the index to ensure there can be no background merges during the subsequent searches that would invalidate the cache
+        forceMerge(client, index);
+
+        QueryBuilder timeoutQueryBuilder = new TermQueryBuilder("k", "hello") {
+            @Override
+            protected Query doToQuery(QueryShardContext context) {
+                return new TermQuery(new Term("k", "hello")) {
+                    @Override
+                    public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
+                        // Create the weight before sleeping. Otherwise, TermStates.build() (in the call to super.createWeight()) will
+                        // sometimes throw an exception on timeout, rather than timing out gracefully.
+                        Weight result = super.createWeight(searcher, scoreMode, boost);
+                        try {
+                            // Pick 500 ms as it's the same duration used in SearchTimeoutIT.testSimpleTimeout() to ensure a timeout
+                            // (We can't directly reuse their ScriptQuery-based logic as it isn't cacheable)
+                            Thread.sleep(500);
+                        } catch (InterruptedException ignored) {}
+                        return result;
+                    }
+                };
+            }
+        };
+
+        SearchResponse resp = client.prepareSearch(index)
+            .setRequestCache(true)
+            .setQuery(timeoutQueryBuilder)
+            .setTimeout(TimeValue.ZERO)
+            .get();
+        assertTrue(resp.isTimedOut());
+        RequestCacheStats requestCacheStats = getRequestCacheStats(client, index);
+        // The cache should be empty as the timed-out query was invalidated
+        assertEquals(0, requestCacheStats.getMemorySizeInBytes());
+    }
+
     private Path[] shardDirectory(String server, Index index, int shard) {
         NodeEnvironment env = internalCluster().getInstance(NodeEnvironment.class, server);
         final Path[] paths = env.availableShardPaths(new ShardId(index, shard));

server/src/main/java/org/opensearch/indices/IndicesRequestCache.java

@@ -310,12 +310,11 @@ BytesReference getOrCompute(
      * @param cacheKey the cache key to invalidate
      */
     void invalidate(IndicesService.IndexShardCacheEntity cacheEntity, DirectoryReader reader, BytesReference cacheKey) {
-        assert reader.getReaderCacheHelper() != null;
-        String readerCacheKeyId = null;
-        if (reader instanceof OpenSearchDirectoryReader) {
-            IndexReader.CacheHelper cacheHelper = ((OpenSearchDirectoryReader) reader).getDelegatingCacheHelper();
-            readerCacheKeyId = ((OpenSearchDirectoryReader.DelegatingCacheHelper) cacheHelper).getDelegatingCacheKey().getId();
-        }
+        assert reader.getReaderCacheHelper() instanceof OpenSearchDirectoryReader.DelegatingCacheHelper;
+        OpenSearchDirectoryReader.DelegatingCacheHelper delegatingCacheHelper = (OpenSearchDirectoryReader.DelegatingCacheHelper) reader
+            .getReaderCacheHelper();
+        String readerCacheKeyId = delegatingCacheHelper.getDelegatingCacheKey().getId();
+
         IndexShard indexShard = (IndexShard) cacheEntity.getCacheIdentity();
         cache.invalidate(getICacheKey(new Key(indexShard.shardId(), cacheKey, readerCacheKeyId, System.identityHashCode(indexShard))));
     }

server/src/main/java/org/opensearch/indices/IndicesService.java

@@ -68,6 +68,7 @@
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.common.lease.Releasable;
 import org.opensearch.common.lifecycle.AbstractLifecycleComponent;
+import org.opensearch.common.lucene.index.OpenSearchDirectoryReader;
 import org.opensearch.common.settings.IndexScopedSettings;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Setting.Property;
@@ -1754,6 +1755,9 @@ public boolean canCache(ShardSearchRequest request, SearchContext context) {
         if (context.getQueryShardContext().isCacheable() == false) {
             return false;
         }
+        if (!(context.searcher().getIndexReader().getReaderCacheHelper() instanceof OpenSearchDirectoryReader.DelegatingCacheHelper)) {
+            return false;
+        }
         return true;
 
     }