Fix query timeout handling by checking searchContext.isSearchTimedOut()

Signed-off-by: Andriy Redko <[email protected]>

Author: reta

Diff for: CHANGELOG.md

@@ -74,6 +74,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Skip remote-repositories validations for node-joins when RepositoriesService is not in sync with cluster-state ([#16763](https://github.com/opensearch-project/OpenSearch/pull/16763))
 - Fix _list/shards API failing when closed indices are present ([#16606](https://github.com/opensearch-project/OpenSearch/pull/16606))
 - Fix remote shards balance ([#15335](https://github.com/opensearch-project/OpenSearch/pull/15335))
+- Fix query timeout handling by checking searchContext.isSearchTimedOut() ([#16882](https://github.com/opensearch-project/OpenSearch/pull/16882))
 
 ### Security
 

Diff for: server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java

@@ -282,7 +282,7 @@ protected void search(List<LeafReaderContext> leaves, Weight weight, Collector c
             // reader order here.
             if (searchContext.shouldUseTimeSeriesDescSortOptimization()) {
                 for (int i = leaves.size() - 1; i >= 0; i--) {
-                    searchLeaf(leaves.get(i), weight, collector);
+                   searchLeaf(leaves.get(i), weight, collector);
                 }
             } else {
                 for (int i = 0; i < leaves.size(); i++) {
@@ -305,6 +305,10 @@ protected void search(List<LeafReaderContext> leaves, Weight weight, Collector c
      */
     @Override
     protected void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collector) throws IOException {
+        // If search is already flagged as timed out, do not continue
+        if (searchContext.isSearchTimedOut()) {
+            return;
+        }
 
         // Check if at all we need to call this leaf for collecting results.
         if (canMatch(ctx) == false) {