From 94b7873b49a83ac769b6b83ddd715b88e0697dbf Mon Sep 17 00:00:00 2001
From: Simon Willnauer <simonw@apache.org>
Date: Tue, 4 Oct 2016 18:05:00 +0200
Subject: [PATCH] Add a #markAsNotCachable() method to context to mark requests
 as not cachable

---
 .../index/query/InnerHitBuilder.java          |  3 ++-
 .../index/query/QueryRewriteContext.java      |  9 ++++++++
 .../index/query/QueryShardContext.java        |  9 ++++++++
 .../index/query/ScriptQueryBuilder.java       |  4 +++-
 .../ScriptScoreFunctionBuilder.java           |  1 +
 .../elasticsearch/indices/IndicesService.java |  2 +-
 .../elasticsearch/search/SearchService.java   |  4 ++--
 .../SignificantTermsAggregationBuilder.java   |  3 +++
 .../heuristics/ScriptHeuristic.java           |  6 ++++++
 .../heuristics/SignificanceHeuristic.java     |  4 ++++
 .../ScriptedMetricAggregationBuilder.java     |  1 +
 .../scripted/ScriptedMetricAggregator.java    |  2 ++
 .../tophits/TopHitsAggregationBuilder.java    |  3 +++
 .../tophits/TopHitsAggregatorFactory.java     |  1 +
 .../ValuesSourceAggregationBuilder.java       |  8 +++++--
 .../search/internal/SearchContext.java        | 16 ++++++++------
 .../search/sort/ScriptSortBuilder.java        |  3 +--
 .../phrase/PhraseSuggestionBuilder.java       |  1 +
 .../index/query/ScriptQueryBuilderTests.java  | 21 ++++++++++++-------
 .../FunctionScoreQueryBuilderTests.java       | 11 ++++++++++
 .../script/mustache/TemplateQueryBuilder.java |  1 +
 .../test/AbstractQueryTestCase.java           | 12 +++++++++++
 22 files changed, 102 insertions(+), 23 deletions(-)

diff --git a/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java b/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java
index 454c808a5d5..0f4a1f3c429 100644
--- a/core/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java
+++ b/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()));
diff --git a/core/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java b/core/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java
index f12605088e6..64329c314dd 100644
--- a/core/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java
+++ b/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;
+    }
 }
diff --git a/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java b/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java
index 377afca2f68..25b7ef35961 100644
--- a/core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java
+++ b/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();
     }
diff --git a/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java b/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
index 3ff924b28db..5a65b4438f0 100644
--- a/core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
+++ b/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);
     }
+
+
 }
diff --git a/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java b/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java
index e2fbc5955d7..d7bff2bdd44 100644
--- a/core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java
+++ b/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());
diff --git a/core/src/main/java/org/elasticsearch/indices/IndicesService.java b/core/src/main/java/org/elasticsearch/indices/IndicesService.java
index abc9873efaf..b658f70ed58 100644
--- a/core/src/main/java/org/elasticsearch/indices/IndicesService.java
+++ b/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;
diff --git a/core/src/main/java/org/elasticsearch/search/SearchService.java b/core/src/main/java/org/elasticsearch/search/SearchService.java
index 4334c5cf541..da3cd4729db 100644
--- a/core/src/main/java/org/elasticsearch/search/SearchService.java
+++ b/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()));
             }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/SignificantTermsAggregationBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/SignificantTermsAggregationBuilder.java
index 245297d72cc..bc210c7f897 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/SignificantTermsAggregationBuilder.java
+++ b/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);
     }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java
index c933f9ef596..1c3f5b7ed56 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/ScriptHeuristic.java
+++ b/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;
+    }
 }
 
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/SignificanceHeuristic.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/SignificanceHeuristic.java
index db9711c1a8d..094f529e92c 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/significant/heuristics/SignificanceHeuristic.java
+++ b/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;
+    }
 }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java
index 244881a5155..b13b6fe3019 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregationBuilder.java
+++ b/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);
     }
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java
index e2d3034fa11..89ec5183c32 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java
+++ b/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;
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java
index 3547db7140c..6c26743d4dc 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregationBuilder.java
+++ b/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);
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregatorFactory.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregatorFactory.java
index 7c6a743a20b..5dff86c5eed 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsAggregatorFactory.java
+++ b/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());
diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceAggregationBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceAggregationBuilder.java
index 8f14a1ffaf9..cbf57aae48a 100644
--- a/core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceAggregationBuilder.java
+++ b/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) {
diff --git a/core/src/main/java/org/elasticsearch/search/internal/SearchContext.java b/core/src/main/java/org/elasticsearch/search/internal/SearchContext.java
index 63a1995b08f..1f90587abb0 100644
--- a/core/src/main/java/org/elasticsearch/search/internal/SearchContext.java
+++ b/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.
      */
diff --git a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
index 0a7cb5e1b36..d41e8c66d6f 100644
--- a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
+++ b/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());
 
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionBuilder.java b/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionBuilder.java
index b0cd6a20499..a17a4f5de80 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggestionBuilder.java
+++ b/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);
diff --git a/core/src/test/java/org/elasticsearch/index/query/ScriptQueryBuilderTests.java b/core/src/test/java/org/elasticsearch/index/query/ScriptQueryBuilderTests.java
index 63b8101515e..b86babb5a31 100644
--- a/core/src/test/java/org/elasticsearch/index/query/ScriptQueryBuilderTests.java
+++ b/core/src/test/java/org/elasticsearch/index/query/ScriptQueryBuilderTests.java
@@ -51,7 +51,7 @@ public class ScriptQueryBuilderTests extends AbstractQueryTestCase<ScriptQueryBu
 
     public void testFromJsonVerbose() throws IOException {
         String json =
-                "{\n" +
+            "{\n" +
                 "  \"script\" : {\n" +
                 "    \"script\" : {\n" +
                 "      \"inline\" : \"5\",\n" +
@@ -71,13 +71,13 @@ public class ScriptQueryBuilderTests extends AbstractQueryTestCase<ScriptQueryBu
 
     public void testFromJson() throws IOException {
         String json =
-                "{\n" +
-                        "  \"script\" : {\n" +
-                        "    \"script\" : \"5\"," +
-                        "    \"boost\" : 1.0,\n" +
-                        "    \"_name\" : \"PcKdEyPOmR\"\n" +
-                        "  }\n" +
-                        "}";
+            "{\n" +
+                "  \"script\" : {\n" +
+                "    \"script\" : \"5\"," +
+                "    \"boost\" : 1.0,\n" +
+                "    \"_name\" : \"PcKdEyPOmR\"\n" +
+                "  }\n" +
+                "}";
 
         ScriptQueryBuilder parsed = (ScriptQueryBuilder) parseQuery(json);
         assertEquals(json, "5", parsed.script().getScript());
@@ -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;
+    }
 }
diff --git a/core/src/test/java/org/elasticsearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java b/core/src/test/java/org/elasticsearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java
index 0eab330e367..bbc7d945046 100644
--- a/core/src/test/java/org/elasticsearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java
+++ b/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;
+    }
 }
diff --git a/modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/TemplateQueryBuilder.java b/modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/TemplateQueryBuilder.java
index 4a19bcb4b93..0ed8b03ede3 100644
--- a/modules/lang-mustache/src/main/java/org/elasticsearch/script/mustache/TemplateQueryBuilder.java
+++ b/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();
diff --git a/test/framework/src/main/java/org/elasticsearch/test/AbstractQueryTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/AbstractQueryTestCase.java
index 6225f5fa5d0..854eb767bd0 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/AbstractQueryTestCase.java
+++ b/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.