Switch executeQueryPhaseAsync and streamNext from block_on to spawn (#20892)

* Switch executeQueryPhaseAsync and streamNext from block_on to spawn

This makes the JVM thread returns immediately after submitting work.

To prevent thread leak in tests, in runtime_manager.rs, split io_runtime
into a cheap Handle (for sharing) and Mutex<Option<Runtime>>
(for owned shutdown). shutdown() now calls io_runtime.shutdown_timeout(10s)
after cpu_executor.join_blocking(), blocking until all datafusion-io-*
threads have fully terminated.

Signed-off-by: Aravind Sagar <sagarara@amazon.com>

* Make some variable names clearer

Signed-off-by: Aravind Sagar <sagarara@amazon.com>

---------

Signed-off-by: Aravind Sagar <sagarara@amazon.com>

Author: aravindsagar

plugins/engine-datafusion/jni/src/lib.rs

@@ -12,10 +12,11 @@ use std::num::NonZeroUsize;
  * compatible open source license.
  */
 use std::ptr::addr_of_mut;
-use jni::objects::{JByteArray, JClass, JMap, JObject};
+use jni::objects::{GlobalRef, JByteArray, JClass, JMap, JObject};
 use jni::objects::JLongArray;
 use jni::sys::{jboolean, jbyteArray, jint, jlong, jstring};
 use jni::{JNIEnv, JavaVM};
+use std::future::Future;
 use std::sync::{Arc, OnceLock};
 use arrow_array::{Array, RecordBatch, StructArray};
 use arrow_array::ffi::FFI_ArrowArray;
@@ -67,7 +68,7 @@ use tokio::runtime::Runtime;
 use std::result;
 use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
-use futures::TryStreamExt;
+use futures::{TryStreamExt, FutureExt};
 
 pub type Result<T, E = DataFusionError> = result::Result<T, E>;
 
@@ -128,6 +129,59 @@ where
     })
 }
 
+/// Extract a human-readable message from a panic payload.
+fn panic_message(payload: &Box<dyn std::any::Any + Send>) -> String {
+    if let Some(s) = payload.downcast_ref::<&str>() {
+        s.to_string()
+    } else if let Some(s) = payload.downcast_ref::<String>() {
+        s.clone()
+    } else {
+        "unknown panic payload".to_string()
+    }
+}
+
+/// Spawn an async task on `runtime` that calls an ActionListener exactly once.
+///
+/// The entire `task` future runs inside `catch_unwind`. Any panic is converted
+/// to a `DataFusionError` and surfaced to the Java caller via `listener_ref`.
+/// This ensures the `CompletableFuture` on the Java side is always completed,
+/// never left hanging.
+///
+/// `on_ok` receives the success value and is responsible for calling the
+/// appropriate `set_action_listener_ok_*` variant. `T` is inferred from
+/// the closure, which in turn pins the `Output` type of `task`.
+fn spawn_jni_task<Fut, T, FOk>(
+    runtime: &tokio::runtime::Handle,
+    task_name: &'static str,
+    listener_ref: GlobalRef,
+    task: Fut,
+    on_ok: FOk,
+)
+where
+    Fut: Future<Output = Result<T, DataFusionError>> + Send + 'static,
+    T: Send + 'static,
+    FOk: FnOnce(&mut JNIEnv, &GlobalRef, T) + Send + 'static,
+{
+    let _ = runtime.spawn(async move {
+        let result = std::panic::AssertUnwindSafe(task)
+            .catch_unwind()
+            .await
+            .unwrap_or_else(|panic| {
+                let msg = panic_message(&panic);
+                log_error!("{} panicked: {}", task_name, msg);
+                Err(DataFusionError::Execution(format!("{} panicked: {}", task_name, msg)))
+            });
+
+        with_jni_env(|env| match result {
+            Ok(value) => on_ok(env, &listener_ref, value),
+            Err(e) => {
+                log_error!("{} failed: {}", task_name, e);
+                set_action_listener_error_global(env, &listener_ref, &e);
+            }
+        });
+    });
+}
+
 /// Initialize the logger for Rust->Java logging bridge.
 /// This should be called once when the native library is loaded.
 #[no_mangle]
@@ -619,9 +673,11 @@ pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_executeQu
     let table_path = shard_view.table_path();
     let files_meta = shard_view.files_metadata();
 
-    io_runtime.block_on(async move {
-
-        let result = query_executor::execute_query_with_cross_rt_stream(
+    spawn_jni_task(
+        &io_runtime,
+        "executeQueryPhaseAsync",
+        listener_ref,
+        query_executor::execute_query_with_cross_rt_stream(
             table_path,
             files_meta,
             table_name,
@@ -630,22 +686,9 @@ pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_executeQu
             target_partitions,
             runtime,
             cpu_executor,
-        ).await;
-
-        match result {
-            Ok(stream_ptr) => {
-                with_jni_env(|env| {
-                    set_action_listener_ok_global(env, &listener_ref, stream_ptr);
-                });
-            }
-            Err(e) => {
-                with_jni_env(|env| {
-                    log_error!("Query execution failed: {}", e);
-                    set_action_listener_error_global(env, &listener_ref, &e);
-                });
-            }
-        }
-    });
+        ),
+        |env, listener_ref, stream_pointer| set_action_listener_ok_global(env, listener_ref, stream_pointer),
+    );
 }
 
 #[no_mangle]
@@ -680,22 +723,13 @@ pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_fetchSegm
     let shard_view = unsafe { &*(shard_view_ptr as *const ShardView) };
     let files_meta = shard_view.files_metadata();
 
-    io_runtime.block_on(async move {
-        let file_stats = util::fetch_segment_statistics(files_meta).await;
-        match file_stats {
-            Ok(map) => {
-                with_jni_env(|env| {
-                    set_action_listener_ok_global_with_map(env, &listener_ref, &map);
-                });
-            }
-            Err(e) => {
-                with_jni_env(|env| {
-                    log_error!("Collecting file stats failed: {}", e);
-                    set_action_listener_error_global(env, &listener_ref, &e);
-                });
-            }
-        }
-    });
+    spawn_jni_task(
+        &io_runtime,
+        "fetchSegmentStats",
+        listener_ref,
+        async move { util::fetch_segment_statistics(files_meta).await },
+        |env, listener_ref, stats_map| set_action_listener_ok_global_with_map(env, listener_ref, &stats_map),
+    );
 }
 
 #[no_mangle]
@@ -732,41 +766,40 @@ pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_streamNex
     let stream_ptr = stream;
     let io_runtime = manager.io_runtime.clone();
 
-    io_runtime.block_on(async move {
-
-        let stream = unsafe { &mut *(stream_ptr as *mut RecordBatchStreamAdapter<CrossRtStream>) };
-        let result = stream.try_next().await;
-
-        // Uncomment for monitoring stream next
-        // let result = STREAM_NEXT_MONITOR.instrument(async {
-        //         stream.try_next().await
-        // }).await;
-
-        // Use thread-local JNI env - auto-attaches!
-        with_jni_env(|env| {
+    // Ensure stream_ptr lifetime is guaranteed beyond the spawn boundary
+    // (e.g., wrap in Arc<Mutex<...>> or ensure sequential access contract)
+    spawn_jni_task(
+        &io_runtime,
+        "streamNext",
+        listener_ref,
+        async move {
+            let stream = unsafe { &mut *(stream_ptr as *mut RecordBatchStreamAdapter<CrossRtStream>) };
+            // Poll the stream with monitoring
+            let result = stream.try_next().await?;
+
+            // Uncomment for monitoring stream next
+            // let result = STREAM_NEXT_MONITOR.instrument(async {
+            //         stream.try_next().await
+            // }).await;
             match result {
-                Ok(Some(batch)) => {
+                Some(batch) => {
                     log_info!("[RUST streamNext] Batch produced: {} rows, {} columns, schema: {:?}",
                         batch.num_rows(), batch.num_columns(), batch.schema().fields().iter().map(|f| f.name().as_str()).collect::<Vec<_>>());
                     // Convert to FFI
                     let struct_array: StructArray = batch.into();
                     let array_data = struct_array.into_data();
                     let ffi_array = FFI_ArrowArray::new(&array_data);
-                    let ffi_array_ptr = Box::into_raw(Box::new(ffi_array));
-                    set_action_listener_ok_global(env, &listener_ref, ffi_array_ptr as jlong);
+                    Ok(Box::into_raw(Box::new(ffi_array)) as jlong)
                 }
-                Ok(None) => {
+                None => {
                     log_info!("[RUST streamNext] End of stream reached");
                     // End of stream
-                    set_action_listener_ok_global(env, &listener_ref, 0);
-                }
-                Err(err) => {
-                    log_error!("Stream next failed: {}", err);
-                    set_action_listener_error_global(env, &listener_ref, &err);
+                    Ok(0)
                 }
             }
-        });
-    });
+        },
+        |env, listener_ref, data_pointer| set_action_listener_ok_global(env, listener_ref, data_pointer),
+    );
     // Function returns immediately to java - async rust work continues in background
 }
 

plugins/engine-datafusion/jni/src/runtime_manager.rs

@@ -2,9 +2,10 @@ use crate::executor::DedicatedExecutor;
 use crate::io::register_io_runtime;
 use vectorized_exec_spi::log_info;
 use log::info;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
 use datafusion::error::DataFusionError;
-use tokio::runtime::{Builder, Runtime};
+use tokio::runtime::{Builder, Handle, Runtime};
 
 #[derive(Debug, Clone)]
 pub struct RuntimeConfig {
@@ -71,7 +72,12 @@ impl RuntimeConfig {
 }
 
 pub struct RuntimeManager {
-    pub io_runtime: Arc<Runtime>,
+    /// Cloneable handle to the IO runtime, used by all JNI entry points to submit
+    /// tasks without taking ownership of the runtime.
+    pub io_runtime: Handle,
+    /// Owned runtime kept behind a Mutex so shutdown() can take() it and call
+    /// shutdown_timeout(), which blocks until all IO threads have fully stopped.
+    io_runtime_owned: Mutex<Option<Runtime>>,
     pub(crate) cpu_executor: DedicatedExecutor,
 }
 
@@ -88,36 +94,37 @@ impl RuntimeManager {
     pub fn with_config(config: RuntimeConfig) -> Self {
         log_info!("Creating RuntimeManager with config: {:?}", config);
 
-        // IO Runtime
-        let io_runtime = Arc::new(
-            Builder::new_multi_thread()
-                .worker_threads(config.effective_io_threads())
-                .thread_name("datafusion-io")
-                .enable_all()
-                .build()
-                .expect("Failed to create IO runtime"),
-        );
+        // IO Runtime — build first so we can extract a Handle for sharing
+        let io_runtime_rt = Builder::new_multi_thread()
+            .worker_threads(config.effective_io_threads())
+            .thread_name("datafusion-io")
+            .enable_all()
+            .build()
+            .expect("Failed to create IO runtime");
+
+        let io_handle = io_runtime_rt.handle().clone();
 
-        // Register IO runtime for current thread
-        register_io_runtime(Some(io_runtime.handle().clone()));
+        // Register IO runtime for the calling (JNI) thread.
+        register_io_runtime(Some(io_handle.clone()));
 
         // CPU Executor with its own runtime
         let mut cpu_runtime_builder = Builder::new_multi_thread();
-        let io_handle = io_runtime.handle().clone();
+        let io_handle_for_cpu = io_handle.clone();
 
         cpu_runtime_builder
             .worker_threads(config.effective_cpu_threads())
             .thread_name("datafusion-cpu")
             .enable_all()
             .on_thread_start(move || {
                 // Register IO runtime for each CPU thread
-                register_io_runtime(Some(io_handle.clone()));
+                register_io_runtime(Some(io_handle_for_cpu.clone()));
             });
 
         let cpu_executor = DedicatedExecutor::new("datafusion-cpu", cpu_runtime_builder);
 
         Self {
-            io_runtime,
+            io_runtime: io_handle,
+            io_runtime_owned: Mutex::new(Some(io_runtime_rt)),
             cpu_executor,
         }
     }
@@ -149,9 +156,14 @@ impl RuntimeManager {
 
     pub fn shutdown(&self) {
         info!("Shutting down RuntimeManager");
+        // Shut down CPU executor first — waits for all in-flight CPU tasks to finish.
         self.cpu_executor.join_blocking();
-         // TODO: io_runtime spawned threads seem to have issue and are leaking
-
+        // Take the owned IO runtime out of the Option and call shutdown_timeout, which
+        // blocks until every datafusion-io-* thread has fully terminated.
+        if let Some(rt) = self.io_runtime_owned.lock().unwrap().take() {
+            rt.shutdown_timeout(Duration::from_secs(10));
+        }
+        info!("RuntimeManager shut down complete");
     }
 }
 

plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/DataFusionReaderManagerTests.java

@@ -52,6 +52,7 @@
 import static org.opensearch.index.engine.Engine.SearcherScope.INTERNAL;
 
 public class DataFusionReaderManagerTests extends OpenSearchTestCase {
+
     private static DataFusionService service;
     Supplier<IndexFileDeleter> noOpFileDeleterSupplier;
 
@@ -468,7 +469,7 @@ private ShardPath createShardPathWithResourceFiles(String indexName, int shardId
 
     private void verifySearchResults(DatafusionSearcher searcher, DatafusionQuery datafusionQuery, Map<String, Long> expectedResults) throws Exception {
         Map<String, Object[]> finalRes = new HashMap<>();
-        searcher.searchAsync(datafusionQuery, service.getRuntimePointer()).whenComplete((streamPointer, error)-> {
+        searcher.searchAsync(datafusionQuery, service.getRuntimePointer()).whenCompleteAsync((streamPointer, error)-> {
             RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
             RecordBatchStream stream = new RecordBatchStream(streamPointer, service.getRuntimePointer(), allocator);