From 489f6cc3a3fcc0d07950c992409c1f65410557b1 Mon Sep 17 00:00:00 2001 From: Arpit Bandejiya Date: Thu, 30 Jul 2026 11:45:05 +0530 Subject: [PATCH] Fix constant ~5s stall in coordinator reduce teardown on limited queries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every PPL query whose plan includes a late-materialization fetch phase (any `sort | head K | fields ` shape) paid a fixed ~5s in the LATE_MATERIALIZATION stage regardless of rows, columns, or bytes scanned, with a "timed out waiting for reduce teardown" WARN per query. Root cause is a teardown cycle in DatafusionReduceSink.closeImpl: close() waited on reduceDone, the drain was parked in stream_next waiting for a pipeline-breaking SortExec/TopK to emit, the TopK was waiting for input EOF, and the input senders were only closed after close() returned. The 5s await timeout was the only thing breaking the cycle. Cancellation could not break it either, for two Rust-side reasons: - QUERY_REGISTRY was keyed one-tracker-per-context_id, but an LM query opens TWO reduce sinks (built at graph-build time) that both register under the same ctx.taskId(). The second insert overwrote the first tracker and the first Drop removed the shared key, so cancel_query found no entry (silent no-op) while the sibling stream still ran. - stream_next re-fetched its token from the registry on every call; a stale registry returned None, degrading cancellable_or to a bare uncancellable await. Fixes: - Rust: QUERY_REGISTRY holds Vec> per id. Register appends; Drop removes only its own tracker (ptr_eq) and drops the key when empty; cancel_query cancels every tracker under the id and logs registry misses instead of silently no-oping. - Rust: stream_next uses its handle's own cancellation token via new QueryTrackingContext::cancellation_token(), immune to registry churn. - Java: closeImpl's REDUCING branch signals end-of-input first — tryClose() on each input sender (dropping a sender closes its mpsc, which the native plan reads as EOF, letting the TopK emit and the drain finish naturally with rows preserved). tryClose, not close(): a feeder parked in send() holds the sender read lock, and a blocking close would deadlock (covered by the existing testCloseWhileFeederParkedOnFullChannelDoesNotDeadlock). If tryClose fails the producer is live, so fall back to cancel; if EOF does not finish the drain within 5s, WARN + cancel so close() can never hang. Cancel-first teardown was not just slow but wrong once cancellation worked: it aborted the TopK before EOF and projections mixing sort keys with fetched columns returned 0 rows. Measured on a 749M-doc 5-shard composite index: 5.15s -> ~0.19s per query (~27x), all shapes row-correct, zero teardown WARNs. Tests: - test_sibling_stream_drop_does_not_orphan_survivor (Rust) - test_cancel_query_cancels_all_siblings (Rust) - testCloseWhileReducingSignalsEofAndPreservesRows (Java) — fails on the old code both ways: 5s close stall without the fix, 0 rows with cancel-first teardown. --- .../rust/src/api.rs | 7 +- .../rust/src/query_tracker.rs | 188 +++++++++++++++--- .../datafusion/DatafusionPartitionSender.java | 23 +++ .../be/datafusion/DatafusionReduceSink.java | 53 ++++- .../datafusion/DatafusionReduceSinkTests.java | 67 +++++++ 5 files changed, 306 insertions(+), 32 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index c884828e71791..3fac9e23e33f7 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs @@ -1361,7 +1361,12 @@ pub unsafe fn stream_get_schema(stream_ptr: i64) -> Result /// on the same stream. pub async unsafe fn stream_next(stream_ptr: i64) -> Result { let handle = &mut *(stream_ptr as *mut QueryStreamHandle); - let token = query_tracker::get_cancellation_token(handle._query_tracking_context.context_id()); + // 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) let result = cancellation::cancellable_or(token.as_ref(), None, async { diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs index 672ac1dcdfad5..59874558f81d5 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs @@ -211,7 +211,19 @@ impl QueryTracker { // Global registry // --------------------------------------------------------------------------- -static QUERY_REGISTRY: Lazy>> = Lazy::new(DashMap::new); +/// Registry of live query trackers, keyed by context_id. +/// +/// The value is a Vec because ONE context_id can have MULTIPLE live native +/// streams: a query with a late-materialization stage opens two coordinator +/// reduce sinks (stage 1 and stage 3), and both call +/// `executeLocalPlan(..., ctx.taskId())` with the same query-level task id. +/// The previous `DashMap>` silently overwrote the +/// first tracker on the second insert, and `Drop` removed the key +/// unconditionally — so the first stream to close deleted the entry for its +/// still-running sibling. `cancel_query` then found nothing, an in-flight +/// `stream_next` could not be interrupted, and the reduce sink's +/// `reduceDone.await(5s)` latch expired on every LM query (~5s stall). +static QUERY_REGISTRY: Lazy>>> = Lazy::new(DashMap::new); // --------------------------------------------------------------------------- // Registry snapshot — top-N FFM export @@ -314,7 +326,7 @@ pub fn snapshot_top_n_by_current(out: &mut [WireQueryMetric]) -> usize { let mut sample_logged = 0usize; for entry in QUERY_REGISTRY.iter() { - let tracker = entry.value(); + for tracker in entry.value() { let bytes_raw = tracker.memory_pool.current_bytes(); let completed = tracker.is_completed(); if sample_logged < 5 { @@ -342,6 +354,7 @@ pub fn snapshot_top_n_by_current(out: &mut [WireQueryMetric]) -> usize { })); } } + } } let mut written = 0usize; @@ -382,12 +395,38 @@ const FLUSH_WORKER_COUNT: usize = 4; /// [`flush_cpu_runtime`] after `stream_close` to ensure deferred drops are /// processed and GroupValues buffers are freed. pub fn cancel_query(context_id: i64) { - if let Some(tracker) = QUERY_REGISTRY.get(&context_id) { + // 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> = 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 = QUERY_REGISTRY.iter().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:?}" + ); + return; + } + native_bridge_common::log_debug!( + "cancel_query: cancelling context_id={context_id} ({} tracker(s))", + trackers.len() + ); + // 0 is the "not cancelled" sentinel in cancelled_at_nanos. The FIRST deref + // of PROCESS_START in a process creates the Instant and measures elapsed + // on it immediately, which can legitimately be 0ns — clamp to 1 so a real + // cancellation is never mistaken for "not cancelled". + let nanos = (PROCESS_START.elapsed().as_nanos() as u64).max(1); + for tracker in trackers { tracker.cancellation_token.cancel(); if let Some(handle) = tracker.abort_handle.get() { handle.abort(); } - let nanos = PROCESS_START.elapsed().as_nanos() as u64; tracker .cancelled_at_nanos .compare_exchange(0, nanos, Ordering::Release, Ordering::Relaxed) @@ -407,7 +446,7 @@ pub fn cancel_query(context_id: i64) { pub fn flush_cpu_runtime(context_id: i64) { let rt_handle = QUERY_REGISTRY .get(&context_id) - .and_then(|tracker| tracker.cpu_runtime_handle.get().cloned()); + .and_then(|trackers| trackers.last().and_then(|t| t.cpu_runtime_handle.get().cloned())); if let Some(handle) = rt_handle { flush_cpu_runtime_with_handle(&handle, context_id); @@ -447,28 +486,41 @@ pub fn flush_cpu_runtime_with_handle(handle: &tokio::runtime::Handle, context_id pub fn take_cpu_runtime_handle(context_id: i64) -> Option { QUERY_REGISTRY .get(&context_id) - .and_then(|tracker| tracker.cpu_runtime_handle.get().cloned()) + .and_then(|trackers| trackers.last().and_then(|t| t.cpu_runtime_handle.get().cloned())) } -/// Clone the cancellation token for the given context_id, if registered. +/// Clone a cancellation token for the given context_id, if registered. +/// +/// With multiple streams under one id this returns the MOST RECENTLY +/// registered tracker's token — callers that own a `QueryTrackingContext` +/// should prefer its own token (see `QueryTrackingContext::cancellation_token`) +/// over this registry lookup, which can go stale mid-stream. pub fn get_cancellation_token(context_id: i64) -> Option { QUERY_REGISTRY .get(&context_id) - .map(|t| t.cancellation_token.clone()) + .and_then(|trackers| trackers.last().map(|t| t.cancellation_token.clone())) } /// Store the CPU task's AbortHandle for the given context_id. +/// +/// Targets the most recently registered tracker for the id — matching the +/// call pattern where a stream registers its context then immediately +/// attaches its handles. pub fn set_abort_handle(context_id: i64, handle: AbortHandle) { - if let Some(tracker) = QUERY_REGISTRY.get(&context_id) { - tracker.abort_handle.set(handle).ok(); + 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(tracker) = QUERY_REGISTRY.get(&context_id) { - tracker.cpu_runtime_handle.set(handle).ok(); + if let Some(trackers) = QUERY_REGISTRY.get(&context_id) { + if let Some(tracker) = trackers.last() { + tracker.cpu_runtime_handle.set(handle).ok(); + } } } @@ -482,12 +534,13 @@ pub fn count_cancelled_running(threshold: Duration) -> (i64, i64) { let threshold_nanos = threshold.as_nanos() as u64; let now_nanos = PROCESS_START.elapsed().as_nanos() as u64; for entry in QUERY_REGISTRY.iter() { - let tracker = entry.value(); - let cancelled_nanos = tracker.cancelled_at_nanos.load(Ordering::Acquire); - if cancelled_nanos > 0 && (now_nanos - cancelled_nanos) >= threshold_nanos { - match tracker.query_type { - QueryType::Shard => shard_count += 1, - QueryType::Coordinator => coordinator_count += 1, + for tracker in entry.value() { + let cancelled_nanos = tracker.cancelled_at_nanos.load(Ordering::Acquire); + if cancelled_nanos > 0 && (now_nanos - cancelled_nanos) >= threshold_nanos { + match tracker.query_type { + QueryType::Shard => shard_count += 1, + QueryType::Coordinator => coordinator_count += 1, + } } } } @@ -535,7 +588,13 @@ impl QueryTrackingContext { completed: AtomicBool::new(false), wall_nanos: std::sync::atomic::AtomicU64::new(0), }); - QUERY_REGISTRY.insert(context_id, Arc::clone(&tracker)); + // Append, don't insert: a second stream registering under the same + // context_id (LM queries have two reduce sinks) must not orphan the + // first stream's tracker. + QUERY_REGISTRY + .entry(context_id) + .or_default() + .push(Arc::clone(&tracker)); Self { tracker: Some(tracker), phantom_reservation: None, @@ -592,6 +651,16 @@ impl QueryTrackingContext { pub fn context_id(&self) -> i64 { self.tracker.as_ref().map_or(0, |t| t.context_id) } + + /// This context's OWN cancellation token, or `None` if tracking is + /// disabled. Prefer this over `get_cancellation_token(context_id)` when a + /// `QueryTrackingContext` is at hand: the registry lookup targets the most + /// recently registered tracker for the id (which may be a sibling stream) + /// and returns `None` once the entry is gone — silently making the + /// caller's await uncancellable. + pub fn cancellation_token(&self) -> Option { + self.tracker.as_ref().map(|t| t.cancellation_token.clone()) + } } impl Drop for QueryTrackingContext { @@ -610,7 +679,19 @@ impl Drop for QueryTrackingContext { // query is never simultaneously visible in both the registry scan // (current) and the total counter — preventing double-counting in // snapshot_cancellation_stats. - QUERY_REGISTRY.remove(&tracker.context_id); + // + // Remove ONLY this context's own tracker (ptr_eq), never the whole + // key blindly: sibling streams under the same context_id may still + // be live. Drop the key when the Vec empties. + 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()); + } + } // If this query was cancelled and ran past the threshold, bump the total counter. let cancelled_nanos = tracker.cancelled_at_nanos.load(Ordering::Acquire); @@ -741,15 +822,64 @@ mod tests { assert!(!QUERY_REGISTRY.contains_key(&ctx_id)); } + /// Regression test for the ~5s late-materialization stall: an LM query + /// registers TWO reduce streams under the SAME context_id (stage-1 and + /// stage-3 sinks both pass ctx.taskId()). The old `DashMap>` + /// registry overwrote the first tracker at insert and removed the whole + /// key at first Drop — so `cancel_query` on the shared id found nothing + /// and the surviving stream could never be interrupted. + #[test] + fn test_sibling_stream_drop_does_not_orphan_survivor() { + let global = make_global_pool(10_000); + let ctx_id = 50_042; + let first = QueryTrackingContext::new(ctx_id, Arc::clone(&global), QueryType::Coordinator); + let second = QueryTrackingContext::new(ctx_id, global, QueryType::Coordinator); + + // Both trackers visible under the shared id. + assert_eq!(QUERY_REGISTRY.get(&ctx_id).unwrap().len(), 2); + + // First stream closes (stage-1 drains to EOF and drops) … + drop(first); + + // … the survivor MUST still be registered and cancellable. + assert_eq!(QUERY_REGISTRY.get(&ctx_id).unwrap().len(), 1); + let survivor_token = second.cancellation_token().unwrap(); + assert!(!survivor_token.is_cancelled()); + cancel_query(ctx_id); + assert!(survivor_token.is_cancelled()); + + drop(second); + assert!(!QUERY_REGISTRY.contains_key(&ctx_id)); + } + + /// cancel_query on a shared id must cancel EVERY live tracker, not just one. + #[test] + fn test_cancel_query_cancels_all_siblings() { + let global = make_global_pool(10_000); + let ctx_id = 50_043; + let a = QueryTrackingContext::new(ctx_id, Arc::clone(&global), QueryType::Coordinator); + let b = QueryTrackingContext::new(ctx_id, global, QueryType::Coordinator); + + let tok_a = a.cancellation_token().unwrap(); + let tok_b = b.cancellation_token().unwrap(); + cancel_query(ctx_id); + assert!(tok_a.is_cancelled()); + assert!(tok_b.is_cancelled()); + + drop(a); + drop(b); + assert!(!QUERY_REGISTRY.contains_key(&ctx_id)); + } + #[test] fn test_wall_secs_ticks_while_running() { let global = make_global_pool(10_000); let ctx_id = 50_002; let _ctx = QueryTrackingContext::new(ctx_id, global, QueryType::Shard); - let t1 = QUERY_REGISTRY.get(&ctx_id).unwrap().wall_secs(); + let t1 = QUERY_REGISTRY.get(&ctx_id).unwrap().last().unwrap().wall_secs(); thread::sleep(Duration::from_millis(50)); - let t2 = QUERY_REGISTRY.get(&ctx_id).unwrap().wall_secs(); + let t2 = QUERY_REGISTRY.get(&ctx_id).unwrap().last().unwrap().wall_secs(); assert!(t2 - t1 >= 0.04); drop(_ctx); @@ -879,7 +1009,7 @@ mod tests { let ctx = QueryTrackingContext::new(ctx_id, global, QueryType::Shard); // Not cancelled yet - let tracker = QUERY_REGISTRY.get(&ctx_id).unwrap(); + let tracker = Arc::clone(QUERY_REGISTRY.get(&ctx_id).unwrap().last().unwrap()); assert!(tracker.cancelled_at_nanos.load(Ordering::Relaxed) == 0); // Cancel @@ -901,6 +1031,8 @@ mod tests { let first = QUERY_REGISTRY .get(&ctx_id) .unwrap() + .last() + .unwrap() .cancelled_at_nanos .load(Ordering::Relaxed); @@ -909,6 +1041,8 @@ mod tests { let second = QUERY_REGISTRY .get(&ctx_id) .unwrap() + .last() + .unwrap() .cancelled_at_nanos .load(Ordering::Relaxed); @@ -924,7 +1058,7 @@ mod tests { let ctx = QueryTrackingContext::new(ctx_id, global, QueryType::Shard); // Not cancelled — cancelled_at_nanos should be 0 - let tracker = QUERY_REGISTRY.get(&ctx_id).unwrap(); + let tracker = Arc::clone(QUERY_REGISTRY.get(&ctx_id).unwrap().last().unwrap()); assert_eq!(tracker.cancelled_at_nanos.load(Ordering::Relaxed), 0); drop(tracker); @@ -932,7 +1066,7 @@ mod tests { cancel_query(ctx_id); // Now cancelled_at_nanos should be > 0 - let tracker = QUERY_REGISTRY.get(&ctx_id).unwrap(); + let tracker = Arc::clone(QUERY_REGISTRY.get(&ctx_id).unwrap().last().unwrap()); assert!(tracker.cancelled_at_nanos.load(Ordering::Relaxed) > 0); drop(tracker); @@ -956,12 +1090,12 @@ mod tests { cancel_query(coord_id); // Verify each query is registered, cancelled, and has correct type - let shard_tracker = QUERY_REGISTRY.get(&shard_id).unwrap(); + let shard_tracker = Arc::clone(QUERY_REGISTRY.get(&shard_id).unwrap().last().unwrap()); assert!(shard_tracker.cancelled_at_nanos.load(Ordering::Relaxed) > 0); assert_eq!(shard_tracker.query_type, QueryType::Shard); drop(shard_tracker); - let coord_tracker = QUERY_REGISTRY.get(&coord_id).unwrap(); + let coord_tracker = Arc::clone(QUERY_REGISTRY.get(&coord_id).unwrap().last().unwrap()); assert!(coord_tracker.cancelled_at_nanos.load(Ordering::Relaxed) > 0); assert_eq!(coord_tracker.query_type, QueryType::Coordinator); drop(coord_tracker); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java index 06d60a01a0ef5..435f68dc3f09b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java @@ -76,6 +76,29 @@ public void close() { } } + /** + * Non-blocking close attempt, for signalling end-of-input from a thread that must + * not park. Returns {@code true} if the sender was closed (or already closed); + * {@code false} if the write lock could not be acquired immediately — i.e. a feeder + * is currently parked inside {@link #send} holding the read lock, and blocking here + * would deadlock (see DatafusionReduceSinkTests + * #testCloseWhileFeederParkedOnFullChannelDoesNotDeadlock). Callers that get + * {@code false} must unblock the producer by other means (e.g. cancel, which drops + * the native receiver and pops the parked send with {@code ReceiverDropped}). + */ + public boolean tryClose() { + if (!lifecycle.writeLock().tryLock()) { + return false; + } + try { + super.close(); + logger.debug("[sender] closed ptr={} (tryClose)", ptr); + return true; + } finally { + lifecycle.writeLock().unlock(); + } + } + @Override protected void doClose() { NativeBridge.senderClose(ptr); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java index 6ece9d4f20a7f..a73c18b9d81cb 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java @@ -363,12 +363,57 @@ protected void feedBatchUnderLock(VectorSchemaRoot batch) { protected Exception closeImpl() { SinkState before = state.compareAndExchange(SinkState.READY, SinkState.DONE); if (before == SinkState.REDUCING) { - // Drain in flight — fire cancel so it unblocks, then wait for reduce's - // finally to complete teardown (releases Arrow batches from the allocator). - fireCancelQuery(); + // Drain in flight. Close the input senders FIRST: dropping a sender closes + // its mpsc channel, which the native plan's streaming input reads as + // end-of-input. Pipeline-breaking operators above it (SortExec/TopK) cannot + // emit until they see EOF, so without this the drain blocks in streamNext + // waiting for output that structurally cannot exist yet, while we block + // here waiting for the drain — a cycle only the await timeout used to + // break (observed as a constant ~5s stall on every LM query, with the + // WARN below firing each time). + // + // EOF (not cancelQuery) is the correct unblocking mechanism on this path: + // close() here means "no more input is coming", not "abort". Cancelling + // instead aborts the TopK before it emits and the drain returns the + // cancellation sentinel with ZERO rows (verified: projections mixing sort + // keys with fetched columns returned 0 rows when this branch cancelled). + // fireCancelQuery remains the mechanism for genuine aborts via cancel(). + // + // MUST be tryClose, not close(): a feeder thread parked inside send() + // (bounded channel full) holds the sender READ lock; a blocking close() + // here would wait forever on the WRITE lock — the exact deadlock covered + // by testCloseWhileFeederParkedOnFullChannelDoesNotDeadlock. When + // tryClose fails, cancel instead: cancel drops the native receiver, + // which pops the parked send with ReceiverDropped, releasing the lock. + boolean allClosed = true; + for (DatafusionPartitionSender sender : sendersByChildStageId.values()) { + try { + allClosed &= sender.tryClose(); + } catch (Exception e) { + logger.warn("[reduce-sink] error closing input sender for EOF: taskId={}", ctx.taskId(), e); + } + } + if (allClosed == false) { + // A producer is parked mid-send: rows are still streaming in, so this + // close is racing a live producer (not the clean drained-early case). + // Cancel is the only lock-free way to unblock it; dropped rows are + // acceptable here because the consumer already stopped listening. + logger.debug("[reduce-sink] input sender busy at close; cancelling to unblock: taskId={}", ctx.taskId()); + fireCancelQuery(); + } try { if (!reduceDone.await(5, java.util.concurrent.TimeUnit.SECONDS)) { - logger.warn("[reduce-sink] timed out waiting for reduce teardown: taskId={}", ctx.taskId()); + // EOF didn't unblock the drain (e.g. the native plan is stuck for + // another reason). Fall back to a hard cancel so close() cannot + // hang, then give teardown a short grace period. + logger.warn( + "[reduce-sink] reduce did not finish after input EOF; falling back to cancel: taskId={}", + ctx.taskId() + ); + fireCancelQuery(); + if (!reduceDone.await(5, java.util.concurrent.TimeUnit.SECONDS)) { + logger.warn("[reduce-sink] timed out waiting for reduce teardown: taskId={}", ctx.taskId()); + } } } catch (InterruptedException e) { Thread.currentThread().interrupt(); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java index b2c8800558d20..67e93b7b09861 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java @@ -150,6 +150,73 @@ public void testFeedDrainsSumToDownstream() throws Exception { } } + /** + * Regression test for the late-materialization ~5s stall: {@code close()} arriving while + * {@code reduce()} is mid-drain (state=REDUCING) must (a) return promptly and (b) NOT lose + * rows the native plan has yet to emit. + * + *

The SUM plan is pipeline-breaking: the aggregate emits nothing until its input reaches + * end-of-input. Closing the sink from outside while the drain is parked used to fire + * {@code cancel_query} and wait on {@code reduceDone} for a hard-coded 5s — but the input + * senders were only closed AFTER that wait, so the aggregate could not emit, the drain could + * not finish, and every such close burned the full timeout (observed as a constant ~5s on + * every LM query, one WARN per query). Worse, once cancellation actually worked, the cancel + * aborted the aggregate BEFORE EOF and the result was 0 rows. + * + *

The fix closes the input senders first (tryClose → EOF → aggregate emits → drain ends + * naturally). This test feeds rows, calls {@code close()} WITHOUT signalling per-child EOF + * (simulating the Stitcher's early close), and asserts the SUM still arrives and close() + * returns in far less than the old 5s timeout. + */ + public void testCloseWhileReducingSignalsEofAndPreservesRows() throws Exception { + NativeBridge.initTokioRuntimeManager(2); + Path spillDir = createTempDir("datafusion-spill"); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); + + try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { + Schema inputSchema = new Schema(List.of(new Field("x", FieldType.nullable(new ArrowType.Int(64, true)), null))); + byte[] substrait = buildSumSubstraitBytes(DatafusionReduceSink.INPUT_ID); + + CapturingSink downstream = new CapturingSink(); + // Non-zero taskId so the cancel fallback path is wired if the test regresses. + ExchangeSinkContext ctx = new ExchangeSinkContext( + "q-close-eof", + 0, + 4242L, + substrait, + alloc, + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstraitBytes(DatafusionReduceSink.INPUT_ID))), + downstream + ); + + DatafusionReduceSink sink = new DatafusionReduceSink(ctx, runtimeHandle); + PlainActionFuture drainDone = PlainActionFuture.newFuture(); + Thread.ofVirtual().start(() -> sink.reduce(drainDone)); + + sink.feed(makeBatch(alloc, inputSchema, new long[] { 1L, 2L, 3L })); + sink.feed(makeBatch(alloc, inputSchema, new long[] { 4L, 5L, 6L })); + // Give the drain time to actually park in stream_next waiting for more input, + // so close() observes state=REDUCING (the stalling branch). + Thread.sleep(200); + + // Close WITHOUT sinkForChild(0).close() — this is the Stitcher-style early close. + long closeStartNanos = System.nanoTime(); + sink.close(); + long closeMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - closeStartNanos); + + drainDone.actionGet(10, TimeUnit.SECONDS); + + assertEquals("EOF-on-close must let the aggregate emit: SUM(1..6)", 21L, downstream.total); + assertTrue("downstream should have received the aggregate row", downstream.totalRows >= 1); + // The old behavior burned the full 5s await; EOF-first should finish in millis. + // Generous bound to stay unflaky on slow CI hosts while still catching a 5s burn. + assertTrue("close() should not burn the teardown timeout, took " + closeMillis + "ms", closeMillis < 4000); + } finally { + runtimeHandle.close(); + } + } + /** * Coordinator reduce running {@code SELECT x FROM "input-0" LIMIT 3} produces exactly the * limited output and tears down cleanly. Feeds far more rows than the limit through a real