searchBack Pressure

Signed-off-by: Pradeep L <spradeel@amazon.com>

Author: Pradeep L

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.
@@ -113,4 +115,18 @@ default FilterDelegationHandle getFilterDelegationHandle(List<DelegatedExpressio
     default void configureFilterDelegation(FilterDelegationHandle handle, BackendExecutionContext backendContext) {
         throw new UnsupportedOperationException("configureFilterDelegation not implemented for [" + name() + "]");
     }
+
+    /**
+     * Returns a snapshot of this backend's currently-tracked queries, keyed by {@code contextId}.
+     *
+     * <p>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.
+     *
+     * <p>Default implementation returns an empty map so backends that do not track per-query
+     * metrics don't have to opt in.
+     */
+    default Map<Long, QueryExecutionMetrics> getActiveQueryMetrics() {
+        return Collections.emptyMap();
+    }
 }

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/QueryExecutionMetrics.java

@@ -0,0 +1,26 @@
+/*
+ * 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.
+ *
+ * <p>The context id (the map key in {@link AnalyticsSearchBackendPlugin#getActiveQueryMetrics()})
+ * is not repeated here — this record holds only the remaining parameters: memory accounting,
+ * wall time, and whether the query has completed but not yet been drained.
+ *
+ * @param currentBytes bytes currently reserved by the query's memory pool
+ * @param peakBytes    high-water mark of bytes reserved during the query's lifetime
+ * @param wallNanos    live or frozen wall-clock duration in nanoseconds
+ * @param completed    {@code true} when the query has finished executing but its entry has not
+ *                     yet been drained from the backend's internal registry
+ *
+ * @opensearch.internal
+ */
+public record QueryExecutionMetrics(long currentBytes, long peakBytes, long wallNanos, boolean completed) {}

sandbox/libs/dataformat-native/rust/Cargo.toml

@@ -80,7 +80,7 @@ opensearch-repository-azure = { path = "../../../plugins/native-repository-azure
 opensearch-repository-fs = { path = "../../../plugins/native-repository-fs/src/main/rust" }
 
 [profile.release]
-lto = true
+lto = "thin"
 codegen-units = 1
 incremental = true
 debug = "line-tables-only"

sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs

@@ -12,6 +12,7 @@ use std::slice;
 use std::str;
 use std::sync::Arc;
 
+use log::info;
 use native_bridge_common::ffm_safe;
 use parking_lot::RwLock;
 
@@ -202,6 +203,52 @@ pub extern "C" fn df_cancel_query(context_id: i64) {
     api::cancel_query(context_id);
 }
 
+// ---------------------------------------------------------------------------
+// Per-query registry snapshot (two-phase)
+//
+// Java calls `df_query_registry_len` to size a buffer, then `df_query_registry_snapshot`
+// to populate it. See `query_tracker::WireQueryMetric` for the wire layout.
+// ---------------------------------------------------------------------------
+
+/// Returns the current number of entries in the query registry.
+/// Value is racy — treat it as a sizing hint only.
+#[no_mangle]
+pub extern "C" fn df_query_registry_len() -> i64 {
+    let len = crate::query_tracker::query_registry_len();
+    info!("[nativemem-bp] ffm.df_query_registry_len -> {}", len);
+    len as i64
+}
+
+/// Copies up to `cap_entries` `WireQueryMetric`s into the caller-provided buffer.
+/// Returns the number of entries actually written.
+///
+/// Safety: `out_ptr` must be non-null, 8-byte aligned, and point to storage for
+/// at least `cap_entries * size_of::<WireQueryMetric>()` bytes.
+#[ffm_safe]
+#[no_mangle]
+pub unsafe extern "C" fn df_query_registry_snapshot(out_ptr: *mut u8, cap_entries: i64) -> i64 {
+    use crate::query_tracker::{snapshot_query_registry, WireQueryMetric};
+
+    if cap_entries < 0 {
+        return Err(format!("negative capacity: {cap_entries}"));
+    }
+    if cap_entries == 0 {
+        info!("[nativemem-bp] ffm.df_query_registry_snapshot: capacity=0, nothing to write");
+        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_query_registry(out);
+    info!(
+        "[nativemem-bp] ffm.df_query_registry_snapshot: wrote {} entries (capacity {})",
+        written, cap_entries
+    );
+    Ok(written as i64)
+}
+
 #[ffm_safe]
 #[no_mangle]
 pub unsafe extern "C" fn df_sql_to_substrait(