Fix constant ~5s stall in coordinator reduce teardown on limited queries

[alchemist51]

Description

Every PPL query whose plan includes a late-materialization fetch phase (any ... | sort <col> | head K | fields <non-sort-key-col> shape) pays a fixed ~5s in the LATE_MATERIALIZATION stage — invariant to row count, column count, and bytes scanned — and logs one WARN per query:

[reduce-sink] timed out waiting for reduce teardown: taskId=...

Profile of an affected query (50 rows out of 749M docs; scan itself is 60–90ms):

Stage	Type	Elapsed
0	SHARD_FRAGMENT	78 ms
1	COORDINATOR_REDUCE	79 ms
2	LATE_MATERIALIZATION	5037 ms
3	COORDINATOR_REDUCE	5038 ms

Root cause: a teardown cycle that only the 5s timeout could break

DatafusionReduceSink.closeImpl (REDUCING branch) fired cancel_query and then waited on reduceDone.await(5, SECONDS). That wait sat inside a cycle:

1. close() waits on reduceDone,
2. reduceDone.countDown() runs only after the drain loop exits,
3. the drain is parked in stream_next waiting for a pipeline-breaking SortExec/TopK to emit,
4. the TopK cannot emit until its input reaches end-of-input,
5. the input senders were only closed after close() returned.

The cancel could not break the cycle either, because of two Rust-side defects:

• Registry overwrite/orphan. QUERY_REGISTRY was DashMap<i64, Arc<QueryTracker>> — one tracker per context_id. But a query with an LM stage opens two coordinator reduce sinks (both constructed at graph-build time, before any stage runs), and both register under the same ctx.taskId(). The second insert silently overwrote the first tracker, and the first stream to close removed the shared key in Drop — so cancel_query found no entry (silent no-op) while the sibling stream was still running. Verified live: cancel_query: NO REGISTRY ENTRY for context_id=943 (no-op). live ids=[] one line after fireCancelQuery: taskId=943.
• Token lookup goes stale. stream_next re-fetched its cancellation token from the registry on every call. Once the entry was gone, it got None, and cancellable_or(None, ...) degrades to a bare, structurally uncancellable await.

jstack during the stall confirms it: the drain thread is RUNNABLE inside the FFM downcall with cpu=0.43ms — waiting, not working.

Why cancel-first teardown was also wrong (not just slow)

With the Rust fixes in place, cancellation actually works — and cancel-first teardown then aborts the TopK before EOF, so projections mixing sort keys with fetched columns returned 0 rows. The close-path semantics are "no more input is coming", not "abort"; EOF is the correct signal.

The fix

Rust (query_tracker.rs, api.rs):

• QUERY_REGISTRY becomes DashMap<i64, Vec<Arc<QueryTracker>>>: register appends; Drop removes only its own tracker (Arc::ptr_eq retain) and drops the key when empty; cancel_query cancels every tracker under the id, and logs registry misses instead of silently no-oping.
• stream_next uses its handle's own token via new QueryTrackingContext::cancellation_token() — immune to registry churn.

Java (DatafusionReduceSink, DatafusionPartitionSender):

• closeImpl's REDUCING branch signals EOF first: tryClose() each input sender (dropping a sender closes its mpsc; the native plan reads that as end-of-input, the TopK emits, the drain finishes naturally, rows preserved).
• tryClose() (new), not close(): a feeder parked in send() holds the sender read lock; a blocking close would deadlock — the exact scenario covered by the existing testCloseWhileFeederParkedOnFullChannelDoesNotDeadlock, which caught the first version of this patch.
• If tryClose fails (live producer mid-send) → fall back to cancel (dropped rows are correct semantics; the consumer already stopped listening). If EOF doesn't finish the drain within 5s → WARN + cancel, so close() can never hang.

Measured impact

On a 749M-doc, 5-shard composite (parquet-primary) index, where <lead-sort-key>=<v> | sort - <second-key> | head 50 | fields <2 non-sort-key cols>:

	before	after
latency (warm)	5.15 s	0.17–0.19 s (~27×)
mixed sort-key+fetched projection	50 rows (only because the timeout expired before EOF)	50 rows, sha-identical output
teardown WARNs	1 per query	0

head K variants, no-sort, and aggregation shapes all row-correct after the fix.

Related Issues

Related to the observability gap in #22601 (the LM stage emits no plan/metrics, which is why this required jstack + DEBUG logs to diagnose).

Check List

• Functionality includes testing.
  • Rust: test_sibling_stream_drop_does_not_orphan_survivor, test_cancel_query_cancels_all_siblings (new), 23/23 query_tracker tests pass.
  • Java: testCloseWhileReducingSignalsEofAndPreservesRows (new) — fails on the old code both ways (5s close stall pre-fix; 0 rows with cancel-first teardown). 12/12 DatafusionReduceSinkTests pass.
• 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.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 1e55371)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Race in send() after closeRequested In send(), the closeRequested flag is checked inside the read lock and throws IllegalStateException if set. However, requestEarlyTermination() sets closeRequested = true before attempting to acquire the write lock (or before signalling native early termination). A send that has already passed the closeRequested check but not yet called NativeBridge.senderSend can proceed, but then the finally block calls closeAfterInFlightSend() which acquires the write lock and calls super.close() — freeing the native sender pointer. If another concurrent send() is between the check and the native call, or if requestEarlyTermination races with an in-flight send that had not yet observed closeRequested, the observable behavior differs from the doc comment ("A send that is already in progress completes"). Also, throwing IllegalStateException from send() when closeRequested is set may break callers that expect the SENDER_SEND_RECEIVER_DROPPED return contract rather than an exception. public long send(long arrayAddr, long schemaAddr) { lifecycle.readLock().lock(); try { if (closeRequested) { throw new IllegalStateException("sender close requested"); } long rc = NativeBridge.senderSend(getPointer(), arrayAddr, schemaAddr); if (rc == NativeBridge.SENDER_SEND_RECEIVER_DROPPED) { receiverDropped = true; } return rc; } finally { lifecycle.readLock().unlock(); closeAfterInFlightSend(); } } Handle attachment ambiguity with siblings set_abort_handle and set_cpu_runtime_handle target trackers.last(), assuming the just-registered tracker is at the tail. This is only safe if no other stream registers under the same context_id between the caller's registration and the handle attachment. For LM queries with two coordinator reduce sinks constructed at graph-build time, if both QueryTrackingContext::new calls interleave with subsequent set_abort_handle calls, the abort handle for the first stream can be attached to the second tracker (or vice versa), leaving one stream without an abort handle and effectively uncancellable. Consider keying by Arc<QueryTracker> identity instead of positional last(). pub fn set_abort_handle(context_id: i64, handle: AbortHandle) { if let Some(trackers) = QUERY_REGISTRY.get(&context_id) { if let Some(tracker) = trackers.last() { tracker.abort_handle.set(handle).ok(); } } } /// Store the CPU runtime handle for the given context_id so that /// `cancel_query` can flush deferred drops on that runtime. pub fn set_cpu_runtime_handle(context_id: i64, handle: tokio::runtime::Handle) { if let Some(trackers) = QUERY_REGISTRY.get(&context_id) { if let Some(tracker) = trackers.last() { tracker.cpu_runtime_handle.set(handle).ok(); } } Cancellation now surfaces as error stream_next was changed from cancellable_or (which mapped to a benign sentinel) to cancellable which returns Err on cancel, then propagates via map_err(DataFusionError::Execution)?. This changes the caller-visible contract: consumers that previously treated cancellation as normal EOS will now observe a hard error. Verify all callers (including the Java drain returning "cancellation sentinel with zero rows") handle the new error path correctly and don't surface spurious query failures for user-initiated cancels or graceful teardown scenarios where the token is fired. pub async unsafe fn stream_next(stream_ptr: i64) -> Result<i64, DataFusionError> { let handle = &mut *(stream_ptr as *mut QueryStreamHandle); // Use the handle's OWN token, not a registry lookup by context_id. The // registry entry can be removed by a sibling stream's Drop (same id) while // this stream is mid-flight; a `None` token here silently degrades // `cancellable_or` to a bare uncancellable await — the reduce sink's // cancel then can't interrupt an in-flight drain (the ~5s LM stall). let token = handle._query_tracking_context.cancellation_token(); // Fetch the next batch (cancellation-aware). Query cancellation is an abort, // not normal end-of-stream: callers must receive an error rather than the // same zero sentinel used for EOF. let result = cancellation::cancellable( token.as_ref(), handle._query_tracking_context.context_id(), async { handle .stream .try_next() .await .map_err(|e: DataFusionError| e) }, ) .await .map_err(DataFusionError::Execution)?; Cancel-then-close ordering On the failure path, cancellable.cancel() is called immediately before parentSink.close(). If cancel() is asynchronous and races with close(), the sink may observe close as graceful EOF before the cancel takes effect, defeating the purpose of the change. Confirm cancel() synchronously marks the sink cancelled (or close() observes cancellation state) so the parent reduce cannot complete successfully in the failure path. } else { // A fetch failure is not normal input completion. Abort a cancellable parent // before closing it so its native stream reports an error instead of treating // close as graceful EOF and completing the top-level query successfully. if (parentSink instanceof CancellableExchangeSink cancellable) { cancellable.cancel(); } try { parentSink.close(); } catch (Exception ignore) {}

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 1e55371

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Prevent double-close of native sender handle send() throws IllegalStateException when closeRequested is set, but the finally block still calls closeAfterInFlightSend(), which will attempt to acquire the write lock and close the native handle. If requestEarlyTermination() was concurrently invoked and already closed the sender under the write lock, this second close via closeUnderWriteLock will call super.close() again on an already-closed handle. Guard closeUnderWriteLock against double-close or short-circuit when the underlying handle is already closed. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java [136-139] -public long send(long arrayAddr, long schemaAddr) { - lifecycle.readLock().lock(); - try { - if (closeRequested) { - throw new IllegalStateException("sender close requested"); - } - long rc = NativeBridge.senderSend(getPointer(), arrayAddr, schemaAddr); - if (rc == NativeBridge.SENDER_SEND_RECEIVER_DROPPED) { - receiverDropped = true; - } - return rc; - } finally { - lifecycle.readLock().unlock(); - closeAfterInFlightSend(); +private void closeUnderWriteLock(String reason) { + if (isClosed()) { + return; } + super.close(); + logger.debug("[sender] closed ptr={} ({})", ptr, reason); } Suggestion importance[1-10]: 6 __ Why: Potential double-close concern is plausible: if requestEarlyTermination() closes under the write lock and then the in-flight send's finally path also calls closeAfterInFlightSend(), super.close() could be invoked twice. However, NativeHandle's base close is typically idempotent, so the actual impact depends on that contract; the suggestion is defensive but not clearly critical.	Low
General	Make shard-lock release ordering explicit QUERY_REGISTRY.get() returns a Ref that holds the DashMap shard read lock for its entire scope. Iterating and cloning Arcs inside the Some(entry) => ... arm still executes while the shard lock is held, so the comment about "cancel outside it" is only satisfied because the match scope ends before the loop below. Bind the snapshot to a let and explicitly drop(entry) (or move the collect into its own block) to make the lock-release ordering unambiguous and prevent future edits from re-introducing the deadlock. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [402-405] pub fn cancel_query(context_id: i64) { - // Snapshot the trackers under the entry lock, then cancel outside it — - // `token.cancel()` can wake tasks whose teardown re-enters the registry - // (Drop → get_mut on the same shard) and would deadlock under DashMap's - // shard lock. - let trackers: Vec<Arc<QueryTracker>> = match QUERY_REGISTRY.get(&context_id) { - Some(entry) => entry.iter().map(Arc::clone).collect(), - None => Vec::new(), + let trackers: Vec<Arc<QueryTracker>> = { + match QUERY_REGISTRY.get(&context_id) { + Some(entry) => { + let snapshot = entry.iter().map(Arc::clone).collect(); + drop(entry); + snapshot + } + None => Vec::new(), + } }; Suggestion importance[1-10]: 3 __ Why: The existing code already drops the Ref at the end of the match arm before the cancellation loop. The suggestion is a minor readability improvement to make the drop explicit but doesn't fix an actual bug.	Low
	Snapshot senders map before iteration Iterating over sendersByChildStageId.values() while another thread may be feeding those senders is safe, but if sendersByChildStageId is a plain HashMap populated during setup and never mutated afterward this is fine — otherwise a ConcurrentModificationException could fire. Verify the map is either a concurrent collection or effectively immutable after construction, and consider snapshotting to a local list before iteration to be defensive. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [386-392] -for (DatafusionPartitionSender sender : sendersByChildStageId.values()) { +for (DatafusionPartitionSender sender : new java.util.ArrayList<>(sendersByChildStageId.values())) { try { sender.requestEarlyTermination(); } catch (Exception e) { logger.warn("[reduce-sink] error signalling input EOF: taskId={}", ctx.taskId(), e); } } -try { - if (!reduceDone.await(5, java.util.concurrent.TimeUnit.SECONDS)) { Suggestion importance[1-10]: 3 __ Why: The suggestion asks the author to verify the map's concurrency semantics. It's a defensive change with limited impact, and the map is likely immutable after setup based on the surrounding code pattern.	Low
	Document single-poller invariant on receiver mutex poll_next now acquires a parking_lot::Mutex on every poll. If terminate_early is called from another thread while the receiver task is being polled on the same async runtime, contention is unlikely to deadlock, but holding a blocking mutex across poll_recv means a concurrent terminate_early will spin/block briefly under contention. More importantly, if any future code path ever polls this stream while already holding the lock (e.g. via re-entrancy from a waker), parking_lot will deadlock silently since it's not reentrant. Consider using try_lock or documenting the single-poller invariant to make the assumption explicit. sandbox/plugins/analytics-backend-datafusion/rust/src/partition_stream.rs [145-147] impl Stream for PartitionStreamReceiver { type Item = Result<RecordBatch, DataFusionError>; + // Invariant: only one task polls this receiver at a time (DataFusion's + // streaming table contract). `terminate_early` briefly contends this + // mutex to close the channel; contention is bounded to a single call. fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { self.rx.lock().poll_recv(cx) } } Suggestion importance[1-10]: 2 __ Why: Documentation-only suggestion about a single-poller invariant; low impact and doesn't address a real bug.	Low

---

Previous suggestions

Suggestions up to commit 32cc99b

Category	Suggestion	Impact
General	Return dropped sentinel instead of throwing Throwing IllegalStateException when closeRequested is observed at entry means a benign race between requestEarlyTermination() and a feeder calling send() surfaces as an error to the feeder instead of the documented SENDER_SEND_RECEIVER_DROPPED return. Return the receiver-dropped sentinel (and set receiverDropped = true) so the feeder cleanly stops feeding, matching the contract in the Javadoc. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java [60-62] public long send(long arrayAddr, long schemaAddr) { lifecycle.readLock().lock(); try { if (closeRequested) { - throw new IllegalStateException("sender close requested"); + receiverDropped = true; + return NativeBridge.SENDER_SEND_RECEIVER_DROPPED; } long rc = NativeBridge.senderSend(getPointer(), arrayAddr, schemaAddr); if (rc == NativeBridge.SENDER_SEND_RECEIVER_DROPPED) { receiverDropped = true; } return rc; } finally { lifecycle.readLock().unlock(); closeAfterInFlightSend(); } } Suggestion importance[1-10]: 6 __ Why: Reasonable observation: throwing IllegalStateException on a benign race between requestEarlyTermination() and send() breaks the documented contract of returning SENDER_SEND_RECEIVER_DROPPED. However, the impact depends on caller expectations and this may be intentional to detect misuse.	Low
	Guard cancel so close still runs cancel() may itself throw (e.g. from native bridge errors during teardown), which would skip the subsequent parentSink.close() and leak the sink. Wrap the cancel call in try/catch (mirroring the swallow-and-log pattern used in the success branch) so close always runs. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/Stitcher.java [213-218] if (parentSink instanceof CancellableExchangeSink cancellable) { - cancellable.cancel(); + try { + cancellable.cancel(); + } catch (Exception e) { + logger.warn("[Stitcher] parentSink.cancel() failed after fetch failure", e); + } } try { parentSink.close(); } catch (Exception ignore) {} Suggestion importance[1-10]: 6 __ Why: Legitimate defensive improvement: if cancel() throws, close() is skipped and the sink leaks. Wrapping in try/catch matches the pattern already used elsewhere in the file.	Low
	Avoid blocking mutex inside poll_next Using std::sync::Mutex inside an async poll_next blocks the executor thread if terminate_early() is holding the lock, and any panic while holding it poisons the mutex for the receiver. Since only the sender's terminate_early and the receiver's poll_next contend, consider parking_lot::Mutex (non-poisoning) or storing the receiver in a tokio::sync::Mutex/only exposing close() via a separate channel to avoid poll-time blocking and poison hazards. sandbox/plugins/analytics-backend-datafusion/rust/src/partition_stream.rs [147-152] fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { - self.rx - .lock() - .expect("partition receiver mutex poisoned") - .poll_recv(cx) + let mut guard = self.rx.lock().expect("partition receiver mutex poisoned"); + guard.poll_recv(cx) } Suggestion importance[1-10]: 4 __ Why: The concern about std::sync::Mutex in async context is valid in principle, but contention is limited (only terminate_early briefly). The improved_code is essentially equivalent to the existing code, not addressing the raised concern.	Low
Possible issue	Release shard lock before cancelling QUERY_REGISTRY.get(...) returns a Ref that holds a DashMap shard read lock; iterating and cloning inside the match arm keeps that lock live across the collect. Bind the Ref to a variable, snapshot into the Vec, and explicitly drop it before cancellation so any Drop path re-entering the same shard cannot deadlock, matching the intent already stated in the surrounding comment. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [402-405] -let trackers: Vec<Arc<QueryTracker>> = match QUERY_REGISTRY.get(&context_id) { - Some(entry) => entry.iter().map(Arc::clone).collect(), - None => Vec::new(), +let trackers: Vec<Arc<QueryTracker>> = { + match QUERY_REGISTRY.get(&context_id) { + Some(entry) => { + let snapshot: Vec<Arc<QueryTracker>> = entry.iter().map(Arc::clone).collect(); + drop(entry); + snapshot + } + None => Vec::new(), + } }; Suggestion importance[1-10]: 5 __ Why: Valid concern about DashMap shard lock lifetime, though the match arm's Ref should be dropped at the end of the match expression before the cancellation loop runs. The explicit drop is defensive but the original code likely already releases the lock correctly.	Low

Suggestions up to commit 8d86838

Category	Suggestion	Impact
General	Treat close exceptions as unclosed senders If sender.tryClose() throws, allClosed is not updated for that sender and the loop treats it as if it succeeded — meaning the code may skip the cancel fallback even though a sender failed to close. Set allClosed = false in the catch block to guarantee the cancel path runs when any sender close fails. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [388-395] boolean allClosed = true; for (DatafusionPartitionSender sender : sendersByChildStageId.values()) { try { allClosed &= sender.tryClose(); } catch (Exception e) { + allClosed = false; logger.warn("[reduce-sink] error closing input sender for EOF: taskId={}", ctx.taskId(), e); } } Suggestion importance[1-10]: 6 __ Why: Valid correctness improvement: if tryClose() throws, allClosed remains true and the cancel fallback is skipped, potentially leaving a producer parked. Setting allClosed = false in the catch ensures the cancel path runs.	Low
Possible issue	Avoid holding shard lock during registry iteration The QUERY_REGISTRY.get(&context_id) returns a Ref that holds a shard read lock; keeping it alive across .iter().map(...).collect() is fine, but the subsequent QUERY_REGISTRY.iter() call in the empty-branch can deadlock if any shard lock is still held. Ensure the get guard is dropped before iterating the registry by binding it in an explicit scope. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [402-410] -let trackers: Vec<Arc<QueryTracker>> = match QUERY_REGISTRY.get(&context_id) { - Some(entry) => entry.iter().map(Arc::clone).collect(), - None => Vec::new(), +let trackers: Vec<Arc<QueryTracker>> = { + match QUERY_REGISTRY.get(&context_id) { + Some(entry) => entry.iter().map(Arc::clone).collect(), + None => Vec::new(), + } }; if trackers.is_empty() { - // A cancel for an id that was never registered (or already finished) - // does nothing. Silent no-ops here hid a real bug (the ~5s LM stall): - // log the miss and the live ids so any recurrence is visible. let live: Vec<i64> = QUERY_REGISTRY.iter().map(|e| *e.key()).collect(); Suggestion importance[1-10]: 3 __ Why: The existing match expression already scopes the Ref guard so it is dropped before the subsequent QUERY_REGISTRY.iter() call; the "improved_code" merely wraps it in an explicit block without changing semantics. Marginal to no impact.	Low

Suggestions up to commit 2959599

Category	Suggestion	Impact
Possible issue	Treat close exceptions as not-closed If sender.tryClose() throws, allClosed is not updated, so a partially-failed close can leave allClosed=true and the cancel-fallback path is skipped — the drain may then hang until the 5s await times out. Treat a thrown exception as "not closed" so the cancel fallback still fires and unblocks the producer. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [388-395] boolean allClosed = true; for (DatafusionPartitionSender sender : sendersByChildStageId.values()) { try { allClosed &= sender.tryClose(); } catch (Exception e) { + allClosed = false; logger.warn("[reduce-sink] error closing input sender for EOF: taskId={}", ctx.taskId(), e); } } Suggestion importance[1-10]: 7 __ Why: Correctly identifies that a thrown exception from tryClose() leaves allClosed unchanged, potentially skipping the cancel fallback and reintroducing the very stall this PR fixes. Setting allClosed = false on exception ensures the cancel path fires.	Medium
	Release shard lock before cancelling The QUERY_REGISTRY.get(&context_id) guard is held across the entry.iter().map(...).collect() call, meaning the DashMap shard read lock is still held while collecting. To truly snapshot outside the shard lock (as the comment claims), bind the guard to a scope that drops before cancellation, or explicitly drop it after collection. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [402-405] -let trackers: Vec<Arc<QueryTracker>> = match QUERY_REGISTRY.get(&context_id) { - Some(entry) => entry.iter().map(Arc::clone).collect(), - None => Vec::new(), +let trackers: Vec<Arc<QueryTracker>> = { + match QUERY_REGISTRY.get(&context_id) { + Some(entry) => { + let v = entry.iter().map(Arc::clone).collect(); + drop(entry); + v + } + None => Vec::new(), + } }; Suggestion importance[1-10]: 6 __ Why: Valid concern: the DashMap get guard remains alive across collect() in the match arm, meaning the shard lock is held during collection. While collection itself doesn't cancel, the explicit drop makes intent clearer and matches the comment's claim. Moderate impact since the actual cancel loop runs after the match expression completes.	Low
General	Bound diagnostic scan cost on miss Iterating the entire QUERY_REGISTRY on every miss to log all live ids is O(N) and can be expensive under load or if cancels are frequent for unknown ids. Guard this diagnostic collection behind a debug-enabled check or limit the number of ids logged to avoid pathological overhead. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [410-413] -let live: Vec<i64> = QUERY_REGISTRY.iter().map(|e| *e.key()).collect(); +let live: Vec<i64> = QUERY_REGISTRY.iter().take(32).map(|e| *e.key()).collect(); native_bridge_common::log_debug!( - "cancel_query: NO REGISTRY ENTRY for context_id={context_id} (no-op). live ids={live:?}" + "cancel_query: NO REGISTRY ENTRY for context_id={context_id} (no-op). live ids (up to 32)={live:?}" ); Suggestion importance[1-10]: 3 __ Why: Minor performance concern for a debug-level log path; the impact is low since log_debug! may already be gated and cancel_query misses should be rare.	Low

Suggestions up to commit b81c88f

Category	Suggestion	Impact
Possible issue	Avoid re-entering shard lock during removal Acquiring get_mut and then calling remove_if on the same key in DashMap can deadlock because both take shard-level write locks; even though drop(entry) is called first, holding a RefMut and then re-entering the same shard is a known DashMap footgun and can also race with a concurrent insert re-populating the Vec. Restructure so the retain and the empty check are atomic under a single shard operation, e.g. using remove_if with a predicate that mutates via entry API, or explicitly re-check emptiness after re-acquisition. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [686-694] -if let Some(mut entry) = QUERY_REGISTRY.get_mut(&tracker.context_id) { - entry.retain(|t| !Arc::ptr_eq(t, tracker)); - let now_empty = entry.is_empty(); - drop(entry); - if now_empty { - QUERY_REGISTRY - .remove_if(&tracker.context_id, |_, v| v.is_empty()); +let should_remove = { + if let Some(mut entry) = QUERY_REGISTRY.get_mut(&tracker.context_id) { + entry.retain(|t| !Arc::ptr_eq(t, tracker)); + entry.is_empty() + } else { + false } +}; +if should_remove { + QUERY_REGISTRY.remove_if(&tracker.context_id, |_, v| v.is_empty()); } Suggestion importance[1-10]: 7 __ Why: Valid concern: holding a RefMut from get_mut on a DashMap shard and then calling remove_if on the same key (even after drop(entry)) risks re-entrancy issues and can race with concurrent inserts. Restructuring to release the mut ref before removal is a reasonable improvement, though the existing code does drop(entry) first.	Medium
	Positional tracker lookup is race-prone Using trackers.last() to target "the most recently registered tracker" is racy: between a stream's new and its set_abort_handle call, a sibling stream can register under the same context_id and become last(), causing the handle to be attached to the wrong tracker (and the intended tracker to never receive its abort handle). Consider passing/using an Arc/identifier of the specific tracker so the setter targets it by identity rather than positional order. The same concern applies to set_cpu_runtime_handle, take_cpu_runtime_handle, and flush_cpu_runtime. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [509-514] pub fn set_abort_handle(context_id: i64, handle: AbortHandle) { + // NOTE: positional `last()` is racy with concurrent sibling registration; + // callers should identify the target tracker explicitly (e.g. by Arc). if let Some(trackers) = QUERY_REGISTRY.get(&context_id) { if let Some(tracker) = trackers.last() { tracker.abort_handle.set(handle).ok(); } } } Suggestion importance[1-10]: 6 __ Why: Valid concern about the race window between tracker registration and set_abort_handle/set_cpu_runtime_handle when siblings register concurrently under the same context_id. However, the suggestion only annotates with a comment rather than fixing it, limiting its impact.	Low
General	Treat sender close errors as not-closed If sender.tryClose() throws, allClosed retains its prior value and the failure is silently swallowed for the fallback decision — the loop may conclude "all closed" while one sender actually errored, skipping the cancel fallback and risking a hang. Set allClosed = false in the catch block so a throwing sender forces the cancel path. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [388-395] boolean allClosed = true; for (DatafusionPartitionSender sender : sendersByChildStageId.values()) { try { allClosed &= sender.tryClose(); } catch (Exception e) { + allClosed = false; logger.warn("[reduce-sink] error closing input sender for EOF: taskId={}", ctx.taskId(), e); } } Suggestion importance[1-10]: 6 __ Why: Correct observation: if tryClose throws, allClosed isn't updated to false, potentially skipping the cancel fallback and risking a hang. Setting allClosed = false in the catch block is a small but meaningful robustness improvement.	Low

Suggestions up to commit 10d1ad4

Category	Suggestion	Impact
Possible issue	Treat tryClose exceptions as not-closed If sender.tryClose() throws, allClosed stays true for that sender, so the fallback cancel path may be skipped even though the sender is neither closed nor known-closed. Treat a throwing tryClose as "not closed" so the cancel fallback still fires to guarantee close() cannot hang. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [388-395] boolean allClosed = true; for (DatafusionPartitionSender sender : sendersByChildStageId.values()) { try { allClosed &= sender.tryClose(); } catch (Exception e) { + allClosed = false; logger.warn("[reduce-sink] error closing input sender for EOF: taskId={}", ctx.taskId(), e); } } Suggestion importance[1-10]: 7 __ Why: Correctly identifies that an exception from tryClose() leaves allClosed unchanged, potentially skipping the cancel fallback and risking a hang — a legitimate defensive improvement.	Medium
	Fix racy tracker selection by context_id Using trackers.last() to attach the abort handle is racy: if a sibling stream registers under the same context_id between the caller's own registration and this call, last() returns the sibling's tracker and the caller's abort handle is attached to the wrong tracker (and its own OnceLock remains unset). Callers should pass their own QueryTrackingContext (or context_id + tracker identity) so the handle is attached to the correct tracker. The same concern applies to set_cpu_runtime_handle. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [509-514] +pub fn set_abort_handle(context_id: i64, handle: AbortHandle) { + if let Some(trackers) = QUERY_REGISTRY.get(&context_id) { + if let Some(tracker) = trackers.last() { + tracker.abort_handle.set(handle).ok(); + } + } +} - Suggestion importance[1-10]: 6 __ Why: Valid concern about a potential race where last() may not point to the caller's tracker if a sibling registers concurrently, but improved_code is identical to existing_code, weakening the actionability.	Low
General	Simplify race-safe registry entry removal Between drop(entry) and remove_if, another thread can register a new sibling tracker under the same context_id, making the Vec non-empty. The remove_if predicate correctly guards against this, but only if remove_if re-checks under the shard lock (which it does). This is fine, but be aware that if a new sibling registers after your is_empty() check and before remove_if, its tracker will correctly be preserved — verify that any new registration path uses entry().or_default() (it does) rather than assuming absence. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [686-694] if let Some(mut entry) = QUERY_REGISTRY.get_mut(&tracker.context_id) { entry.retain(|t| !Arc::ptr_eq(t, tracker)); - let now_empty = entry.is_empty(); drop(entry); - if now_empty { - QUERY_REGISTRY - .remove_if(&tracker.context_id, |_, v| v.is_empty()); - } + QUERY_REGISTRY.remove_if(&tracker.context_id, |_, v| v.is_empty()); } Suggestion importance[1-10]: 3 __ Why: The simplification is minor and the original code is already correct. The suggestion mostly asks the author to verify behavior rather than fix a real bug.	Low
	Ensure shard guard is released before cancel The QUERY_REGISTRY.get(&context_id) guard is held across the entry.iter()...collect() call, keeping the DashMap shard read-locked while iterating. The comment above states the intent is to snapshot under the entry lock then release before cancelling — make that explicit by scoping the guard into a block so it is definitively dropped before the cancellation loop runs. sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs [402-405] -let trackers: Vec<Arc<QueryTracker>> = match QUERY_REGISTRY.get(&context_id) { - Some(entry) => entry.iter().map(Arc::clone).collect(), - None => Vec::new(), +let trackers: Vec<Arc<QueryTracker>> = { + match QUERY_REGISTRY.get(&context_id) { + Some(entry) => entry.iter().map(Arc::clone).collect(), + None => Vec::new(), + } }; Suggestion importance[1-10]: 2 __ Why: The suggested change is essentially cosmetic; the original code already drops the guard at the end of the statement before the cancellation loop runs. Wrapping in a block adds negligible clarity.	Low

[github-actions]

✅ Gradle check result for 489f6cc: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.53%. Comparing base (8cac2f5) to head (1e55371).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22609      +/-   ##
============================================
+ Coverage     71.50%   71.53%   +0.02%     
+ Complexity    77037    77002      -35     
============================================
  Files          6156     6156              
  Lines        358417   358417              
  Branches      52243    52243              
============================================
+ Hits         256285   256380      +95     
+ Misses        81808    81643     -165     
- Partials      20324    20394      +70     

☔ View full report in Codecov by Harness.
📢 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.

[mch2]

thanks @alchemist51 lgtm

[github-actions]

Persistent review updated to latest commit 10d1ad4

[github-actions]

Persistent review updated to latest commit b81c88f

[github-actions]

Persistent review updated to latest commit 2959599

[github-actions]

✅ Gradle check result for 2959599: SUCCESS

[github-actions]

Persistent review updated to latest commit 8d86838

[github-actions]

✅ Gradle check result for 8d86838: SUCCESS

[github-actions]

Persistent review updated to latest commit 32cc99b

[github-actions]

Persistent review updated to latest commit 1e55371

[github-actions]

✅ Gradle check result for 1e55371: SUCCESS