Early terminate the slices for concurrent search if CircuitBreakingEx…

…ception is thrown

Signed-off-by: Neetika Singhal <[email protected]>

server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java

@@ -67,6 +67,7 @@
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.common.lease.Releasable;
 import org.opensearch.common.lucene.search.TopDocsAndMaxScore;
+import org.opensearch.core.common.breaker.CircuitBreakingException;
 import org.opensearch.search.DocValueFormat;
 import org.opensearch.search.SearchService;
 import org.opensearch.search.dfs.AggregatedDfs;
@@ -269,6 +270,11 @@ public void search(
 
     @Override
     protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
+
+        if (searchContext.isCircuitBreakerTripped()) {
+            return;
+        }
+
         // Time series based workload by default traverses segments in desc order i.e. latest to the oldest order.
         // This is actually beneficial for search queries to start search on latest segments first for time series workload.
         // That can slow down ASC order queries on timestamp workload. So to avoid that slowdown, we will reverse leaf
@@ -297,6 +303,9 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto
         if (canMatch(ctx) == false) {
             return;
         }
+        if (searchContext.isCircuitBreakerTripped()) {
+            return;
+        }
 
         final LeafCollector leafCollector;
         try {
@@ -315,6 +324,9 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto
         } catch (QueryPhase.TimeExceededException e) {
             searchContext.setSearchTimedOut(true);
             return;
+        } catch (CircuitBreakingException e) {
+            searchContext.setCircuitBreakerTripped(true);
+            return;
         }
         // catch early terminated exception and rethrow?
         Bits liveDocs = ctx.reader().getLiveDocs();
@@ -330,6 +342,9 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto
                 } catch (QueryPhase.TimeExceededException e) {
                     searchContext.setSearchTimedOut(true);
                     return;
+                } catch (CircuitBreakingException e) {
+                    searchContext.setCircuitBreakerTripped(true);
+                    return;
                 }
             }
         } else {
@@ -349,6 +364,9 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto
                 } catch (QueryPhase.TimeExceededException e) {
                     searchContext.setSearchTimedOut(true);
                     return;
+                } catch (CircuitBreakingException e) {
+                    searchContext.setCircuitBreakerTripped(true);
+                    return;
                 }
             }
         }

server/src/main/java/org/opensearch/search/internal/SearchContext.java

@@ -122,6 +122,7 @@ public List<Aggregator> toAggregators(Collection<Collector> collectors) {
     private InnerHitsContext innerHitsContext;
 
     private volatile boolean searchTimedOut;
+    private volatile boolean circuitBreakerTripped;
 
     protected SearchContext() {}
 
@@ -139,6 +140,14 @@ public void setSearchTimedOut(boolean searchTimedOut) {
         this.searchTimedOut = searchTimedOut;
     }
 
+    public boolean isCircuitBreakerTripped() {
+        return circuitBreakerTripped;
+    }
+
+    public void setCircuitBreakerTripped(boolean circuitBreakerTripped) {
+        this.circuitBreakerTripped = circuitBreakerTripped;
+    }
+
     @Override
     public final void close() {
         if (closed.compareAndSet(false, true)) {