Added queryTimeout to the IndexSearcher to ensure that queries that uses this timeout during search exit when queries times out

[navneet1v]

Description

Added queryTimeout to the IndexSearcher to ensure that queries that uses this timeout during search exit when queries times out

Problem

OpenSearch implements its own query timeout/cancellation mechanism via MutableQueryTimeout and ExitableDirectoryReader, which hooks into Lucene at the reader level (terms, points, bulk scoring). However, it never calls IndexSearcher.setTimeout() to set the timeout on Lucene's IndexSearcher.

This is a problem because Lucene's AbstractKnnVectorQuery relies on IndexSearcher.getTimeout() to enforce timeouts during KNN vector search. Specifically, in AbstractKnnVectorQuery.rewrite():

KnnCollectorManager manager = getKnnCollectorManager(k, searcher);
manager = new TimeLimitingKnnCollectorManager(manager, searcher.getTimeout());

Since searcher.getTimeout() returns null, TimeLimitingKnnCollectorManager gets a null QueryTimeout, and KNN vector searches completely ignore the query timeout. This means a slow KNN query can run indefinitely even when a user has set a timeout on their search request or the cluster has search.default_search_timeout configured.

Additionally, ExitableDirectoryReader.ExitableLeafReader does not wrap FloatVectorValues, ByteVectorValues, or any vector-related APIs, so there are no cancellation checks during HNSW graph traversal either.

What's NOT covered by the existing timeout mechanism

Lucene API	Wrapped by ExitableLeafReader?	Timeout enforced?
Terms / TermsEnum	✅ Yes (ExitableTerms, ExitableTermsEnum)	✅ Yes
PointValues / PointTree	✅ Yes (ExitablePointValues, ExitablePointTree)	✅ Yes
HNSW graph traversal (KNN search)	❌ No	❌ No
TimeLimitingKnnCollectorManager	N/A (uses IndexSearcher.getTimeout())	❌ No — getTimeout() returns null

Fix

Two changes in ContextIndexSearcher.java:

1. MutableQueryTimeout now implements both interfaces:

   • ExitableDirectoryReader.QueryCancellation (OpenSearch's interface — existing)
   • org.apache.lucene.index.QueryTimeout (Lucene's interface — new)

   The new shouldExit() method delegates to checkCancelled() — if any cancellation runnable throws, it returns true.

2. setTimeout(cancellable) is called in the ContextIndexSearcher constructor so that searcher.getTimeout() returns the MutableQueryTimeout instance instead of null.

Impact

• KNN vector queries now respect the query timeout via Lucene's TimeLimitingKnnCollectorManager.
• Any future Lucene feature that uses IndexSearcher.getTimeout() will automatically work.
• Zero impact on existing behavior — the existing ExitableDirectoryReader + addQueryCancellation mechanism continues to work exactly as before for terms, points, and bulk scoring.

Testing

Added 5 unit tests in ContextIndexSearcherTests:

Test	What it verifies
testTimeoutIsSetOnSearcher	getTimeout() returns non-null after construction — confirms setTimeout(cancellable) is called
testTimeoutShouldExitReturnsFalseWhenNoCancellations	shouldExit() returns false when no cancellation runnables are registered
testTimeoutShouldExitReturnsFalseWhenCancellationDoesNotThrow	shouldExit() returns false when a no-op cancellation is registered (query still within time)
testTimeoutShouldExitReturnsTrueWhenCancellationThrows	shouldExit() returns true when a cancellation runnable throws (timeout exceeded)
testTimeoutShouldExitReflectsRemoval	shouldExit() transitions from true → false after the throwing cancellation is removed via removeQueryCancellation()

Related Issues

NA

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 350affc)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 Multiple PR themes Sub-PR theme: Bridge MutableQueryTimeout with Lucene QueryTimeout for KNN search timeout support Relevant files: server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java Sub-PR theme: Add tests for MutableQueryTimeout shouldExit behavior Relevant files: server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java

⚡ Recommended focus areas for review Null Timeout The setTimeout(cancellable) is only called when cancellable != null. However, when cancellable is null (e.g., when wrapWithExitableDirectoryReader is false or no cancellation is registered), searcher.getTimeout() will still return null, meaning KNN queries will still ignore timeouts in those cases. It should be verified whether there are valid execution paths where a timeout should be enforced but cancellable is null. if (cancellable != null) { setTimeout(cancellable); } Exception Swallowing In shouldExit(), only QueryPhase.TimeExceededException and TaskCancelledException are caught and converted to true. Any other RuntimeException thrown by a cancellation runnable will propagate unexpectedly to Lucene's TimeLimitingKnnCollectorManager, which may not handle arbitrary exceptions gracefully. Consider whether other exception types (e.g., SearchContextMissingException) should also be caught and treated as a signal to exit. public boolean shouldExit() { try { checkCancelled(); } catch (QueryPhase.TimeExceededException | TaskCancelledException e) { return true; } return false; } Resource Leak In withContextIndexSearcher, the ContextIndexSearcher is created but never closed. If ContextIndexSearcher implements Closeable or holds resources, this could cause resource leaks in tests. The searcher should be closed in a try-with-resources or finally block. ContextIndexSearcher searcher = new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, null, searchContext ); test.accept(searcher); } } Thread Safety MutableQueryTimeout uses a HashSet for runnables, which is not thread-safe. The shouldExit() method iterates over runnables (via checkCancelled()) while add/remove/clear can be called concurrently. If shouldExit() is called from Lucene's KNN search thread while the main thread modifies the set, a ConcurrentModificationException or data race could occur. Consider using a thread-safe collection. private final Set<Runnable> runnables = new HashSet<>(); private Runnable add(Runnable action) { Objects.requireNonNull(action, "cancellation runnable should not be null"); if (runnables.add(action) == false) { throw new IllegalArgumentException("Cancellation runnable already added"); } return action; } private void remove(Runnable action) { runnables.remove(action); } @Override public void checkCancelled() { for (Runnable timeout : runnables) { timeout.run(); } } @Override public boolean isEnabled() { return runnables.isEmpty() == false; } /** * Implements {@link QueryTimeout#shouldExit()} by delegating to {@link #checkCancelled()}. * Returns {@code true} if a registered cancellation runnable throws a * {@link org.opensearch.search.query.QueryPhase.TimeExceededException} (timeout) or * {@link org.opensearch.core.tasks.TaskCancelledException} (task cancellation), * indicating that the query should be terminated early. * <p> * This is called by Lucene's {@link org.apache.lucene.search.TimeLimitingKnnCollectorManager} * during KNN vector search to check whether the search should be terminated early. */ @Override public boolean shouldExit() { try { checkCancelled(); } catch (QueryPhase.TimeExceededException | TaskCancelledException e) { return true; } return false; } public void clear() { runnables.clear(); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 350affc
Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Guard against unsafe cast to QueryTimeout The cancellable field is of type ExitableDirectoryReader.QueryCancellation, but setTimeout expects a QueryTimeout. Since MutableQueryTimeout now implements both interfaces, this cast is safe — but only if cancellable is actually a MutableQueryTimeout. If a different implementation of QueryCancellation is ever passed, this will fail at runtime with a ClassCastException. Add an explicit type check or cast with a clear error message to guard against this. server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java [166-168] if (cancellable != null) { - setTimeout(cancellable); + if (cancellable instanceof QueryTimeout) { + setTimeout((QueryTimeout) cancellable); + } else { + throw new IllegalArgumentException( + "cancellable must implement QueryTimeout to support searcher-level timeout, got: " + cancellable.getClass() + ); + } } Suggestion importance[1-10]: 5 __ Why: The concern is valid — if a non-MutableQueryTimeout implementation of QueryCancellation is passed, the cast would fail. However, looking at the codebase, cancellable is always a MutableQueryTimeout instance (it's a private static inner class), making this a theoretical rather than practical risk. The suggestion adds defensive programming value but is not critical.	Low
General	Ensure searcher resources are properly closed The ContextIndexSearcher is never closed in withContextIndexSearcher, which may leak resources (e.g., the underlying reader or executor). Wrap it in a try-with-resources block to ensure proper cleanup. server/src/test/java/org/opensearch/search/internal/ContextIndexSearcherTests.java [665-674] -ContextIndexSearcher searcher = new ContextIndexSearcher( +try (ContextIndexSearcher searcher = new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, null, searchContext -); -test.accept(searcher); +)) { + test.accept(searcher); +} Suggestion importance[1-10]: 4 __ Why: The ContextIndexSearcher is not closed after use in the test helper, which could leak resources. Wrapping it in a try-with-resources block is a valid improvement for test hygiene, though the impact is limited to test code.	Low
	Clarify exception handling in shouldExit The shouldExit() method silently swallows QueryPhase.TimeExceededException and TaskCancelledException and returns true, but any other RuntimeException propagates unchecked. This asymmetry means that if a new cancellation exception type is introduced in the future, it will propagate rather than signal exit. Consider documenting this explicitly or catching RuntimeException broadly and re-throwing only non-cancellation exceptions. server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java [671-678] +@Override public boolean shouldExit() { try { checkCancelled(); } catch (QueryPhase.TimeExceededException | TaskCancelledException e) { return true; + } catch (RuntimeException e) { + // Re-throw unexpected exceptions so they are not silently swallowed + throw e; } return false; } Suggestion importance[1-10]: 1 __ Why: The improved_code is functionally identical to the existing_code — the added catch (RuntimeException e) { throw e; } block does exactly what already happens implicitly when an uncaught exception propagates. This suggestion adds no behavioral change and only adds noise.	Low

---

Previous suggestions

Suggestions up to commit 73bb8c9

Category	Suggestion	Impact
Possible issue	Guard against null timeout being set setTimeout is called unconditionally with cancellable, but when wrapWithExitableDirectoryReader is false, the cancellable may not be properly wired to the reader. More critically, if cancellable is null (as seen in tests passing null), calling setTimeout(null) could clear any previously set timeout or cause unexpected behavior. A null check should be added before calling setTimeout. server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java [161] -setTimeout(cancellable); +if (cancellable != null) { + setTimeout(cancellable); +} Suggestion importance[1-10]: 6 __ Why: The tests explicitly pass null as the cancellable parameter, and calling setTimeout(null) could clear any previously set timeout or cause unexpected behavior. Adding a null check before calling setTimeout is a valid defensive measure to prevent unintended side effects.	Low
	Narrow caught exception type in timeout check The shouldExit() method catches all RuntimeException types, which is overly broad and could mask unrelated runtime errors (e.g., NullPointerException, IllegalStateException). It should only catch the specific cancellation exception types (e.g., TaskCancelledException or SearchTimeoutException) that checkCancelled() is expected to throw to signal a timeout or cancellation. server/src/main/java/org/opensearch/search/internal/ContextIndexSearcher.java [637-645] @Override public boolean shouldExit() { try { checkCancelled(); - } catch (RuntimeException e) { + } catch (TaskCancelledException | TimeExceededException e) { return true; } return false; } Suggestion importance[1-10]: 5 __ Why: The concern about catching all RuntimeException types is valid from a correctness standpoint, but checkCancelled() in OpenSearch typically throws specific runtime exceptions. The suggestion to use TaskCancelledException | TimeExceededException may not be accurate since the actual exception types thrown by checkCancelled() need to be verified, and the shouldExit() contract may intentionally use broad exception catching to handle any cancellation signal.	Low

[navneet1v]

Updated code from the suggestion added by PR bot

[github-actions]

Persistent review updated to latest commit 350affc

[github-actions]

✅ Gradle check result for 350affc: SUCCESS

[codecov]

Codecov Report

❌ Patch coverage is 87.50000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 73.37%. Comparing base (75c87b7) to head (350affc).
⚠️ Report is 27 commits behind head on main.

Files with missing lines	Patch %	Lines
...ensearch/search/internal/ContextIndexSearcher.java	87.50%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21316      +/-   ##
============================================
- Coverage     73.40%   73.37%   -0.04%     
+ Complexity    74092    74030      -62     
============================================
  Files          5948     5948              
  Lines        336527   336534       +7     
  Branches      48552    48553       +1     
============================================
- Hits         247035   246931     -104     
- Misses        69836    69907      +71     
- Partials      19656    19696      +40     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.