[AURON #2189] Periodically publish native metrics for long-running Flink tasks

Author: jennychen

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<Long> 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 <T> Optional<T> getOptional(ConfigOption<T> option);
 
     public <T> T get(ConfigOption<T> option) {

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

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<RowData
     private volatile boolean isRunning;
     private transient String auronOperatorIdWithSubtaskIndex;
     private transient MetricNode nativeMetric;
-    private transient MetricGroup metricGroup;
     private transient ObjectMapper mapper;
 
     // Kafka Consumer for partition metadata discovery only (does NOT consume data)
@@ -338,22 +337,10 @@ PhysicalPlanNode applyMergedCalcPlan(PhysicalPlanNode sourcePlan) {
 
     @Override
     public void run(SourceContext<RowData> sourceContext) throws Exception {
-        metricGroup = getRuntimeContext().getMetricGroup();
-        final Map<String, Counter> 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

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<MetricNode> 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.
+     *
+     * <p>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<MetricNode> 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.

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<org.apache.auron.metric.MetricNode> 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
     // ====================================================================

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.<RowData>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));
+    }
 }

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;