Make tokio concurrency settings dynamic (opensearch-project#21817)

* Make tokio concurrency settings dynamic

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

Author: AjayRajNelapudi

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

@@ -88,6 +88,54 @@ pub extern "C" fn df_shutdown_runtime_manager() {
     }
 }
 
+/// Updates the effective permit count of a named concurrency gate.
+/// Gate names: "fragment_executor" (targets DedicatedExecutor gate) or
+///             "reduce" (targets RuntimeManager coordinator gate).
+///
+/// Scale-up is synchronous. Scale-down spawns an async task on the IO
+/// runtime to acquire poison permits (may need to wait for in-flight
+/// queries to release).
+///
+/// Java: NativeBridge.updateConcurrencyGate(String, int)
+#[ffm_safe]
+#[no_mangle]
+pub unsafe extern "C" fn df_update_concurrency_gate(
+    gate_name_ptr: *const u8,
+    gate_name_len: i64,
+    new_max_permits: u32,
+) -> i64 {
+    let gate_name = str_from_raw(gate_name_ptr, gate_name_len)
+        .map_err(|e| format!("df_update_concurrency_gate: {}", e))?;
+
+    let mgr = match get_rt_manager() {
+        Ok(m) => m,
+        Err(_) => {
+            warn!("df_update_concurrency_gate called before runtime init");
+            return Ok(0);
+        }
+    };
+
+    let gate = match gate_name {
+        "fragment_executor" => mgr.cpu_executor().concurrency_gate().clone(),
+        "reduce" => mgr.coordinator_gate().clone(),
+        other => {
+            warn!("df_update_concurrency_gate: unknown gate '{}'", other);
+            return Ok(0);
+        }
+    };
+
+    let io_runtime = mgr.io_runtime.clone();
+    let gate_name_owned = gate_name.to_string();
+
+    // Spawn the resize on the IO runtime. Scale-up completes immediately;
+    // scale-down may need to wait for permits to become available.
+    io_runtime.spawn(async move {
+        gate.resize(new_max_permits, &gate_name_owned).await;
+    });
+
+    Ok(0)
+}
+
 #[ffm_safe]
 #[no_mangle]
 pub unsafe extern "C" fn df_create_global_runtime(
@@ -1149,8 +1197,8 @@ pub unsafe extern "C" fn df_stats(out_ptr: *mut u8, out_cap: i64) -> i64 {
         query_execution: pack_task_monitor(query_execution_monitor()),
         stream_next: pack_task_monitor(stream_next_monitor()),
         plan_setup: pack_task_monitor(plan_setup_monitor()),
-        datanode_gate: pack_partition_gate(mgr.cpu_executor.concurrency_gate()),
-        coordinator_gate: pack_partition_gate(mgr.coordinator_gate()),
+        fragment_executor_gate: pack_partition_gate(mgr.cpu_executor.concurrency_gate()),
+        reduce_executor_gate: pack_partition_gate(mgr.coordinator_gate()),
     };
 
     // Copy struct bytes to caller buffer
@@ -1262,4 +1310,121 @@ mod tests {
         assert_eq!(send_outcome_to_code(SendOutcome::ReceiverDropped), SENDER_SEND_RECEIVER_DROPPED);
         assert_eq!(SENDER_SEND_RECEIVER_DROPPED, 1);
     }
+
+    /// Initialize the global runtime manager for tests.
+    /// Uses 2 CPU threads and 1.5 multiplier (default) for both gates.
+    fn init_test_runtime() {
+        df_init_runtime_manager(2, 1.5, 1.5);
+    }
+
+    /// Shutdown and clear the global runtime manager after tests.
+    /// Must be called from a blocking context (not inside an async runtime).
+    fn shutdown_test_runtime() {
+        df_shutdown_runtime_manager();
+    }
+
+    /// Helper: call df_update_concurrency_gate with a Rust string.
+    /// Returns the i64 result (0 = success for the outer call).
+    unsafe fn call_update_gate(gate_name: &str, new_max: u32) -> i64 {
+        df_update_concurrency_gate(
+            gate_name.as_ptr(),
+            gate_name.len() as i64,
+            new_max,
+        )
+    }
+
+    /// Validates: Requirements 2.2, 2.3, 2.4, 2.6
+    ///
+    /// Combined test for FFI gate routing to avoid global state conflicts
+    /// between parallel test threads. Tests are run sequentially within this
+    /// function since they all share the TOKIO_RUNTIME_MANAGER global.
+    ///
+    /// Covers:
+    /// - "fragment_executor" routes to the DedicatedExecutor's gate (Req 2.2)
+    /// - "reduce" routes to the RuntimeManager's coordinator gate (Req 2.3)
+    /// - Unknown gate name logs warning and returns success (Req 2.4)
+    /// - Calling update before runtime init returns success (Req 2.6)
+    #[test]
+    fn test_ffi_gate_routing() {
+        // ── Test 1: update before runtime init returns success (Req 2.6) ──
+        shutdown_test_runtime(); // ensure clean state
+        let result = unsafe { call_update_gate("fragment_executor", 10) };
+        assert_eq!(result, 0, "FFI call should return success even before runtime init");
+
+        // ── Initialize runtime for remaining tests ──
+        init_test_runtime();
+        let mgr = get_rt_manager().expect("runtime should be initialized");
+
+        // ── Test 2: "fragment_executor" routes to CPU executor gate (Req 2.2) ──
+        {
+            let gate = mgr.cpu_executor().concurrency_gate().clone();
+            let initial_max = gate.max_permits();
+            let new_max = initial_max + 4;
+
+            let result = unsafe { call_update_gate("fragment_executor", new_max) };
+            assert_eq!(result, 0, "FFI call should return success for 'fragment_executor'");
+
+            // The resize is spawned on the IO runtime asynchronously.
+            // Wait briefly for it to complete.
+            std::thread::sleep(std::time::Duration::from_millis(200));
+
+            assert_eq!(
+                gate.max_permits(),
+                new_max,
+                "fragment_executor gate max_permits should be updated to {}",
+                new_max
+            );
+        }
+
+        // ── Test 3: "reduce" routes to coordinator gate (Req 2.3) ──
+        {
+            let gate = mgr.coordinator_gate().clone();
+            let initial_max = gate.max_permits();
+            let new_max = initial_max + 4;
+
+            let result = unsafe { call_update_gate("reduce", new_max) };
+            assert_eq!(result, 0, "FFI call should return success for 'reduce'");
+
+            // The resize is spawned on the IO runtime asynchronously.
+            // Wait briefly for it to complete.
+            std::thread::sleep(std::time::Duration::from_millis(200));
+
+            assert_eq!(
+                gate.max_permits(),
+                new_max,
+                "reduce gate max_permits should be updated to {}",
+                new_max
+            );
+        }
+
+        // ── Test 4: unknown gate name returns success without modifying gates (Req 2.4) ──
+        {
+            let fragment_executor_gate = mgr.cpu_executor().concurrency_gate().clone();
+            let reduce_executor_gate = mgr.coordinator_gate().clone();
+
+            let fragment_executor_max_before = fragment_executor_gate.max_permits();
+            let reduce_max_before = reduce_executor_gate.max_permits();
+
+            let result = unsafe { call_update_gate("unknown_gate", 99) };
+            assert_eq!(result, 0, "FFI call should return success even for unknown gate");
+
+            // Wait briefly to ensure no async resize was spawned
+            std::thread::sleep(std::time::Duration::from_millis(100));
+
+            // Neither gate should have been modified
+            assert_eq!(
+                fragment_executor_gate.max_permits(),
+                fragment_executor_max_before,
+                "fragment_executor gate should not be modified for unknown gate name"
+            );
+            assert_eq!(
+                reduce_executor_gate.max_permits(),
+                reduce_max_before,
+                "reduce gate should not be modified for unknown gate name"
+            );
+        }
+
+        // ── Cleanup ──
+        shutdown_test_runtime();
+    }
 }

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

@@ -19,7 +19,7 @@ pub struct RuntimeManager {
     pub io_monitor: RuntimeMonitor,
     pub cpu_monitor: Option<RuntimeMonitor>,
     /// Separate concurrency gate for coordinator-reduce execution.
-    /// Independent from the datanode gate (on DedicatedExecutor) to avoid
+    /// Independent from the fragment executor gate (on DedicatedExecutor) to avoid
     /// deadlock when shard streams and coordinator reduce run concurrently.
     coordinator_gate: Arc<ConcurrencyGate>,
 }
@@ -51,16 +51,16 @@ impl RuntimeManager {
                 register_io_runtime(Some(io_handle.clone()));
             });
 
-        // Datanode concurrency gate: limits concurrent partition tasks from shard scans.
+        // Fragment executor concurrency gate: limits concurrent partition tasks from shard scans.
         let datanode_max_concurrent = (cpu_threads as f64 * datanode_multiplier).max(1.0) as usize;
         let cpu_executor = DedicatedExecutor::new("datafusion-cpu", cpu_runtime_builder, datanode_max_concurrent);
 
         let cpu_monitor = cpu_executor
             .handle()
             .map(|h| RuntimeMonitor::new(&h));
 
-        // Coordinator concurrency gate: limits concurrent partition tasks from reduce execution.
-        // Separate from datanode gate to avoid deadlock (shard streams hold datanode permits
+        // Reduce concurrency gate: limits concurrent partition tasks from reduce execution.
+        // Separate from fragment executor gate to avoid deadlock (shard streams hold fragment executor permits
         // while coordinator reduce runs concurrently on single-node clusters).
         let coordinator_max_concurrent = (cpu_threads as f64 * coordinator_multiplier).max(1.0) as usize;
         let coordinator_gate = Arc::new(ConcurrencyGate::new(coordinator_max_concurrent));

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

@@ -5,7 +5,7 @@
 //! Stats packing helpers for the FFM `df_stats()` function.
 //!
 //! Packs Tokio runtime metrics and per-operation task monitor metrics
-//! into a `#[repr(C)]` `DfStatsBuffer` struct (304 bytes) for efficient
+//! into a `#[repr(C)]` `DfStatsBuffer` struct (336 bytes) for efficient
 //! transfer across the FFM boundary.
 //!
 //! ## Struct layout
@@ -18,8 +18,8 @@
 //! | `query_execution`    | `TaskMonitorRepr`    | 3 × i64 |
 //! | `stream_next`        | `TaskMonitorRepr`    | 3 × i64 |
 //! | `plan_setup`         | `TaskMonitorRepr`    | 3 × i64 |
-//! | `datanode_gate`      | `PartitionGateRepr`  | 4 × i64 |
-//! | `coordinator_gate`   | `PartitionGateRepr`  | 4 × i64 |
+//! | `fragment_executor_gate` | `PartitionGateRepr`  | 6 × i64 |
+//! | `reduce_executor_gate`            | `PartitionGateRepr`  | 6 × i64 |
 
 use tokio::runtime::Handle;
 use tokio_metrics::{RuntimeMonitor, TaskMonitor};
@@ -68,6 +68,8 @@ pub struct PartitionGateRepr {
     pub active_permits: i64,
     pub total_wait_duration_ms: i64,
     pub total_batches_started: i64,
+    pub poison_permits: i64,
+    pub target_max_permits: i64,
 }
 
 #[repr(C)]
@@ -78,19 +80,19 @@ pub struct DfStatsBuffer {
     pub query_execution: TaskMonitorRepr,
     pub stream_next: TaskMonitorRepr,
     pub plan_setup: TaskMonitorRepr,
-    pub datanode_gate: PartitionGateRepr,
-    pub coordinator_gate: PartitionGateRepr,
+    pub fragment_executor_gate: PartitionGateRepr,
+    pub reduce_executor_gate: PartitionGateRepr,
 }
 
 const _: () = assert!(std::mem::size_of::<RuntimeMetricsRepr>() == 9 * 8);
 const _: () = assert!(std::mem::size_of::<TaskMonitorRepr>() == 3 * 8);
-const _: () = assert!(std::mem::size_of::<PartitionGateRepr>() == 4 * 8);
-const _: () = assert!(std::mem::size_of::<DfStatsBuffer>() == 38 * 8);
+const _: () = assert!(std::mem::size_of::<PartitionGateRepr>() == 6 * 8);
+const _: () = assert!(std::mem::size_of::<DfStatsBuffer>() == 42 * 8);
 
 pub mod layout {
     use super::*;
     pub const BUFFER_BYTE_SIZE: usize = std::mem::size_of::<DfStatsBuffer>();
-    const _: () = assert!(BUFFER_BYTE_SIZE == 304);
+    const _: () = assert!(BUFFER_BYTE_SIZE == 336);
 }
 
 /// Snapshot a `RuntimeMonitor` and return a populated `RuntimeMetricsRepr`.
@@ -163,12 +165,16 @@ pub fn pack_task_monitor(monitor: &TaskMonitor) -> TaskMonitorRepr {
 /// | active_permits         | `gate.active_permits()`             |
 /// | total_wait_duration_ms | `gate.total_wait_ms()`              |
 /// | total_batches_started  | `gate.total_queries_admitted()`     |
+/// | poison_permits         | `gate.poison_permits_held()`        |
+/// | target_max_permits     | `gate.target_max_permits()`         |
 pub fn pack_partition_gate(gate: &ConcurrencyGate) -> PartitionGateRepr {
     PartitionGateRepr {
         max_permits: gate.max_permits() as i64,
         active_permits: gate.active_permits() as i64,
         total_wait_duration_ms: gate.total_wait_ms() as i64,
         total_batches_started: gate.total_queries_admitted() as i64,
+        poison_permits: gate.poison_permits_held() as i64,
+        target_max_permits: gate.target_max_permits() as i64,
     }
 }
 
@@ -253,14 +259,14 @@ mod tests {
             query_execution: pack_task_monitor(query_execution_monitor()),
             stream_next: pack_task_monitor(stream_next_monitor()),
             plan_setup: pack_task_monitor(plan_setup_monitor()),
-            datanode_gate: pack_partition_gate(mgr.cpu_executor.concurrency_gate()),
-            coordinator_gate: pack_partition_gate(mgr.coordinator_gate()),
+            fragment_executor_gate: pack_partition_gate(mgr.cpu_executor.concurrency_gate()),
+            reduce_executor_gate: pack_partition_gate(mgr.coordinator_gate()),
         };
 
-        assert_eq!(layout::BUFFER_BYTE_SIZE, 304);
+        assert_eq!(layout::BUFFER_BYTE_SIZE, 336);
         assert!(buf.io_runtime.workers_count > 0, "IO runtime workers_count should be > 0, got {}", buf.io_runtime.workers_count);
-        assert!(buf.datanode_gate.max_permits > 0, "datanode_gate max_permits should be > 0, got {}", buf.datanode_gate.max_permits);
-        assert!(buf.coordinator_gate.max_permits > 0, "coordinator_gate max_permits should be > 0, got {}", buf.coordinator_gate.max_permits);
+        assert!(buf.fragment_executor_gate.max_permits > 0, "fragment_executor_gate max_permits should be > 0, got {}", buf.fragment_executor_gate.max_permits);
+        assert!(buf.reduce_executor_gate.max_permits > 0, "reduce_executor_gate max_permits should be > 0, got {}", buf.reduce_executor_gate.max_permits);
 
         if mgr.cpu_monitor.is_some() {
             assert!(buf.cpu_runtime.workers_count > 0, "CPU runtime workers_count should be > 0, got {}", buf.cpu_runtime.workers_count);
@@ -273,9 +279,9 @@ mod tests {
     #[test]
     fn test_df_stats_buffer_too_small() {
         // Verify that the buffer size assertion holds
-        assert_eq!(std::mem::size_of::<DfStatsBuffer>(), 304);
-        assert_eq!(layout::BUFFER_BYTE_SIZE, 304);
-        // A buffer smaller than 304 bytes should be rejected by df_stats.
+        assert_eq!(std::mem::size_of::<DfStatsBuffer>(), 336);
+        assert_eq!(layout::BUFFER_BYTE_SIZE, 336);
+        // A buffer smaller than 336 bytes should be rejected by df_stats.
         // We can't call df_stats directly without a runtime manager,
         // but we verify the constant is correct.
         assert!(layout::BUFFER_BYTE_SIZE > 0);

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

@@ -420,6 +420,19 @@ public Collection<Object> createComponents(
         clusterService.getClusterSettings()
             .addSettingsUpdateConsumer(DATAFUSION_MEMORY_GUARD_EXECUTION_CRITICAL_THRESHOLD, v -> updateMemoryGuardThresholds());
 
+        // Wire dynamic concurrency gate multiplier settings
+        int cpuThreads = DataFusionService.cpuThreadCount();
+
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(DatafusionSettings.CONCURRENCY_DATANODE_MULTIPLIER, multiplier -> {
+            int newMax = Math.max(1, (int) (cpuThreads * multiplier));
+            NativeBridge.updateConcurrencyGate("fragment_executor", newMax);
+        });
+
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(DatafusionSettings.CONCURRENCY_COORDINATOR_MULTIPLIER, multiplier -> {
+            int newMax = Math.max(1, (int) (cpuThreads * multiplier));
+            NativeBridge.updateConcurrencyGate("reduce", newMax);
+        });
+
         // Apply initial values
         NativeBridge.setMinTargetPartitions(DATAFUSION_MIN_TARGET_PARTITIONS.get(settings));
         NativeBridge.setReduceTargetPartitions(DATAFUSION_REDUCE_TARGET_PARTITIONS.get(settings));

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java

@@ -63,6 +63,14 @@ public static Builder builder() {
         return new Builder();
     }
 
+    /**
+     * Returns the number of available CPU threads used for the dedicated executor.
+     * This is the value used to compute concurrency gate permit counts (cpu_threads × multiplier).
+     */
+    public static int cpuThreadCount() {
+        return Runtime.getRuntime().availableProcessors();
+    }
+
     @Override
     protected void doStart() {
         logger.debug("Starting DataFusion service");

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionService.java

@@ -191,22 +191,24 @@ public final class DatafusionSettings {
 
     // ── Concurrency gate settings ──
 
-    /** Datanode concurrency gate multiplier: max concurrent partition-equivalents = cpu_threads × multiplier. */
+    /** Fragment executor concurrency gate multiplier: max concurrent partition-equivalents = cpu_threads × multiplier. */
     public static final Setting<Double> CONCURRENCY_DATANODE_MULTIPLIER = Setting.doubleSetting(
-        "datafusion.concurrency.datanode_multiplier",
+        "datafusion.concurrency.fragment_executor_multiplier",
         1.5,
         0.1,
         10.0,
-        Setting.Property.NodeScope
+        Setting.Property.NodeScope,
+        Setting.Property.Dynamic
     );
 
-    /** Coordinator concurrency gate multiplier: max concurrent partition-equivalents = cpu_threads × multiplier. */
+    /** Reduce concurrency gate multiplier: max concurrent partition-equivalents = cpu_threads × multiplier. */
     public static final Setting<Double> CONCURRENCY_COORDINATOR_MULTIPLIER = Setting.doubleSetting(
-        "datafusion.concurrency.coordinator_multiplier",
+        "datafusion.concurrency.reduce_executor_multiplier",
         1.5,
         0.1,
         10.0,
-        Setting.Property.NodeScope
+        Setting.Property.NodeScope,
+        Setting.Property.Dynamic
     );
 
     // Query strategy constants