| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you under the Apache License, Version 2.0 (the |
| // "License"); you may not use this file except in compliance |
| // with the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, |
| // software distributed under the License is distributed on an |
| // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| // KIND, either express or implied. See the License for the |
| // specific language governing permissions and limitations |
| // under the License. |
| |
| //! Tokio runtime metrics, gated behind the `runtime-metrics` Cargo feature. |
| //! |
| //! `tokio-metrics` only compiles when `--cfg tokio_unstable` is set, so the |
| //! whole module is conditional on the feature flag. The feature-off variant |
| //! still exists; `runtime_stats()` returns an `Err` whose message points at |
| //! the rebuild command. The Java surface (`SessionContext.runtimeStats()`) |
| //! is therefore stable across both build configurations. |
| //! |
| //! Field layout returned across the FFI boundary (in this exact order, all |
| //! `i64`): |
| //! 0 numWorkers |
| //! 1 liveTasksCount |
| //! 2 globalQueueDepth |
| //! 3 elapsedNanos |
| //! 4 totalBusyNanos |
| //! 5 totalParkCount |
| //! 6 totalPollsCount |
| //! 7 totalNoopCount |
| //! 8 totalStealCount |
| //! 9 totalLocalScheduleCount |
| //! 10 totalOverflowCount |
| |
| #[cfg(not(feature = "runtime-metrics"))] |
| use datafusion_jni_common::errors::JniResult; |
| |
| /// Number of i64 values in the snapshot array; kept here so the Java side and |
| /// the feature-off stub agree on the layout. |
| pub const STATS_FIELD_COUNT: usize = 11; |
| |
| #[cfg(feature = "runtime-metrics")] |
| mod imp { |
| use std::sync::{Mutex, OnceLock}; |
| use std::time::Duration; |
| use tokio_metrics::{RuntimeIntervals, RuntimeMonitor}; |
| |
| use super::STATS_FIELD_COUNT; |
| use datafusion_jni_common::errors::JniResult; |
| |
| /// `RuntimeMonitor::intervals().next()` returns *delta* metrics covering |
| /// the period since the previous call (or, on the very first call, since |
| /// the iterator was constructed). To present the documented |
| /// "totals since runtime start" semantic to Java callers, we own the |
| /// iterator process-wide and accumulate every field that's documented as |
| /// monotonic. |
| /// |
| /// Snapshot fields (`workers_count`, `live_tasks_count`, |
| /// `global_queue_depth`) are point-in-time and pass through unchanged. |
| struct RuntimeAccumulator { |
| intervals: RuntimeIntervals, |
| elapsed: Duration, |
| total_busy: Duration, |
| total_park_count: u64, |
| total_polls_count: u64, |
| total_noop_count: u64, |
| total_steal_count: u64, |
| total_local_schedule_count: u64, |
| total_overflow_count: u64, |
| } |
| |
| static ACC: OnceLock<Mutex<RuntimeAccumulator>> = OnceLock::new(); |
| |
| fn accumulator() -> &'static Mutex<RuntimeAccumulator> { |
| // Warm `crate::runtime()` first so the shared Tokio runtime exists |
| // and its OnceLock initializer has already populated ACC via |
| // `runtime_metrics::init`. We must NOT call `runtime()` from inside |
| // `ACC.get_or_init`: that would acquire ACC's once-init slot, then |
| // (on the very first call) recurse through `RT.get_or_init` -> |
| // `init()` -> `ACC.get_or_init`, which is same-thread reentrancy on |
| // the same OnceLock and deadlocks indefinitely. |
| let _ = crate::runtime(); |
| // After `runtime()` returns, ACC is guaranteed populated. The |
| // defensive `get_or_init` covers a future build path that constructs |
| // the runtime without going through `crate::runtime` (none today); |
| // it is safe to invoke from here because no init lock is held. |
| ACC.get_or_init(build_accumulator) |
| } |
| |
| fn build_accumulator() -> Mutex<RuntimeAccumulator> { |
| let handle = crate::runtime().handle().clone(); |
| let monitor = RuntimeMonitor::new(&handle); |
| Mutex::new(RuntimeAccumulator { |
| intervals: monitor.intervals(), |
| elapsed: Duration::ZERO, |
| total_busy: Duration::ZERO, |
| total_park_count: 0, |
| total_polls_count: 0, |
| total_noop_count: 0, |
| total_steal_count: 0, |
| total_local_schedule_count: 0, |
| total_overflow_count: 0, |
| }) |
| } |
| |
| /// Eagerly construct the runtime-metrics accumulator. Called from |
| /// `crate::runtime()` immediately after the shared Tokio runtime is |
| /// created so the accumulator's interval baseline matches runtime start. |
| /// Without this, `RuntimeMonitor` would only start sampling at the first |
| /// `runtimeStats()` call, silently dropping every poll/busy nanosecond |
| /// that happened before that. |
| pub fn init(handle: &tokio::runtime::Handle) { |
| ACC.get_or_init(|| { |
| let monitor = RuntimeMonitor::new(handle); |
| Mutex::new(RuntimeAccumulator { |
| intervals: monitor.intervals(), |
| elapsed: Duration::ZERO, |
| total_busy: Duration::ZERO, |
| total_park_count: 0, |
| total_polls_count: 0, |
| total_noop_count: 0, |
| total_steal_count: 0, |
| total_local_schedule_count: 0, |
| total_overflow_count: 0, |
| }) |
| }); |
| } |
| |
| pub fn runtime_stats() -> JniResult<[i64; STATS_FIELD_COUNT]> { |
| let mut acc = accumulator() |
| .lock() |
| .map_err(|_| "runtime-metrics accumulator lock poisoned")?; |
| let delta = acc |
| .intervals |
| .next() |
| .ok_or("tokio-metrics RuntimeMonitor returned no interval")?; |
| // Workers count and the point-in-time fields are not deltas; they are |
| // the runtime's current state. Pass through directly. |
| let workers_count = delta.workers_count; |
| let live_tasks_count = delta.live_tasks_count; |
| let global_queue_depth = delta.global_queue_depth; |
| // Accumulate the deltas. |
| acc.elapsed = acc.elapsed.saturating_add(delta.elapsed); |
| acc.total_busy = acc.total_busy.saturating_add(delta.total_busy_duration); |
| acc.total_park_count = acc.total_park_count.saturating_add(delta.total_park_count); |
| acc.total_polls_count = acc |
| .total_polls_count |
| .saturating_add(delta.total_polls_count); |
| acc.total_noop_count = acc.total_noop_count.saturating_add(delta.total_noop_count); |
| acc.total_steal_count = acc |
| .total_steal_count |
| .saturating_add(delta.total_steal_count); |
| acc.total_local_schedule_count = acc |
| .total_local_schedule_count |
| .saturating_add(delta.total_local_schedule_count); |
| acc.total_overflow_count = acc |
| .total_overflow_count |
| .saturating_add(delta.total_overflow_count); |
| |
| Ok([ |
| workers_count as i64, |
| live_tasks_count as i64, |
| global_queue_depth as i64, |
| i128_to_i64_saturating(acc.elapsed.as_nanos() as i128), |
| i128_to_i64_saturating(acc.total_busy.as_nanos() as i128), |
| acc.total_park_count as i64, |
| acc.total_polls_count as i64, |
| acc.total_noop_count as i64, |
| acc.total_steal_count as i64, |
| acc.total_local_schedule_count as i64, |
| acc.total_overflow_count as i64, |
| ]) |
| } |
| |
| fn i128_to_i64_saturating(v: i128) -> i64 { |
| v.clamp(i64::MIN as i128, i64::MAX as i128) as i64 |
| } |
| } |
| |
| #[cfg(feature = "runtime-metrics")] |
| pub use imp::{init, runtime_stats}; |
| |
| /// No-op when the `runtime-metrics` feature is off: there is no monitor to |
| /// initialise. Kept as a regular function (rather than a `cfg` annotation at |
| /// every call site) so the runtime construction path stays unconditional. |
| #[cfg(not(feature = "runtime-metrics"))] |
| pub fn init(_handle: &tokio::runtime::Handle) {} |
| |
| #[cfg(not(feature = "runtime-metrics"))] |
| pub fn runtime_stats() -> JniResult<[i64; STATS_FIELD_COUNT]> { |
| Err( |
| "datafusion-jni was built without the `runtime-metrics` Cargo feature; \ |
| rebuild the native crate with \ |
| `RUSTFLAGS=\"--cfg tokio_unstable\" cargo build -p datafusion-jni --features runtime-metrics` \ |
| to enable SessionContext.runtimeStats" |
| .into(), |
| ) |
| } |
