Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alchemist51 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1361,7 +1361,12 @@ pub unsafe fn stream_get_schema(stream_ptr: i64) -> Result<i64, DataFusionError>
/// on the same stream.
pub async unsafe fn stream_next(stream_ptr: i64) -> Result<i64, DataFusionError> {
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 {
Expand Down
188 changes: 161 additions & 27 deletions sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,19 @@ impl QueryTracker {
// Global registry
// ---------------------------------------------------------------------------

static QUERY_REGISTRY: Lazy<DashMap<i64, Arc<QueryTracker>>> = 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<i64, Arc<QueryTracker>>` 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<DashMap<i64, Vec<Arc<QueryTracker>>>> = Lazy::new(DashMap::new);

// ---------------------------------------------------------------------------
// Registry snapshot — top-N FFM export
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -342,6 +354,7 @@ pub fn snapshot_top_n_by_current(out: &mut [WireQueryMetric]) -> usize {
}));
}
}
}
}

let mut written = 0usize;
Expand Down Expand Up @@ -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<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();
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)
Expand All @@ -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);
Expand Down Expand Up @@ -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<tokio::runtime::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()))
}

/// 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<CancellationToken> {
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();
}
}
}

Expand All @@ -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,
}
}
}
}
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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<CancellationToken> {
self.tracker.as_ref().map(|t| t.cancellation_token.clone())
}
}

impl Drop for QueryTrackingContext {
Expand All @@ -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);
Expand Down Expand Up @@ -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<i64, Arc<_>>`
/// 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);
Expand Down Expand Up @@ -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
Expand All @@ -901,6 +1031,8 @@ mod tests {
let first = QUERY_REGISTRY
.get(&ctx_id)
.unwrap()
.last()
.unwrap()
.cancelled_at_nanos
.load(Ordering::Relaxed);

Expand All @@ -909,6 +1041,8 @@ mod tests {
let second = QUERY_REGISTRY
.get(&ctx_id)
.unwrap()
.last()
.unwrap()
.cancelled_at_nanos
.load(Ordering::Relaxed);

Expand All @@ -924,15 +1058,15 @@ 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);

// Cancel it
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);

Expand All @@ -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);
Expand Down
23 changes: 23 additions & 0 deletions ...kend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Loading
Loading