Search & Count: Add option to early terminate doc collection

Allow users to control document collection termination, if a specified terminate_after number is
set. Upon setting the newly added parameter, the response will include a boolean terminated_early
flag, indicating if the document collection for any shard terminated early.

closes #6876

docs/reference/search/count.asciidoc

@@ -63,6 +63,12 @@ query.
 
 |default_operator |The default operator to be used, can be `AND` or
 `OR`. Defaults to `OR`.
+
+|coming[1.4.0] terminate_after |The maximum count for each shard, upon
+reaching which the query execution will terminate early.
+If set, the response will have a boolean field `terminated_early` to
+indicate whether the query execution has actually terminated_early.
+Defaults to no terminate_after.
 |=======================================================================
 
 [float]

docs/reference/search/request-body.asciidoc

@@ -62,6 +62,12 @@ that point when expired. Defaults to no timeout.
 `query_and_fetch`. Defaults to `query_then_fetch`. See
 <<search-request-search-type,_Search Type_>> for
 more details on the different types of search that can be performed.
+
+|coming[1.4.0] `terminate_after` |The maximum number of documents to collect for
+each shard, upon reaching which the query execution will terminate early.
+If set, the response will have a boolean field `terminated_early` to
+indicate whether the query execution has actually terminated_early.
+Defaults to no terminate_after.
 |=======================================================================
 
 Out of the above, the `search_type` is the one that can not be passed

docs/reference/search/uri-request.asciidoc

@@ -82,6 +82,12 @@ scores and return them as part of each hit.
 within the specified time value and bail with the hits accumulated up to
 that point when expired. Defaults to no timeout.
 
+|coming[1.4.0] `terminate_after` |The maximum number of documents to collect for
+each shard, upon reaching which the query execution will terminate early.
+If set, the response will have a boolean field `terminated_early` to
+indicate whether the query execution has actually terminated_early.
+Defaults to no terminate_after.
+
 |`from` |The starting from index of the hits to return. Defaults to `0`.
 
 |`size` |The number of hits to return. Defaults to `10`.

src/main/java/org/elasticsearch/action/count/CountRequest.java

@@ -20,6 +20,8 @@
 package org.elasticsearch.action.count;
 
 import org.elasticsearch.ElasticsearchGenerationException;
+import org.elasticsearch.ElasticsearchIllegalArgumentException;
+import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.support.QuerySourceBuilder;
 import org.elasticsearch.action.support.broadcast.BroadcastOperationRequest;
@@ -34,6 +36,8 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentHelper;
 
+import static org.elasticsearch.search.internal.SearchContext.DEFAULT_TERMINATE_AFTER;
+
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Map;
@@ -67,6 +71,7 @@ public class CountRequest extends BroadcastOperationRequest<CountRequest> {
     private String[] types = Strings.EMPTY_ARRAY;
 
     long nowInMillis;
+    private int terminateAfter = DEFAULT_TERMINATE_AFTER;
 
     CountRequest() {
     }
@@ -221,6 +226,21 @@ public class CountRequest extends BroadcastOperationRequest<CountRequest> {
         return this.preference;
     }
 
+    /**
+     * Upon reaching <code>terminateAfter</code> counts, the count request will early terminate
+     */
+    public CountRequest terminateAfter(int terminateAfterCount) {
+        if (terminateAfterCount <= 0) {
+            throw new ElasticsearchIllegalArgumentException("terminateAfter must be > 0");
+        }
+        this.terminateAfter = terminateAfterCount;
+        return this;
+    }
+
+    public int terminateAfter() {
+        return this.terminateAfter;
+    }
+
     @Override
     public void readFrom(StreamInput in) throws IOException {
         super.readFrom(in);
@@ -230,6 +250,10 @@ public class CountRequest extends BroadcastOperationRequest<CountRequest> {
         sourceUnsafe = false;
         source = in.readBytesReference();
         types = in.readStringArray();
+
+        if (in.getVersion().onOrAfter(Version.V_1_4_0)) {
+            terminateAfter = in.readVInt();
+        }
     }
 
     @Override
@@ -240,6 +264,10 @@ public class CountRequest extends BroadcastOperationRequest<CountRequest> {
         out.writeOptionalString(preference);
         out.writeBytesReference(source);
         out.writeStringArray(types);
+
+        if (out.getVersion().onOrAfter(Version.V_1_4_0)) {
+            out.writeVInt(terminateAfter);
+        }
     }
 
     @Override

src/main/java/org/elasticsearch/action/count/CountRequestBuilder.java

@@ -91,6 +91,14 @@ public class CountRequestBuilder extends BroadcastOperationRequestBuilder<CountR
         return this;
     }
 
+    /**
+     * The query binary to execute
+     */
+    public CountRequestBuilder setQuery(BytesReference queryBinary) {
+        sourceBuilder().setQuery(queryBinary);
+        return this;
+    }
+
     /**
      * The source to execute.
      */
@@ -115,6 +123,11 @@ public class CountRequestBuilder extends BroadcastOperationRequestBuilder<CountR
         return this;
     }
 
+    public CountRequestBuilder setTerminateAfter(int terminateAfter) {
+        request().terminateAfter(terminateAfter);
+        return this;
+    }
+
     @Override
     protected void doExecute(ActionListener<CountResponse> listener) {
         if (sourceBuilder != null) {

src/main/java/org/elasticsearch/action/count/CountResponse.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.action.count;
 
+import org.elasticsearch.Version;
 import org.elasticsearch.action.ShardOperationFailedException;
 import org.elasticsearch.action.support.broadcast.BroadcastOperationResponse;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -33,15 +34,17 @@ import java.util.List;
  */
 public class CountResponse extends BroadcastOperationResponse {
 
+    private boolean terminatedEarly;
     private long count;
 
     CountResponse() {
 
     }
 
-    CountResponse(long count, int totalShards, int successfulShards, int failedShards, List<ShardOperationFailedException> shardFailures) {
+    CountResponse(long count, boolean hasTerminatedEarly, int totalShards, int successfulShards, int failedShards, List<ShardOperationFailedException> shardFailures) {
         super(totalShards, successfulShards, failedShards, shardFailures);
         this.count = count;
+        this.terminatedEarly = hasTerminatedEarly;
     }
 
     /**
@@ -51,6 +54,13 @@ public class CountResponse extends BroadcastOperationResponse {
         return count;
     }
 
+    /**
+     * True if the request has been terminated early due to enough count
+     */
+    public boolean terminatedEarly() {
+        return this.terminatedEarly;
+    }
+
     public RestStatus status() {
         if (getFailedShards() == 0) {
             if (getSuccessfulShards() == 0 && getTotalShards() > 0) {
@@ -76,11 +86,17 @@ public class CountResponse extends BroadcastOperationResponse {
     public void readFrom(StreamInput in) throws IOException {
         super.readFrom(in);
         count = in.readVLong();
+        if (in.getVersion().onOrAfter(Version.V_1_4_0)) {
+            terminatedEarly = in.readBoolean();
+        }
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeVLong(count);
+        if (out.getVersion().onOrAfter(Version.V_1_4_0)) {
+            out.writeBoolean(terminatedEarly);
+        }
     }
 }

src/main/java/org/elasticsearch/action/count/ShardCountRequest.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.action.count;
 
+import org.elasticsearch.Version;
 import org.elasticsearch.action.support.broadcast.BroadcastShardOperationRequest;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.Strings;
@@ -28,12 +29,15 @@ import org.elasticsearch.common.io.stream.StreamOutput;
 
 import java.io.IOException;
 
+import static org.elasticsearch.search.internal.SearchContext.DEFAULT_TERMINATE_AFTER;
+
 /**
  * Internal count request executed directly against a specific index shard.
  */
 class ShardCountRequest extends BroadcastShardOperationRequest {
 
     private float minScore;
+    private int terminateAfter;
 
     private BytesReference querySource;
 
@@ -55,6 +59,7 @@ class ShardCountRequest extends BroadcastShardOperationRequest {
         this.types = request.types();
         this.filteringAliases = filteringAliases;
         this.nowInMillis = request.nowInMillis;
+        this.terminateAfter = request.terminateAfter();
     }
 
     public float minScore() {
@@ -77,6 +82,10 @@ class ShardCountRequest extends BroadcastShardOperationRequest {
         return this.nowInMillis;
     }
 
+    public int terminateAfter() {
+        return this.terminateAfter;
+    }
+
     @Override
     public void readFrom(StreamInput in) throws IOException {
         super.readFrom(in);
@@ -99,6 +108,12 @@ class ShardCountRequest extends BroadcastShardOperationRequest {
             }
         }
         nowInMillis = in.readVLong();
+
+        if (in.getVersion().onOrAfter(Version.V_1_4_0)) {
+            terminateAfter = in.readVInt();
+        } else {
+            terminateAfter = DEFAULT_TERMINATE_AFTER;
+        }
     }
 
     @Override
@@ -121,5 +136,9 @@ class ShardCountRequest extends BroadcastShardOperationRequest {
             out.writeVInt(0);
         }
         out.writeVLong(nowInMillis);
+
+        if (out.getVersion().onOrAfter(Version.V_1_4_0)) {
+            out.writeVInt(terminateAfter);
+        }
     }
 }

src/main/java/org/elasticsearch/action/count/ShardCountResponse.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.action.count;
 
+import org.elasticsearch.Version;
 import org.elasticsearch.action.support.broadcast.BroadcastShardOperationResponse;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -33,29 +34,41 @@ import java.io.IOException;
 class ShardCountResponse extends BroadcastShardOperationResponse {
 
     private long count;
+    private boolean terminatedEarly;
 
     ShardCountResponse() {
 
     }
 
-    public ShardCountResponse(String index, int shardId, long count) {
+    public ShardCountResponse(String index, int shardId, long count, boolean terminatedEarly) {
         super(index, shardId);
         this.count = count;
+        this.terminatedEarly = terminatedEarly;
     }
 
     public long getCount() {
         return this.count;
     }
 
+    public boolean terminatedEarly() {
+        return this.terminatedEarly;
+    }
+
     @Override
     public void readFrom(StreamInput in) throws IOException {
         super.readFrom(in);
         count = in.readVLong();
+        if (in.getVersion().onOrAfter(Version.V_1_4_0)) {
+            terminatedEarly = in.readBoolean();
+        }
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeVLong(count);
+        if (out.getVersion().onOrAfter(Version.V_1_4_0)) {
+            out.writeBoolean(terminatedEarly);
+        }
     }
 }

src/main/java/org/elasticsearch/action/count/TransportCountAction.java

@@ -58,6 +58,7 @@ import java.util.Set;
 import java.util.concurrent.atomic.AtomicReferenceArray;
 
 import static com.google.common.collect.Lists.newArrayList;
+import static org.elasticsearch.search.internal.SearchContext.DEFAULT_TERMINATE_AFTER;
 
 /**
  *
@@ -139,6 +140,7 @@ public class TransportCountAction extends TransportBroadcastOperationAction<Coun
         int successfulShards = 0;
         int failedShards = 0;
         long count = 0;
+        boolean terminatedEarly = false;
         List<ShardOperationFailedException> shardFailures = null;
         for (int i = 0; i < shardsResponses.length(); i++) {
             Object shardResponse = shardsResponses.get(i);
@@ -152,10 +154,13 @@ public class TransportCountAction extends TransportBroadcastOperationAction<Coun
                 shardFailures.add(new DefaultShardOperationFailedException((BroadcastShardOperationFailedException) shardResponse));
             } else {
                 count += ((ShardCountResponse) shardResponse).getCount();
+                if (((ShardCountResponse) shardResponse).terminatedEarly()) {
+                    terminatedEarly = true;
+                }
                 successfulShards++;
             }
         }
-        return new CountResponse(count, shardsResponses.length(), successfulShards, failedShards, shardFailures);
+        return new CountResponse(count, terminatedEarly, shardsResponses.length(), successfulShards, failedShards, shardFailures);
     }
 
     @Override
@@ -186,10 +191,20 @@ public class TransportCountAction extends TransportBroadcastOperationAction<Coun
                     QueryParseContext.removeTypes();
                 }
             }
+            final boolean hasTerminateAfterCount = request.terminateAfter() != DEFAULT_TERMINATE_AFTER;
+            boolean terminatedEarly = false;
             context.preProcess();
             try {
-                long count = Lucene.count(context.searcher(), context.query());
-                return new ShardCountResponse(request.index(), request.shardId(), count);
+                long count;
+                if (hasTerminateAfterCount) {
+                    final Lucene.EarlyTerminatingCollector countCollector =
+                            Lucene.createCountBasedEarlyTerminatingCollector(request.terminateAfter());
+                    terminatedEarly = Lucene.countWithEarlyTermination(context.searcher(), context.query(), countCollector);
+                    count = countCollector.count();
+                } else {
+                    count = Lucene.count(context.searcher(), context.query());
+                }
+                return new ShardCountResponse(request.index(), request.shardId(), count, terminatedEarly);
             } catch (Exception e) {
                 throw new QueryPhaseExecutionException(context, "failed to execute count", e);
             }

src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java

@@ -129,6 +129,15 @@ public class SearchRequestBuilder extends ActionRequestBuilder<SearchRequest, Se
         return this;
     }
 
+    /**
+     * An optional document count, upon collecting which the search
+     * query will early terminate
+     */
+    public SearchRequestBuilder setTerminateAfter(int terminateAfter) {
+        sourceBuilder().terminateAfter(terminateAfter);
+        return this;
+    }
+
     /**
      * A comma separated list of routing values to control the shards the search will be executed on.
      */

src/main/java/org/elasticsearch/action/search/SearchResponse.java

@@ -116,6 +116,14 @@ public class SearchResponse extends ActionResponse implements StatusToXContent {
         return internalResponse.timedOut();
     }
 
+    /**
+     * Has the search operation terminated early due to reaching
+     * <code>terminateAfter</code>
+     */
+    public Boolean isTerminatedEarly() {
+        return internalResponse.terminatedEarly();
+    }
+
     /**
      * How long the search took.
      */
@@ -181,6 +189,7 @@ public class SearchResponse extends ActionResponse implements StatusToXContent {
         static final XContentBuilderString REASON = new XContentBuilderString("reason");
         static final XContentBuilderString TOOK = new XContentBuilderString("took");
         static final XContentBuilderString TIMED_OUT = new XContentBuilderString("timed_out");
+        static final XContentBuilderString TERMINATED_EARLY = new XContentBuilderString("terminated_early");
     }
 
     @Override
@@ -190,6 +199,9 @@ public class SearchResponse extends ActionResponse implements StatusToXContent {
         }
         builder.field(Fields.TOOK, tookInMillis);
         builder.field(Fields.TIMED_OUT, isTimedOut());
+        if (isTerminatedEarly() != null) {
+            builder.field(Fields.TERMINATED_EARLY, isTerminatedEarly());
+        }
         builder.startObject(Fields._SHARDS);
         builder.field(Fields.TOTAL, getTotalShards());
         builder.field(Fields.SUCCESSFUL, getSuccessfulShards());

src/main/java/org/elasticsearch/action/search/type/TransportSearchScrollScanAction.java

@@ -119,7 +119,7 @@ public class TransportSearchScrollScanAction extends AbstractComponent {
 
         public void start() {
             if (scrollId.getContext().length == 0) {
-                final InternalSearchResponse internalResponse = new InternalSearchResponse(new InternalSearchHits(InternalSearchHits.EMPTY, Long.parseLong(this.scrollId.getAttributes().get("total_hits")), 0.0f), null, null, null, false);
+                final InternalSearchResponse internalResponse = new InternalSearchResponse(new InternalSearchHits(InternalSearchHits.EMPTY, Long.parseLong(this.scrollId.getAttributes().get("total_hits")), 0.0f), null, null, null, false, null);
                 listener.onResponse(new SearchResponse(internalResponse, request.scrollId(), 0, 0, 0l, buildShardFailures()));
                 return;
             }

src/main/java/org/elasticsearch/common/lucene/Lucene.java

@@ -29,6 +29,7 @@ import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.Version;
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ElasticsearchIllegalArgumentException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.common.Nullable;
@@ -39,6 +40,7 @@ import org.elasticsearch.index.analysis.AnalyzerScope;
 import org.elasticsearch.index.analysis.NamedAnalyzer;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 
+import static org.elasticsearch.common.lucene.search.NoopCollector.NOOP_COLLECTOR;
 import java.io.IOException;
 
 /**
@@ -121,12 +123,110 @@ public class Lucene {
 
     public static long count(IndexSearcher searcher, Query query) throws IOException {
         TotalHitCountCollector countCollector = new TotalHitCountCollector();
+        query = wrapCountQuery(query);
+        searcher.search(query, countCollector);
+        return countCollector.getTotalHits();
+    }
+
+    /**
+     * Performs a count on the <code>searcher</code> for <code>query</code>. Terminates
+     * early when the count has reached <code>terminateAfter</code>
+     */
+    public static long count(IndexSearcher searcher, Query query, int terminateAfterCount) throws IOException {
+        EarlyTerminatingCollector countCollector = createCountBasedEarlyTerminatingCollector(terminateAfterCount);
+        countWithEarlyTermination(searcher, query, countCollector);
+        return countCollector.count();
+    }
+
+    /**
+     * Creates count based early termination collector with a threshold of <code>maxCountHits</code>
+     */
+    public final static EarlyTerminatingCollector createCountBasedEarlyTerminatingCollector(int maxCountHits) {
+        return new EarlyTerminatingCollector(maxCountHits);
+    }
+
+    /**
+     * Wraps <code>delegate</code> with count based early termination collector with a threshold of <code>maxCountHits</code>
+     */
+    public final static EarlyTerminatingCollector wrapCountBasedEarlyTerminatingCollector(final Collector delegate, int maxCountHits) {
+        return new EarlyTerminatingCollector(delegate, maxCountHits);
+    }
+
+    /**
+     * Wraps <code>delegate</code> with a time limited collector with a timeout of <code>timeoutInMillis</code>
+     */
+    public final static TimeLimitingCollector wrapTimeLimitingCollector(final Collector delegate, long timeoutInMillis) {
+        return new TimeLimitingCollector(delegate, TimeLimitingCollector.getGlobalCounter(), timeoutInMillis);
+    }
+
+    /**
+     * Performs an exists (count > 0) query on the <code>searcher</code> for <code>query</code>
+     * with <code>filter</code> using the given <code>collector</code>
+     *
+     * The <code>collector</code> can be instantiated using <code>Lucene.createExistsCollector()</code>
+     */
+    public static boolean exists(IndexSearcher searcher, Query query, Filter filter,
+                                 EarlyTerminatingCollector collector) throws IOException {
+        collector.reset();
+        countWithEarlyTermination(searcher, filter, query, collector);
+        return collector.exists();
+    }
+
+
+    /**
+     * Performs an exists (count > 0) query on the <code>searcher</code> for <code>query</code>
+     * using the given <code>collector</code>
+     *
+     * The <code>collector</code> can be instantiated using <code>Lucene.createExistsCollector()</code>
+     */
+    public static boolean exists(IndexSearcher searcher, Query query, EarlyTerminatingCollector collector) throws IOException {
+        collector.reset();
+        countWithEarlyTermination(searcher, query, collector);
+        return collector.exists();
+    }
+
+    /**
+     * Calls <code>countWithEarlyTermination(searcher, null, query, collector)</code>
+     */
+    public static boolean countWithEarlyTermination(IndexSearcher searcher, Query query,
+                                                  EarlyTerminatingCollector collector) throws IOException {
+        return countWithEarlyTermination(searcher, null, query, collector);
+    }
+
+    /**
+     * Performs a count on <code>query</code> and <code>filter</code> with early termination using <code>searcher</code>.
+     * The early termination threshold is specified by the provided <code>collector</code>
+     */
+    public static boolean countWithEarlyTermination(IndexSearcher searcher, Filter filter, Query query,
+                                                        EarlyTerminatingCollector collector) throws IOException {
+        query = wrapCountQuery(query);
+        try {
+            if (filter == null) {
+                searcher.search(query, collector);
+            } else {
+                searcher.search(query, filter, collector);
+            }
+        } catch (EarlyTerminationException e) {
+            // early termination
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * Creates an {@link org.elasticsearch.common.lucene.Lucene.EarlyTerminatingCollector}
+     * with a threshold of <code>1</code>
+     */
+    public final static EarlyTerminatingCollector createExistsCollector() {
+        return createCountBasedEarlyTerminatingCollector(1);
+    }
+
+    private final static Query wrapCountQuery(Query query) {
         // we don't need scores, so wrap it in a constant score query
         if (!(query instanceof ConstantScoreQuery)) {
             query = new ConstantScoreQuery(query);
         }
-        searcher.search(query, countCollector);
-        return countCollector.getTotalHits();
+        return query;
     }
 
     /**
@@ -355,35 +455,70 @@ public class Lucene {
         }
     }
 
-    public static class ExistsCollector extends Collector {
+    /**
+     * This exception is thrown when {@link org.elasticsearch.common.lucene.Lucene.EarlyTerminatingCollector}
+     * reaches early termination
+     * */
+    public final static class EarlyTerminationException extends ElasticsearchException {
+
+        public EarlyTerminationException(String msg) {
+            super(msg);
+        }
+    }
 
-        private boolean exists;
+    /**
+     * A collector that terminates early by throwing {@link org.elasticsearch.common.lucene.Lucene.EarlyTerminationException}
+     * when count of matched documents has reached <code>maxCountHits</code>
+     */
+    public final static class EarlyTerminatingCollector extends Collector {
+
+        private final int maxCountHits;
+        private final Collector delegate;
+        private int count = 0;
+
+        EarlyTerminatingCollector(int maxCountHits) {
+            this.maxCountHits = maxCountHits;
+            this.delegate = NOOP_COLLECTOR;
+        }
+
+        EarlyTerminatingCollector(final Collector delegate, int maxCountHits) {
+            this.maxCountHits = maxCountHits;
+            this.delegate = (delegate == null) ? NOOP_COLLECTOR : delegate;
+        }
 
         public void reset() {
-            exists = false;
+            count = 0;
+        }
+        public int count() {
+            return count;
         }
 
         public boolean exists() {
-            return exists;
+            return count > 0;
         }
 
         @Override
         public void setScorer(Scorer scorer) throws IOException {
-            this.exists = false;
+            delegate.setScorer(scorer);
         }
 
         @Override
         public void collect(int doc) throws IOException {
-            exists = true;
+            delegate.collect(doc);
+
+            if (++count >= maxCountHits) {
+                throw new EarlyTerminationException("early termination [CountBased]");
+            }
         }
 
         @Override
-        public void setNextReader(AtomicReaderContext context) throws IOException {
+        public void setNextReader(AtomicReaderContext atomicReaderContext) throws IOException {
+            delegate.setNextReader(atomicReaderContext);
         }
 
         @Override
         public boolean acceptsDocsOutOfOrder() {
-            return true;
+            return delegate.acceptsDocsOutOfOrder();
         }
     }
 

src/main/java/org/elasticsearch/percolator/PercolateContext.java

@@ -485,6 +485,16 @@ public class PercolateContext extends SearchContext {
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public int terminateAfter() {
+        return DEFAULT_TERMINATE_AFTER;
+    }
+
+    @Override
+    public void terminateAfter(int terminateAfter) {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public SearchContext minimumScore(float minimumScore) {
         throw new UnsupportedOperationException();

src/main/java/org/elasticsearch/percolator/PercolatorService.java

@@ -446,14 +446,13 @@ public class PercolatorService extends AbstractComponent {
         @Override
         public PercolateShardResponse doPercolate(PercolateShardRequest request, PercolateContext context, boolean isNested) {
             long count = 0;
-            Lucene.ExistsCollector collector = new Lucene.ExistsCollector();
+            Lucene.EarlyTerminatingCollector collector = Lucene.createExistsCollector();
             for (Map.Entry<BytesRef, Query> entry : context.percolateQueries().entrySet()) {
-                collector.reset();
                 try {
                     if (isNested) {
-                        context.docSearcher().search(entry.getValue(), NonNestedDocsFilter.INSTANCE, collector);
+                        Lucene.exists(context.docSearcher(), entry.getValue(), NonNestedDocsFilter.INSTANCE, collector);
                     } else {
-                        context.docSearcher().search(entry.getValue(), collector);
+                        Lucene.exists(context.docSearcher(), entry.getValue(), collector);
                     }
                 } catch (Throwable e) {
                     logger.debug("[" + entry.getKey() + "] failed to execute query", e);
@@ -543,19 +542,18 @@ public class PercolatorService extends AbstractComponent {
             long count = 0;
             List<BytesRef> matches = new ArrayList<>();
             List<Map<String, HighlightField>> hls = new ArrayList<>();
-            Lucene.ExistsCollector collector = new Lucene.ExistsCollector();
+            Lucene.EarlyTerminatingCollector collector = Lucene.createExistsCollector();
 
             for (Map.Entry<BytesRef, Query> entry : context.percolateQueries().entrySet()) {
-                collector.reset();
                 if (context.highlight() != null) {
                     context.parsedQuery(new ParsedQuery(entry.getValue(), ImmutableMap.<String, Filter>of()));
                     context.hitContext().cache().clear();
                 }
                 try {
                     if (isNested) {
-                        context.docSearcher().search(entry.getValue(), NonNestedDocsFilter.INSTANCE, collector);
+                        Lucene.exists(context.docSearcher(), entry.getValue(), NonNestedDocsFilter.INSTANCE, collector);
                     } else {
-                        context.docSearcher().search(entry.getValue(), collector);
+                        Lucene.exists(context.docSearcher(), entry.getValue(), collector);
                     }
                 } catch (Throwable e) {
                     logger.debug("[" + entry.getKey() + "] failed to execute query", e);

src/main/java/org/elasticsearch/percolator/QueryCollector.java

@@ -58,7 +58,7 @@ abstract class QueryCollector extends Collector {
     final ESLogger logger;
     boolean isNestedDoc = false;
 
-    final Lucene.ExistsCollector collector = new Lucene.ExistsCollector();
+    final Lucene.EarlyTerminatingCollector collector = Lucene.createExistsCollector();
     BytesRef current;
 
     SortedBinaryDocValues values;
@@ -200,16 +200,15 @@ abstract class QueryCollector extends Collector {
             }
             // run the query
             try {
-                collector.reset();
                 if (context.highlight() != null) {
                     context.parsedQuery(new ParsedQuery(query, ImmutableMap.<String, Filter>of()));
                     context.hitContext().cache().clear();
                 }
 
                 if (isNestedDoc) {
-                    searcher.search(query, NonNestedDocsFilter.INSTANCE, collector);
+                    Lucene.exists(searcher, query, NonNestedDocsFilter.INSTANCE, collector);
                 } else {
-                    searcher.search(query, collector);
+                    Lucene.exists(searcher, query, collector);
                 }
                 if (collector.exists()) {
                     if (!limit || counter < size) {
@@ -259,11 +258,10 @@ abstract class QueryCollector extends Collector {
             }
             // run the query
             try {
-                collector.reset();
                 if (isNestedDoc) {
-                    searcher.search(query, NonNestedDocsFilter.INSTANCE, collector);
+                    Lucene.exists(searcher, query, NonNestedDocsFilter.INSTANCE, collector);
                 } else {
-                    searcher.search(query, collector);
+                    Lucene.exists(searcher, query, collector);
                 }
                 if (collector.exists()) {
                     topDocsCollector.collect(doc);
@@ -323,15 +321,14 @@ abstract class QueryCollector extends Collector {
             }
             // run the query
             try {
-                collector.reset();
                 if (context.highlight() != null) {
                     context.parsedQuery(new ParsedQuery(query, ImmutableMap.<String, Filter>of()));
                     context.hitContext().cache().clear();
                 }
                 if (isNestedDoc) {
-                    searcher.search(query, NonNestedDocsFilter.INSTANCE, collector);
+                    Lucene.exists(searcher, query, NonNestedDocsFilter.INSTANCE, collector);
                 } else {
-                    searcher.search(query, collector);
+                    Lucene.exists(searcher, query, collector);
                 }
                 if (collector.exists()) {
                     if (!limit || counter < size) {
@@ -389,11 +386,10 @@ abstract class QueryCollector extends Collector {
             }
             // run the query
             try {
-                collector.reset();
                 if (isNestedDoc) {
-                    searcher.search(query, NonNestedDocsFilter.INSTANCE, collector);
+                    Lucene.exists(searcher, query, NonNestedDocsFilter.INSTANCE, collector);
                 } else {
-                    searcher.search(query, collector);
+                    Lucene.exists(searcher, query, collector);
                 }
                 if (collector.exists()) {
                     counter++;

src/main/java/org/elasticsearch/rest/action/count/RestCountAction.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.rest.action.count;
 
+import org.elasticsearch.ElasticsearchIllegalArgumentException;
 import org.elasticsearch.action.count.CountRequest;
 import org.elasticsearch.action.count.CountResponse;
 import org.elasticsearch.action.support.IndicesOptions;
@@ -33,6 +34,7 @@ import org.elasticsearch.rest.action.support.RestActions;
 import org.elasticsearch.rest.action.support.RestBuilderListener;
 
 import static org.elasticsearch.action.count.CountRequest.DEFAULT_MIN_SCORE;
+import static org.elasticsearch.search.internal.SearchContext.DEFAULT_TERMINATE_AFTER;
 import static org.elasticsearch.rest.RestRequest.Method.GET;
 import static org.elasticsearch.rest.RestRequest.Method.POST;
 import static org.elasticsearch.rest.action.support.RestActions.buildBroadcastShardsHeader;
@@ -76,12 +78,20 @@ public class RestCountAction extends BaseRestHandler {
         countRequest.types(Strings.splitStringByCommaToArray(request.param("type")));
         countRequest.preference(request.param("preference"));
 
+        final int terminateAfter = request.paramAsInt("terminate_after", DEFAULT_TERMINATE_AFTER);
+        if (terminateAfter < 0) {
+            throw new ElasticsearchIllegalArgumentException("terminateAfter must be > 0");
+        } else if (terminateAfter > 0) {
+            countRequest.terminateAfter(terminateAfter);
+        }
         client.count(countRequest, new RestBuilderListener<CountResponse>(channel) {
             @Override
             public RestResponse buildResponse(CountResponse response, XContentBuilder builder) throws Exception {
                 builder.startObject();
+                if (terminateAfter != DEFAULT_TERMINATE_AFTER) {
+                    builder.field("terminated_early", response.terminatedEarly());
+                }
                 builder.field("count", response.getCount());
-
                 buildBroadcastShardsHeader(builder, response);
 
                 builder.endObject();

src/main/java/org/elasticsearch/rest/action/search/RestSearchAction.java

@@ -37,6 +37,7 @@ import org.elasticsearch.rest.action.support.RestStatusToXContentListener;
 import org.elasticsearch.search.Scroll;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.fetch.source.FetchSourceContext;
+import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.sort.SortOrder;
 
 import static org.elasticsearch.common.unit.TimeValue.parseTimeValue;
@@ -172,6 +173,18 @@ public class RestSearchAction extends BaseRestHandler {
             }
             searchSourceBuilder.timeout(request.paramAsTime("timeout", null));
         }
+        if (request.hasParam("terminate_after")) {
+            if (searchSourceBuilder == null) {
+                searchSourceBuilder = new SearchSourceBuilder();
+            }
+            int terminateAfter = request.paramAsInt("terminate_after",
+                    SearchContext.DEFAULT_TERMINATE_AFTER);
+            if (terminateAfter < 0) {
+                throw new ElasticsearchIllegalArgumentException("terminateAfter must be > 0");
+            } else if (terminateAfter > 0) {
+                searchSourceBuilder.terminateAfter(terminateAfter);
+            }
+        }
 
         String sField = request.param("fields");
         if (sField != null) {

src/main/java/org/elasticsearch/search/aggregations/metrics/tophits/TopHitsContext.java

@@ -367,6 +367,16 @@ public class TopHitsContext extends SearchContext {
         throw new UnsupportedOperationException("Not supported");
     }
 
+    @Override
+    public int terminateAfter() {
+        return context.terminateAfter();
+    }
+
+    @Override
+    public void terminateAfter(int terminateAfter) {
+        throw new UnsupportedOperationException("Not supported");
+    }
+
     @Override
     public SearchContext minimumScore(float minimumScore) {
         throw new UnsupportedOperationException("Not supported");

src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -24,6 +24,7 @@ import com.google.common.base.Charsets;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import org.elasticsearch.ElasticsearchGenerationException;
+import org.elasticsearch.ElasticsearchIllegalArgumentException;
 import org.elasticsearch.client.Requests;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.Strings;
@@ -40,6 +41,7 @@ import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.facet.FacetBuilder;
 import org.elasticsearch.search.fetch.source.FetchSourceContext;
 import org.elasticsearch.search.highlight.HighlightBuilder;
+import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.rescore.RescoreBuilder;
 import org.elasticsearch.search.sort.SortBuilder;
 import org.elasticsearch.search.sort.SortBuilders;
@@ -97,6 +99,7 @@ public class SearchSourceBuilder implements ToXContent {
     private Float minScore;
 
     private long timeoutInMillis = -1;
+    private int terminateAfter = SearchContext.DEFAULT_TERMINATE_AFTER;
 
     private List<String> fieldNames;
     private List<String> fieldDataFields;
@@ -308,6 +311,17 @@ public class SearchSourceBuilder implements ToXContent {
         return this;
     }
 
+    /**
+     * An optional terminate_after to terminate the search after
+     * collecting <code>terminateAfter</code> documents
+     */
+    public  SearchSourceBuilder terminateAfter(int terminateAfter) {
+        if (terminateAfter <= 0) {
+            throw new ElasticsearchIllegalArgumentException("terminateAfter must be > 0");
+        }
+        this.terminateAfter = terminateAfter;
+        return this;
+    }
     /**
      * Adds a sort against the given field name and the sort ordering.
      *
@@ -738,6 +752,10 @@ public class SearchSourceBuilder implements ToXContent {
             builder.field("timeout", timeoutInMillis);
         }
 
+        if (terminateAfter != SearchContext.DEFAULT_TERMINATE_AFTER) {
+            builder.field("terminate_after", terminateAfter);
+        }
+
         if (queryBuilder != null) {
             builder.field("query");
             queryBuilder.toXContent(builder, params);

src/main/java/org/elasticsearch/search/controller/SearchPhaseController.java

@@ -318,11 +318,19 @@ public class SearchPhaseController extends AbstractComponent {
         long totalHits = 0;
         float maxScore = Float.NEGATIVE_INFINITY;
         boolean timedOut = false;
+        Boolean terminatedEarly = null;
         for (AtomicArray.Entry<? extends QuerySearchResultProvider> entry : queryResults) {
             QuerySearchResult result = entry.value.queryResult();
             if (result.searchTimedOut()) {
                 timedOut = true;
             }
+            if (result.terminatedEarly() != null) {
+                if (terminatedEarly == null) {
+                    terminatedEarly = result.terminatedEarly();
+                } else if (result.terminatedEarly()) {
+                    terminatedEarly = true;
+                }
+            }
             totalHits += result.topDocs().totalHits;
             if (!Float.isNaN(result.topDocs().getMaxScore())) {
                 maxScore = Math.max(maxScore, result.topDocs().getMaxScore());
@@ -397,7 +405,7 @@ public class SearchPhaseController extends AbstractComponent {
 
         InternalSearchHits searchHits = new InternalSearchHits(hits.toArray(new InternalSearchHit[hits.size()]), totalHits, maxScore);
 
-        return new InternalSearchResponse(searchHits, facets, aggregations, suggest, timedOut);
+        return new InternalSearchResponse(searchHits, facets, aggregations, suggest, timedOut, terminatedEarly);
     }
 
 }

src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java

@@ -23,6 +23,7 @@ import org.apache.lucene.index.AtomicReaderContext;
 import org.apache.lucene.search.*;
 import org.elasticsearch.common.lease.Releasable;
 import org.elasticsearch.common.lease.Releasables;
+import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.lucene.MinimumScoreCollector;
 import org.elasticsearch.common.lucene.MultiCollector;
 import org.elasticsearch.common.lucene.search.FilteredCollector;
@@ -138,9 +139,17 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
 
     @Override
     public void search(List<AtomicReaderContext> leaves, Weight weight, Collector collector) throws IOException {
-        if (searchContext.timeoutInMillis() != -1) {
+        final boolean timeoutSet = searchContext.timeoutInMillis() != -1;
+        final boolean terminateAfterSet = searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER;
+
+        if (timeoutSet) {
             // TODO: change to use our own counter that uses the scheduler in ThreadPool
-            collector = new TimeLimitingCollector(collector, TimeLimitingCollector.getGlobalCounter(), searchContext.timeoutInMillis());
+            // throws TimeLimitingCollector.TimeExceededException when timeout has reached
+            collector = Lucene.wrapTimeLimitingCollector(collector, searchContext.timeoutInMillis());
+        }
+        if (terminateAfterSet) {
+            // throws Lucene.EarlyTerminationException when given count is reached
+            collector = Lucene.wrapCountBasedEarlyTerminatingCollector(collector, searchContext.terminateAfter());
         }
         if (currentState == Stage.MAIN_QUERY) {
             if (enableMainDocIdSetCollector) {
@@ -165,11 +174,18 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
 
         // we only compute the doc id set once since within a context, we execute the same query always...
         try {
-            if (searchContext.timeoutInMillis() != -1) {
+            if (timeoutSet || terminateAfterSet) {
                 try {
                     super.search(leaves, weight, collector);
                 } catch (TimeLimitingCollector.TimeExceededException e) {
+                    assert timeoutSet : "TimeExceededException thrown even though timeout wasn't set";
                     searchContext.queryResult().searchTimedOut(true);
+                } catch (Lucene.EarlyTerminationException e) {
+                    assert terminateAfterSet : "EarlyTerminationException thrown even though terminateAfter wasn't set";
+                    searchContext.queryResult().terminatedEarly(true);
+                }
+                if (terminateAfterSet && searchContext.queryResult().terminatedEarly() == null) {
+                    searchContext.queryResult().terminatedEarly(false);
                 }
             } else {
                 super.search(leaves, weight, collector);
@@ -205,4 +221,4 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
             searchContext.clearReleasables(Lifetime.COLLECTION);
         }
     }
-}
+}

src/main/java/org/elasticsearch/search/internal/DefaultSearchContext.java

@@ -117,6 +117,9 @@ public class DefaultSearchContext extends SearchContext {
     // timeout in millis
     private long timeoutInMillis = -1;
 
+    // terminate after count
+    private int terminateAfter = DEFAULT_TERMINATE_AFTER;
+
 
     private List<String> groupStats;
 
@@ -471,6 +474,16 @@ public class DefaultSearchContext extends SearchContext {
         this.timeoutInMillis = timeoutInMillis;
     }
 
+    @Override
+    public int terminateAfter() {
+        return terminateAfter;
+    }
+
+    @Override
+    public void terminateAfter(int terminateAfter) {
+        this.terminateAfter = terminateAfter;
+    }
+
     public SearchContext minimumScore(float minimumScore) {
         this.minimumScore = minimumScore;
         return this;

src/main/java/org/elasticsearch/search/internal/InternalSearchResponse.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.search.internal;
 
+import org.elasticsearch.Version;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Streamable;
@@ -41,7 +42,7 @@ import static org.elasticsearch.search.internal.InternalSearchHits.readSearchHit
 public class InternalSearchResponse implements Streamable, ToXContent {
 
     public static InternalSearchResponse empty() {
-        return new InternalSearchResponse(InternalSearchHits.empty(), null, null, null, false);
+        return new InternalSearchResponse(InternalSearchHits.empty(), null, null, null, false, null);
     }
 
     private InternalSearchHits hits;
@@ -54,21 +55,28 @@ public class InternalSearchResponse implements Streamable, ToXContent {
 
     private boolean timedOut;
 
+    private Boolean terminatedEarly = null;
+
     private InternalSearchResponse() {
     }
 
-    public InternalSearchResponse(InternalSearchHits hits, InternalFacets facets, InternalAggregations aggregations, Suggest suggest, boolean timedOut) {
+    public InternalSearchResponse(InternalSearchHits hits, InternalFacets facets, InternalAggregations aggregations, Suggest suggest, boolean timedOut, Boolean terminatedEarly) {
         this.hits = hits;
         this.facets = facets;
         this.aggregations = aggregations;
         this.suggest = suggest;
         this.timedOut = timedOut;
+        this.terminatedEarly = terminatedEarly;
     }
 
     public boolean timedOut() {
         return this.timedOut;
     }
 
+    public Boolean terminatedEarly() {
+        return this.terminatedEarly;
+    }
+
     public SearchHits hits() {
         return hits;
     }
@@ -119,6 +127,10 @@ public class InternalSearchResponse implements Streamable, ToXContent {
             suggest = Suggest.readSuggest(Suggest.Fields.SUGGEST, in);
         }
         timedOut = in.readBoolean();
+
+        if (in.getVersion().onOrAfter(Version.V_1_4_0)) {
+            terminatedEarly = in.readOptionalBoolean();
+        }
     }
 
     @Override
@@ -143,5 +155,10 @@ public class InternalSearchResponse implements Streamable, ToXContent {
             suggest.writeTo(out);
         }
         out.writeBoolean(timedOut);
+
+        if (out.getVersion().onOrAfter(Version.V_1_4_0)) {
+            out.writeOptionalBoolean(terminatedEarly);
+
+        }
     }
 }

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

@@ -72,6 +72,7 @@ import java.util.List;
 public abstract class SearchContext implements Releasable {
 
     private static ThreadLocal<SearchContext> current = new ThreadLocal<>();
+    public final static int DEFAULT_TERMINATE_AFTER = 0;
 
     public static void setCurrent(SearchContext value) {
         current.set(value);
@@ -212,6 +213,10 @@ public abstract class SearchContext implements Releasable {
 
     public abstract void timeoutInMillis(long timeoutInMillis);
 
+    public abstract int terminateAfter();
+
+    public abstract void terminateAfter(int terminateAfter);
+
     public abstract SearchContext minimumScore(float minimumScore);
 
     public abstract Float minimumScore();

src/main/java/org/elasticsearch/search/query/QueryPhase.java

@@ -79,6 +79,7 @@ public class QueryPhase implements SearchPhase {
                 .put("min_score", new MinScoreParseElement())
                 .put("minScore", new MinScoreParseElement())
                 .put("timeout", new TimeoutParseElement())
+                .put("terminate_after", new TerminateAfterParseElement())
                 .putAll(facetPhase.parseElements())
                 .putAll(aggregationPhase.parseElements())
                 .putAll(suggestPhase.parseElements())

src/main/java/org/elasticsearch/search/query/QuerySearchResult.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.search.query;
 
 import org.apache.lucene.search.TopDocs;
+import org.elasticsearch.Version;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.search.SearchShardTarget;
@@ -49,6 +50,7 @@ public class QuerySearchResult extends TransportResponse implements QuerySearchR
     private InternalAggregations aggregations;
     private Suggest suggest;
     private boolean searchTimedOut;
+    private Boolean terminatedEarly = null;
 
     public QuerySearchResult() {
 
@@ -90,6 +92,14 @@ public class QuerySearchResult extends TransportResponse implements QuerySearchR
         return searchTimedOut;
     }
 
+    public void terminatedEarly(boolean terminatedEarly) {
+        this.terminatedEarly = terminatedEarly;
+    }
+
+    public Boolean terminatedEarly() {
+        return this.terminatedEarly;
+    }
+
     public TopDocs topDocs() {
         return topDocs;
     }
@@ -164,6 +174,9 @@ public class QuerySearchResult extends TransportResponse implements QuerySearchR
             suggest = Suggest.readSuggest(Suggest.Fields.SUGGEST, in);
         }
         searchTimedOut = in.readBoolean();
+        if (in.getVersion().onOrAfter(Version.V_1_4_0)) {
+            terminatedEarly = in.readOptionalBoolean();
+        }
     }
 
     @Override
@@ -193,5 +206,8 @@ public class QuerySearchResult extends TransportResponse implements QuerySearchR
             suggest.writeTo(out);
         }
         out.writeBoolean(searchTimedOut);
+        if (out.getVersion().onOrAfter(Version.V_1_4_0)) {
+            out.writeOptionalBoolean(terminatedEarly);
+        }
     }
 }

src/main/java/org/elasticsearch/search/query/TerminateAfterParseElement.java

@@ -0,0 +1,43 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search.query;
+
+import org.elasticsearch.ElasticsearchIllegalArgumentException;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.search.SearchParseElement;
+import org.elasticsearch.search.internal.SearchContext;
+
+/**
+ * Parser element for 'terminate_after'
+ */
+public class TerminateAfterParseElement implements SearchParseElement {
+
+    @Override
+    public void parse(XContentParser parser, SearchContext context) throws Exception {
+        XContentParser.Token token = parser.currentToken();
+        if (token == XContentParser.Token.VALUE_NUMBER) {
+            int terminateAfterCount = parser.intValue();
+            if (terminateAfterCount <= 0) {
+                throw new ElasticsearchIllegalArgumentException("terminateAfter must be > 0");
+            }
+            context.terminateAfter(parser.intValue());
+        }
+    }
+}

src/main/java/org/elasticsearch/search/suggest/phrase/PhraseSuggester.java

@@ -169,12 +169,14 @@ public final class PhraseSuggester extends Suggester<PhraseSuggestionContext> {
                 req = client.prepareSearch()
                         .setPreference(suggestions.getPreference())
                         .setQuery(QueryBuilders.constantScoreQuery(FilterBuilders.bytesFilter(querySource)))
-                        .setSearchType(SearchType.COUNT);
+                        .setSearchType(SearchType.COUNT)
+                        .setTerminateAfter(1);
             } else {
                 req = client.prepareSearch()
                         .setPreference(suggestions.getPreference())
                         .setQuery(querySource)
-                        .setSearchType(SearchType.COUNT);
+                        .setSearchType(SearchType.COUNT)
+                        .setTerminateAfter(1);
             }
             multiSearchRequestBuilder.add(req);
         }

src/test/java/org/elasticsearch/count/simple/SimpleCountTests.java

@@ -20,13 +20,18 @@
 package org.elasticsearch.count.simple;
 
 import org.elasticsearch.action.count.CountResponse;
+import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.test.ElasticsearchIntegrationTest;
 import org.junit.Test;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.ExecutionException;
 
+import static org.elasticsearch.cluster.metadata.IndexMetaData.SETTING_NUMBER_OF_REPLICAS;
+import static org.elasticsearch.cluster.metadata.IndexMetaData.SETTING_NUMBER_OF_SHARDS;
 import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
 import static org.elasticsearch.index.query.QueryBuilders.rangeQuery;
@@ -113,6 +118,45 @@ public class SimpleCountTests extends ElasticsearchIntegrationTest {
         assertHitCount(countResponse, 2l);
     }
 
+    @Test
+    public void simpleCountEarlyTerminationTests() throws Exception {
+        // set up one shard only to test early termination
+        prepareCreate("test").setSettings(
+                SETTING_NUMBER_OF_SHARDS, 1,
+                SETTING_NUMBER_OF_REPLICAS, 0).get();
+        ensureGreen();
+        int max = randomIntBetween(3, 29);
+        List<IndexRequestBuilder> docbuilders = new ArrayList<>(max);
+
+        for (int i = 1; i <= max; i++) {
+            String id = String.valueOf(i);
+            docbuilders.add(client().prepareIndex("test", "type1", id).setSource("field", "2010-01-"+ id +"T02:00"));
+        }
+
+        indexRandom(true, docbuilders);
+        ensureGreen();
+        refresh();
+
+        String upperBound = "2010-01-" + String.valueOf(max+1) + "||+2d";
+        String lowerBound = "2009-12-01||+2d";
+
+        // sanity check
+        CountResponse countResponse = client().prepareCount("test").setQuery(QueryBuilders.rangeQuery("field").gte(lowerBound).lte(upperBound)).execute().actionGet();
+        assertHitCount(countResponse, max);
+
+        // threshold <= actual count
+        for (int i = 1; i <= max; i++) {
+            countResponse = client().prepareCount("test").setQuery(QueryBuilders.rangeQuery("field").gte(lowerBound).lte(upperBound)).setTerminateAfter(i).execute().actionGet();
+            assertHitCount(countResponse, i);
+            assertTrue(countResponse.terminatedEarly());
+        }
+
+        // threshold > actual count
+        countResponse = client().prepareCount("test").setQuery(QueryBuilders.rangeQuery("field").gte(lowerBound).lte(upperBound)).setTerminateAfter(max + randomIntBetween(1, max)).execute().actionGet();
+        assertHitCount(countResponse, max);
+        assertFalse(countResponse.terminatedEarly());
+    }
+
     @Test
     public void localDependentDateTests() throws Exception {
         assertAcked(prepareCreate("test")

src/test/java/org/elasticsearch/search/simple/SimpleSearchTests.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.search.simple;
 
 import org.elasticsearch.ElasticsearchIllegalArgumentException;
+import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.common.settings.ImmutableSettings;
 import org.elasticsearch.common.xcontent.XContentFactory;
@@ -28,8 +29,12 @@ import org.elasticsearch.test.ElasticsearchIntegrationTest;
 import org.elasticsearch.test.junit.annotations.TestLogging;
 import org.junit.Test;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.ExecutionException;
 
+import static org.elasticsearch.cluster.metadata.IndexMetaData.SETTING_NUMBER_OF_REPLICAS;
+import static org.elasticsearch.cluster.metadata.IndexMetaData.SETTING_NUMBER_OF_SHARDS;
 import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
 import static org.elasticsearch.index.query.QueryBuilders.rangeQuery;
@@ -200,4 +205,43 @@ public class SimpleSearchTests extends ElasticsearchIntegrationTest {
 
         }
     }
+
+    @Test
+    public void simpleTerminateAfterCountTests() throws Exception {
+        prepareCreate("test").setSettings(
+                SETTING_NUMBER_OF_SHARDS, 1,
+                SETTING_NUMBER_OF_REPLICAS, 0).get();
+        ensureGreen();
+        int max = randomIntBetween(3, 29);
+        List<IndexRequestBuilder> docbuilders = new ArrayList<>(max);
+
+        for (int i = 1; i <= max; i++) {
+            String id = String.valueOf(i);
+            docbuilders.add(client().prepareIndex("test", "type1", id).setSource("field", "2010-01-"+ id +"T02:00"));
+        }
+
+        indexRandom(true, docbuilders);
+        ensureGreen();
+        refresh();
+
+        String upperBound = "2010-01-" + String.valueOf(max+1) + "||+2d";
+        String lowerBound = "2009-12-01||+2d";
+
+        SearchResponse searchResponse;
+
+        for (int i = 1; i <= max; i++) {
+            searchResponse = client().prepareSearch("test")
+                    .setQuery(QueryBuilders.rangeQuery("field").gte(lowerBound).lte(upperBound))
+                    .setTerminateAfter(i).execute().actionGet();
+            assertHitCount(searchResponse, (long)i);
+            assertTrue(searchResponse.isTerminatedEarly());
+        }
+
+        searchResponse = client().prepareSearch("test")
+                .setQuery(QueryBuilders.rangeQuery("field").gte(lowerBound).lte(upperBound))
+                .setTerminateAfter(2 * max).execute().actionGet();
+
+        assertHitCount(searchResponse, max);
+        assertFalse(searchResponse.isTerminatedEarly());
+    }
 }

src/test/java/org/elasticsearch/test/TestSearchContext.java

@@ -76,6 +76,7 @@ public class TestSearchContext extends SearchContext {
 
     ContextIndexSearcher searcher;
     int size;
+    private int terminateAfter = DEFAULT_TERMINATE_AFTER;
 
     public TestSearchContext(ThreadPool threadPool, CacheRecycler cacheRecycler, PageCacheRecycler pageCacheRecycler, BigArrays bigArrays, IndexService indexService, FilterCache filterCache, IndexFieldDataService indexFieldDataService) {
         this.cacheRecycler = cacheRecycler;
@@ -351,6 +352,16 @@ public class TestSearchContext extends SearchContext {
     public void timeoutInMillis(long timeoutInMillis) {
     }
 
+    @Override
+    public int terminateAfter() {
+        return terminateAfter;
+    }
+
+    @Override
+    public void terminateAfter(int terminateAfter) {
+        this.terminateAfter = terminateAfter;
+    }
+
     @Override
     public SearchContext minimumScore(float minimumScore) {
         return null;