diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java
index 59f9b1f899f92..42d5fbe386ed5 100644
--- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java
@@ -8,7 +8,9 @@
package org.opensearch.analytics.spi;
+import java.util.Collections;
import java.util.List;
+import java.util.Map;
/**
* SPI extension point for backend query engine plugins.
@@ -114,6 +116,24 @@ default void configureFilterDelegation(FilterDelegationHandle handle, BackendExe
throw new UnsupportedOperationException("configureFilterDelegation not implemented for [" + name() + "]");
}
+ /**
+ * Returns a snapshot of this backend's currently-tracked queries, keyed by {@code contextId}.
+ *
+ *
The map is a point-in-time view — entries can register or drain concurrently on the
+ * backend side. Implementations MUST return a non-null map (empty when nothing is tracked)
+ * and SHOULD make it unmodifiable so callers cannot mutate backend state.
+ *
+ *
Implementations MAY cap the result to a top-N subset by current memory usage to bound
+ * the FFI cost (the DataFusion backend caps at the heaviest 10 live queries). Callers that
+ * need a complete enumeration should not rely on this method.
+ *
+ *
Default implementation returns an empty map so backends that do not track per-query
+ * metrics don't have to opt in.
+ */
+ default Map getTopQueriesByMemory() {
+ return Collections.emptyMap();
+ }
+
/**
* Install a thread tracker for attribution of delegation callbacks executing on foreign threads.
* Called after {@link #configureFilterDelegation}. Pass {@code null} to clear.
diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/QueryExecutionMetrics.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/QueryExecutionMetrics.java
new file mode 100644
index 0000000000000..0e03287d3ead7
--- /dev/null
+++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/QueryExecutionMetrics.java
@@ -0,0 +1,21 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.analytics.spi;
+
+/**
+ * Snapshot of per-query execution metrics reported by an analytics backend.
+ *
+ *
The context id (the map key in {@link AnalyticsSearchBackendPlugin#getTopQueriesByMemory()})
+ * is not repeated here — this record holds only the memory accounting fields.
+ *
+ * @param currentBytes bytes currently reserved by the query's memory pool
+ *
+ * @opensearch.internal
+ */
+public record QueryExecutionMetrics(long currentBytes) {}
diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs
index 140f3e8bbc15a..b8fc8ce406b0e 100644
--- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs
+++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs
@@ -206,6 +206,47 @@ pub extern "C" fn df_cancel_query(context_id: i64) {
api::cancel_query(context_id);
}
+// ---------------------------------------------------------------------------
+// Per-query registry top-N snapshot
+//
+// One FFM call: Java allocates a buffer sized for `N` entries, Rust selects
+// the heaviest live queries by `current_bytes` (bounded min-heap of size N)
+// and writes them back-to-back. See `query_tracker::WireQueryMetric` for the
+// wire layout.
+// ---------------------------------------------------------------------------
+
+/// Copies up to `cap_entries` of the heaviest live queries (by
+/// `current_bytes` desc) as `WireQueryMetric`s into the caller-provided buffer.
+/// Returns the number of entries actually written.
+///
+/// Order of entries within the buffer is unspecified. Completed and zero-byte
+/// trackers are filtered out.
+///
+/// Safety: `out_ptr` must be non-null, 8-byte aligned, and point to storage
+/// for at least `cap_entries * size_of::()` bytes.
+#[ffm_safe]
+#[no_mangle]
+pub unsafe extern "C" fn df_query_registry_top_n_by_current(
+ out_ptr: *mut u8,
+ cap_entries: i64,
+) -> i64 {
+ use crate::query_tracker::{snapshot_top_n_by_current, WireQueryMetric};
+
+ if cap_entries < 0 {
+ return Err(format!("negative capacity: {cap_entries}"));
+ }
+ if cap_entries == 0 {
+ return Ok(0);
+ }
+ if out_ptr.is_null() {
+ return Err("null snapshot buffer".to_string());
+ }
+ let out: &mut [WireQueryMetric] =
+ slice::from_raw_parts_mut(out_ptr as *mut WireQueryMetric, cap_entries as usize);
+ let written = snapshot_top_n_by_current(out);
+ Ok(written as i64)
+}
+
#[ffm_safe]
#[no_mangle]
pub unsafe extern "C" fn df_sql_to_substrait(
diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs
index fb522332de9a4..066aa8214ac82 100644
--- a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs
+++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs
@@ -35,4 +35,4 @@ pub mod session_context;
pub mod statistics_cache;
pub mod udf;
pub mod stats;
-pub mod task_monitors;
\ No newline at end of file
+pub mod task_monitors;
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 90a59a64f0aa8..84d29376b6a1f 100644
--- a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs
+++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_tracker.rs
@@ -150,6 +150,138 @@ impl QueryTracker {
static QUERY_REGISTRY: Lazy>> = Lazy::new(DashMap::new);
+// ---------------------------------------------------------------------------
+// Registry snapshot — top-N FFM export
+// ---------------------------------------------------------------------------
+
+/// Wire representation of a single query's tracker, used by the top-N
+/// snapshot API. Fields are `i64` so the struct crosses the FFM boundary with
+/// a stable, alignment-safe layout (8-byte aligned on every target we support).
+///
+/// | Field | Meaning |
+/// |---------------|-----------------------------------------------------------|
+/// | context_id | `QueryTracker::context_id` |
+/// | current_bytes | `QueryMemoryPool::current_bytes`, clamped to `i64::MAX` |
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct WireQueryMetric {
+ pub context_id: i64,
+ pub current_bytes: i64,
+}
+
+const _: () = assert!(std::mem::size_of::() == 2 * 8);
+
+fn usize_to_i64_saturating(value: usize) -> i64 {
+ if value > i64::MAX as usize {
+ i64::MAX
+ } else {
+ value as i64
+ }
+}
+
+/// Top-N snapshot: copy at most `out.len()` of the heaviest live queries
+/// (ranked by `current_bytes` descending) into `out`. Returns the number of
+/// entries actually written.
+///
+/// Selection is done with a bounded min-heap of capacity `out.len()`, so a
+/// single pass over the registry runs in `O(M log N)` where `M` is the
+/// registry size and `N = out.len()`. The heap holds the *top N candidates*
+/// seen so far; the smallest of those sits at the root, and any new entry
+/// heavier than that root replaces it.
+///
+/// Order of entries written into `out` is unspecified — heap order, not
+/// sorted. Callers that need a sorted view sort the prefix client-side.
+///
+/// Filters applied on the Rust side, in this order:
+/// - completed trackers are skipped (a completed query's `current_bytes`
+/// drops to zero on `Drop` anyway, so this is mostly a fast path)
+/// - `current_bytes == 0` is skipped (registered but un-allocated trackers
+/// are not "heavy" candidates)
+///
+/// `out` may be a raw slice pointing at a Java-owned buffer; it must be valid
+/// for writes of `out.len() * size_of::()` bytes and properly
+/// aligned for `WireQueryMetric` (8-byte aligned).
+pub fn snapshot_top_n_by_current(out: &mut [WireQueryMetric]) -> usize {
+ use std::cmp::{Ordering, Reverse};
+ use std::collections::BinaryHeap;
+
+ /// Heap entry — POD copy of the tracker fields sampled at iteration
+ /// time. We deliberately avoid carrying an `Arc` here:
+ /// re-reading the live tracker at materialization time would let the
+ /// reported `current_bytes` drift away from the value the entry was
+ /// ranked on (e.g. a query finishes between rank and report and ships
+ /// out as zero), and would also defeat the `current_bytes == 0` filter
+ /// applied during iteration. Sampling once keeps rank and report
+ /// consistent and removes the per-survivor `Arc::clone` from the hot
+ /// path.
+ #[derive(Clone, Copy)]
+ struct HeapEntry {
+ context_id: i64,
+ current_bytes: i64,
+ }
+
+ impl Eq for HeapEntry {}
+ impl PartialEq for HeapEntry {
+ fn eq(&self, other: &Self) -> bool {
+ self.current_bytes == other.current_bytes
+ }
+ }
+ impl Ord for HeapEntry {
+ fn cmp(&self, other: &Self) -> Ordering {
+ self.current_bytes.cmp(&other.current_bytes)
+ }
+ }
+ impl PartialOrd for HeapEntry {
+ fn partial_cmp(&self, other: &Self) -> Option {
+ Some(self.cmp(other))
+ }
+ }
+
+ let n = out.len();
+ if n == 0 {
+ return 0;
+ }
+
+ // BinaryHeap is a max-heap; Reverse flips it so the smallest current_bytes
+ // sits at the root. That root is the weakest member of the current top-N
+ // set — the one a heavier candidate displaces.
+ let mut heap: BinaryHeap> = BinaryHeap::with_capacity(n);
+
+ for entry in QUERY_REGISTRY.iter() {
+ let tracker = entry.value();
+ if tracker.is_completed() {
+ continue;
+ }
+ let current_raw = tracker.memory_pool.current_bytes();
+ let current_bytes = usize_to_i64_saturating(current_raw);
+ if current_bytes == 0 {
+ continue;
+ }
+ let candidate = HeapEntry {
+ context_id: tracker.context_id,
+ current_bytes,
+ };
+ if heap.len() < n {
+ heap.push(Reverse(candidate));
+ } else if let Some(Reverse(min)) = heap.peek() {
+ if candidate.current_bytes > min.current_bytes {
+ heap.pop();
+ heap.push(Reverse(candidate));
+ }
+ }
+ }
+
+ let mut written = 0usize;
+ for Reverse(entry) in heap.into_iter() {
+ out[written] = WireQueryMetric {
+ context_id: entry.context_id,
+ current_bytes: entry.current_bytes,
+ };
+ written += 1;
+ }
+ written
+}
+
/// Fire the cancellation token for the given context_id.
/// No-op for unknown or already-completed queries.
pub fn cancel_query(context_id: i64) {
@@ -220,8 +352,8 @@ impl Drop for QueryTrackingContext {
fn drop(&mut self) {
if let Some(tracker) = &self.tracker {
tracker.mark_completed();
- debug!(
- "Query completed ctx={}: wall={:.3}s, mem_current={}B, mem_peak={}B",
+ debug!(
+ "Query completed ctx={}: wall={:.3}s, mem_current={}B, mem_peak={}B",
tracker.context_id,
tracker.wall_secs(),
tracker.memory_pool.current_bytes(),
@@ -467,4 +599,178 @@ mod tests {
assert!(!QUERY_REGISTRY.contains_key(&ctx_id));
}
+
+ // -----------------------------------------------------------------------
+ // Top-N snapshot tests
+ // -----------------------------------------------------------------------
+
+ /// Helper: drain all live tracker contexts so each top-N test starts from
+ /// a clean registry. Tests in this module are tagged with unique id ranges
+ /// but the registry is process-wide, so we only assert on contexts we own.
+ fn fresh_buf(n: usize) -> Vec {
+ vec![
+ WireQueryMetric {
+ context_id: 0,
+ current_bytes: 0,
+ };
+ n
+ ]
+ }
+
+ #[test]
+ fn test_top_n_zero_capacity_returns_zero() {
+ let mut buf: Vec = Vec::new();
+ let written = snapshot_top_n_by_current(&mut buf);
+ assert_eq!(written, 0);
+ }
+
+ #[test]
+ fn test_top_n_picks_highest_current_bytes() {
+ let global = make_global_pool(10_000_000);
+ let id_lo = 70_000;
+ let id_md = 70_001;
+ let id_hi = 70_002;
+
+ let ctx_lo = QueryTrackingContext::new(id_lo, Arc::clone(&global));
+ let ctx_md = QueryTrackingContext::new(id_md, Arc::clone(&global));
+ let ctx_hi = QueryTrackingContext::new(id_hi, Arc::clone(&global));
+
+ let pool_lo: Arc = ctx_lo.memory_pool().unwrap();
+ let pool_md: Arc = ctx_md.memory_pool().unwrap();
+ let pool_hi: Arc = ctx_hi.memory_pool().unwrap();
+
+ let mut r_lo = make_reservation(&pool_lo, "lo");
+ let mut r_md = make_reservation(&pool_md, "md");
+ let mut r_hi = make_reservation(&pool_hi, "hi");
+ r_lo.try_grow(1_000).unwrap();
+ r_md.try_grow(5_000).unwrap();
+ r_hi.try_grow(9_000).unwrap();
+
+ // Top-2 must contain ids hi and md but not lo. Other parallel tests
+ // may push entries into the registry, so we filter to our id range.
+ let mut buf = fresh_buf(2);
+ let written = snapshot_top_n_by_current(&mut buf);
+ assert!(written >= 1 && written <= 2);
+
+ let our: Vec<&WireQueryMetric> = buf[..written]
+ .iter()
+ .filter(|m| (id_lo..=id_hi).contains(&m.context_id))
+ .collect();
+ // At minimum, id_hi (the heaviest) must be in our slice if any of
+ // ours made it in.
+ if !our.is_empty() {
+ assert!(our.iter().any(|m| m.context_id == id_hi));
+ assert!(!our.iter().any(|m| m.context_id == id_lo));
+ }
+
+ drop(r_lo);
+ drop(r_md);
+ drop(r_hi);
+ drop(ctx_lo);
+ drop(ctx_md);
+ drop(ctx_hi);
+ QUERY_REGISTRY.remove(&id_lo);
+ QUERY_REGISTRY.remove(&id_md);
+ QUERY_REGISTRY.remove(&id_hi);
+ }
+
+ #[test]
+ fn test_top_n_skips_completed_and_zero_byte_trackers() {
+ let global = make_global_pool(1_000_000);
+ let live_id = 70_100;
+ let zero_id = 70_101;
+ let done_id = 70_102;
+
+ let live_ctx = QueryTrackingContext::new(live_id, Arc::clone(&global));
+ let live_pool: Arc = live_ctx.memory_pool().unwrap();
+ let mut r_live = make_reservation(&live_pool, "live");
+ r_live.try_grow(4_096).unwrap();
+
+ // Registered but never reserved — current_bytes stays 0.
+ let _zero_ctx = QueryTrackingContext::new(zero_id, Arc::clone(&global));
+
+ // Completed before snapshot. Drop reservation first so QueryMemoryPool
+ // is settled, then drop the context to flip the completed flag.
+ let done_ctx = QueryTrackingContext::new(done_id, Arc::clone(&global));
+ let done_pool: Arc = done_ctx.memory_pool().unwrap();
+ let mut r_done = make_reservation(&done_pool, "done");
+ r_done.try_grow(8_192).unwrap();
+ drop(r_done);
+ drop(done_ctx);
+
+ let mut buf = fresh_buf(8);
+ let written = snapshot_top_n_by_current(&mut buf);
+
+ let our: Vec<&WireQueryMetric> = buf[..written]
+ .iter()
+ .filter(|m| [live_id, zero_id, done_id].contains(&m.context_id))
+ .collect();
+ assert!(our.iter().any(|m| m.context_id == live_id));
+ assert!(!our.iter().any(|m| m.context_id == zero_id));
+ assert!(!our.iter().any(|m| m.context_id == done_id));
+
+ drop(r_live);
+ drop(live_ctx);
+ QUERY_REGISTRY.remove(&live_id);
+ QUERY_REGISTRY.remove(&zero_id);
+ QUERY_REGISTRY.remove(&done_id);
+ }
+
+ #[test]
+ fn test_top_n_with_buffer_larger_than_live_set() {
+ let global = make_global_pool(1_000_000);
+ let id = 70_200;
+ let ctx = QueryTrackingContext::new(id, global);
+ let pool: Arc = ctx.memory_pool().unwrap();
+ let mut r = make_reservation(&pool, "only");
+ r.try_grow(2_048).unwrap();
+
+ let sentinel = WireQueryMetric {
+ context_id: -1,
+ current_bytes: -1,
+ };
+ let mut buf = vec![sentinel; 16];
+ let written = snapshot_top_n_by_current(&mut buf);
+ assert!(written >= 1);
+ assert!(written < buf.len());
+ // Tail past `written` keeps the sentinel — proves the snapshot did not
+ // touch it.
+ for slot in &buf[written..] {
+ assert_eq!(slot.context_id, -1);
+ }
+
+ drop(r);
+ drop(ctx);
+ QUERY_REGISTRY.remove(&id);
+ }
+
+ #[test]
+ fn test_top_n_caps_at_buffer_capacity() {
+ let global = make_global_pool(10_000_000);
+ let ids: Vec = (70_300..70_310).collect();
+ let mut contexts = Vec::with_capacity(ids.len());
+ let mut reservations = Vec::with_capacity(ids.len());
+ for (i, id) in ids.iter().enumerate() {
+ let ctx = QueryTrackingContext::new(*id, Arc::clone(&global));
+ let pool: Arc = ctx.memory_pool().unwrap();
+ let mut r = make_reservation(&pool, "cap");
+ r.try_grow((i as usize + 1) * 1_000).unwrap();
+ reservations.push(r);
+ contexts.push(ctx);
+ }
+
+ let mut buf = fresh_buf(3);
+ let written = snapshot_top_n_by_current(&mut buf);
+ assert!(written <= 3, "must not exceed buffer capacity");
+
+ for r in reservations.drain(..) {
+ drop(r);
+ }
+ for ctx in contexts.drain(..) {
+ drop(ctx);
+ }
+ for id in &ids {
+ QUERY_REGISTRY.remove(id);
+ }
+ }
}
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java
index d677195e8c2aa..7112c8246dac8 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java
@@ -30,6 +30,7 @@
import org.opensearch.analytics.spi.JoinCapability;
import org.opensearch.analytics.spi.NumericToDoubleAdapter;
import org.opensearch.analytics.spi.ProjectCapability;
+import org.opensearch.analytics.spi.QueryExecutionMetrics;
import org.opensearch.analytics.spi.ScalarFunction;
import org.opensearch.analytics.spi.ScalarFunctionAdapter;
import org.opensearch.analytics.spi.ScanCapability;
@@ -704,6 +705,14 @@ public void configureFilterDelegation(FilterDelegationHandle handle, BackendExec
FilterTreeCallbacks.setHandle(handle);
}
+ @Override
+ public Map getTopQueriesByMemory() {
+ // Delegate to the plugin that owns the DataFusionService and native runtime.
+ // Keeping the real implementation on DataFusionPlugin lets operators call it
+ // directly (e.g., from a REST action) without going through the SPI.
+ return plugin.getTopQueriesByMemory();
+ }
+
@Override
public void setDelegationThreadTracker(org.opensearch.analytics.spi.DelegationThreadTracker tracker) {
FilterTreeCallbacks.setThreadTracker(tracker);
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java
index bad7879a96322..cf06ec8aff2b9 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java
@@ -11,7 +11,9 @@
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin;
+import org.opensearch.analytics.spi.QueryExecutionMetrics;
import org.opensearch.be.datafusion.action.DataFusionStatsAction;
+import org.opensearch.be.datafusion.nativelib.NativeBridge;
import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
import org.opensearch.cluster.node.DiscoveryNodes;
import org.opensearch.cluster.service.ClusterService;
@@ -35,6 +37,7 @@
import org.opensearch.rest.RestController;
import org.opensearch.rest.RestHandler;
import org.opensearch.script.ScriptService;
+import org.opensearch.search.backpressure.trackers.NativeMemoryUsageTracker;
import org.opensearch.threadpool.ThreadPool;
import org.opensearch.transport.client.Client;
import org.opensearch.watcher.ResourceWatcherService;
@@ -42,7 +45,9 @@
import java.io.IOException;
import java.util.Collection;
import java.util.Collections;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
import java.util.function.Supplier;
import io.substrait.extension.DefaultExtensionCatalog;
@@ -105,6 +110,14 @@ public class DataFusionPlugin extends Plugin implements SearchBackEndPlugin createComponents(
this.datafusionSettings = new DatafusionSettings(clusterService);
+ // Expose per-task native-memory usage to search backpressure. The tracker calls
+ // this supplier once per refresh (invoked by the backpressure service at the top of
+ // doRun() and nodeStats()), snapshotting all live queries in one FFM call. Per-task
+ // evaluation then reads from the tracker's cached map — no FFM call per task.
+ //
+ // The OpenSearch task id is used as the DataFusion context_id at query launch
+ // (see ShardScanInstructionHandler / DatafusionSearchExecEngine), so the map is
+ // already keyed by Task#getId on the consumer side.
+ NativeMemoryUsageTracker.setSnapshotSupplier(this::currentBytesByTaskId);
+ NativeMemoryUsageTracker.setNativeMemoryBudgetSupplier(() -> DATAFUSION_MEMORY_POOL_LIMIT.get(clusterService.getSettings()));
+
this.substraitExtensions = loadSubstraitExtensions();
return Collections.singletonList(dataFusionService);
}
+ /**
+ * Project the active-query metrics map down to {@code taskId -> currentBytes} for the
+ * backpressure native-memory tracker. One FFM snapshot per call, capped at
+ * {@link #ACTIVE_QUERY_METRICS_TOP_N} entries (the heaviest live queries).
+ * Returns an empty map when the service isn't running, so startup/shutdown races
+ * don't surface bad data.
+ */
+ private Map currentBytesByTaskId() {
+ if (dataFusionService == null) {
+ return Collections.emptyMap();
+ }
+ Map metrics = getTopQueriesByMemory();
+ if (metrics.isEmpty()) {
+ return Collections.emptyMap();
+ }
+ Map out = new HashMap<>(metrics.size());
+ for (Map.Entry e : metrics.entrySet()) {
+ out.put(e.getKey(), e.getValue().currentBytes());
+ }
+ if (logger.isDebugEnabled()) {
+ logger.debug("native memory snapshot: {} active queries", out.size());
+ }
+ return out;
+ }
+
/**
* Loads the Substrait default extension catalog with the plugin's classloader as the
* thread context classloader. Jackson polymorphic deserialization (used by Substrait
@@ -284,4 +333,27 @@ public void close() throws IOException {
dataFusionService.close();
}
}
+
+ /**
+ * Snapshot the native DataFusion per-query registry and return a {@code contextId -> metrics}
+ * map. Returns an empty map when the service is not yet running (startup) or has been stopped
+ * (shutdown), so callers never see a half-initialized view.
+ *
+ *
Each entry mirrors one {@code QueryTracker} on the Rust side — current and peak memory
+ * reservation, wall time, and whether the query has completed but not yet been drained.
+ * The map contains at most {@link #ACTIVE_QUERY_METRICS_TOP_N} entries — the heaviest live
+ * queries by {@code current_bytes}, selected on the Rust side. Iteration order matches the
+ * order Rust drained the bounded min-heap (unspecified but stable per snapshot).
+ */
+ @Override
+ public Map getTopQueriesByMemory() {
+ if (dataFusionService == null) {
+ return Collections.emptyMap();
+ }
+ Map result = NativeBridge.getTopNQueriesByMemory(ACTIVE_QUERY_METRICS_TOP_N);
+ if (logger.isDebugEnabled()) {
+ logger.debug("getTopQueriesByMemory: {} entries from native registry", result.size());
+ }
+ return result;
+ }
}
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java
index 45440e3c6f3d5..d261cb9514deb 100644
--- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java
@@ -9,6 +9,7 @@
package org.opensearch.be.datafusion.nativelib;
import org.opensearch.analytics.backend.jni.NativeHandle;
+import org.opensearch.analytics.spi.QueryExecutionMetrics;
import org.opensearch.be.datafusion.stats.DataFusionStats;
import org.opensearch.be.datafusion.stats.NativeExecutorsStats;
import org.opensearch.be.datafusion.stats.TaskMonitorStats;
@@ -22,8 +23,11 @@
import java.lang.foreign.SymbolLookup;
import java.lang.foreign.ValueLayout;
import java.lang.invoke.MethodHandle;
+import java.util.Collections;
+import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
+import java.util.Map;
/**
* FFM bridge to native DataFusion library.
@@ -84,6 +88,7 @@ public final class NativeBridge {
private static final MethodHandle EXECUTE_WITH_CONTEXT;
private static final MethodHandle CANCEL_QUERY;
private static final MethodHandle STATS;
+ private static final MethodHandle QUERY_REGISTRY_TOP_N_BY_CURRENT;
private static final MethodHandle PREPARE_PARTIAL_PLAN;
private static final MethodHandle PREPARE_FINAL_PLAN;
private static final MethodHandle EXECUTE_LOCAL_PREPARED_PLAN;
@@ -407,6 +412,12 @@ public final class NativeBridge {
FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG)
);
+ // i64 df_query_registry_top_n_by_current(out_ptr, cap_entries)
+ QUERY_REGISTRY_TOP_N_BY_CURRENT = linker.downcallHandle(
+ lib.find("df_query_registry_top_n_by_current").orElseThrow(),
+ FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG)
+ );
+
// ── Distributed aggregate: prepare partial/final plans ──
// i64 df_prepare_partial_plan(handle_ptr, bytes_ptr, bytes_len)
PREPARE_PARTIAL_PLAN = linker.downcallHandle(
@@ -709,6 +720,55 @@ public static DataFusionStats stats() {
}
}
+ // ---- Per-query registry top-N snapshot ----
+
+ /**
+ * Snapshot the {@code n} heaviest live queries by {@code current_bytes}.
+ *
+ *
One FFM call: Java allocates an {@code n × ENTRY_BYTES} segment, Rust
+ * runs a bounded min-heap of size {@code n} over its registry and writes
+ * the survivors as {@link QueryRegistryLayout} entries. Completed and
+ * zero-byte trackers are filtered on the Rust side. Order within the
+ * buffer is unspecified.
+ *
+ *
Mirrors the {@link #stats()} pattern: the layout class decodes
+ * directly into the caller's final type ({@link QueryExecutionMetrics})
+ * with no transport-only intermediate.
+ *
+ * @param n maximum entries to return; must be non-negative. Zero short-circuits
+ * without crossing the FFM boundary.
+ * @return a non-null map of {@code contextId → metrics}, possibly empty,
+ * with at most {@code n} entries.
+ * @throws IllegalArgumentException if {@code n} is negative or implies a
+ * buffer larger than {@link Integer#MAX_VALUE} bytes
+ */
+ public static Map getTopNQueriesByMemory(int n) {
+ if (n < 0) {
+ throw new IllegalArgumentException("n must be non-negative: " + n);
+ }
+ if (n == 0) {
+ return Collections.emptyMap();
+ }
+ long bytes = (long) n * QueryRegistryLayout.ENTRY_BYTES;
+ if (bytes > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException("n too large for a single snapshot: " + n);
+ }
+ try (var call = new NativeCall()) {
+ var seg = call.buf((int) bytes);
+ long written = call.invoke(QUERY_REGISTRY_TOP_N_BY_CURRENT, seg, (long) n);
+ if (written == 0L) {
+ return Collections.emptyMap();
+ }
+ int rows = (int) Math.min(written, (long) n);
+ Map out = new HashMap<>(rows);
+ for (int i = 0; i < rows; i++) {
+ long ctxId = QueryRegistryLayout.readContextId(seg, i);
+ out.put(ctxId, QueryRegistryLayout.readMetrics(seg, i));
+ }
+ return Collections.unmodifiableMap(out);
+ }
+ }
+
// ---- Stubs ----
public static byte[] sqlToSubstrait(long readerPtr, String tableName, String sql, long runtimePtr) {
diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/QueryRegistryLayout.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/QueryRegistryLayout.java
new file mode 100644
index 0000000000000..568d57dc4851f
--- /dev/null
+++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/QueryRegistryLayout.java
@@ -0,0 +1,84 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.be.datafusion.nativelib;
+
+import org.opensearch.analytics.spi.QueryExecutionMetrics;
+
+import java.lang.foreign.MemoryLayout;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.StructLayout;
+import java.lang.foreign.ValueLayout;
+
+/**
+ * Mirrors the Rust {@code query_tracker::WireQueryMetric} {@code #[repr(C)]}
+ * struct (2 × i64 = 16 bytes) and decodes a strided buffer of those structs
+ * directly into the SPI types consumed by
+ * {@link org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin#getTopQueriesByMemory}.
+ *
+ *
Decoder reads each row by slicing the segment to that row's 16-byte
+ * window. We deliberately avoid {@code SequenceLayout}-derived
+ * {@code VarHandle}s because their bounds check enforces the segment span
+ * the entire sequence layout, not just up to the requested row. With a
+ * generic {@code Long.MAX_VALUE / ENTRY_BYTES} sequence layout the read
+ * always fails an {@link IndexOutOfBoundsException}.
+ *
+ *
Buffer shape (populated by {@code df_query_registry_top_n_by_current}):
+ *
+ */
+public final class QueryRegistryLayout {
+
+ /** Layout of a single wire entry. */
+ public static final StructLayout ENTRY_LAYOUT = MemoryLayout.structLayout(
+ ValueLayout.JAVA_LONG.withName("context_id"),
+ ValueLayout.JAVA_LONG.withName("current_bytes")
+ );
+
+ /** Byte size of one wire entry. Matches {@code size_of::()} on the Rust side. */
+ public static final long ENTRY_BYTES = ENTRY_LAYOUT.byteSize();
+
+ // Byte offsets within a single entry. Field order must match
+ // ENTRY_LAYOUT and the Rust #[repr(C)] WireQueryMetric struct.
+ private static final long OFF_CONTEXT_ID = 0L;
+ private static final long OFF_CURRENT_BYTES = 8L;
+
+ static {
+ long expected = 2L * Long.BYTES;
+ if (ENTRY_BYTES != expected) {
+ throw new AssertionError("QueryRegistryLayout entry size mismatch: expected " + expected + " but got " + ENTRY_BYTES);
+ }
+ }
+
+ private QueryRegistryLayout() {}
+
+ /**
+ * Read the {@code context_id} of the entry at zero-based index {@code i}.
+ *
+ * @param seg buffer populated by {@code df_query_registry_top_n_by_current}
+ * @param i row index, must be {@code < written} returned by the FFI call
+ */
+ public static long readContextId(MemorySegment seg, int i) {
+ long rowOffset = (long) i * ENTRY_BYTES;
+ return seg.get(ValueLayout.JAVA_LONG, rowOffset + OFF_CONTEXT_ID);
+ }
+
+ /**
+ * Read the metrics fields of the entry at zero-based index {@code i} into
+ * a fresh {@link QueryExecutionMetrics} record.
+ *
+ * @param seg buffer populated by {@code df_query_registry_top_n_by_current}
+ * @param i row index, must be {@code < written} returned by the FFI call
+ */
+ public static QueryExecutionMetrics readMetrics(MemorySegment seg, int i) {
+ long rowOffset = (long) i * ENTRY_BYTES;
+ return new QueryExecutionMetrics(seg.get(ValueLayout.JAVA_LONG, rowOffset + OFF_CURRENT_BYTES));
+ }
+}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsQueryTask.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsQueryTask.java
index 8fcb1b911048f..bb9ac23642f78 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsQueryTask.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsQueryTask.java
@@ -15,11 +15,10 @@
import org.opensearch.common.Nullable;
import org.opensearch.common.unit.TimeValue;
import org.opensearch.core.tasks.TaskId;
-import org.opensearch.tasks.CancellableTask;
-import org.opensearch.tasks.SearchBackpressureTask;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
/**
* Coordinator-level cancellable task representing a running analytics query.
@@ -28,7 +27,7 @@
*
* @opensearch.internal
*/
-public class AnalyticsQueryTask extends CancellableTask implements SearchBackpressureTask {
+public class AnalyticsQueryTask extends SearchTask {
private static final Logger logger = LogManager.getLogger(AnalyticsQueryTask.class);
@@ -49,7 +48,7 @@ public AnalyticsQueryTask(
id,
type,
action,
- "queryId[" + queryId + "]",
+ (Supplier) () -> "queryId[" + queryId + "]",
parentTaskId,
headers,
cancelAfterTimeInterval != null ? cancelAfterTimeInterval : TimeValue.MINUS_ONE
@@ -58,15 +57,15 @@ public AnalyticsQueryTask(
this.cancelAfterTimeInterval = cancelAfterTimeInterval;
}
- public AnalyticsQueryTask(long id, String type, String action, String queryId, TaskId parentTaskId, Map headers) {
- this(id, type, action, queryId, parentTaskId, headers, null);
- }
-
@Override
public boolean shouldCancelChildrenOnCancellation() {
return true;
}
+ public AnalyticsQueryTask(long id, String type, String action, String queryId, TaskId parentTaskId, Map headers) {
+ this(id, type, action, queryId, parentTaskId, headers, null);
+ }
+
public String getQueryId() {
return queryId;
}
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsShardTask.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsShardTask.java
index 8bed052216aa2..8155091958a94 100644
--- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsShardTask.java
+++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsShardTask.java
@@ -8,19 +8,26 @@
package org.opensearch.analytics.exec.task;
+import org.opensearch.action.search.SearchShardTask;
import org.opensearch.core.tasks.TaskId;
-import org.opensearch.tasks.CancellableTask;
import java.util.Map;
/**
* Data-node shard task representing a single shard fragment execution.
- * Analogous to {@link org.opensearch.action.search.SearchShardTask}.
- * Cancelling this task does not cascade to children.
+ *
Extends {@link SearchShardTask} so that search backpressure and
+ * {@code TaskResourceTrackingService} observe and track analytics shard
+ * tasks alongside regular search shard tasks. The native-memory tracker
+ * in {@code SearchBackpressureService} filters by {@code SearchShardTask},
+ * so inheriting from it is the integration point that exposes per-query
+ * DataFusion memory to cancellation.
+ *
+ *
{@link NativeMemoryTestPlugin} satisfies the second condition by installing fake snapshot
+ * and budget suppliers from {@code createComponents}, mirroring the production wiring path
+ * {@code DataFusionPlugin} uses. The first and third are checked via {@code assumeTrue} so the
+ * suite is skipped on macOS/Windows CI rather than failing.
+ */
+@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 1)
+public class NativeMemorySearchBackpressureIT extends OpenSearchIntegTestCase {
+
+ private static final TimeValue TIMEOUT = new TimeValue(10, TimeUnit.SECONDS);
+ private static final long DEFAULT_BUDGET_BYTES = 1024L * 1024L * 1024L; // 1 GiB
+
+ /**
+ * Per-task-id native bytes published into {@link NativeMemoryUsageTracker}'s snapshot
+ * supplier. Tests mutate this map between requests; the supplier installed by
+ * {@link NativeMemoryTestPlugin} reads the live reference on every refresh.
+ */
+ private static final Map SNAPSHOT = new ConcurrentHashMap<>();
+ private static final AtomicLong BUDGET_BYTES = new AtomicLong(DEFAULT_BUDGET_BYTES);
+
+ @Override
+ protected Collection> nodePlugins() {
+ final List> plugins = new ArrayList<>(super.nodePlugins());
+ plugins.add(NativeMemoryTestPlugin.class);
+ return plugins;
+ }
+
+ @Before
+ public final void setupNodeSettings() {
+ // Native-memory duress + tracking only engages on Linux; skip otherwise.
+ assumeTrue("native-memory backpressure path is Linux-only (OsProbe.getProcessNativeMemoryBytes)", Constants.LINUX);
+
+ // Force native-memory duress to trip after a single observation by setting the threshold
+ // to 0.0 — any usedFraction above zero counts as a breach. Combined with a tiny
+ // numSuccessiveBreaches the duress flips on the next tick deterministically.
+ Settings request = Settings.builder()
+ .put(NodeDuressSettings.SETTING_CPU_THRESHOLD.getKey(), 1.0) // disable CPU duress
+ .put(NodeDuressSettings.SETTING_HEAP_THRESHOLD.getKey(), 1.0) // disable heap duress
+ .put(NodeDuressSettings.SETTING_NATIVE_MEMORY_THRESHOLD.getKey(), 0.0)
+ .put(NodeDuressSettings.SETTING_NUM_SUCCESSIVE_BREACHES.getKey(), 1)
+ // search_backpressure.node_duress.native_memory_limit must be > 0 or the duress
+ // lambda short-circuits ("totalNative is zero"). 1 byte is enough — the probe
+ // reports a non-zero RssAnon.
+ .put(NodeDuressSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.getKey(), "1b")
+ .build();
+ assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(request).get());
+
+ SNAPSHOT.clear();
+ BUDGET_BYTES.set(DEFAULT_BUDGET_BYTES);
+ }
+
+ @After
+ public final void cleanupNodeSettings() {
+ SNAPSHOT.clear();
+ BUDGET_BYTES.set(DEFAULT_BUDGET_BYTES);
+ assertAcked(
+ client().admin()
+ .cluster()
+ .prepareUpdateSettings()
+ .setPersistentSettings(Settings.builder().putNull("*"))
+ .setTransientSettings(Settings.builder().putNull("*"))
+ );
+ }
+
+ /**
+ * The shard task's reported native bytes exceed {@code budget * fraction}; SBP must cancel
+ * with a "native memory usage exceeded" reason when running in {@code enforced} mode.
+ */
+ public void testSearchShardTaskCancellationWithHighNativeMemory() throws InterruptedException {
+ Settings request = Settings.builder()
+ .put(SearchBackpressureSettings.SETTING_MODE.getKey(), "enforced")
+ // Threshold = 0.5 of 1 GiB = 512 MiB. Snapshot publishes 768 MiB for the task —
+ // strictly above, so the evaluator must produce a cancellation reason.
+ .put(SearchShardTaskSettings.SETTING_NATIVE_MEMORY_PERCENT_THRESHOLD.getKey(), 0.5)
+ .build();
+ assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(request).get());
+
+ ExceptionCatchingListener listener = new ExceptionCatchingListener();
+ client().execute(TestTransportAction.ACTION, new TestRequest(taskId -> SNAPSHOT.put(taskId, 768L * 1024L * 1024L)), listener);
+ assertTrue("task should have been cancelled within " + TIMEOUT, listener.latch.await(TIMEOUT.getSeconds() * 2, TimeUnit.SECONDS));
+
+ Exception caughtException = listener.getException();
+ assertNotNull("SearchShardTask should have been cancelled with TaskCancelledException", caughtException);
+ MatcherAssert.assertThat(caughtException, instanceOf(TaskCancelledException.class));
+ MatcherAssert.assertThat(caughtException.getMessage(), containsString("native memory usage exceeded"));
+ }
+
+ /**
+ * Per-task threshold is set, but the published bytes are below the evaluator's cutoff. SBP
+ * must NOT cancel even though native-memory duress is active. This pins the contract that
+ * duress alone doesn't kill tasks — the evaluator decides per task.
+ */
+ public void testSearchShardTaskNotCancelledBelowThreshold() throws InterruptedException {
+ Settings request = Settings.builder()
+ .put(SearchBackpressureSettings.SETTING_MODE.getKey(), "enforced")
+ // Threshold = 0.9 of 1 GiB ≈ 922 MiB. Snapshot publishes 100 MiB — well under.
+ .put(SearchShardTaskSettings.SETTING_NATIVE_MEMORY_PERCENT_THRESHOLD.getKey(), 0.9)
+ .build();
+ assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(request).get());
+
+ ExceptionCatchingListener listener = new ExceptionCatchingListener();
+ client().execute(TestTransportAction.ACTION, new TestRequest(taskId -> SNAPSHOT.put(taskId, 100L * 1024L * 1024L)), listener);
+ // Task self-completes after TIMEOUT; allow extra room.
+ assertTrue(
+ "task should have completed within " + (TIMEOUT.getSeconds() * 3) + "s",
+ listener.latch.await(TIMEOUT.getSeconds() * 3, TimeUnit.SECONDS)
+ );
+ assertNull("SearchShardTask under threshold must not be cancelled by the native-memory tracker", listener.getException());
+ }
+
+ /**
+ * Native-memory threshold setting explicitly set to {@code 0.0} (the rollout-gate inert state).
+ * Even with heavy reported bytes the tracker must stay silent, because {@code fraction == 0}
+ * disables evaluation. Pins the spec's "operator can opt out" contract.
+ */
+ public void testTrackerInertWhenPerTaskFractionIsZero() throws InterruptedException {
+ Settings request = Settings.builder()
+ .put(SearchBackpressureSettings.SETTING_MODE.getKey(), "enforced")
+ // Explicitly set to 0.0 to override the production default and force the inert path.
+ .put(SearchShardTaskSettings.SETTING_NATIVE_MEMORY_PERCENT_THRESHOLD.getKey(), 0.0)
+ .build();
+ assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(request).get());
+
+ ExceptionCatchingListener listener = new ExceptionCatchingListener();
+ client().execute(
+ TestTransportAction.ACTION,
+ new TestRequest(taskId -> SNAPSHOT.put(taskId, 4L * 1024L * 1024L * 1024L)), // 4 GiB
+ listener
+ );
+ assertTrue("task should have completed without cancellation", listener.latch.await(TIMEOUT.getSeconds() * 3, TimeUnit.SECONDS));
+ assertNull("Per-task fraction explicitly set to 0 — the tracker must not cancel", listener.getException());
+ }
+
+ /**
+ * In {@code monitor_only} mode, SBP must log breaches but never invoke
+ * {@code TaskManager.cancelTaskAndDescendants}. The task therefore self-completes.
+ */
+ public void testMonitorOnlyModeDoesNotCancel() throws InterruptedException {
+ Settings request = Settings.builder()
+ .put(SearchBackpressureSettings.SETTING_MODE.getKey(), "monitor_only")
+ .put(SearchShardTaskSettings.SETTING_NATIVE_MEMORY_PERCENT_THRESHOLD.getKey(), 0.5)
+ .build();
+ assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(request).get());
+
+ ExceptionCatchingListener listener = new ExceptionCatchingListener();
+ client().execute(TestTransportAction.ACTION, new TestRequest(taskId -> SNAPSHOT.put(taskId, 768L * 1024L * 1024L)), listener);
+ assertTrue(
+ "task should have completed without cancellation in monitor_only",
+ listener.latch.await(TIMEOUT.getSeconds() * 3, TimeUnit.SECONDS)
+ );
+ assertNull("monitor_only must not cancel", listener.getException());
+ }
+
+ /**
+ * {@code search_backpressure.mode = disabled} short-circuits the entire {@code doRun} loop
+ * before any tracker is consulted. The task must complete normally even with a heavy
+ * snapshot reading.
+ */
+ public void testDisabledModeShortCircuits() throws InterruptedException {
+ Settings request = Settings.builder()
+ .put(SearchBackpressureSettings.SETTING_MODE.getKey(), "disabled")
+ .put(SearchShardTaskSettings.SETTING_NATIVE_MEMORY_PERCENT_THRESHOLD.getKey(), 0.5)
+ .build();
+ assertAcked(client().admin().cluster().prepareUpdateSettings().setPersistentSettings(request).get());
+
+ ExceptionCatchingListener listener = new ExceptionCatchingListener();
+ client().execute(TestTransportAction.ACTION, new TestRequest(taskId -> SNAPSHOT.put(taskId, 999L * 1024L * 1024L)), listener);
+ assertTrue(
+ "task should have completed without cancellation in disabled mode",
+ listener.latch.await(TIMEOUT.getSeconds() * 3, TimeUnit.SECONDS)
+ );
+ assertNull("disabled mode must not cancel", listener.getException());
+ }
+
+ // -----------------------------------------------------------------------
+ // Test plugin + transport action
+ // -----------------------------------------------------------------------
+
+ /**
+ * Installs the snapshot + budget suppliers on {@link NativeMemoryUsageTracker} during
+ * {@code createComponents} — same boot-time hook the production DataFusion plugin uses.
+ * SBP is constructed in {@code Node} after {@code createComponents} runs, so the
+ * suppliers are visible when {@code isNativeTrackingSupported()} is evaluated.
+ */
+ public static class NativeMemoryTestPlugin extends Plugin implements ActionPlugin {
+ @Override
+ public Collection