From db0c4ce4cbfc51c19a2568553c58774740ebc653 Mon Sep 17 00:00:00 2001 From: jennychen Date: Mon, 31 Aug 2026 18:30:46 +0800 Subject: [PATCH 1/2] [AURON #2189] Periodically publish native metrics for long-running Flink tasks --- .../configuration/AuronConfiguration.java | 14 ++ .../apache/auron/jni/AuronAdaptorTest.java | 1 + .../kafka/AuronKafkaSourceFunction.java | 23 +-- .../auron/flink/metric/FlinkMetricNode.java | 36 ++++ .../operator/FlinkAuronCalcOperator.java | 27 +-- .../AuronKafkaSourceFunctionMergeTest.java | 54 +++++ native-engine/auron-jni-bridge/src/conf.rs | 1 + native-engine/auron/src/metrics.rs | 189 ++++++++++++++++-- native-engine/auron/src/rt.rs | 83 +++++++- 9 files changed, 355 insertions(+), 73 deletions(-) diff --git a/auron-core/src/main/java/org/apache/auron/configuration/AuronConfiguration.java b/auron-core/src/main/java/org/apache/auron/configuration/AuronConfiguration.java index a2bd08d4b..391ed0085 100644 --- a/auron-core/src/main/java/org/apache/auron/configuration/AuronConfiguration.java +++ b/auron-core/src/main/java/org/apache/auron/configuration/AuronConfiguration.java @@ -65,6 +65,20 @@ public abstract class AuronConfiguration { + "if not configured, the default value of 1 is used.") .withDefaultValue(1); + /** + * How often the native runtime publishes DataFusion metrics to {@code MetricNode} while a task + * is running. Java {@code MetricNode.add} is incremental, so the native side sends positive + * deltas. {@code 0} disables the timer and publishes only when the native runtime finalizes + * (the historical batch behavior). + */ + public static final ConfigOption METRICS_UPDATE_INTERVAL_MS = new ConfigOption<>(Long.class) + .withKey("auron.metrics.update.interval.ms") + .withCategory("Runtime Configuration") + .withDescription("Interval in milliseconds for publishing native execution metrics to MetricNode " + + "during a running task. Set to 0 to publish only when the native runtime finalizes. " + + "Default is 1000ms so long-running Flink tasks report live counters.") + .withDefaultValue(1000L); + public abstract Optional getOptional(ConfigOption option); public T get(ConfigOption option) { diff --git a/auron-core/src/test/java/org/apache/auron/jni/AuronAdaptorTest.java b/auron-core/src/test/java/org/apache/auron/jni/AuronAdaptorTest.java index edf317c1d..8c11fa49a 100644 --- a/auron-core/src/test/java/org/apache/auron/jni/AuronAdaptorTest.java +++ b/auron-core/src/test/java/org/apache/auron/jni/AuronAdaptorTest.java @@ -36,6 +36,7 @@ public void testRetrieveConfigWithAuronAdaptor() { assertEquals(auronConfig.getInteger(AuronConfiguration.BATCH_SIZE), 10000); assertEquals(auronConfig.getDouble(AuronConfiguration.MEMORY_FRACTION), 0.6, 0.0); assertEquals(auronConfig.getString(AuronConfiguration.NATIVE_LOG_LEVEL), "info"); + assertEquals(auronConfig.getLong(AuronConfiguration.METRICS_UPDATE_INTERVAL_MS), 1000L); } @Test diff --git a/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunction.java b/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunction.java index b4d4dfa39..b15079e26 100644 --- a/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunction.java +++ b/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunction.java @@ -27,6 +27,7 @@ import org.apache.auron.flink.arrow.FlinkArrowReader; import org.apache.auron.flink.arrow.FlinkArrowUtils; import org.apache.auron.flink.configuration.FlinkAuronConfiguration; +import org.apache.auron.flink.metric.FlinkMetricNode; import org.apache.auron.flink.runtime.operator.AuronPlanTreeRewriter; import org.apache.auron.flink.runtime.operator.FlinkAuronFunction; import org.apache.auron.flink.table.data.AuronColumnarRowData; @@ -56,7 +57,6 @@ import org.apache.flink.api.common.typeinfo.TypeInformation; import org.apache.flink.api.java.tuple.Tuple2; import org.apache.flink.configuration.Configuration; -import org.apache.flink.metrics.Counter; import org.apache.flink.metrics.MetricGroup; import org.apache.flink.runtime.state.FunctionInitializationContext; import org.apache.flink.runtime.state.FunctionSnapshotContext; @@ -130,7 +130,6 @@ public class AuronKafkaSourceFunction extends RichParallelSourceFunction sourceContext) throws Exception { - metricGroup = getRuntimeContext().getMetricGroup(); - final Map flinkCounters = new HashMap<>(); - - nativeMetric = new MetricNode(new ArrayList<>()) { - @Override - public void add(String name, long value) { - // Integration with Flink metrics - Counter counter = flinkCounters.get(name); - if (counter == null) { - counter = metricGroup.counter(name); - flinkCounters.put(name, counter); - } - counter.inc(value); - LOG.debug("Metric Auron Source: {} = {}", name, value); - } - }; + // Mirror physicalPlanNode (KafkaScan, or fused Project[Filter?[KafkaScan]]) so native + // periodic metric walks via MetricNode.getChild(i) do not IndexOutOfBounds. + nativeMetric = + FlinkMetricNode.fromPlan(physicalPlanNode, getRuntimeContext().getMetricGroup()); // The native output carries [meta, logical], where logical is the projected output when a // merged Calc plan is active and the original output otherwise. The metadata column count // and per-field positions both derive from KAFKA_AURON_META_FIELDS so adding or reordering diff --git a/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/metric/FlinkMetricNode.java b/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/metric/FlinkMetricNode.java index 361f52f78..93519c454 100644 --- a/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/metric/FlinkMetricNode.java +++ b/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/metric/FlinkMetricNode.java @@ -16,10 +16,12 @@ */ package org.apache.auron.flink.metric; +import java.util.Collections; import java.util.List; import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; import org.apache.auron.metric.MetricNode; +import org.apache.auron.protobuf.PhysicalPlanNode; import org.apache.flink.metrics.Counter; import org.apache.flink.metrics.MetricGroup; @@ -54,6 +56,40 @@ public FlinkMetricNode(MetricGroup metricGroup, List children) { this.metricGroup = Objects.requireNonNull(metricGroup, "metricGroup"); } + /** + * Recursively constructs a {@link FlinkMetricNode} whose shape mirrors {@code node}'s plan + * tree. Native code walks the execution plan during periodic metric updates and at + * finalization, indexing into the parallel metric tree via {@link MetricNode#getChild(int)}; + * an empty children list would throw {@link IndexOutOfBoundsException} on the first {@code + * getChild(0)} call. All levels share the same Flink {@link MetricGroup} so named counters + * aggregate at the operator scope. + * + *

Supported leaves are {@code FFIReader} (standalone Calc) and {@code KafkaScan} (Kafka + * source, including after Calc fusion). Supported unary nodes are {@code Projection} and + * {@code Filter}. + * + * @param node the physical plan whose shape the metric tree must match; must not be null + * @param metricGroup the Flink metric group all tree levels register counters against; must + * not be null + * @return the root of the mirrored metric tree + */ + public static FlinkMetricNode fromPlan(PhysicalPlanNode node, MetricGroup metricGroup) { + Objects.requireNonNull(node, "node"); + Objects.requireNonNull(metricGroup, "metricGroup"); + final List children; + if (node.hasFfiReader() || node.hasKafkaScan()) { + children = Collections.emptyList(); + } else if (node.hasProjection()) { + children = Collections.singletonList(fromPlan(node.getProjection().getInput(), metricGroup)); + } else if (node.hasFilter()) { + children = Collections.singletonList(fromPlan(node.getFilter().getInput(), metricGroup)); + } else { + throw new IllegalArgumentException( + "Unexpected plan node type for metric tree: " + node.getPhysicalPlanTypeCase()); + } + return new FlinkMetricNode(metricGroup, children); + } + /** * Adds {@code value} to the Flink {@link Counter} named {@code name}, creating the counter on * first use. Non-positive values are ignored. diff --git a/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/runtime/operator/FlinkAuronCalcOperator.java b/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/runtime/operator/FlinkAuronCalcOperator.java index 371418379..770951b21 100644 --- a/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/runtime/operator/FlinkAuronCalcOperator.java +++ b/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/runtime/operator/FlinkAuronCalcOperator.java @@ -153,7 +153,7 @@ public void open() throws Exception { String opIdWithSubtask = rc.getOperatorUniqueID() + "-" + rc.getIndexOfThisSubtask(); this.childAllocator = FlinkArrowUtils.createChildAllocator("FlinkAuronCalc-" + opIdWithSubtask); - this.metricNode = buildMetricTree(plan, getMetricGroup()); + this.metricNode = FlinkMetricNode.fromPlan(plan, getMetricGroup()); this.exporter = new FlinkArrowFFIExporter(childAllocator, inputRowType, BATCH_ROW_LIMIT); // UUID disambiguates re-runs after operator restart so a stale registration cannot // collide with the new one. @@ -318,31 +318,6 @@ static PhysicalPlanNode injectFfiReaderLeaf(PhysicalPlanNode node, String resour + "Project[FFIReader] / Filter[FFIReader] / FFIReader shape; got: "); } - /** - * Recursively constructs a {@link FlinkMetricNode} whose shape mirrors {@code node}'s plan - * tree. Native code walks the plan tree at finalization time and indexes into the parallel - * metric tree via {@link org.apache.auron.metric.MetricNode#getChild(int)}; an empty children - * list would throw {@link IndexOutOfBoundsException} on the first {@code getChild(0)} call. - * All levels share the same Flink {@link org.apache.flink.metrics.MetricGroup} so named - * counters aggregate at the operator scope. - */ - private static FlinkMetricNode buildMetricTree(PhysicalPlanNode node, org.apache.flink.metrics.MetricGroup mg) { - final List children; - if (node.hasFfiReader()) { - children = Collections.emptyList(); - } else if (node.hasProjection()) { - children = Collections.singletonList( - buildMetricTree(node.getProjection().getInput(), mg)); - } else if (node.hasFilter()) { - children = - Collections.singletonList(buildMetricTree(node.getFilter().getInput(), mg)); - } else { - throw new IllegalArgumentException( - "Unexpected plan node type for metric tree: " + node.getPhysicalPlanTypeCase()); - } - return new FlinkMetricNode(mg, children); - } - // ==================================================================== // SupportsAuronNative // ==================================================================== diff --git a/auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunctionMergeTest.java b/auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunctionMergeTest.java index 86110328c..485f92012 100644 --- a/auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunctionMergeTest.java +++ b/auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/connector/kafka/AuronKafkaSourceFunctionMergeTest.java @@ -18,11 +18,14 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.Properties; +import org.apache.auron.flink.metric.FlinkMetricNode; import org.apache.auron.flink.utils.SchemaConverters; +import org.apache.auron.metric.MetricNode; import org.apache.auron.protobuf.FFIReaderExecNode; import org.apache.auron.protobuf.FilterExecNode; import org.apache.auron.protobuf.KafkaScanExecNode; @@ -31,6 +34,7 @@ import org.apache.auron.protobuf.PhysicalPlanNode; import org.apache.auron.protobuf.ProjectionExecNode; import org.apache.flink.api.common.eventtime.WatermarkStrategy; +import org.apache.flink.metrics.groups.UnregisteredMetricsGroup; import org.apache.flink.table.data.RowData; import org.apache.flink.table.types.logical.IntType; import org.apache.flink.table.types.logical.LogicalType; @@ -228,4 +232,54 @@ void testApplyMergedCalcPlanReturnsSourcePlanWhenNoPlanStaged() { fn.setWatermarkStrategy(WatermarkStrategy.forMonotonousTimestamps()); assertEquals(source, fn.applyMergedCalcPlan(source)); } + + @Test + void testMetricTreeForKafkaScanLeafHasNoChildren() { + FlinkMetricNode root = FlinkMetricNode.fromPlan(kafkaScan(), new UnregisteredMetricsGroup()); + assertLeaf(root); + } + + @Test + void testMetricTreeForFusedProjectKafkaScan() { + PhysicalPlanNode fused = + AuronKafkaSourceFunction.buildMergedPlan(logicalProjection(ffiReaderPlaceholder()), kafkaScan()); + + FlinkMetricNode root = FlinkMetricNode.fromPlan(fused, new UnregisteredMetricsGroup()); + MetricNode scan = root.getChild(0); + assertLeaf(scan); + assertThrows(IndexOutOfBoundsException.class, () -> root.getChild(1)); + } + + @Test + void testMetricTreeForFusedProjectFilterKafkaScan() { + PhysicalPlanNode filter = PhysicalPlanNode.newBuilder() + .setFilter(FilterExecNode.newBuilder() + .setInput(ffiReaderPlaceholder()) + .build()) + .build(); + PhysicalPlanNode fused = AuronKafkaSourceFunction.buildMergedPlan(logicalProjection(filter), kafkaScan()); + + FlinkMetricNode root = FlinkMetricNode.fromPlan(fused, new UnregisteredMetricsGroup()); + MetricNode filterNode = root.getChild(0); + MetricNode scan = filterNode.getChild(0); + assertLeaf(scan); + assertThrows(IndexOutOfBoundsException.class, () -> root.getChild(1)); + assertThrows(IndexOutOfBoundsException.class, () -> filterNode.getChild(1)); + } + + @Test + void testMetricTreeForApplyMergedCalcPlanOutput() { + AuronKafkaSourceFunction fn = newFunction(); + RowType projected = RowType.of(new LogicalType[] {new IntType()}, new String[] {"int"}); + fn.setMergedCalcPlan(logicalProjection(ffiReaderPlaceholder()), projected); + + PhysicalPlanNode fused = fn.applyMergedCalcPlan(kafkaScan()); + FlinkMetricNode root = FlinkMetricNode.fromPlan(fused, new UnregisteredMetricsGroup()); + assertLeaf(root.getChild(0)); + } + + private static void assertLeaf(MetricNode node) { + assertNotNull(node); + assertThrows(IndexOutOfBoundsException.class, () -> node.getChild(0)); + } } diff --git a/native-engine/auron-jni-bridge/src/conf.rs b/native-engine/auron-jni-bridge/src/conf.rs index 9f160beb1..a68f94a2e 100644 --- a/native-engine/auron-jni-bridge/src/conf.rs +++ b/native-engine/auron-jni-bridge/src/conf.rs @@ -63,6 +63,7 @@ define_conf!(BooleanConf, ORC_SCHEMA_CASE_SENSITIVE); define_conf!(IntConf, UDAF_FALLBACK_NUM_UDAFS_TRIGGER_SORT_AGG); define_conf!(BooleanConf, PARSE_JSON_ERROR_FALLBACK); define_conf!(StringConf, NATIVE_LOG_LEVEL); +define_conf!(LongConf, METRICS_UPDATE_INTERVAL_MS); pub trait BooleanConf { fn key(&self) -> &'static str; diff --git a/native-engine/auron/src/metrics.rs b/native-engine/auron/src/metrics.rs index 30d957a39..d59edbc0e 100644 --- a/native-engine/auron/src/metrics.rs +++ b/native-engine/auron/src/metrics.rs @@ -13,46 +13,199 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc}; use auron_jni_bridge::{jni_call, jni_new_string}; use datafusion::{common::Result, physical_plan::ExecutionPlan}; use jni::objects::JObject; +/// Last-published absolute metric values, keyed by plan-node path + metric name +/// (for example `output_rows` at the root or `0.1.elapsed_compute`). +pub type MetricSnapshot = HashMap; + pub fn update_metric_node( metric_node: JObject, execution_plan: Arc, + snapshot: &mut MetricSnapshot, +) -> Result<()> { + update_metric_node_at(metric_node, execution_plan, snapshot, "") +} + +fn update_metric_node_at( + metric_node: JObject, + execution_plan: Arc, + snapshot: &mut MetricSnapshot, + node_path: &str, ) -> Result<()> { if metric_node.is_null() { return Ok(()); } - // update current node - update_metrics( - metric_node, - &execution_plan - .metrics() - .unwrap_or_default() - .iter() - .map(|m| m.value()) - .map(|m| (m.name(), m.as_usize() as i64)) - .collect::>(), - )?; + // Bind MetricsSet so metric name &str values outlive this call. + let metrics_set = execution_plan.metrics().unwrap_or_default(); + let metric_values: Vec<(&str, i64)> = metrics_set + .iter() + .map(|m| m.value()) + .map(|m| (m.name(), m.as_usize() as i64)) + .collect(); + let deltas = compute_positive_deltas(snapshot, node_path, &metric_values); + update_metrics(metric_node, &deltas)?; - // update children nodes for (i, &child_plan) in execution_plan.children().iter().enumerate() { let child_metric_node = jni_call!( MetricNode(metric_node).getChild(i as i32) -> JObject )?; - update_metric_node(child_metric_node.as_obj(), child_plan.clone())?; + update_metric_node_at( + child_metric_node.as_obj(), + child_plan.clone(), + snapshot, + &child_node_path(node_path, i), + )?; } Ok(()) } -fn update_metrics(metric_node: JObject, metric_values: &[(&str, i64)]) -> Result<()> { - for &(name, value) in metric_values { - let jname = jni_new_string!(&name)?; - jni_call!(MetricNode(metric_node).add(jname.as_obj(), value) -> ())?; +fn update_metrics(metric_node: JObject, metric_values: &[(String, i64)]) -> Result<()> { + for (name, value) in metric_values { + let jname = jni_new_string!(name)?; + jni_call!(MetricNode(metric_node).add(jname.as_obj(), *value) -> ())?; } Ok(()) } + +fn snapshot_key(node_path: &str, name: &str) -> String { + if node_path.is_empty() { + name.to_string() + } else { + format!("{node_path}.{name}") + } +} + +fn child_node_path(node_path: &str, child_index: usize) -> String { + if node_path.is_empty() { + child_index.to_string() + } else { + format!("{node_path}.{child_index}") + } +} + +/// Aggregates same-name metrics on one node, then returns JNI deltas for values +/// that increased since the last publish. Non-positive deltas are skipped and +/// do not update the snapshot (so a later rebound is not over-counted). +fn compute_positive_deltas( + snapshot: &mut MetricSnapshot, + node_path: &str, + metric_values: &[(&str, i64)], +) -> Vec<(String, i64)> { + let mut current_by_name: HashMap<&str, i64> = HashMap::new(); + for &(name, value) in metric_values { + *current_by_name.entry(name).or_insert(0) += value; + } + + let mut deltas = Vec::new(); + for (name, current) in current_by_name { + let key = snapshot_key(node_path, name); + let last = snapshot.get(&key).copied().unwrap_or(0); + let delta = current - last; + if delta > 0 { + snapshot.insert(key, current); + deltas.push((name.to_string(), delta)); + } + } + deltas +} + +#[cfg(test)] +mod tests { + use super::*; + + fn sorted(mut deltas: Vec<(String, i64)>) -> Vec<(String, i64)> { + deltas.sort_by(|a, b| a.0.cmp(&b.0)); + deltas + } + + #[test] + fn first_publish_sends_absolute_as_delta_from_zero() { + let mut snapshot = MetricSnapshot::new(); + let deltas = compute_positive_deltas( + &mut snapshot, + "", + &[("output_rows", 10), ("elapsed_compute", 3)], + ); + assert_eq!( + sorted(deltas), + vec![ + ("elapsed_compute".to_string(), 3), + ("output_rows".to_string(), 10), + ] + ); + assert_eq!(snapshot.get("output_rows"), Some(&10)); + assert_eq!(snapshot.get("elapsed_compute"), Some(&3)); + } + + #[test] + fn second_publish_sends_only_increase() { + let mut snapshot = MetricSnapshot::new(); + let _ = compute_positive_deltas(&mut snapshot, "", &[("output_rows", 10)]); + let deltas = compute_positive_deltas(&mut snapshot, "", &[("output_rows", 15)]); + assert_eq!(deltas, vec![("output_rows".to_string(), 5)]); + assert_eq!(snapshot.get("output_rows"), Some(&15)); + } + + #[test] + fn unchanged_or_non_positive_is_skipped() { + let mut snapshot = MetricSnapshot::new(); + let _ = compute_positive_deltas(&mut snapshot, "", &[("output_rows", 10)]); + let deltas = compute_positive_deltas( + &mut snapshot, + "", + &[("output_rows", 10), ("spilled", 0), ("gauge", -1)], + ); + assert!(deltas.is_empty()); + assert_eq!(snapshot.get("output_rows"), Some(&10)); + assert!(!snapshot.contains_key("spilled")); + assert!(!snapshot.contains_key("gauge")); + } + + #[test] + fn gauge_dip_does_not_change_last_published() { + let mut snapshot = MetricSnapshot::new(); + let _ = compute_positive_deltas(&mut snapshot, "", &[("mem", 100)]); + let dip = compute_positive_deltas(&mut snapshot, "", &[("mem", 40)]); + assert!(dip.is_empty()); + assert_eq!(snapshot.get("mem"), Some(&100)); + let rebound = compute_positive_deltas(&mut snapshot, "", &[("mem", 120)]); + assert_eq!(rebound, vec![("mem".to_string(), 20)]); + assert_eq!(snapshot.get("mem"), Some(&120)); + } + + #[test] + fn same_name_on_one_node_is_summed() { + let mut snapshot = MetricSnapshot::new(); + let deltas = + compute_positive_deltas(&mut snapshot, "", &[("output_rows", 4), ("output_rows", 6)]); + assert_eq!(deltas, vec![("output_rows".to_string(), 10)]); + let again = + compute_positive_deltas(&mut snapshot, "", &[("output_rows", 4), ("output_rows", 7)]); + assert_eq!(again, vec![("output_rows".to_string(), 1)]); + } + + #[test] + fn node_paths_are_independent() { + let mut snapshot = MetricSnapshot::new(); + let root = compute_positive_deltas(&mut snapshot, "", &[("output_rows", 10)]); + let child = compute_positive_deltas(&mut snapshot, "0", &[("output_rows", 3)]); + assert_eq!(root, vec![("output_rows".to_string(), 10)]); + assert_eq!(child, vec![("output_rows".to_string(), 3)]); + assert_eq!(snapshot.get("output_rows"), Some(&10)); + assert_eq!(snapshot.get("0.output_rows"), Some(&3)); + } + + #[test] + fn snapshot_key_and_child_path() { + assert_eq!(snapshot_key("", "output_rows"), "output_rows"); + assert_eq!(snapshot_key("0.1", "output_rows"), "0.1.output_rows"); + assert_eq!(child_node_path("", 0), "0"); + assert_eq!(child_node_path("0", 1), "0.1"); + } +} diff --git a/native-engine/auron/src/rt.rs b/native-engine/auron/src/rt.rs index 932ae3136..e23c85f90 100644 --- a/native-engine/auron/src/rt.rs +++ b/native-engine/auron/src/rt.rs @@ -21,6 +21,7 @@ use std::{ atomic::{AtomicBool, Ordering}, mpsc::Receiver, }, + time::Duration, }; use arrow::{ @@ -29,7 +30,9 @@ use arrow::{ record_batch::RecordBatch, }; use auron_jni_bridge::{ - conf::{IntConf, TASK_CPUS, TOKIO_WORKER_THREADS_PER_CPU}, + conf::{ + IntConf, LongConf, METRICS_UPDATE_INTERVAL_MS, TASK_CPUS, TOKIO_WORKER_THREADS_PER_CPU, + }, is_task_running, jni_call, jni_call_static, jni_convert_byte_array, jni_exception_check, jni_exception_occurred, jni_new_global_ref, jni_new_object, jni_new_string, }; @@ -52,13 +55,14 @@ use datafusion_ext_plans::{ }; use futures::{FutureExt, StreamExt}; use jni::objects::{GlobalRef, JObject}; +use parking_lot::Mutex; use prost::Message; use tokio::{runtime::Runtime, task::JoinHandle}; use crate::{ handle_unwinded_scope, logging::{THREAD_PARTITION_ID, THREAD_STAGE_ID, THREAD_TID}, - metrics::update_metric_node, + metrics::{MetricSnapshot, update_metric_node}, }; pub struct NativeExecutionRuntime { @@ -68,6 +72,8 @@ pub struct NativeExecutionRuntime { batch_receiver: Receiver>>, tokio_runtime: Runtime, join_handle: JoinHandle<()>, + metrics_ticker: Option>, + metric_state: Arc>, // Flag to indicate runtime is being finalized - used to gracefully handle SendError is_finalizing: Arc, } @@ -119,6 +125,7 @@ impl NativeExecutionRuntime { let classloader_global = jni_new_global_ref!(classloader.as_obj())?; let mut tokio_runtime_builder = tokio::runtime::Builder::new_multi_thread(); tokio_runtime_builder + .enable_time() .thread_name(format!( "auron-native-stage-{stage_id}-part-{partition_id}-tid-{tid}" )) @@ -238,6 +245,15 @@ impl NativeExecutionRuntime { }); }); + let metric_state = Arc::new(Mutex::new(MetricSnapshot::new())); + let metrics_ticker = spawn_metrics_ticker( + &tokio_runtime, + native_wrapper.clone(), + execution_plan.clone(), + metric_state.clone(), + is_finalizing.clone(), + ); + let native_execution_runtime = Self { exec_ctx: exec_ctx.clone(), native_wrapper: native_wrapper.clone(), @@ -245,6 +261,8 @@ impl NativeExecutionRuntime { tokio_runtime, batch_receiver, join_handle, + metrics_ticker, + metric_state, is_finalizing, }; Ok(native_execution_runtime) @@ -288,13 +306,18 @@ impl NativeExecutionRuntime { let partition = self.exec_ctx.partition_id(); log::info!("(partition={partition}) native execution finalizing"); + // Stop the ticker before the last flush so it cannot JNI into MetricNode + // after we drop native_wrapper. Abort does not wait; the mutex serializes + // an in-flight publish with this final update. + self.is_finalizing.store(true, Ordering::Release); + if let Some(metrics_ticker) = &self.metrics_ticker { + metrics_ticker.abort(); + } self.update_metrics().unwrap_or_default(); drop(self.plan); - // Set finalizing flag before dropping receiver and native_wrapper to prevent - // concurrent set_error calls from next_batch/tokio workers from accessing a - // freed GlobalRef after finalize completes. - self.is_finalizing.store(true, Ordering::Release); + // Drop receiver after is_finalizing so concurrent next_batch/tokio workers + // skip set_error instead of touching a freed GlobalRef. drop(self.batch_receiver); cancel_all_tasks(&self.exec_ctx.task_ctx()); // cancel all pending streams @@ -304,12 +327,50 @@ impl NativeExecutionRuntime { } fn update_metrics(&self) -> Result<()> { - let metrics = jni_call!( - AuronCallNativeWrapper(self.native_wrapper.as_obj()).getMetrics() -> JObject - )?; - update_metric_node(metrics.as_obj(), self.plan.clone())?; - Ok(()) + publish_plan_metrics(&self.native_wrapper, &self.plan, &self.metric_state) + } +} + +fn spawn_metrics_ticker( + tokio_runtime: &Runtime, + native_wrapper: GlobalRef, + plan: Arc, + metric_state: Arc>, + is_finalizing: Arc, +) -> Option> { + let interval_ms = METRICS_UPDATE_INTERVAL_MS.value().unwrap_or(1000); + if interval_ms <= 0 { + return None; } + + Some(tokio_runtime.spawn(async move { + let mut interval = tokio::time::interval(Duration::from_millis(interval_ms as u64)); + // The first tick completes immediately; skip it so we do not JNI all-zero + // metrics. + interval.tick().await; + loop { + interval.tick().await; + if is_finalizing.load(Ordering::Acquire) { + break; + } + if let Err(err) = publish_plan_metrics(&native_wrapper, &plan, &metric_state) { + log::warn!("periodic metric update failed: {err}"); + } + } + })) +} + +fn publish_plan_metrics( + native_wrapper: &GlobalRef, + plan: &Arc, + metric_state: &Mutex, +) -> Result<()> { + let metrics = jni_call!( + AuronCallNativeWrapper(native_wrapper.as_obj()).getMetrics() -> JObject + )?; + let mut snapshot = metric_state.lock(); + update_metric_node(metrics.as_obj(), plan.clone(), &mut snapshot)?; + Ok(()) } fn set_error(native_wrapper: &GlobalRef, message: &str, cause: Option) -> Result<()> { From 38088ffe8b60632283891ef84e894e76a1f2d6b8 Mon Sep 17 00:00:00 2001 From: jennychen Date: Wed, 2 Sep 2026 16:57:33 +0800 Subject: [PATCH 2/2] [AURON #2189] introduce auron.metrics.update.enabled to only enable metrics update for flink --- .../configuration/AuronConfiguration.java | 25 +++++++++++++++---- .../apache/auron/jni/AuronAdaptorTest.java | 1 + .../FlinkAuronConfiguration.java | 13 ++++++++++ .../FlinkAuronConfigurationTest.java | 8 ++++++ native-engine/auron-jni-bridge/src/conf.rs | 1 + native-engine/auron/src/rt.rs | 7 +++++- 6 files changed, 49 insertions(+), 6 deletions(-) diff --git a/auron-core/src/main/java/org/apache/auron/configuration/AuronConfiguration.java b/auron-core/src/main/java/org/apache/auron/configuration/AuronConfiguration.java index 391ed0085..49655b92f 100644 --- a/auron-core/src/main/java/org/apache/auron/configuration/AuronConfiguration.java +++ b/auron-core/src/main/java/org/apache/auron/configuration/AuronConfiguration.java @@ -65,18 +65,33 @@ public abstract class AuronConfiguration { + "if not configured, the default value of 1 is used.") .withDefaultValue(1); + /** + * Whether the native runtime periodically publishes DataFusion metrics to {@code MetricNode} + * while a task is running. Disabled by default so Spark batch jobs keep finalize-only + * reporting. Flink overrides this option to {@code true} because long-running tasks never + * finalize. + */ + public static final ConfigOption METRICS_UPDATE_ENABLED = new ConfigOption<>(Boolean.class) + .withKey("auron.metrics.update.enabled") + .withCategory("Runtime Configuration") + .withDescription("Enable periodic native metric publishing to MetricNode while a task is running. " + + "Disabled by default (Spark finalize-only). Flink enables this so long-running tasks " + + "report live counters. When enabled, auron.metrics.update.interval.ms controls the " + + "publish interval.") + .withDefaultValue(false); + /** * How often the native runtime publishes DataFusion metrics to {@code MetricNode} while a task - * is running. Java {@code MetricNode.add} is incremental, so the native side sends positive - * deltas. {@code 0} disables the timer and publishes only when the native runtime finalizes - * (the historical batch behavior). + * is running. Used only when {@link #METRICS_UPDATE_ENABLED} is true. Java {@code + * MetricNode.add} is incremental, so the native side sends positive deltas. {@code 0} disables + * the timer and publishes only when the native runtime finalizes. */ public static final ConfigOption METRICS_UPDATE_INTERVAL_MS = new ConfigOption<>(Long.class) .withKey("auron.metrics.update.interval.ms") .withCategory("Runtime Configuration") .withDescription("Interval in milliseconds for publishing native execution metrics to MetricNode " - + "during a running task. Set to 0 to publish only when the native runtime finalizes. " - + "Default is 1000ms so long-running Flink tasks report live counters.") + + "during a running task when auron.metrics.update.enabled is true. Set to 0 to publish " + + "only when the native runtime finalizes. Default is 1000ms.") .withDefaultValue(1000L); public abstract Optional getOptional(ConfigOption option); diff --git a/auron-core/src/test/java/org/apache/auron/jni/AuronAdaptorTest.java b/auron-core/src/test/java/org/apache/auron/jni/AuronAdaptorTest.java index 8c11fa49a..6d92e1022 100644 --- a/auron-core/src/test/java/org/apache/auron/jni/AuronAdaptorTest.java +++ b/auron-core/src/test/java/org/apache/auron/jni/AuronAdaptorTest.java @@ -37,6 +37,7 @@ public void testRetrieveConfigWithAuronAdaptor() { assertEquals(auronConfig.getDouble(AuronConfiguration.MEMORY_FRACTION), 0.6, 0.0); assertEquals(auronConfig.getString(AuronConfiguration.NATIVE_LOG_LEVEL), "info"); assertEquals(auronConfig.getLong(AuronConfiguration.METRICS_UPDATE_INTERVAL_MS), 1000L); + assertEquals(auronConfig.getBoolean(AuronConfiguration.METRICS_UPDATE_ENABLED), false); } @Test diff --git a/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/configuration/FlinkAuronConfiguration.java b/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/configuration/FlinkAuronConfiguration.java index bfbe151b1..b841603c1 100644 --- a/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/configuration/FlinkAuronConfiguration.java +++ b/auron-flink-extension/auron-flink-runtime/src/main/java/org/apache/auron/flink/configuration/FlinkAuronConfiguration.java @@ -63,6 +63,19 @@ public class FlinkAuronConfiguration extends AuronConfiguration { .withDescription("Enable collection of additional metrics for input batch statistics.") .withDefaultValue(false); + /** + * Periodic native metric publishing for long-running Flink tasks. Same JNI field name as + * {@link AuronConfiguration#METRICS_UPDATE_ENABLED} so native lookup on this class does not + * fall back to the Spark/shared default of {@code false}. + */ + public static final ConfigOption METRICS_UPDATE_ENABLED = new ConfigOption<>(Boolean.class) + .withKey("auron.metrics.update.enabled") + .withCategory("Runtime Configuration") + .withDescription("Enable periodic native metric publishing to MetricNode while a Flink task is " + + "running. Default is true because Flink tasks are long-lived and do not finalize. " + + "Set flink.auron.metrics.update.enabled=false to restore finalize-only reporting.") + .withDefaultValue(true); + private final Configuration flinkConfig; public FlinkAuronConfiguration() { diff --git a/auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/configuration/FlinkAuronConfigurationTest.java b/auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/configuration/FlinkAuronConfigurationTest.java index ea198a914..4a16c8ca1 100644 --- a/auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/configuration/FlinkAuronConfigurationTest.java +++ b/auron-flink-extension/auron-flink-runtime/src/test/java/org/apache/auron/flink/configuration/FlinkAuronConfigurationTest.java @@ -47,6 +47,7 @@ public void testGetConfigFromFlinkConfig() { assertEquals(config.get(AuronConfiguration.NATIVE_LOG_LEVEL), "DEBUG"); assertEquals(config.get(AuronConfiguration.MEMORY_FRACTION), 0.6); // default value assertEquals(true, config.get(FlinkAuronConfiguration.FAIL_BACK_FLINK_ENGINE_ENABLED)); + assertEquals(true, config.get(FlinkAuronConfiguration.METRICS_UPDATE_ENABLED)); } @Test @@ -55,4 +56,11 @@ public void testFailBackFlinkEngineEnabledDefaultsToTrue() { "auron.failback.flink.engine.enabled", FlinkAuronConfiguration.FAIL_BACK_FLINK_ENGINE_ENABLED.key()); assertEquals(Boolean.TRUE, FlinkAuronConfiguration.FAIL_BACK_FLINK_ENGINE_ENABLED.defaultValue()); } + + @Test + public void testMetricsUpdateEnabledDefaultsToTrueOnFlink() { + assertEquals("auron.metrics.update.enabled", FlinkAuronConfiguration.METRICS_UPDATE_ENABLED.key()); + assertEquals(Boolean.TRUE, FlinkAuronConfiguration.METRICS_UPDATE_ENABLED.defaultValue()); + assertEquals(Boolean.FALSE, AuronConfiguration.METRICS_UPDATE_ENABLED.defaultValue()); + } } diff --git a/native-engine/auron-jni-bridge/src/conf.rs b/native-engine/auron-jni-bridge/src/conf.rs index a68f94a2e..55bb3e3d4 100644 --- a/native-engine/auron-jni-bridge/src/conf.rs +++ b/native-engine/auron-jni-bridge/src/conf.rs @@ -63,6 +63,7 @@ define_conf!(BooleanConf, ORC_SCHEMA_CASE_SENSITIVE); define_conf!(IntConf, UDAF_FALLBACK_NUM_UDAFS_TRIGGER_SORT_AGG); define_conf!(BooleanConf, PARSE_JSON_ERROR_FALLBACK); define_conf!(StringConf, NATIVE_LOG_LEVEL); +define_conf!(BooleanConf, METRICS_UPDATE_ENABLED); define_conf!(LongConf, METRICS_UPDATE_INTERVAL_MS); pub trait BooleanConf { diff --git a/native-engine/auron/src/rt.rs b/native-engine/auron/src/rt.rs index e23c85f90..0ae8ad633 100644 --- a/native-engine/auron/src/rt.rs +++ b/native-engine/auron/src/rt.rs @@ -31,7 +31,8 @@ use arrow::{ }; use auron_jni_bridge::{ conf::{ - IntConf, LongConf, METRICS_UPDATE_INTERVAL_MS, TASK_CPUS, TOKIO_WORKER_THREADS_PER_CPU, + BooleanConf, IntConf, LongConf, METRICS_UPDATE_ENABLED, METRICS_UPDATE_INTERVAL_MS, + TASK_CPUS, TOKIO_WORKER_THREADS_PER_CPU, }, is_task_running, jni_call, jni_call_static, jni_convert_byte_array, jni_exception_check, jni_exception_occurred, jni_new_global_ref, jni_new_object, jni_new_string, @@ -338,6 +339,10 @@ fn spawn_metrics_ticker( metric_state: Arc>, is_finalizing: Arc, ) -> Option> { + let enabled = METRICS_UPDATE_ENABLED.value().unwrap_or(false); + if !enabled { + return None; + } let interval_ms = METRICS_UPDATE_INTERVAL_MS.value().unwrap_or(1000); if interval_ms <= 0 { return None;