diff --git a/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollector.java b/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollector.java index 5906c2dfc94a4..5cf41ce5bdab8 100644 --- a/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollector.java +++ b/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollector.java @@ -59,7 +59,7 @@ public static final class EarlyTerminationException extends RuntimeException { } private final int maxCountHits; - private final AtomicLong numCollected; + private long numCollected; private final boolean forceTermination; private boolean earlyTerminated; @@ -74,19 +74,11 @@ public static final class EarlyTerminationException extends RuntimeException { super(delegate); this.maxCountHits = maxCountHits; this.forceTermination = forceTermination; - this.numCollected = new AtomicLong(); - } - - EarlyTerminatingCollector(final Collector delegate, int maxCountHits, boolean forceTermination, AtomicLong numCollected) { - super(delegate); - this.maxCountHits = maxCountHits; - this.forceTermination = forceTermination; - this.numCollected = numCollected; } @Override public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException { - if (numCollected.get() >= maxCountHits) { + if (numCollected >= maxCountHits) { earlyTerminated = true; if (forceTermination) { throw new EarlyTerminationException("early termination [CountBased]"); @@ -97,7 +89,7 @@ public LeafCollector getLeafCollector(LeafReaderContext context) throws IOExcept return new FilterLeafCollector(super.getLeafCollector(context)) { @Override public void collect(int doc) throws IOException { - if (numCollected.incrementAndGet() > maxCountHits) { + if (++numCollected > maxCountHits) { earlyTerminated = true; if (forceTermination) { throw new EarlyTerminationException("early termination [CountBased]"); diff --git a/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollectorManager.java b/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollectorManager.java index 5db7d82206a48..a912e4c76c0ec 100644 --- a/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollectorManager.java +++ b/server/src/main/java/org/opensearch/search/query/EarlyTerminatingCollectorManager.java @@ -30,13 +30,11 @@ public class EarlyTerminatingCollectorManager private final CollectorManager manager; private final int maxCountHits; private boolean forceTermination; - private final AtomicLong numCollected; EarlyTerminatingCollectorManager(CollectorManager manager, int maxCountHits, boolean forceTermination) { this.manager = manager; this.maxCountHits = maxCountHits; this.forceTermination = forceTermination; - this.numCollected = new AtomicLong(); } @Override @@ -44,8 +42,7 @@ public EarlyTerminatingCollector newCollector() throws IOException { return new EarlyTerminatingCollector( manager.newCollector(), maxCountHits, - forceTermination /* forced termination is not supported */, - numCollected + forceTermination /* forced termination is not supported */ ); }