diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 2ba55e618308..a794299cb608 100644 --- a/runners/kafka-streams/build.gradle +++ b/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") diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java index 3e97638695e1..cdb59f67cce9 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java +++ b/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) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java index 817746bf002f..972afa15c358 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java +++ b/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 diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java index 38d3601e814f..fb26e9e71f6b 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java +++ b/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 upstreamTransformIds) { + Set 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 FnDataReceiver 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> mainInputReceiver() { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java index a2e6ed837c7d..59c233a78bb1 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java +++ b/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()) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java index 0a045ff7395b..89a2d9825cbc 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java +++ b/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 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. diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java index 4d73eee9581a..8ab29182c281 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java +++ b/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()); } /** diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java index 74169acb81cb..de3a23911adc 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java +++ b/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. */ diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/MetricsTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/MetricsTest.java new file mode 100644 index 000000000000..558a1bab91a7 --- /dev/null +++ b/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}. + * + *

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 { + private final Counter counter = Metrics.counter(NAMESPACE, COUNTER_NAME); + + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver 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 counter = Iterables.getOnlyElement(query.getCounters()); + // Create.of(1, 2, 3) feeds the DoFn exactly three elements. + assertThat(counter.getAttempted(), is(3L)); + } +}