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