Expose tokio-metrics as datafusion plugin stats (opensearch-project#21303)

* Expose tokio-metrics as datafusion plugin stats

Signed-off-by: Ajay Raj Nelapudi <ajnelapu@amazon.com>

* refactor: Extract stats SPI to libs/plugin-stats-spi

Signed-off-by: Ajay Raj Nelapudi <ajnelapu@amazon.com>

* Move DF specific constructs to DF plugin

Signed-off-by: Ajay Raj Nelapudi <ajnelapu@amazon.com>

* Remove RUNTIME_METRICS_COUNT and resolve sandbox failures

Signed-off-by: Ajay Raj Nelapudi <ajnelapu@amazon.com>

* Add local_queue_depth to runtime metrics

Signed-off-by: Ajay Raj Nelapudi <ajnelapu@amazon.com>

* Use Handle::metrics() for point-in-time stats instead of interval deltas

Signed-off-by: Ajay Raj Nelapudi <ajnelapu@amazon.com>

---------

Signed-off-by: Ajay Raj Nelapudi <ajnelapu@amazon.com>
Co-authored-by: Ajay Raj Nelapudi <ajnelapu@amazon.com>

Author: AjayRajNelapudi

gradle/libs.versions.toml

@@ -109,6 +109,10 @@ flatbuffers           = "2.0.0"
 calcite               = "1.41.0"
 calcite_os_rev        = "1"
 
+# property-based testing
+jqwik                 = "1.9.2"
+junit_jupiter         = "5.11.3"
+junit_platform        = "1.11.3"
 [libraries]
 antlr4-runtime = { group = "org.antlr", name = "antlr4-runtime", version.ref = "antlr4" }
 asm-analysis = { group = "org.ow2.asm", name = "asm-analysis", version.ref = "asm" }

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

@@ -70,6 +70,7 @@ thiserror = "1.0"
 async-trait = "0.1"
 bytes = "1"
 criterion = { version = "0.5", features = ["async_tokio"] }
+tokio-metrics = { version = "0.5", features = ["rt"] }
 
 # Internal
 native-bridge-common = { path = "common" }

sandbox/libs/plugin-stats-spi/build.gradle

@@ -0,0 +1,24 @@
+/*
+ * 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.
+ */
+
+/*
+ * SPI interfaces for plugin stats collection.
+ * Contains only PluginStats (marker interface) and BackendStatsProvider.
+ * Consumed by sandbox plugins that report backend statistics.
+ */
+
+dependencies {
+    api project(':libs:opensearch-core')
+    api project(':libs:opensearch-common')
+}
+
+testingConventions.enabled = false
+
+tasks.named('forbiddenApisMain').configure {
+    replaceSignatureFiles 'jdk-signatures'
+}

sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/BackendStatsProvider.java

@@ -0,0 +1,32 @@
+/*
+ * 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.plugin.stats;
+
+/**
+ * Each backend (DataFusion, Parquet, future engines) implements this interface
+ * to provide its stats to the Mustang Stats Framework. The Analytics Plugin
+ * discovers {@code BackendStatsProvider} implementations and iterates over them
+ * to collect stats from all registered backends.
+ */
+public interface BackendStatsProvider {
+
+    /**
+     * Returns the backend's identifier, e.g. {@code "datafusion"}, {@code "parquet"}.
+     *
+     * @return a non-null backend name
+     */
+    String name();
+
+    /**
+     * Returns the backend's stats object.
+     *
+     * @return a non-null {@link PluginStats} instance
+     */
+    PluginStats getBackendStats();
+}

sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/PluginStats.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.plugin.stats;
+
+/**
+ * Marker interface for all backend stats types in the Mustang Stats Framework.
+ *
+ * <p>Intentionally empty — serves as the common type for
+ * {@link BackendStatsProvider#getBackendStats()}. Each backend's top-level stats
+ * class (e.g. {@code DataFusionStats}) implements this interface so the Analytics
+ * Plugin can discover and iterate over them.
+ */
+public interface PluginStats {
+    // marker — no methods
+}

sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/package-info.java

@@ -0,0 +1,22 @@
+/*
+ * 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.
+ */
+
+/**
+ * SPI stats types for the Mustang Stats Framework.
+ *
+ * <p>This package contains the stats interfaces shared between
+ * the OpenSearch server and native backend plugins. Types here are visible to
+ * both sides without requiring a plugin dependency.
+ *
+ * <p>Key types:
+ * <ul>
+ *   <li>{@link org.opensearch.plugin.stats.PluginStats} — marker interface for all backend stats</li>
+ *   <li>{@link org.opensearch.plugin.stats.BackendStatsProvider} — interface for backends to provide stats</li>
+ * </ul>
+ */
+package org.opensearch.plugin.stats;

sandbox/plugins/analytics-backend-datafusion/build.gradle

@@ -39,6 +39,9 @@ dependencies {
   // Shared native bridge lib (provides the unified .so and FFM SymbolLookup)
   implementation project(':sandbox:libs:dataformat-native')
 
+  // Canonical stats SPI classes (PluginStats, BackendStatsProvider)
+  implementation project(':sandbox:libs:plugin-stats-spi')
+
   // Provided at runtime by the parent analytics-engine plugin; compile-only to avoid jar hell.
   compileOnly project(':sandbox:libs:analytics-framework')
   // analytics-engine's RelNode types (OpenSearchStageInputScan) are referenced by the
@@ -192,6 +195,47 @@ testingConventions.enabled = false
 
 tasks.named('forbiddenPatterns').configure {
     exclude '**/*.parquet'
+    exclude '**/*.dylib'
+    exclude '**/*.so'
+    exclude '**/*.dll'
+}
+
+// ---- Property-based tests (jqwik / JUnit 5 Platform) ----
+
+sourceSets {
+  propertyTest {
+    java {
+      srcDir 'src/propertyTest/java'
+    }
+    compileClasspath += sourceSets.main.output
+    runtimeClasspath += sourceSets.main.output
+  }
+}
+
+configurations {
+  propertyTestImplementation.extendsFrom implementation, compileOnly
+  propertyTestRuntimeOnly.extendsFrom runtimeOnly
+}
+
+dependencies {
+  propertyTestImplementation "net.jqwik:jqwik:${versions.jqwik}"
+  propertyTestImplementation "org.junit.jupiter:junit-jupiter-api:${versions.junit_jupiter}"
+  propertyTestRuntimeOnly "org.junit.platform:junit-platform-launcher:${versions.junit_platform}"
+  // Jackson for JSON parsing in property tests
+  propertyTestImplementation "com.fasterxml.jackson.core:jackson-databind:${versions.jackson_databind}"
+}
+
+tasks.register('propertyTest', Test) {
+  description = 'Run jqwik property-based tests'
+  group = 'verification'
+  useJUnitPlatform {
+    includeEngines 'jqwik'
+  }
+  testClassesDirs = sourceSets.propertyTest.output.classesDirs
+  classpath = sourceSets.propertyTest.runtimeClasspath
+  jvmArgs += ["--add-opens", "java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED"]
+  // Disable the security manager for property tests (jqwik is not compatible)
+  systemProperty 'tests.security.manager', 'false'
 }
 
 tasks.matching { it.name == 'missingJavadoc' }.configureEach {
@@ -206,3 +250,7 @@ tasks.named('thirdPartyAudit').configure {
         'org.apache.calcite.server.ServerDdlExecutor'
     )
 }
+
+// jqwik property tests don't ship with the randomized-testing framework that
+// forbiddenApis signatures reference — skip the check for this source set.
+tasks.matching { it.name == 'forbiddenApisPropertyTest' }.configureEach { enabled = false }

sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml

@@ -1,4 +1,5 @@
 [package]
+# DataFusion analytics backend native library
 name = "opensearch-datafusion"
 version = "0.1.0"
 edition = "2021"
@@ -46,6 +47,8 @@ thiserror = { workspace = true }
 # convert_tz UDF
 chrono-tz = "0.10"
 
+tokio-metrics = { workspace = true }
+
 [dev-dependencies]
 criterion = { workspace = true }
 tempfile = { workspace = true }

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

@@ -154,6 +154,14 @@ impl DedicatedExecutor {
         }
     }
 
+    /// Returns a clone of the underlying Tokio runtime `Handle`, if the
+    /// executor has not been shut down. Used to create a
+    /// `tokio_metrics::RuntimeMonitor` for the CPU runtime.
+    pub fn handle(&self) -> Option<Handle> {
+        let state = self.state.read();
+        state.handle.clone()
+    }
+
     pub fn shutdown(&self) {
         let mut state = self.state.write();
         state.handle = None;

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

@@ -651,3 +651,59 @@ pub unsafe extern "C" fn df_execute_with_context(
             .map_err(|e| e.to_string())
     }
 }
+
+// ---- Stats collection ----
+
+/// Collects all native executor metrics into a caller-provided byte buffer.
+///
+/// The buffer must have capacity for at least `size_of::<DfStatsBuffer>()` bytes (224).
+/// Returns 0 on success.
+#[ffm_safe]
+#[no_mangle]
+pub unsafe extern "C" fn df_stats(out_ptr: *mut u8, out_cap: i64) -> i64 {
+    use crate::stats::{layout, pack_runtime_metrics, pack_task_monitor, DfStatsBuffer, RuntimeMetricsRepr};
+    use crate::task_monitors::{
+        query_execution_monitor, stream_next_monitor,
+        fetch_phase_monitor, segment_stats_monitor,
+    };
+
+    if out_cap < 0 || (out_cap as usize) < layout::BUFFER_BYTE_SIZE {
+        return Err(format!(
+            "stats buffer too small: need {} but got {}",
+            layout::BUFFER_BYTE_SIZE, out_cap
+        ));
+    }
+
+    let mgr = get_rt_manager()?;
+
+    // IO runtime (always present)
+    let io_runtime = pack_runtime_metrics(&mgr.io_monitor, mgr.io_runtime.handle());
+
+    // CPU runtime (optional — zeroed when absent)
+    let cpu_runtime = if let Some(ref cpu_mon) = mgr.cpu_monitor {
+        if let Some(cpu_handle) = mgr.cpu_executor.handle() {
+            pack_runtime_metrics(cpu_mon, &cpu_handle)
+        } else {
+            RuntimeMetricsRepr::zeroed()
+        }
+    } else {
+        RuntimeMetricsRepr::zeroed()
+    };
+
+    let buf = DfStatsBuffer {
+        io_runtime,
+        cpu_runtime,
+        query_execution: pack_task_monitor(query_execution_monitor()),
+        stream_next: pack_task_monitor(stream_next_monitor()),
+        fetch_phase: pack_task_monitor(fetch_phase_monitor()),
+        segment_stats: pack_task_monitor(segment_stats_monitor()),
+    };
+
+    // Copy struct bytes to caller buffer
+    std::ptr::copy_nonoverlapping(
+        &buf as *const DfStatsBuffer as *const u8,
+        out_ptr,
+        std::mem::size_of::<DfStatsBuffer>(),
+    );
+    Ok(0)
+}