Add a #markAsNotCachable() method to context to mark requests as not cachable

core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java

@@ -574,8 +574,9 @@ public final class InnerHitBuilder extends ToXContentToBytes implements Writeabl
             innerHitsContext.docValueFieldsContext(new DocValueFieldsContext(docValueFields));
         }
         if (scriptFields != null) {
+            context.markAsNotCachable();
             for (ScriptField field : scriptFields) {
-                SearchScript searchScript = innerHitsContext.scriptService().search(innerHitsContext.lookup(), field.script(),
+                SearchScript searchScript = innerHitsContext.getQueryShardContext().getScriptService().search(innerHitsContext.lookup(), field.script(),
                         ScriptContext.Standard.SEARCH, Collections.emptyMap());
                 innerHitsContext.scriptFields().add(new org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(
                         field.fieldName(), searchScript, field.ignoreFailure()));

core/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java

@@ -41,6 +41,7 @@ public class QueryRewriteContext implements ParseFieldMatcherSupplier {
     protected final Client client;
     protected final IndexReader reader;
     protected final ClusterState clusterState;
+    protected boolean cachable;
 
     public QueryRewriteContext(IndexSettings indexSettings, MapperService mapperService, ScriptService scriptService,
                                IndicesQueriesRegistry indicesQueriesRegistry, Client client, IndexReader reader,
@@ -116,4 +117,12 @@ public class QueryRewriteContext implements ParseFieldMatcherSupplier {
         String defaultScriptLanguage = ScriptSettings.getLegacyDefaultLang(indexSettings.getNodeSettings());
         return new QueryParseContext(defaultScriptLanguage, indicesQueriesRegistry, parser, indexSettings.getParseFieldMatcher());
     }
+
+    public void markAsNotCachable() {
+        this.cachable = false;
+    }
+
+    public boolean isCachable() {
+        return cachable;
+    }
 }

core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java

@@ -324,6 +324,15 @@ public class QueryShardContext extends QueryRewriteContext {
         }
     }
 
+    @Override
+    public void markAsNotCachable() {
+        super.markAsNotCachable();
+        SearchContext current = SearchContext.current();
+        if (current != null) {
+           current.markAsNotCachable();
+        }
+    }
+
     public final Index index() {
         return indexSettings.getIndex();
     }

core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java

@@ -41,7 +41,6 @@ import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
 import java.util.Collections;
-import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 
@@ -134,6 +133,7 @@ public class ScriptQueryBuilder extends AbstractQueryBuilder<ScriptQueryBuilder>
 
     @Override
     protected Query doToQuery(QueryShardContext context) throws IOException {
+        context.markAsNotCachable();
         return new ScriptQuery(script, context.getScriptService(), context.lookup());
     }
 
@@ -216,4 +216,6 @@ public class ScriptQueryBuilder extends AbstractQueryBuilder<ScriptQueryBuilder>
     protected boolean doEquals(ScriptQueryBuilder other) {
         return Objects.equals(script, other.script);
     }
+
+
 }

core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java

@@ -95,6 +95,7 @@ public class ScriptScoreFunctionBuilder extends ScoreFunctionBuilder<ScriptScore
 
     @Override
     protected ScoreFunction doToFunction(QueryShardContext context) {
+        context.markAsNotCachable();
         try {
             SearchScript searchScript = context.getScriptService().search(context.lookup(), script, ScriptContext.Standard.SEARCH,
                     Collections.emptyMap());

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

@@ -1082,7 +1082,7 @@ public class IndicesService extends AbstractLifecycleComponent
         }
         // if now in millis is used (or in the future, a more generic "isDeterministic" flag
         // then we can't cache based on "now" key within the search request, as it is not deterministic
-        if (context.nowInMillisUsed()) {
+        if (context.isCachable() == false) {
             return false;
         }
         return true;

core/src/main/java/org/elasticsearch/search/SearchService.java

@@ -526,7 +526,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
             // be cached. If it is still present in the request (e.g. in a range
             // aggregation) it will still be caught when the aggregation is
             // evaluated.
-            context.resetNowInMillisUsed();
+            context.resetCanCache();
             if (request.scroll() != null) {
                 context.scrollContext(new ScrollContext());
                 context.scrollContext().scroll = request.scroll();
@@ -735,7 +735,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         }
         if (source.scriptFields() != null) {
             for (org.elasticsearch.search.builder.SearchSourceBuilder.ScriptField field : source.scriptFields()) {
-                SearchScript searchScript = context.scriptService().search(context.lookup(), field.script(), ScriptContext.Standard.SEARCH,
+                SearchScript searchScript = scriptService.search(context.lookup(), field.script(), ScriptContext.Standard.SEARCH,
                         Collections.emptyMap());
                 context.scriptFields().add(new ScriptField(field.fieldName(), searchScript, field.ignoreFailure()));
             }

core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/SignificantTermsAggregationBuilder.java

@@ -220,6 +220,9 @@ public class SignificantTermsAggregationBuilder extends ValuesSourceAggregationB
     @Override
     protected ValuesSourceAggregatorFactory<ValuesSource, ?> innerBuild(AggregationContext context, ValuesSourceConfig<ValuesSource> config,
             AggregatorFactory<?> parent, Builder subFactoriesBuilder) throws IOException {
+        if (significanceHeuristic.canCache()) {
+            context.searchContext().markAsNotCachable();
+        }
         return new SignificantTermsAggregatorFactory(name, type, config, includeExclude, executionHint, filterBuilder,
                 bucketCountThresholds, significanceHeuristic, context, parent, subFactoriesBuilder, metaData);
     }

core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java

@@ -78,6 +78,7 @@ public class ScriptHeuristic extends SignificanceHeuristic {
 
     @Override
     public void initialize(SearchContext context) {
+        context.markAsNotCachable();
         initialize(context.scriptService());
     }
 
@@ -217,5 +218,10 @@ public class ScriptHeuristic extends SignificanceHeuristic {
             return Long.toString(value);
         }
     }
+
+    @Override
+    public boolean canCache() {
+        return false;
+    }
 }
 

core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/SignificanceHeuristic.java

@@ -57,4 +57,8 @@ public abstract class SignificanceHeuristic implements NamedWriteable, ToXConten
     public void initialize(SearchContext context) {
 
     }
+
+    public boolean canCache() {
+        return true;
+    }
 }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java

@@ -182,6 +182,7 @@ public class ScriptedMetricAggregationBuilder extends AbstractAggregationBuilder
     @Override
     protected ScriptedMetricAggregatorFactory doBuild(AggregationContext context, AggregatorFactory<?> parent,
             Builder subfactoriesBuilder) throws IOException {
+        context.searchContext().markAsNotCachable();
         return new ScriptedMetricAggregatorFactory(name, type, initScript, mapScript, combineScript, reduceScript, params, context,
                 parent, subfactoriesBuilder, metaData);
     }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java

@@ -53,10 +53,12 @@ public class ScriptedMetricAggregator extends MetricsAggregator {
         this.params = params;
         ScriptService scriptService = context.searchContext().scriptService();
         if (initScript != null) {
+            context.searchContext().markAsNotCachable();
             scriptService.executable(initScript, ScriptContext.Standard.AGGS, Collections.emptyMap()).run();
         }
         this.mapScript = scriptService.search(context.searchContext().lookup(), mapScript, ScriptContext.Standard.AGGS, Collections.emptyMap());
         if (combineScript != null) {
+            context.searchContext().markAsNotCachable();
             this.combineScript = scriptService.executable(combineScript, ScriptContext.Standard.AGGS, Collections.emptyMap());
         } else {
             this.combineScript = null;

core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java

@@ -527,6 +527,9 @@ public class TopHitsAggregationBuilder extends AbstractAggregationBuilder<TopHit
     @Override
     protected TopHitsAggregatorFactory doBuild(AggregationContext context, AggregatorFactory<?> parent, Builder subfactoriesBuilder)
             throws IOException {
+        if (scriptFields != null && scriptFields.isEmpty() == false) {
+            context.searchContext().markAsNotCachable();
+        }
         return new TopHitsAggregatorFactory(name, type, from, size, explain, version, trackScores, sorts, highlightBuilder,
             storedFieldsContext, fieldDataFields, scriptFields, fetchSourceContext, context,
             parent, subfactoriesBuilder, metaData);

core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregatorFactory.java

@@ -99,6 +99,7 @@ public class TopHitsAggregatorFactory extends AggregatorFactory<TopHitsAggregato
             subSearchContext.docValueFieldsContext(new DocValueFieldsContext(docValueFields));
         }
         if (scriptFields != null) {
+            subSearchContext.markAsNotCachable();
             for (ScriptField field : scriptFields) {
                 SearchScript searchScript = subSearchContext.scriptService().search(subSearchContext.lookup(), field.script(),
                         ScriptContext.Standard.SEARCH, Collections.emptyMap());

core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceAggregationBuilder.java

@@ -376,8 +376,12 @@ public abstract class ValuesSourceAggregationBuilder<VS extends ValuesSource, AB
     }
 
     private SearchScript createScript(Script script, SearchContext context) {
-        return script == null ? null
-                : context.scriptService().search(context.lookup(), script, ScriptContext.Standard.AGGS, Collections.emptyMap());
+        if (script == null) {
+            return null;
+        } else {
+            context.markAsNotCachable();
+            return context.scriptService().search(context.lookup(), script, ScriptContext.Standard.AGGS, Collections.emptyMap());
+        }
     }
 
     private static DocValueFormat resolveFormat(@Nullable String format, @Nullable ValueType valueType) {

core/src/main/java/org/elasticsearch/search/internal/SearchContext.java

@@ -118,7 +118,7 @@ public abstract class SearchContext extends AbstractRefCounted implements Releas
         }
     }
 
-    private boolean nowInMillisUsed;
+    private boolean canCache = true;
 
     @Override
     protected final void closeInternal() {
@@ -163,16 +163,16 @@ public abstract class SearchContext extends AbstractRefCounted implements Releas
     public abstract long getOriginNanoTime();
 
     public final long nowInMillis() {
-        nowInMillisUsed = true;
+        markAsNotCachable();
         return nowInMillisImpl();
     }
 
-    public final boolean nowInMillisUsed() {
-        return nowInMillisUsed;
+    public final boolean isCachable() {
+        return canCache;
     }
 
-    public final void resetNowInMillisUsed() {
-        this.nowInMillisUsed = false;
+    public final void resetCanCache() {
+        this.canCache = true;
     }
 
     protected abstract long nowInMillisImpl();
@@ -402,6 +402,10 @@ public abstract class SearchContext extends AbstractRefCounted implements Releas
     /** Return a view of the additional query collectors that should be run for this context. */
     public abstract Map<Class<?>, Collector> queryCollectors();
 
+    public final void markAsNotCachable() {
+        this.canCache = false;
+    }
+
     /**
      * The life time of an object that is used during search execution.
      */

core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java

@@ -55,9 +55,7 @@ import org.elasticsearch.search.MultiValueMode;
 
 import java.io.IOException;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Locale;
-import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 
@@ -286,6 +284,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
 
     @Override
     public SortFieldAndFormat build(QueryShardContext context) throws IOException {
+        context.markAsNotCachable();
         final SearchScript searchScript = context.getScriptService().search(
                 context.lookup(), script, ScriptContext.Standard.SEARCH, Collections.emptyMap());
 

core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionBuilder.java

@@ -633,6 +633,7 @@ public class PhraseSuggestionBuilder extends SuggestionBuilder<PhraseSuggestionB
         }
 
         if (this.collateQuery != null) {
+            context.markAsNotCachable();
             CompiledScript compiledScript = context.getScriptService().compile(this.collateQuery, ScriptContext.Standard.SEARCH,
                     Collections.emptyMap());
             suggestionContext.setCollateQueryScript(compiledScript);

core/src/test/java/org/elasticsearch/index/query/ScriptQueryBuilderTests.java

@@ -89,4 +89,9 @@ public class ScriptQueryBuilderTests extends AbstractQueryTestCase<ScriptQueryBu
         //adding additional objects within the params object.
         return Collections.singleton(Script.ScriptField.PARAMS.getPreferredName());
     }
+
+    @Override
+    protected boolean isCachable(ScriptQueryBuilder queryBuilder) {
+        return false;
+    }
 }

core/src/test/java/org/elasticsearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java

@@ -789,4 +789,15 @@ public class FunctionScoreQueryBuilderTests extends AbstractQueryTestCase<Functi
                     RandomScoreFunctionBuilderWithFixedSeed::new, RandomScoreFunctionBuilderWithFixedSeed::fromXContent));
         }
     }
+
+    @Override
+    protected boolean isCachable(FunctionScoreQueryBuilder queryBuilder) {
+        FilterFunctionBuilder[] filterFunctionBuilders = queryBuilder.filterFunctionBuilders();
+        for (FilterFunctionBuilder builder : filterFunctionBuilders) {
+            if (builder.getScoreFunction() instanceof ScriptScoreFunctionBuilder) {
+                return false;
+            }
+        }
+        return true;
+    }
 }

modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/TemplateQueryBuilder.java

@@ -120,6 +120,7 @@ public class TemplateQueryBuilder extends AbstractQueryBuilder<TemplateQueryBuil
 
     @Override
     protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
+        queryRewriteContext.markAsNotCachable();
         ExecutableScript executable = queryRewriteContext.getScriptService().executable(template,
             ScriptContext.Standard.SEARCH, Collections.emptyMap());
         BytesReference querySource = (BytesReference) executable.run();

test/framework/src/main/java/org/elasticsearch/test/AbstractQueryTestCase.java

@@ -630,12 +630,24 @@ public abstract class AbstractQueryTestCase<QB extends AbstractQueryBuilder<QB>>
     }
 
     private QueryBuilder rewriteQuery(QB queryBuilder, QueryRewriteContext rewriteContext) throws IOException {
+        final boolean wasCachable = rewriteContext.isCachable();
         QueryBuilder rewritten = QueryBuilder.rewriteQuery(queryBuilder, rewriteContext);
+        if (wasCachable == true) { // it's not resettable so we can only check it once
+            if (isCachable(queryBuilder)) {
+                assert rewriteContext.isCachable() : queryBuilder.toString();
+            } else {
+                assert rewriteContext.isCachable() == false;
+            }
+        }
         // extra safety to fail fast - serialize the rewritten version to ensure it's serializable.
         assertSerialization(rewritten);
         return rewritten;
     }
 
+    protected boolean isCachable(QB queryBuilder) {
+        return true;
+    }
+
     /**
      * Few queries allow you to set the boost and queryName on the java api, although the corresponding parser
      * doesn't parse them as they are not supported. This method allows to disable boost and queryName related tests for those queries.