Fix inaccurate total hit count in _search/template api

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

@@ -72,7 +72,6 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
             searchTemplateRequest = SearchTemplateRequest.fromXContent(parser);
         }
         searchTemplateRequest.setRequest(searchRequest);
-        RestSearchAction.checkRestTotalHits(request, searchRequest);
 
         return channel -> client.execute(SearchTemplateAction.INSTANCE, searchTemplateRequest, new RestStatusToXContentListener<>(channel));
     }

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

@@ -32,11 +32,13 @@
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.rest.action.search.RestSearchAction;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.script.TemplateScript;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.transport.TransportService;
 
@@ -110,8 +112,27 @@ static SearchRequest convert(SearchTemplateRequest searchTemplateRequest, Search
             builder.parseXContent(parser, false);
             builder.explain(searchTemplateRequest.isExplain());
             builder.profile(searchTemplateRequest.isProfile());
+            checkRestTotalHitsAsInt(searchRequest, builder);
             searchRequest.source(builder);
         }
         return searchRequest;
     }
+
+    private static void checkRestTotalHitsAsInt(SearchRequest searchRequest, SearchSourceBuilder searchSourceBuilder) {
+        if (searchRequest.source() == null) {
+            searchRequest.source(new SearchSourceBuilder());
+        }
+        Integer trackTotalHitsUpTo = searchRequest.source().trackTotalHitsUpTo();
+        // trackTotalHitsUpTo is set to Integer.MAX_VALUE when `rest_total_hits_as_int` is true
+        if (trackTotalHitsUpTo != null) {
+            if (searchSourceBuilder.trackTotalHitsUpTo() == null) {
+                // trackTotalHitsUpTo should be set here, ensure that we can get an accurate total hits count
+                searchSourceBuilder.trackTotalHitsUpTo(trackTotalHitsUpTo);
+            } else if (searchSourceBuilder.trackTotalHitsUpTo() != SearchContext.TRACK_TOTAL_HITS_ACCURATE
+                && searchSourceBuilder.trackTotalHitsUpTo() != SearchContext.TRACK_TOTAL_HITS_DISABLED) {
+                throw new IllegalArgumentException("[" + RestSearchAction.TOTAL_HITS_AS_INT_PARAM + "] cannot be used " +
+                    "if the tracking of total hits is not accurate, got " + searchSourceBuilder.trackTotalHitsUpTo());
+            }
+        }
+    }
 }

modules/lang-mustache/src/test/resources/rest-api-spec/test/lang_mustache/30_search_template.yml

@@ -40,6 +40,19 @@
 
   - match: { hits.total: 2 }
 
+---
+"Test with invalid track_total_hits":
+
+  - do:
+      catch:              bad_request
+      search_template:
+        rest_total_hits_as_int: true
+        body: { "source" : { "query": { "match_{{template}}": {} }, "track_total_hits": "{{trackTotalHits}}" }, "params" : { "template" : "all", "trackTotalHits" : 1 } }
+
+  - match: { status: 400 }
+  - match: { error.type: illegal_argument_exception }
+  - match: { error.reason: "[rest_total_hits_as_int] cannot be used if the tracking of total hits is not accurate, got 1" }
+
 ---
 "Missing template search request":
 

modules/lang-mustache/src/test/resources/rest-api-spec/test/lang_mustache/50_multi_search_template.yml

@@ -94,13 +94,21 @@ setup:
           - source: '{"query": {"{{query_type}}": {} }' # Unknown query type
             params:
               query_type: "unknown"
+          # Search 4 has an unsupported track_total_hits
+          - index: index_*
+          - source: '{"query": {"match_all": {} }, "track_total_hits": "{{trackTotalHits}}" }'
+            params:
+              trackTotalHits: 1
 
   - match:  { responses.0.hits.total:     2  }
   - match:  { responses.1.error.root_cause.0.type: json_e_o_f_exception }
   - match:  { responses.1.error.root_cause.0.reason: "/Unexpected.end.of.input/" }
   - match:  { responses.2.hits.total:     1  }
   - match:  { responses.3.error.root_cause.0.type: parsing_exception }
   - match:  { responses.3.error.root_cause.0.reason: "/unknown.query.\\[unknown\\]/" }
+  - match:  { responses.4.error.root_cause.0.type: illegal_argument_exception }
+  - match:  { responses.4.error.root_cause.0.reason: "[rest_total_hits_as_int] cannot be used if the tracking of total hits is not accurate, got 1" }
+
 
 ---
 "Multi-search template with invalid request":