Add cancellation checkpoints in field data loading and aggregation paths

[kaushalmahi12]

Description

Cancelled search tasks continue consuming CPU indefinitely because key code paths in field data loading and global ordinals building have zero cancellation checkpoints. This causes "zombie threads" — tasks cancelled but consuming 100% CPU for days/weeks.

Root Cause

The call chain AggregationPhase.preProcess → AggregatorFactories.createTopLevelAggregators → TermsAggregatorFactory → GlobalOrdinalsBuilder.build → AbstractIndexOrdinalsFieldData.loadGlobalDirect → PagedBytesIndexFieldData.loadDirect → OrdinalsBuilder.addDoc → GrowableWriter.get/set has no cancellation checkpoints.

While ExitableDirectoryReader already wraps TermsEnum.next() with cancellation checks, it does not wrap PostingsEnum. The inner loop in PagedBytesIndexFieldData.loadDirect() iterates docsEnum.nextDoc() → OrdinalsBuilder.addDoc() in a tight CPU loop with no way to interrupt it once a task is cancelled.

Changes

This PR adds cancellation checks at three levels:

1. ExitablePostingsEnum in ExitableDirectoryReader (the critical fix)

• New ExitablePostingsEnum wraps PostingsEnum.nextDoc() and advance() with sampled cancellation checks (every 8191 calls, matching the existing ExitableIntersectVisitor pattern)
• Wired from ExitableTermsEnum.postings() so any code path that iterates postings through the exitableReader gets coverage automatically
• This directly covers the tight CPU loop: docsEnum.nextDoc() → OrdinalsBuilder.addDoc() → GrowableWriter.get/set
• The wrapping survives through FrequencyFilter and RamAccountingTermsEnum because FilteredTermsEnum.postings() delegates to the underlying tenum

2. GlobalOrdinalsBuilder.build() — per-segment cancellation check

• New overloaded build() accepts a Runnable cancellationCheck, called between segment iterations
• Old method delegates with () -> {} (no-op) — fully backward compatible
• AbstractIndexOrdinalsFieldData.loadGlobalDirect(DirectoryReader, Runnable) overload threads the check through

3. AggregatorFactories.createTopLevelAggregators() — per-factory cancellation check

• Checks searchContext.isCancelled() before each factories[i].create() call
• Allows deeply nested aggregation trees (4-8 levels observed in production) to be interrupted between levels

Why PostingsEnum is the only Lucene construct that needed wrapping

Other LeafReader iteration constructs (NumericDocValues, SortedSetDocValues, etc.) are used during aggregation collection, which is driven by the query scorer that already has cancellation checks. PostingsEnum is the only one used in an unbounded tight loop (loadDirect) that runs outside the query scorer's cancellation-aware iteration.

Backward Compatibility

• Zero interface changes (IndexFieldData, IndexOrdinalsFieldData, IndexFieldDataCache)
• Zero method signature changes on existing methods
• All new methods are overloads with backward-compatible defaults
• Warmers, recovery, and non-search paths are unaffected
• Plugins/modules that implement IndexFieldData don't need any changes
• When lowLevelCancellation is disabled, ExitableDirectoryReader doesn't wrap the reader, so there is zero overhead

Cache Safety

IndicesFieldDataCache uses computeIfAbsent — if loadDirect() or loadGlobalDirect() throws TaskCancelledException, the entry is simply not cached. No partial/corrupt cache entries.

Testing

• SearchCancellationTests.testExitablePostingsEnum — verifies ExitablePostingsEnum throws TaskCancelledException on nextDoc() and advance() when cancelled
• SearchCancellationTests.testExitablePostingsEnumNoOpWhenCancellationDisabled — verifies no wrapping overhead when cancellation is disabled
• GlobalOrdinalsBuilderTests — verifies per-segment cancellation check fires, delayed cancellation works, and original method backward compat
• AggregatorFactoriesCancellationTests — verifies createTopLevelAggregators throws when isCancelled() is true

Check List

• New functionality includes testing
• New functionality has been documented
• API changes companion pull request - NA
• Commits are signed per the DCO using --signoff

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 Commits, please see the Contribution Guidelines.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit b68ca57)

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: Add ExitablePostingsEnum to ExitableDirectoryReader Relevant files: server/src/main/java/org/opensearch/search/internal/ExitableDirectoryReader.java server/src/test/java/org/opensearch/search/SearchCancellationTests.java Sub-PR theme: Add cancellation checkpoints to GlobalOrdinalsBuilder Relevant files: server/src/main/java/org/opensearch/index/fielddata/ordinals/GlobalOrdinalsBuilder.java server/src/main/java/org/opensearch/index/fielddata/plain/AbstractIndexOrdinalsFieldData.java server/src/test/java/org/opensearch/index/fielddata/ordinals/GlobalOrdinalsBuilderTests.java Sub-PR theme: Add cancellation check in AggregatorFactories.createTopLevelAggregators Relevant files: server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java server/src/test/java/org/opensearch/search/aggregations/AggregatorFactoriesCancellationTests.java

⚡ Recommended focus areas for review Reuse Parameter Ignored In ExitableTermsEnum.postings(), the reuse parameter is intentionally discarded and null is passed to in.postings(). While the comment explains this is because the wrapper type differs, this means that callers who pass a reuse buffer for performance optimization will silently get no reuse benefit. This could cause increased GC pressure in hot paths. Consider documenting this trade-off more explicitly or finding a way to support reuse. public PostingsEnum postings(PostingsEnum reuse, int flags) throws IOException { // Don't reuse when wrapping, since the wrapper type differs from the delegate type final PostingsEnum postings = in.postings(null, flags); return new ExitablePostingsEnum(postings, queryCancellation); } Inconsistent Sampling advance() in ExitablePostingsEnum always checks cancellation (no sampling), while nextDoc() uses sampled checks every 8191 calls. This inconsistency means advance() has higher overhead than nextDoc(). If advance() is called frequently (e.g., in conjunction with a DISI), this could introduce measurable latency. Consider applying the same sampling strategy to advance() for consistency, or document why advance() warrants an unsampled check. public int advance(int target) throws IOException { queryCancellation.checkCancelled(); return in.advance(target); } Resource Leak In testBuildWithCancellationBetweenSegments() and testBuildWithDelayedCancellation(), RandomIndexWriter w is not closed in a try-with-resources block. If an exception occurs before w.close() is called (e.g., during w.getReader()), the writer will not be closed. The writer is closed inside the inner try block via w.close(), but this is only reached if w.getReader() succeeds. RandomIndexWriter w = new RandomIndexWriter(random(), dir); w.w.getConfig().setMergePolicy(NoMergePolicy.INSTANCE); // Create 3 segments with distinct terms for (int seg = 0; seg < 3; seg++) { for (int i = 0; i < 10; i++) { Document doc = new Document(); doc.add(new StringField("field", "seg" + seg + "_term" + i, Field.Store.NO)); w.addDocument(doc); } w.flush(); } try (IndexReader reader = w.getReader()) { w.close(); assertTrue("Need multiple segments for global ordinals", reader.leaves().size() > 1); IndexOrdinalsFieldData fieldData = mockFieldData("field", reader); // Build without cancellation — should succeed assertNotNull(GlobalOrdinalsBuilder.build( reader, fieldData, new NoneCircuitBreakerService(), logger, AbstractLeafOrdinalsFieldData.DEFAULT_SCRIPT_FUNCTION, () -> {} )); // Build with immediate cancellation — should throw between segments expectThrows(TaskCancelledException.class, () -> GlobalOrdinalsBuilder.build( reader, fieldData, new NoneCircuitBreakerService(), logger, AbstractLeafOrdinalsFieldData.DEFAULT_SCRIPT_FUNCTION, () -> { throw new TaskCancelledException("cancelled"); } )); } } Mock Behavior The test calls createSearchContext(...) which likely returns a real or partially-mocked SearchContext, and then uses when(searchContext.isCancelled()).thenReturn(...). If SearchContext is not a mock, this will fail at runtime. Verify that createSearchContext returns a Mockito mock or spy that supports stubbing isCancelled(). SearchContext searchContext = createSearchContext( searcher, createIndexSettings(), new MatchAllDocsQuery(), bucketConsumer, keywordField("field") ); // Build AggregatorFactories from a builder with an actual aggregation TermsAggregationBuilder aggBuilder = new TermsAggregationBuilder("terms").field("field").size(10); QueryShardContext qsc = searchContext.getQueryShardContext(); AggregatorFactories.Builder factoriesBuilder = new AggregatorFactories.Builder().addAggregator(aggBuilder); AggregatorFactories factories = factoriesBuilder.build(qsc, null); // Verify it works when not cancelled when(searchContext.isCancelled()).thenReturn(false); List<Aggregator> aggregators = factories.createTopLevelAggregators(searchContext); assertFalse(aggregators.isEmpty()); // Now mark as cancelled — should throw TaskCancelledException when(searchContext.isCancelled()).thenReturn(true);

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to b68ca57

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Ensure SearchContext is a mock for stubbing The createSearchContext method in AggregatorTestCase typically returns a real or partially-real SearchContext, not a Mockito mock. Calling when(searchContext.isCancelled()).thenReturn(false) on a non-mock object will throw a MissingMethodInvocationException or silently fail at runtime. You should verify that createSearchContext returns a mock, or explicitly create a mock SearchContext for this test. server/src/test/java/org/opensearch/search/aggregations/AggregatorFactoriesCancellationTests.java [56-67] -SearchContext searchContext = createSearchContext( - searcher, createIndexSettings(), new MatchAllDocsQuery(), bucketConsumer, keywordField("field") -); -... +SearchContext searchContext = mock(SearchContext.class); +// configure required SearchContext methods used by createTopLevelAggregators when(searchContext.isCancelled()).thenReturn(false); -List<Aggregator> aggregators = factories.createTopLevelAggregators(searchContext); +// ... set up other necessary mock behaviors Suggestion importance[1-10]: 6 __ Why: This is a valid concern — if createSearchContext in AggregatorTestCase returns a real object rather than a Mockito mock, calling when(...).thenReturn(...) on it would fail. However, the improved_code is overly simplified and doesn't reflect the actual test setup needed, making it hard to apply directly. The issue is worth investigating but the suggested fix is incomplete.	Low
General	Enable delegate reuse in postings wrapping The reuse parameter is completely ignored and always passed as null to the delegate. While the comment explains why the wrapper itself can't be reused, the underlying delegate could still benefit from reuse if the passed-in reuse is already an ExitablePostingsEnum wrapping the same type. More importantly, passing null unconditionally prevents any reuse optimization in the underlying reader, which can increase GC pressure in hot paths like field data loading. server/src/main/java/org/opensearch/search/internal/ExitableDirectoryReader.java [215-220] @Override public PostingsEnum postings(PostingsEnum reuse, int flags) throws IOException { - // Don't reuse when wrapping, since the wrapper type differs from the delegate type - final PostingsEnum postings = in.postings(null, flags); + // Unwrap reuse if it's an ExitablePostingsEnum so the delegate can reuse its inner enum + PostingsEnum reuseInner = (reuse instanceof ExitablePostingsEnum) ? ((ExitablePostingsEnum) reuse).in : null; + final PostingsEnum postings = in.postings(reuseInner, flags); return new ExitablePostingsEnum(postings, queryCancellation); } Suggestion importance[1-10]: 5 __ Why: The suggestion is technically valid — unwrapping the reuse parameter to pass the inner PostingsEnum to the delegate could reduce GC pressure in hot paths. However, the in field in ExitablePostingsEnum is private, so the cast access shown in improved_code would require making it package-private or adding a getter. The impact is moderate as it's an optimization rather than a correctness fix.	Low
	Fix misleading cancellation test comment The cancellation logic arms after the first check passes, but the comment says "cancel after first segment." However, the cancellationCheck is called at the start of each loop iteration (before loading the segment), so the first call sets cancelled=true and the second call (before segment index 1) throws. This means segment 0 is loaded but segment 1 is never loaded — the test actually cancels before the second segment loads, not after. The test comment and variable name are misleading and could cause confusion about the actual behavior being tested. server/src/test/java/org/opensearch/index/fielddata/ordinals/GlobalOrdinalsBuilderTests.java [98-103] () -> { if (cancelled.get()) { - throw new TaskCancelledException("cancelled after first segment"); + throw new TaskCancelledException("cancelled before second segment"); } - cancelled.set(true); // arm cancellation after first check passes + cancelled.set(true); // first check passes; subsequent checks will cancel } Suggestion importance[1-10]: 3 __ Why: The suggestion correctly identifies that the comment "cancel after first segment" is misleading since the cancellation actually fires before the second segment loads. This is a minor documentation/comment accuracy issue with low functional impact.	Low

---

Previous suggestions

Suggestions up to commit 6f5e248

Category	Suggestion	Impact
General	Reuse inner delegate to avoid unnecessary allocations The reuse parameter is completely ignored and always passed as null to the delegate. While the comment explains this is intentional to avoid type mismatch, you should check if the reuse argument is already an ExitablePostingsEnum and unwrap its inner delegate to pass for reuse, improving performance by avoiding unnecessary object allocation. server/src/main/java/org/opensearch/search/internal/ExitableDirectoryReader.java [215-220] @Override public PostingsEnum postings(PostingsEnum reuse, int flags) throws IOException { - // Don't reuse when wrapping, since the wrapper type differs from the delegate type - final PostingsEnum postings = in.postings(null, flags); + // Unwrap ExitablePostingsEnum to reuse the inner delegate + PostingsEnum innerReuse = (reuse instanceof ExitablePostingsEnum) ? ((ExitablePostingsEnum) reuse).in : null; + final PostingsEnum postings = in.postings(innerReuse, flags); return new ExitablePostingsEnum(postings, queryCancellation); } Suggestion importance[1-10]: 5 __ Why: This is a valid performance optimization that unwraps the ExitablePostingsEnum to reuse its inner delegate, avoiding unnecessary object allocations. The improvement is moderate as it reduces GC pressure in hot paths.	Low
	Fix confusing cancellation logic in test The cancellation logic has a bug: segmentCount[0]++ is post-increment, so when segmentCount[0] is 0, it evaluates to 0 (false for > 0) and then increments to 1. Then if (segmentCount[0] == 1) sets cancelled to true. On the second call, segmentCount[0]++ evaluates to 1 (true for > 0) and cancelled.get() is true, so it throws. This works but the logic is confusing and fragile — the second if block will never execute after the first if throws, but the ordering of the two if statements means cancelled is set after the count check, which could cause issues if the logic were slightly different. Consider simplifying with a clearer counter approach. server/src/test/java/org/opensearch/index/fielddata/ordinals/GlobalOrdinalsBuilderTests.java [94-105] AtomicBoolean cancelled = new AtomicBoolean(false); int[] segmentCount = { 0 }; expectThrows(TaskCancelledException.class, () -> GlobalOrdinalsBuilder.build( reader, fieldData, new NoneCircuitBreakerService(), logger, AbstractLeafOrdinalsFieldData.DEFAULT_SCRIPT_FUNCTION, () -> { - if (segmentCount[0]++ > 0 && cancelled.get()) { + int count = segmentCount[0]++; + if (count == 0) { + cancelled.set(true); + } else if (cancelled.get()) { throw new TaskCancelledException("cancelled after first segment"); } - if (segmentCount[0] == 1) cancelled.set(true); } )); Suggestion importance[1-10]: 4 __ Why: The existing logic does work correctly as analyzed in the suggestion, but the improved version with a captured count variable is clearer and less error-prone. This is a minor readability/maintainability improvement in test code.	Low
Possible issue	Ensure search context supports mock stubbing in test The test relies on createSearchContext returning a mock or spy that supports when(...).thenReturn(...), but if createSearchContext returns a real SearchContext object, the when stubbing will fail at runtime. Verify that createSearchContext returns a Mockito mock/spy, or explicitly create a mock SearchContext for the cancellation behavior being tested. server/src/test/java/org/opensearch/search/aggregations/AggregatorFactoriesCancellationTests.java [56-76] -SearchContext searchContext = createSearchContext( - searcher, createIndexSettings(), new MatchAllDocsQuery(), bucketConsumer, keywordField("field") -); -... +SearchContext searchContext = mock(SearchContext.class); +// configure necessary mock behavior for createTopLevelAggregators when(searchContext.isCancelled()).thenReturn(false); List<Aggregator> aggregators = factories.createTopLevelAggregators(searchContext); assertFalse(aggregators.isEmpty()); // Now mark as cancelled — should throw TaskCancelledException when(searchContext.isCancelled()).thenReturn(true); expectThrows( TaskCancelledException.class, () -> factories.createTopLevelAggregators(searchContext) ); Suggestion importance[1-10]: 3 __ Why: The suggestion raises a valid concern about whether createSearchContext returns a mock, but the improved_code is incomplete and doesn't properly configure the mock for createTopLevelAggregators. Since AggregatorTestCase.createSearchContext likely already returns a mock/spy, this is more of a verification concern than a definitive bug.	Low

Suggestions up to commit 0a67de9

Category	Suggestion	Impact
Possible issue	Fix incorrect post-increment logic in cancellation lambda The logic in the cancellation lambda has a bug: segmentCount[0] is post-incremented before the check, so when segmentCount[0] is 0 it becomes 1 after ++, and the cancelled.set(true) branch (segmentCount[0] == 1) is never reached because the increment already happened. The cancellation will never be set, so TaskCancelledException will never be thrown and the test will fail. server/src/test/java/org/opensearch/index/fielddata/ordinals/GlobalOrdinalsBuilderTests.java [94-105] AtomicBoolean cancelled = new AtomicBoolean(false); int[] segmentCount = { 0 }; expectThrows(TaskCancelledException.class, () -> GlobalOrdinalsBuilder.build( reader, fieldData, new NoneCircuitBreakerService(), logger, AbstractLeafOrdinalsFieldData.DEFAULT_SCRIPT_FUNCTION, () -> { - if (segmentCount[0]++ > 0 && cancelled.get()) { + int count = segmentCount[0]++; + if (count == 0) { + cancelled.set(true); + } else if (cancelled.get()) { throw new TaskCancelledException("cancelled after first segment"); } - if (segmentCount[0] == 1) cancelled.set(true); } )); Suggestion importance[1-10]: 7 __ Why: The analysis of the post-increment bug is correct: after segmentCount[0]++ executes when segmentCount[0] is 0, the value becomes 1, so the if (segmentCount[0] == 1) branch is never reached and cancelled is never set to true, meaning the test would never throw TaskCancelledException. The improved code correctly separates the increment from the checks.	Medium
	Fix mocking on non-mock SearchContext instance The test calls when(searchContext.isCancelled()).thenReturn(false) but createSearchContext likely returns a real SearchContext object, not a mock, so Mockito stubbing will fail at runtime with a NotAMockException. The SearchContext should be explicitly mocked, or the test should use a mechanism supported by the real SearchContext to control the cancellation state. server/src/test/java/org/opensearch/search/aggregations/AggregatorFactoriesCancellationTests.java [56-76] -SearchContext searchContext = createSearchContext( - searcher, createIndexSettings(), new MatchAllDocsQuery(), bucketConsumer, keywordField("field") -); -... +SearchContext searchContext = mock(SearchContext.class); +// configure required SearchContext methods used by createTopLevelAggregators when(searchContext.isCancelled()).thenReturn(false); List<Aggregator> aggregators = factories.createTopLevelAggregators(searchContext); assertFalse(aggregators.isEmpty()); // Now mark as cancelled — should throw TaskCancelledException when(searchContext.isCancelled()).thenReturn(true); expectThrows( TaskCancelledException.class, () -> factories.createTopLevelAggregators(searchContext) ); Suggestion importance[1-10]: 6 __ Why: The suggestion raises a valid concern about calling when(searchContext.isCancelled()) on a potentially non-mock object. However, since this test extends AggregatorTestCase, createSearchContext may return a mock or a test-specific implementation that supports Mockito stubbing. The actual behavior depends on the test framework internals, so this is a potential but not certain issue.	Low
General	Reuse existing wrapper to avoid unnecessary allocations The postings method always creates a new ExitablePostingsEnum even when reuse is already an ExitablePostingsEnum wrapping the same delegate. If reuse is an ExitablePostingsEnum, its inner delegate could be reused to avoid unnecessary object allocation. Additionally, passing null to in.postings() unconditionally discards any valid reuse opportunity from the underlying delegate. server/src/main/java/org/opensearch/search/internal/ExitableDirectoryReader.java [215-220] @Override public PostingsEnum postings(PostingsEnum reuse, int flags) throws IOException { - // Don't reuse when wrapping, since the wrapper type differs from the delegate type - final PostingsEnum postings = in.postings(null, flags); + PostingsEnum reuseInner = null; + if (reuse instanceof ExitablePostingsEnum) { + reuseInner = ((ExitablePostingsEnum) reuse).in; + } + final PostingsEnum postings = in.postings(reuseInner, flags); + if (reuse instanceof ExitablePostingsEnum) { + ((ExitablePostingsEnum) reuse).resetDelegate(postings); + return reuse; + } return new ExitablePostingsEnum(postings, queryCancellation); } Suggestion importance[1-10]: 3 __ Why: The suggestion proposes reusing ExitablePostingsEnum instances to reduce allocations, but the improved_code references a resetDelegate method that doesn't exist in the PR code, making it non-applicable as written. The optimization is minor and the implementation is incomplete.	Low

[github-actions]

❌ Gradle check result for 0a67de9: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit 6f5e248

[github-actions]

Persistent review updated to latest commit b68ca57

[github-actions]

❌ Gradle check result for b68ca57: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[kkhatua]

@jainankitk , @bowenlan-amzn , @sgup432

Could you please help review this PR?