force termination with no synchronization between threads

jed326 · Sep 25, 2023 · ef167c0 · ef167c0
1 parent d2bf2f5
commit ef167c0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 15 deletions.
diff --git 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
@@ -30,22 +30,19 @@ public class EarlyTerminatingCollectorManager<C extends Collector>
     private final CollectorManager<C, ReduceableSearchResult> manager;
     private final int maxCountHits;
     private boolean forceTermination;
-    private final AtomicLong numCollected;
 
     EarlyTerminatingCollectorManager(CollectorManager<C, ReduceableSearchResult> manager, int maxCountHits, boolean forceTermination) {
         this.manager = manager;
         this.maxCountHits = maxCountHits;
         this.forceTermination = forceTermination;
-        this.numCollected = new AtomicLong();
     }
 
     @Override
     public EarlyTerminatingCollector newCollector() throws IOException {
         return new EarlyTerminatingCollector(
             manager.newCollector(),
             maxCountHits,
-            forceTermination /* forced termination is not supported */,
-            numCollected
+            forceTermination /* forced termination is not supported */
         );
     }