[GSoC 2026] Kafka Streams runner: surface SDK-harness metrics as MetricResults

runners/kafka-streams/build.gradle

@@ -45,6 +45,7 @@ dependencies {
   implementation project(path: ":sdks:java:core", configuration: "shadow")
   implementation project(path: ":runners:kafka-streams:proto", configuration: "shadow")
   implementation project(path: ":model:pipeline", configuration: "shadow")
+  implementation project(path: ":model:fn-execution", configuration: "shadow")
   implementation project(path: ":model:job-management", configuration: "shadow")
   implementation project(":runners:core-java")
   permitUnusedDeclared project(":runners:core-java")

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java

@@ -62,7 +62,8 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo)
 
     KafkaStreams kafkaStreams = new KafkaStreams(topology, streamsConfig(jobInfo));
     kafkaStreams.start();
-    return new KafkaStreamsPortablePipelineResult(kafkaStreams);
+    return new KafkaStreamsPortablePipelineResult(
+        kafkaStreams, context.getMetricsContainerStepMap());
   }
 
   private Properties streamsConfig(JobInfo jobInfo) {

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java

@@ -21,6 +21,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import org.apache.beam.model.jobmanagement.v1.JobApi;
+import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
 import org.apache.beam.runners.jobsubmission.PortablePipelineResult;
 import org.apache.beam.sdk.metrics.MetricResults;
 import org.apache.kafka.streams.KafkaStreams;
@@ -41,11 +42,16 @@ class KafkaStreamsPortablePipelineResult implements PortablePipelineResult {
       LoggerFactory.getLogger(KafkaStreamsPortablePipelineResult.class);
 
   private final KafkaStreams kafkaStreams;
+  // The job's metrics accumulator, shared by reference with the topology's stage processors, which
+  // update it as the SDK harness reports bundle metrics.
+  private final MetricsContainerStepMap metricsContainerStepMap;
   private final CountDownLatch terminated = new CountDownLatch(1);
   private volatile boolean cancelled = false;
 
-  KafkaStreamsPortablePipelineResult(KafkaStreams kafkaStreams) {
+  KafkaStreamsPortablePipelineResult(
+      KafkaStreams kafkaStreams, MetricsContainerStepMap metricsContainerStepMap) {
     this.kafkaStreams = kafkaStreams;
+    this.metricsContainerStepMap = metricsContainerStepMap;
     kafkaStreams.setStateListener(
         (newState, oldState) -> {
           if (newState == KafkaStreams.State.NOT_RUNNING || newState == KafkaStreams.State.ERROR) {
@@ -104,8 +110,9 @@ public State waitUntilFinish() {
 
   @Override
   public MetricResults metrics() {
-    throw new UnsupportedOperationException(
-        "Metrics are not yet implemented in the Kafka Streams runner.");
+    // Attempted values only: the runner does not distinguish committed results yet (that needs
+    // metrics to be folded into the exactly-once commit, which lands with the durability work).
+    return MetricsContainerStepMap.asAttemptedOnlyMetricResults(metricsContainerStepMap);
   }
 
   @Override

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java

@@ -20,7 +20,10 @@
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleProgressResponse;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse;
 import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
 import org.apache.beam.runners.fnexecution.control.BundleProgressHandler;
 import org.apache.beam.runners.fnexecution.control.ExecutableStageContext;
 import org.apache.beam.runners.fnexecution.control.OutputReceiverFactory;
@@ -32,6 +35,7 @@
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
 import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.apache.kafka.streams.processor.api.Processor;
 import org.apache.kafka.streams.processor.api.ProcessorContext;
@@ -74,6 +78,10 @@ class ExecutableStageProcessor
   // This stage's own transform id, stamped on every watermark it forwards so downstream watermark
   // aggregators know which transform the report came from — regardless of who consumes it.
   private final String transformId;
+  // This stage's Beam metrics container, updated from the final MonitoringInfos the SDK harness
+  // reports as each bundle completes. The pipeline result reads the containing step map as
+  // MetricResults.
+  private final MetricsContainerImpl metricsContainer;
 
   // pendingOutputs is enqueued by SDK harness threads (inside the OutputReceiverFactory callback)
   // and drained by the Kafka Streams processing thread on bundle close; needs to be thread-safe.
@@ -98,16 +106,20 @@ class ExecutableStageProcessor
    * @param transformId this stage's own transform id, stamped on the watermarks it emits
    * @param upstreamTransformIds the transform ids feeding this stage (known from the pipeline
    *     graph), whose reports the {@link WatermarkAggregator} waits for
+   * @param metricsContainer this stage's container in the job's metrics step map, updated with the
+   *     harness's per-bundle MonitoringInfos
    */
   ExecutableStageProcessor(
       RunnerApi.ExecutableStagePayload stagePayload,
       JobInfo jobInfo,
       String transformId,
-      Set<String> upstreamTransformIds) {
+      Set<String> upstreamTransformIds,
+      MetricsContainerImpl metricsContainer) {
     this.stagePayload = stagePayload;
     this.jobInfo = jobInfo;
     this.transformId = transformId;
     this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
+    this.metricsContainer = metricsContainer;
   }
 
   @Override
@@ -179,11 +191,37 @@ public <OutputT> FnDataReceiver<OutputT> create(String pCollectionId) {
             };
           }
         };
+    // Fold the harness's reported metrics into this stage's container when each bundle completes.
+    // Only the completion response is applied: it carries the bundle's final cumulative values, and
+    // the container's update() adds counter values, so also applying mid-bundle progress snapshots
+    // would double-count them. Live mid-bundle metrics can come later if a use appears.
+    BundleProgressHandler progressHandler =
+        new BundleProgressHandler() {
+          @Override
+          public void onProgress(ProcessBundleProgressResponse progress) {
+            // Deliberately not folded into the container; see comment above.
+            if (LOG.isDebugEnabled()) {
+              LOG.debug(
+                  "Stage {} bundle progress: {}",
+                  transformId,
+                  TextFormat.printer().printToString(progress));
+            }
+          }
+
+          @Override
+          public void onCompleted(ProcessBundleResponse response) {
+            if (LOG.isDebugEnabled()) {
+              LOG.debug(
+                  "Stage {} bundle completed: {}",
+                  transformId,
+                  TextFormat.printer().printToString(response));
+            }
+            metricsContainer.update(response.getMonitoringInfosList());
+          }
+        };
     currentBundle =
         factory.getBundle(
-            outputReceiverFactory,
-            StateRequestHandler.unsupported(),
-            BundleProgressHandler.ignored());
+            outputReceiverFactory, StateRequestHandler.unsupported(), progressHandler);
   }
 
   private FnDataReceiver<WindowedValue<?>> mainInputReceiver() {

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java

@@ -86,12 +86,17 @@ public void translate(
     Topology topology = context.getTopology();
     // The stage stamps its own transform id on the watermarks it emits, and aggregates its input
     // watermark from the reports of its single upstream transform (the producer of its input
-    // PCollection, whose node name is the upstream transform id).
+    // PCollection, whose node name is the upstream transform id). Harness-reported metrics land in
+    // this stage's container of the job's metrics step map.
     topology.addProcessor(
         transformId,
         () ->
             new ExecutableStageProcessor(
-                stagePayload, context.getJobInfo(), transformId, ImmutableSet.of(parentProcessor)),
+                stagePayload,
+                context.getJobInfo(),
+                transformId,
+                ImmutableSet.of(parentProcessor),
+                context.getMetricsContainerStepMap().getContainer(transformId)),
         parentProcessor);
 
     if (!transform.getOutputsMap().isEmpty()) {

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java

@@ -20,6 +20,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.regex.Pattern;
+import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
 import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
 import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions;
 import org.apache.kafka.streams.Topology;
@@ -45,6 +46,14 @@ public class KafkaStreamsTranslationContext {
   private final KafkaStreamsPipelineOptions pipelineOptions;
   private final Topology topology;
   private final Map<String, String> pCollectionIdToProcessorName;
+  // Accumulates the Beam metrics reported by the SDK harness, one container per executable stage.
+  // Processors update it as bundles complete (in-JVM reference sharing); the pipeline result
+  // exposes it as MetricResults. Sharing one container across a stage's parallel tasks is safe and
+  // correct: the metric cells are thread-safe (atomic cells in concurrent maps) and the updates are
+  // per-bundle final values applied with add semantics, so concurrent tasks accumulate rather than
+  // overwrite. Aggregation across multiple runner JVMs is out of scope until the multi-instance
+  // work.
+  private final MetricsContainerStepMap metricsContainerStepMap = new MetricsContainerStepMap();
 
   public static KafkaStreamsTranslationContext create(
       JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) {
@@ -77,6 +86,16 @@ public Topology getTopology() {
     return topology;
   }
 
+  /**
+   * Returns the job's metrics accumulator: one {@link
+   * org.apache.beam.runners.core.metrics.MetricsContainerImpl container} per executable stage,
+   * updated by the stage processors as the SDK harness reports bundle metrics, and read by the
+   * pipeline result via {@link MetricsContainerStepMap#asAttemptedOnlyMetricResults}.
+   */
+  public MetricsContainerStepMap getMetricsContainerStepMap() {
+    return metricsContainerStepMap;
+  }
+
   /**
    * Registers the processor node that produces the given Beam PCollection. Downstream translators
    * resolve their parent processor names by looking up the input PCollection id.

runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java

@@ -25,10 +25,12 @@
 import java.util.Set;
 import java.util.UUID;
 import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.core.metrics.MetricsContainerStepMap;
 import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
 import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsPipelineTranslator;
 import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsTranslationContext;
 import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.metrics.MetricResults;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.PortablePipelineOptions;
@@ -107,8 +109,12 @@ public static KafkaStreamsTranslationContext translate(Pipeline pipeline) {
     return context;
   }
 
-  /** Translates and drives the pipeline to quiescence through a {@link TopologyTestDriver}. */
-  public static void run(Pipeline pipeline) {
+  /**
+   * Translates and drives the pipeline to quiescence through a {@link TopologyTestDriver}. Returns
+   * the metrics the SDK harness reported while running (attempted values), so tests can assert on
+   * user counters — the same surface PAssert uses to verify its assertions ran.
+   */
+  public static MetricResults run(Pipeline pipeline) {
     KafkaStreamsTranslationContext context = translate(pipeline);
     Topology topology = context.getTopology();
     try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig(pipeline))) {
@@ -117,6 +123,8 @@ public static void run(Pipeline pipeline) {
       driver.advanceWallClockTime(Duration.ofSeconds(1));
       roundTripInternalTopics(driver, internalTopics(topology));
     }
+    return MetricsContainerStepMap.asAttemptedOnlyMetricResults(
+        context.getMetricsContainerStepMap());
   }
 
   /**

runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java

@@ -21,6 +21,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 
 import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
 import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
@@ -57,7 +58,8 @@ private static ExecutableStageProcessor newProcessor() {
         RunnerApi.ExecutableStagePayload.getDefaultInstance(),
         jobInfo,
         STAGE_ID,
-        ImmutableSet.of(UPSTREAM_ID));
+        ImmutableSet.of(UPSTREAM_ID),
+        new MetricsContainerImpl(STAGE_ID));
   }
 
   /** A report from the upstream transform's given partition. */

runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/MetricsTest.java

@@ -0,0 +1,79 @@
+/*
+ * 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.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.MetricNameFilter;
+import org.apache.beam.sdk.metrics.MetricQueryResults;
+import org.apache.beam.sdk.metrics.MetricResult;
+import org.apache.beam.sdk.metrics.MetricResults;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.metrics.MetricsFilter;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.junit.Test;
+
+/**
+ * End-to-end test that user metrics reported by a {@link DoFn} in the SDK harness surface through
+ * the runner as {@link MetricResults}.
+ *
+ * <p>The harness reports metrics as Fn API MonitoringInfos with each bundle; the stage processor
+ * folds them into the job's metrics step map, and the runner exposes them as attempted {@link
+ * MetricResults}. This is the surface {@code PAssert} uses to verify its assertions actually ran,
+ * so it is a prerequisite for the {@code @ValidatesRunner} suite.
+ */
+public class MetricsTest {
+
+  private static final String NAMESPACE = "MetricsTest";
+  private static final String COUNTER_NAME = "elements";
+
+  /** Increments a user counter for every element the harness feeds it. */
+  private static class CountingFn extends DoFn<Integer, Integer> {
+    private final Counter counter = Metrics.counter(NAMESPACE, COUNTER_NAME);
+
+    @ProcessElement
+    public void processElement(@Element Integer input, OutputReceiver<Integer> out) {
+      counter.inc();
+      out.output(input);
+    }
+  }
+
+  @Test
+  public void userCounterFromHarnessSurfacesInMetricResults() {
+    Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions());
+    pipeline.apply("create", Create.of(1, 2, 3)).apply("count", ParDo.of(new CountingFn()));
+
+    MetricResults metrics = KafkaStreamsTestRunner.run(pipeline);
+
+    MetricQueryResults query =
+        metrics.queryMetrics(
+            MetricsFilter.builder()
+                .addNameFilter(MetricNameFilter.named(NAMESPACE, COUNTER_NAME))
+                .build());
+    MetricResult<Long> counter = Iterables.getOnlyElement(query.getCounters());
+    // Create.of(1, 2, 3) feeds the DoFn exactly three elements.
+    assertThat(counter.getAttempted(), is(3L));
+  }
+}