From 2e4453071537e3a756b79dc668cc3ad920ca80e1 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Sat, 16 May 2026 14:42:46 +0500 Subject: [PATCH 01/37] Add Kafka Streams runner skeleton module and portable entry points --- CHANGES.md | 1 + build.gradle.kts | 1 + runners/kafka-streams/build.gradle | 62 ++++++++++ .../kafka/streams/KafkaStreamsJobInvoker.java | 99 ++++++++++++++++ .../streams/KafkaStreamsJobServerDriver.java | 109 ++++++++++++++++++ .../streams/KafkaStreamsPipelineOptions.java | 76 ++++++++++++ .../streams/KafkaStreamsPipelineResult.java | 69 +++++++++++ .../streams/KafkaStreamsPipelineRunner.java | 48 ++++++++ .../kafka/streams/KafkaStreamsRunner.java | 101 ++++++++++++++++ .../streams/KafkaStreamsRunnerRegistrar.java | 48 ++++++++ .../runners/kafka/streams/package-info.java | 20 ++++ .../KafkaStreamsPipelineTranslator.java | 70 +++++++++++ .../KafkaStreamsTranslationContext.java | 47 ++++++++ .../streams/translation/package-info.java | 20 ++++ settings.gradle.kts | 1 + 15 files changed, 772 insertions(+) create mode 100644 runners/kafka-streams/build.gradle create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriver.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineResult.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerRegistrar.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/package-info.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/package-info.java diff --git a/CHANGES.md b/CHANGES.md index 52475a99d8e1..18d2698a2875 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -71,6 +71,7 @@ ## New Features / Improvements +* (Java) Added a `runners/kafka-streams` Gradle module with portable job server and runner entry points; translation fails fast with an explicit unsupported-URN message until transforms are implemented ([#38465](https://github.com/apache/beam/issues/38465)). * Capability introduces an indicator for aggregations and timers firing during a pipeline drain, allowing users and sinks to recognize and appropriately handle potentially incomplete or partial data ([#36884](https://github.com/apache/beam/issues/36884)). * Added support for setting disk provisioned IOPS and throughput in Dataflow runner via `--diskProvisionedIops` and `--diskProvisionedThroughputMibps` pipeline options (Java/Go/Python) ([#38349](https://github.com/apache/beam/issues/38349)). * TriggerStateMachineRunner changes from BitSetCoder to SentinelBitSetCoder to diff --git a/build.gradle.kts b/build.gradle.kts index 4af8fa3f1ab4..b3b9fdd7fdf0 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -278,6 +278,7 @@ tasks.register("javaPreCommit") { dependsOn(":runners:java-fn-execution:build") dependsOn(":runners:java-job-service:build") dependsOn(":runners:jet:build") + dependsOn(":runners:kafka-streams:build") dependsOn(":runners:local-java:build") dependsOn(":runners:portability:java:build") dependsOn(":runners:prism:java:build") diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle new file mode 100644 index 000000000000..54474502ad7b --- /dev/null +++ b/runners/kafka-streams/build.gradle @@ -0,0 +1,62 @@ +/* + * 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. + */ + +plugins { id 'org.apache.beam.module' } + +def kafka_version = '3.9.0' + +applyJavaNature( + automaticModuleName: 'org.apache.beam.runners.kafka.streams', +) + +description = "Apache Beam :: Runners :: Kafka Streams" + +evaluationDependsOn(":sdks:java:core") +evaluationDependsOn(":runners:core-java") + +configurations.configureEach { + resolutionStrategy.eachDependency { details -> + if (details.requested.group == "org.apache.kafka") { + details.useVersion(kafka_version) + details.because("Kafka Streams runner is developed against Kafka ${kafka_version}.") + } + } +} + +dependencies { + compileOnly project(":sdks:java:build-tools") + permitUnusedDeclared project(":sdks:java:build-tools") + + implementation project(path: ":sdks:java:core", configuration: "shadow") + implementation project(path: ":model:pipeline", configuration: "shadow") + implementation project(":runners:core-java") + permitUnusedDeclared project(":runners:core-java") + implementation project(":runners:java-fn-execution") + implementation project(":runners:java-job-service") + implementation project(":runners:portability:java") + implementation project(path: ":sdks:java:extensions:google-cloud-platform-core") + implementation library.java.args4j + implementation library.java.joda_time + implementation library.java.slf4j_api + implementation library.java.vendored_grpc_1_69_0 + implementation library.java.vendored_guava_32_1_2_jre + implementation "org.apache.kafka:kafka-clients:$kafka_version" + implementation "org.apache.kafka:kafka-streams:$kafka_version" + permitUnusedDeclared "org.apache.kafka:kafka-clients:$kafka_version" + permitUnusedDeclared "org.apache.kafka:kafka-streams:$kafka_version" +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java new file mode 100644 index 000000000000..bd1ed4c1bcf6 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java @@ -0,0 +1,99 @@ +/* + * 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; + +import java.util.UUID; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.jobsubmission.JobInvocation; +import org.apache.beam.runners.jobsubmission.JobInvoker; +import org.apache.beam.runners.jobsubmission.PortablePipelineRunner; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Struct; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ListeningExecutorService; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Job invoker for the Kafka Streams portable runner. */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class KafkaStreamsJobInvoker extends JobInvoker { + + private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsJobInvoker.class); + + public static KafkaStreamsJobInvoker create( + KafkaStreamsJobServerDriver.KafkaStreamsServerConfiguration serverConfig) { + return new KafkaStreamsJobInvoker(serverConfig); + } + + private final KafkaStreamsJobServerDriver.KafkaStreamsServerConfiguration serverConfig; + + protected KafkaStreamsJobInvoker( + KafkaStreamsJobServerDriver.KafkaStreamsServerConfiguration serverConfig) { + super("kafka-streams-runner-job-invoker-%d"); + this.serverConfig = serverConfig; + } + + @Override + protected JobInvocation invokeWithExecutor( + RunnerApi.Pipeline pipeline, + Struct options, + @Nullable String retrievalToken, + ListeningExecutorService executorService) { + + LOG.trace( + "Parsing pipeline options (job server {}:{})", + serverConfig.getHost(), + serverConfig.getPort()); + KafkaStreamsPipelineOptions kafkaStreamsOptions = + PipelineOptionsTranslation.fromProto(options).as(KafkaStreamsPipelineOptions.class); + + String invocationId = + String.format("%s_%s", kafkaStreamsOptions.getJobName(), UUID.randomUUID().toString()); + + PortablePipelineRunner pipelineRunner = new KafkaStreamsPipelineRunner(kafkaStreamsOptions); + kafkaStreamsOptions.setRunner(null); + + LOG.info("Invoking job {} with pipeline runner {}", invocationId, pipelineRunner); + return createJobInvocation( + invocationId, + retrievalToken, + executorService, + pipeline, + kafkaStreamsOptions, + pipelineRunner); + } + + protected JobInvocation createJobInvocation( + String invocationId, + String retrievalToken, + ListeningExecutorService executorService, + RunnerApi.Pipeline pipeline, + KafkaStreamsPipelineOptions kafkaStreamsOptions, + PortablePipelineRunner pipelineRunner) { + JobInfo jobInfo = + JobInfo.create( + invocationId, + kafkaStreamsOptions.getJobName(), + retrievalToken, + PipelineOptionsTranslation.toProto(kafkaStreamsOptions)); + return new JobInvocation(jobInfo, executorService, pipeline, pipelineRunner); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriver.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriver.java new file mode 100644 index 000000000000..ddeac8b7c959 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriver.java @@ -0,0 +1,109 @@ +/* + * 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; + +import org.apache.beam.runners.jobsubmission.JobServerDriver; +import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; +import org.apache.beam.sdk.fn.server.ServerFactory; +import org.apache.beam.sdk.io.FileSystems; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.kohsuke.args4j.CmdLineException; +import org.kohsuke.args4j.CmdLineParser; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Driver that starts a Beam job server for the Kafka Streams portable runner. */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class KafkaStreamsJobServerDriver extends JobServerDriver { + + private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsJobServerDriver.class); + + /** Runner-specific configuration for the job server process. */ + public static class KafkaStreamsServerConfiguration extends ServerConfiguration {} + + public static void main(String[] args) throws Exception { + PipelineOptions options = PipelineOptionsFactory.create(); + options.as(GcsOptions.class).setGcsUploadBufferSizeBytes(1024 * 1024); + FileSystems.setDefaultPipelineOptions(options); + fromParams(args).run(); + } + + private static void printUsage(CmdLineParser parser) { + System.err.println( + String.format( + "Usage: java %s arguments...", KafkaStreamsJobServerDriver.class.getSimpleName())); + parser.printUsage(System.err); + System.err.println(); + } + + public static KafkaStreamsServerConfiguration parseArgs(String[] args) { + KafkaStreamsServerConfiguration configuration = new KafkaStreamsServerConfiguration(); + CmdLineParser parser = new CmdLineParser(configuration); + try { + parser.parseArgument(args); + } catch (CmdLineException e) { + LOG.error("Unable to parse command line arguments.", e); + printUsage(parser); + throw new IllegalArgumentException("Unable to parse command line arguments.", e); + } + return configuration; + } + + /** Used by tests and tooling to construct a driver from command-line parameters. */ + public static KafkaStreamsJobServerDriver fromParams(String[] args) { + return fromConfig(parseArgs(args)); + } + + public static KafkaStreamsJobServerDriver fromConfig( + KafkaStreamsServerConfiguration configuration) { + return create( + configuration, + createJobServerFactory(configuration), + createArtifactServerFactory(configuration), + () -> KafkaStreamsJobInvoker.create(configuration)); + } + + public static KafkaStreamsJobServerDriver fromConfig( + KafkaStreamsServerConfiguration configuration, JobInvokerFactory jobInvokerFactory) { + return create( + configuration, + createJobServerFactory(configuration), + createArtifactServerFactory(configuration), + jobInvokerFactory); + } + + private static KafkaStreamsJobServerDriver create( + KafkaStreamsServerConfiguration configuration, + ServerFactory jobServerFactory, + ServerFactory artifactServerFactory, + JobInvokerFactory jobInvokerFactory) { + return new KafkaStreamsJobServerDriver( + configuration, jobServerFactory, artifactServerFactory, jobInvokerFactory); + } + + private KafkaStreamsJobServerDriver( + KafkaStreamsServerConfiguration configuration, + ServerFactory jobServerFactory, + ServerFactory artifactServerFactory, + JobInvokerFactory jobInvokerFactory) { + super(configuration, jobServerFactory, artifactServerFactory, jobInvokerFactory); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java new file mode 100644 index 000000000000..019b37cba770 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java @@ -0,0 +1,76 @@ +/* + * 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; + +import java.nio.file.Paths; +import org.apache.beam.sdk.options.Default; +import org.apache.beam.sdk.options.DefaultValueFactory; +import org.apache.beam.sdk.options.Description; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PortablePipelineOptions; + +/** Pipeline options for the Kafka Streams runner. */ +public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { + + @Description("Comma-separated list of host:port Kafka brokers used by the Kafka Streams client.") + @Default.String("localhost:9092") + String getBootstrapServers(); + + void setBootstrapServers(String bootstrapServers); + + @Description( + "Kafka Streams application.id (must be unique for each distinct topology using the same " + + "input topics in a Kafka cluster).") + @Default.String("beam-kafka-streams-runner") + String getApplicationId(); + + void setApplicationId(String applicationId); + + @Description( + "Kafka Streams processing.guarantee setting, for example at_least_once or exactly_once_v2.") + @Default.String("exactly_once_v2") + String getProcessingGuarantee(); + + void setProcessingGuarantee(String processingGuarantee); + + @Description("Soft cap on the number of elements per bundle.") + @Default.Integer(1000) + int getMaxBundleSize(); + + void setMaxBundleSize(int maxBundleSize); + + @Description("Soft cap on bundle wall-clock duration in milliseconds.") + @Default.Integer(1000) + int getMaxBundleTimeMs(); + + void setMaxBundleTimeMs(int maxBundleTimeMs); + + @Description("Directory where Kafka Streams stores local state.") + @Default.InstanceFactory(StateDirDefaultFactory.class) + String getStateDir(); + + void setStateDir(String stateDir); + + /** Default {@link #getStateDir()} under the JVM temp directory. */ + class StateDirDefaultFactory implements DefaultValueFactory { + @Override + public String create(PipelineOptions options) { + return Paths.get(System.getProperty("java.io.tmpdir"), "beam-kafka-streams-state").toString(); + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineResult.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineResult.java new file mode 100644 index 000000000000..776eaa6629f5 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineResult.java @@ -0,0 +1,69 @@ +/* + * 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; + +import java.io.IOException; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.metrics.MetricResults; +import org.joda.time.Duration; + +/** + * Forwards {@link PipelineResult} calls to a delegate and stops an embedded job server when the + * pipeline reaches a terminal state. + */ +class KafkaStreamsPipelineResult implements PipelineResult { + + private final PipelineResult delegate; + private final Runnable stopJobServer; + + KafkaStreamsPipelineResult(PipelineResult delegate, Runnable stopJobServer) { + this.delegate = delegate; + this.stopJobServer = stopJobServer; + } + + @Override + public State getState() { + return delegate.getState(); + } + + @Override + public State cancel() throws IOException { + State state = delegate.cancel(); + stopJobServer.run(); + return state; + } + + @Override + public State waitUntilFinish(Duration duration) { + State state = delegate.waitUntilFinish(duration); + stopJobServer.run(); + return state; + } + + @Override + public State waitUntilFinish() { + State state = delegate.waitUntilFinish(); + stopJobServer.run(); + return state; + } + + @Override + public MetricResults metrics() { + return delegate.metrics(); + } +} 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 new file mode 100644 index 000000000000..cc7464e22786 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java @@ -0,0 +1,48 @@ +/* + * 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; + +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.jobsubmission.PortablePipelineResult; +import org.apache.beam.runners.jobsubmission.PortablePipelineRunner; +import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsPipelineTranslator; +import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsTranslationContext; + +/** Executes a portable pipeline by translating it to Kafka Streams. */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class KafkaStreamsPipelineRunner implements PortablePipelineRunner { + + private final KafkaStreamsPipelineOptions pipelineOptions; + + public KafkaStreamsPipelineRunner(KafkaStreamsPipelineOptions pipelineOptions) { + this.pipelineOptions = pipelineOptions; + } + + @Override + public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) throws Exception { + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + KafkaStreamsTranslationContext context = + translator.createTranslationContext(jobInfo, pipelineOptions); + RunnerApi.Pipeline prepared = translator.prepareForTranslation(pipeline); + translator.translate(context, prepared); + throw new IllegalStateException("Translation unexpectedly completed without an executor"); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java new file mode 100644 index 000000000000..ce8f0544b681 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java @@ -0,0 +1,101 @@ +/* + * 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; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.portability.PortableRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.PipelineRunner; +import org.apache.beam.sdk.options.ExperimentalOptions; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A {@link PipelineRunner} that submits portable jobs to an in-process or external Beam job service + * backed by the Kafka Streams translation path. + */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class KafkaStreamsRunner extends PipelineRunner { + + private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsRunner.class); + + private final KafkaStreamsPipelineOptions pipelineOptions; + + public static KafkaStreamsRunner fromOptions(PipelineOptions options) { + return new KafkaStreamsRunner(options.as(KafkaStreamsPipelineOptions.class)); + } + + protected KafkaStreamsRunner(KafkaStreamsPipelineOptions pipelineOptions) { + this.pipelineOptions = pipelineOptions; + } + + @Override + public PipelineResult run(Pipeline pipeline) { + assignPortableDefaults(pipelineOptions); + KafkaStreamsJobServerDriver jobServerDriver = null; + try { + if (Strings.isNullOrEmpty(pipelineOptions.getJobEndpoint())) { + LOG.info("No job endpoint configured; starting an embedded Kafka Streams job server."); + KafkaStreamsJobServerDriver.KafkaStreamsServerConfiguration configuration = + new KafkaStreamsJobServerDriver.KafkaStreamsServerConfiguration(); + configuration.setPort(0); + jobServerDriver = KafkaStreamsJobServerDriver.fromConfig(configuration); + pipelineOptions.setJobEndpoint(jobServerDriver.start()); + } + PortableRunner portableRunner = PortableRunner.fromOptions(pipelineOptions); + PipelineResult result = portableRunner.run(pipeline); + if (jobServerDriver != null) { + return new KafkaStreamsPipelineResult(result, jobServerDriver::stop); + } + return result; + } catch (IOException e) { + if (jobServerDriver != null) { + jobServerDriver.stop(); + } + throw new RuntimeException(e); + } + } + + private static void assignPortableDefaults(KafkaStreamsPipelineOptions pipelineOptions) { + if (Strings.isNullOrEmpty(pipelineOptions.getDefaultEnvironmentType())) { + pipelineOptions.setDefaultEnvironmentType(Environments.ENVIRONMENT_LOOPBACK); + } + ExperimentalOptions experimentalOptions = pipelineOptions.as(ExperimentalOptions.class); + List experiments = + experimentalOptions.getExperiments() == null + ? new ArrayList<>() + : new ArrayList<>(experimentalOptions.getExperiments()); + if (!experiments.contains("beam_fn_api")) { + experiments.add("beam_fn_api"); + experimentalOptions.setExperiments(experiments); + } + } + + @Override + public String toString() { + return "KafkaStreamsRunner#" + hashCode(); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerRegistrar.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerRegistrar.java new file mode 100644 index 000000000000..ac3c64b97bb0 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerRegistrar.java @@ -0,0 +1,48 @@ +/* + * 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; + +import com.google.auto.service.AutoService; +import org.apache.beam.sdk.PipelineRunner; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsRegistrar; +import org.apache.beam.sdk.runners.PipelineRunnerRegistrar; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; + +/** {@link com.google.auto.service.AutoService} registrations for the Kafka Streams runner. */ +public class KafkaStreamsRunnerRegistrar { + private KafkaStreamsRunnerRegistrar() {} + + /** Registers {@link KafkaStreamsRunner}. */ + @AutoService(PipelineRunnerRegistrar.class) + public static class Runner implements PipelineRunnerRegistrar { + @Override + public Iterable>> getPipelineRunners() { + return ImmutableList.of(KafkaStreamsRunner.class); + } + } + + /** Registers {@link KafkaStreamsPipelineOptions}. */ + @AutoService(PipelineOptionsRegistrar.class) + public static class Options implements PipelineOptionsRegistrar { + @Override + public Iterable> getPipelineOptions() { + return ImmutableList.of(KafkaStreamsPipelineOptions.class); + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/package-info.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/package-info.java new file mode 100644 index 000000000000..c9def4d1a4d7 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/package-info.java @@ -0,0 +1,20 @@ +/* + * 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. + */ + +/** Kafka Streams runner: portable pipeline execution backed by Apache Kafka Streams. */ +package org.apache.beam.runners.kafka.streams; diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java new file mode 100644 index 000000000000..cc915d604b68 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -0,0 +1,70 @@ +/* + * 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 java.util.Map; +import java.util.TreeMap; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; + +/** + * Translates a portable Beam pipeline into a Kafka Streams {@code Topology}. + * + *

The initial implementation only validates the graph and fails fast with an explicit message + * for transforms that are not yet supported. + */ +public class KafkaStreamsPipelineTranslator { + + public KafkaStreamsTranslationContext createTranslationContext( + JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { + return KafkaStreamsTranslationContext.create(jobInfo, pipelineOptions); + } + + /** Returns the pipeline to translate (placeholder for future fusion / expansion steps). */ + public RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) { + return pipeline; + } + + /** + * Translates the pipeline. Throws {@link UnsupportedOperationException} with a clear URN message + * for the first unsupported primitive encountered. + */ + public void translate(KafkaStreamsTranslationContext context, RunnerApi.Pipeline pipeline) { + Map transforms = pipeline.getComponents().getTransformsMap(); + TreeMap ordered = new TreeMap<>(transforms); + for (Map.Entry entry : ordered.entrySet()) { + RunnerApi.PTransform transform = entry.getValue(); + if (!transform.hasSpec()) { + continue; + } + String urn = transform.getSpec().getUrn(); + if (urn.isEmpty()) { + continue; + } + throw new UnsupportedOperationException( + "No translator registered for URN " + + urn + + " (jobId=" + + context.getJobInfo().jobId() + + ")"); + } + throw new UnsupportedOperationException( + "No translator registered for pipeline (no transform URNs found)"); + } +} 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 new file mode 100644 index 000000000000..7c6d3d079159 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java @@ -0,0 +1,47 @@ +/* + * 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 org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; + +/** Mutable state shared while translating a portable pipeline into a Kafka Streams topology. */ +public class KafkaStreamsTranslationContext { + + private final JobInfo jobInfo; + private final KafkaStreamsPipelineOptions pipelineOptions; + + public static KafkaStreamsTranslationContext create( + JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { + return new KafkaStreamsTranslationContext(jobInfo, pipelineOptions); + } + + private KafkaStreamsTranslationContext( + JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { + this.jobInfo = jobInfo; + this.pipelineOptions = pipelineOptions; + } + + public JobInfo getJobInfo() { + return jobInfo; + } + + public KafkaStreamsPipelineOptions getPipelineOptions() { + return pipelineOptions; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/package-info.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/package-info.java new file mode 100644 index 000000000000..9c09b9ceeb0b --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/package-info.java @@ -0,0 +1,20 @@ +/* + * 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. + */ + +/** Portable pipeline translation to Kafka Streams topologies. */ +package org.apache.beam.runners.kafka.streams.translation; diff --git a/settings.gradle.kts b/settings.gradle.kts index fc5f40c23d17..b92b254981fe 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -143,6 +143,7 @@ include(":runners:google-cloud-dataflow-java:examples-streaming") include(":runners:java-fn-execution") include(":runners:java-job-service") include(":runners:jet") +include(":runners:kafka-streams") include(":runners:local-java") include(":runners:portability:java") include(":runners:prism") From 61c891a69cad50e1fd4b2e4a336843efe4d5e6bb Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Tue, 19 May 2026 16:10:28 +0500 Subject: [PATCH 02/37] Address review notes on KafkaStreamsPipelineResult and state dir - cancel(): wrap delegate.cancel() in try/finally so the embedded job server is always stopped, even if cancellation throws IOException. - waitUntilFinish(Duration): only stop the job server when the returned state is terminal, so a timed-out wait does not prematurely kill the job server while the pipeline is still running. - waitUntilFinish(): wrap in try/finally for the same defensive cleanup reason as cancel(). - KafkaStreamsPipelineOptions.StateDirDefaultFactory: include the job name in the default Kafka Streams state directory so that multiple pipelines on the same host (e.g. parallel tests) do not collide and hit a LockException. --- .../streams/KafkaStreamsPipelineOptions.java | 14 ++++++++++-- .../streams/KafkaStreamsPipelineResult.java | 22 +++++++++++++------ 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java index 019b37cba770..911ffce2eea3 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java @@ -66,11 +66,21 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { void setStateDir(String stateDir); - /** Default {@link #getStateDir()} under the JVM temp directory. */ + /** + * Default {@link #getStateDir()} under the JVM temp directory. + * + *

The job name is included in the path so that multiple pipelines running on the same host + * (e.g. parallel tests) do not collide on the same Kafka Streams state directory and trigger a + * {@code LockException}. + */ class StateDirDefaultFactory implements DefaultValueFactory { @Override public String create(PipelineOptions options) { - return Paths.get(System.getProperty("java.io.tmpdir"), "beam-kafka-streams-state").toString(); + return Paths.get( + System.getProperty("java.io.tmpdir"), + "beam-kafka-streams-state", + options.getJobName()) + .toString(); } } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineResult.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineResult.java index 776eaa6629f5..65ff8b77b180 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineResult.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineResult.java @@ -43,23 +43,31 @@ public State getState() { @Override public State cancel() throws IOException { - State state = delegate.cancel(); - stopJobServer.run(); - return state; + try { + return delegate.cancel(); + } finally { + stopJobServer.run(); + } } @Override public State waitUntilFinish(Duration duration) { State state = delegate.waitUntilFinish(duration); - stopJobServer.run(); + // A null/non-terminal state means the wait timed out and the pipeline is still running; + // keep the job server alive so the caller can continue to interact with the job. + if (state != null && state.isTerminal()) { + stopJobServer.run(); + } return state; } @Override public State waitUntilFinish() { - State state = delegate.waitUntilFinish(); - stopJobServer.run(); - return state; + try { + return delegate.waitUntilFinish(); + } finally { + stopJobServer.run(); + } } @Override From cef6544e7920219e7c02e85f9d9902f6cc33ea7e Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Thu, 21 May 2026 20:33:28 +0500 Subject: [PATCH 03/37] Address review feedback on Kafka Streams Runner skeleton --- ...m_PreCommit_Java_Kafka_Streams_Runner.json | 4 + .github/workflows/README.md | 1 + ...am_PreCommit_Java_Kafka_Streams_Runner.yml | 118 ++++++++++++++++++ runners/kafka-streams/build.gradle | 5 + .../kafka/streams/KafkaStreamsJobInvoker.java | 22 +++- .../streams/KafkaStreamsJobServerDriver.java | 3 - .../streams/KafkaStreamsPipelineOptions.java | 7 -- .../streams/KafkaStreamsPipelineRunner.java | 3 - .../kafka/streams/KafkaStreamsRunner.java | 14 +-- .../KafkaStreamsJobServerDriverTest.java | 71 +++++++++++ .../KafkaStreamsPipelineOptionsTest.java | 74 +++++++++++ .../KafkaStreamsPipelineTranslatorTest.java | 113 +++++++++++++++++ 12 files changed, 408 insertions(+), 27 deletions(-) create mode 100644 .github/trigger_files/beam_PreCommit_Java_Kafka_Streams_Runner.json create mode 100644 .github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriverTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptionsTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java diff --git a/.github/trigger_files/beam_PreCommit_Java_Kafka_Streams_Runner.json b/.github/trigger_files/beam_PreCommit_Java_Kafka_Streams_Runner.json new file mode 100644 index 000000000000..5abe02fc09c7 --- /dev/null +++ b/.github/trigger_files/beam_PreCommit_Java_Kafka_Streams_Runner.json @@ -0,0 +1,4 @@ +{ + "comment": "Modify this file in a trivial way to cause this test suite to run.", + "modification": 1 +} diff --git a/.github/workflows/README.md b/.github/workflows/README.md index c6a95b29b4c0..e70d5e17d7a7 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -249,6 +249,7 @@ PreCommit Jobs run in a schedule and also get triggered in a PR if relevant sour | [ PreCommit Java HBase IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_HBase_IO_Direct.yml) | N/A |`Run Java_HBase_IO_Direct PreCommit`| [![.github/workflows/beam_PreCommit_Java_HBase_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_HBase_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_HBase_IO_Direct.yml?query=event%3Aschedule) | | [ PreCommit Java HCatalog IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_HCatalog_IO_Direct.yml) | N/A |`Run Java_HCatalog_IO_Direct PreCommit`| [![.github/workflows/beam_PreCommit_Java_HCatalog_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_HCatalog_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_HCatalog_IO_Direct.yml?query=event%3Aschedule) | | [ PreCommit Java Kafka IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Kafka_IO_Direct.yml) | N/A |`Run Java_Kafka_IO_Direct PreCommit`| [![.github/workflows/beam_PreCommit_Java_Kafka_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Kafka_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Kafka_IO_Direct.yml?query=event%3Aschedule) | +| [ PreCommit Java Kafka Streams Runner ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml) | N/A |`Run Java_Kafka_Streams_Runner PreCommit`| [![.github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml?query=event%3Aschedule) | | [ PreCommit Java InfluxDb IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_InfluxDb_IO_Direct.yml) | N/A |`Run Java_InfluxDb_IO_Direct PreCommit`| [![.github/workflows/beam_PreCommit_Java_InfluxDb_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_InfluxDb_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_InfluxDb_IO_Direct.yml?query=event%3Aschedule) | | [ PreCommit Java IOs Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_IOs_Direct.yml) | N/A |`Run Java_IOs_Direct PreCommit`| N/A | | [ PreCommit Java JDBC IO Direct ](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_JDBC_IO_Direct.yml) | N/A |`Run Java_JDBC_IO_Direct PreCommit`| [![.github/workflows/beam_PreCommit_Java_JDBC_IO_Direct.yml](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_JDBC_IO_Direct.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_PreCommit_Java_JDBC_IO_Direct.yml?query=event%3Aschedule) | diff --git a/.github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml b/.github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml new file mode 100644 index 000000000000..564b2bbc4bc8 --- /dev/null +++ b/.github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml @@ -0,0 +1,118 @@ +# 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. + +name: PreCommit Java Kafka Streams Runner + +on: + push: + tags: ['v*'] + branches: ['master', 'release-*'] + paths: + - "runners/kafka-streams/**" + - ".github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml" + pull_request_target: + branches: ['master', 'release-*'] + paths: + - "runners/kafka-streams/**" + - 'release/trigger_all_tests.json' + - '.github/trigger_files/beam_PreCommit_Java_Kafka_Streams_Runner.json' + issue_comment: + types: [created] + schedule: + - cron: '15 2/6 * * *' + workflow_dispatch: + +# Setting explicit permissions for the action to avoid the default permissions which are `write-all` in case of pull_request_target event +permissions: + actions: write + pull-requests: write + checks: write + contents: read + deployments: read + id-token: none + issues: write + discussions: read + packages: read + pages: read + repository-projects: read + security-events: read + statuses: read + +# This allows a subsequently queued workflow run to interrupt previous runs +concurrency: + group: '${{ github.workflow }} @ ${{ github.event.pull_request.number || github.event.pull_request.head.label || github.sha || github.head_ref || github.ref }}-${{ github.event.schedule || github.event.comment.id || github.event.sender.login }}' + cancel-in-progress: true + +env: + DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }} + GRADLE_ENTERPRISE_CACHE_USERNAME: ${{ secrets.GE_CACHE_USERNAME }} + GRADLE_ENTERPRISE_CACHE_PASSWORD: ${{ secrets.GE_CACHE_PASSWORD }} + +jobs: + beam_PreCommit_Java_Kafka_Streams_Runner: + name: ${{ matrix.job_name }} (${{ matrix.job_phrase }}) + strategy: + matrix: + job_name: ["beam_PreCommit_Java_Kafka_Streams_Runner"] + job_phrase: ["Run Java_Kafka_Streams_Runner PreCommit"] + timeout-minutes: 60 + if: | + github.event_name == 'push' || + github.event_name == 'pull_request_target' || + (github.event_name == 'schedule' && github.repository == 'apache/beam') || + github.event_name == 'workflow_dispatch' || + github.event.comment.body == 'Run Java_Kafka_Streams_Runner PreCommit' + runs-on: [self-hosted, ubuntu-24.04, main] + steps: + - uses: actions/checkout@v6 + - name: Setup repository + uses: ./.github/actions/setup-action + with: + comment_phrase: ${{ matrix.job_phrase }} + github_token: ${{ secrets.GITHUB_TOKEN }} + github_job: ${{ matrix.job_name }} (${{ matrix.job_phrase }}) + - name: Setup environment + uses: ./.github/actions/setup-environment-action + - name: run Kafka Streams runner build script + uses: ./.github/actions/gradle-command-self-hosted-action + with: + gradle-command: :runners:kafka-streams:build + max-workers: 4 + - name: Archive JUnit Test Results + uses: actions/upload-artifact@v7 + if: ${{ !success() }} + with: + name: JUnit Test Results + path: "**/build/reports/tests/" + - name: Publish JUnit Test Results + uses: EnricoMi/publish-unit-test-result-action@v2 + if: always() + with: + commit: '${{ env.prsha || env.GITHUB_SHA }}' + comment_mode: ${{ github.event_name == 'issue_comment' && 'always' || 'off' }} + files: '**/build/test-results/**/*.xml' + large_files: true + - name: Archive SpotBugs Results + uses: actions/upload-artifact@v7 + if: always() + with: + name: SpotBugs Results + path: '**/build/reports/spotbugs/*.html' + - name: Publish SpotBugs Results + uses: jwgmeligmeyling/spotbugs-github-action@v1.2 + if: always() + with: + name: Publish SpotBugs + path: '**/build/reports/spotbugs/*.html' diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 54474502ad7b..9204fef4e768 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -59,4 +59,9 @@ dependencies { implementation "org.apache.kafka:kafka-streams:$kafka_version" permitUnusedDeclared "org.apache.kafka:kafka-clients:$kafka_version" permitUnusedDeclared "org.apache.kafka:kafka-streams:$kafka_version" + + testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") + testImplementation library.java.hamcrest + testImplementation library.java.junit + testImplementation library.java.mockito_core } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java index bd1ed4c1bcf6..ad12e17da5af 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java @@ -23,17 +23,16 @@ import org.apache.beam.runners.jobsubmission.JobInvocation; import org.apache.beam.runners.jobsubmission.JobInvoker; import org.apache.beam.runners.jobsubmission.PortablePipelineRunner; +import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Struct; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ListeningExecutorService; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** Job invoker for the Kafka Streams portable runner. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public class KafkaStreamsJobInvoker extends JobInvoker { private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsJobInvoker.class); @@ -69,7 +68,7 @@ protected JobInvocation invokeWithExecutor( String.format("%s_%s", kafkaStreamsOptions.getJobName(), UUID.randomUUID().toString()); PortablePipelineRunner pipelineRunner = new KafkaStreamsPipelineRunner(kafkaStreamsOptions); - kafkaStreamsOptions.setRunner(null); + clearRunner(kafkaStreamsOptions); LOG.info("Invoking job {} with pipeline runner {}", invocationId, pipelineRunner); return createJobInvocation( @@ -83,7 +82,7 @@ protected JobInvocation invokeWithExecutor( protected JobInvocation createJobInvocation( String invocationId, - String retrievalToken, + @Nullable String retrievalToken, ListeningExecutorService executorService, RunnerApi.Pipeline pipeline, KafkaStreamsPipelineOptions kafkaStreamsOptions, @@ -92,8 +91,19 @@ protected JobInvocation createJobInvocation( JobInfo.create( invocationId, kafkaStreamsOptions.getJobName(), - retrievalToken, + Strings.nullToEmpty(retrievalToken), PipelineOptionsTranslation.toProto(kafkaStreamsOptions)); return new JobInvocation(jobInfo, executorService, pipeline, pipelineRunner); } + + /** + * Clears the runner class on the pipeline options before serialization. The Beam {@link + * PipelineOptions#setRunner} API accepts {@code null} to clear the runner — this mirrors the same + * pattern used by Flink/Spark portable runners so the runner class is not required on the SDK + * harness classpath. + */ + @SuppressWarnings("nullness") + private static void clearRunner(PipelineOptions options) { + options.setRunner(null); + } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriver.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriver.java index ddeac8b7c959..4f2134689253 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriver.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriver.java @@ -29,9 +29,6 @@ import org.slf4j.LoggerFactory; /** Driver that starts a Beam job server for the Kafka Streams portable runner. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public class KafkaStreamsJobServerDriver extends JobServerDriver { private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsJobServerDriver.class); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java index 911ffce2eea3..f46e0a024ed0 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java @@ -41,13 +41,6 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { void setApplicationId(String applicationId); - @Description( - "Kafka Streams processing.guarantee setting, for example at_least_once or exactly_once_v2.") - @Default.String("exactly_once_v2") - String getProcessingGuarantee(); - - void setProcessingGuarantee(String processingGuarantee); - @Description("Soft cap on the number of elements per bundle.") @Default.Integer(1000) int getMaxBundleSize(); 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 cc7464e22786..7c78d3a751b4 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 @@ -25,9 +25,6 @@ import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsTranslationContext; /** Executes a portable pipeline by translating it to Kafka Streams. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public class KafkaStreamsPipelineRunner implements PortablePipelineRunner { private final KafkaStreamsPipelineOptions pipelineOptions; diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java index ce8f0544b681..a6f3621d1ffd 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java @@ -28,6 +28,7 @@ import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.util.construction.Environments; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -35,9 +36,6 @@ * A {@link PipelineRunner} that submits portable jobs to an in-process or external Beam job service * backed by the Kafka Streams translation path. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public class KafkaStreamsRunner extends PipelineRunner { private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsRunner.class); @@ -55,7 +53,7 @@ protected KafkaStreamsRunner(KafkaStreamsPipelineOptions pipelineOptions) { @Override public PipelineResult run(Pipeline pipeline) { assignPortableDefaults(pipelineOptions); - KafkaStreamsJobServerDriver jobServerDriver = null; + @Nullable KafkaStreamsJobServerDriver jobServerDriver = null; try { if (Strings.isNullOrEmpty(pipelineOptions.getJobEndpoint())) { LOG.info("No job endpoint configured; starting an embedded Kafka Streams job server."); @@ -68,7 +66,8 @@ public PipelineResult run(Pipeline pipeline) { PortableRunner portableRunner = PortableRunner.fromOptions(pipelineOptions); PipelineResult result = portableRunner.run(pipeline); if (jobServerDriver != null) { - return new KafkaStreamsPipelineResult(result, jobServerDriver::stop); + KafkaStreamsJobServerDriver driverForStop = jobServerDriver; + return new KafkaStreamsPipelineResult(result, driverForStop::stop); } return result; } catch (IOException e) { @@ -84,10 +83,9 @@ private static void assignPortableDefaults(KafkaStreamsPipelineOptions pipelineO pipelineOptions.setDefaultEnvironmentType(Environments.ENVIRONMENT_LOOPBACK); } ExperimentalOptions experimentalOptions = pipelineOptions.as(ExperimentalOptions.class); + @Nullable List existingExperiments = experimentalOptions.getExperiments(); List experiments = - experimentalOptions.getExperiments() == null - ? new ArrayList<>() - : new ArrayList<>(experimentalOptions.getExperiments()); + existingExperiments == null ? new ArrayList<>() : new ArrayList<>(existingExperiments); if (!experiments.contains("beam_fn_api")) { experiments.add("beam_fn_api"); experimentalOptions.setExperiments(experiments); diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriverTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriverTest.java new file mode 100644 index 000000000000..8eecc708904d --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobServerDriverTest.java @@ -0,0 +1,71 @@ +/* + * 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; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.not; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; + +import org.junit.Test; + +/** Tests for {@link KafkaStreamsJobServerDriver}. */ +public class KafkaStreamsJobServerDriverTest { + + @Test + public void testConfigurationDefaults() { + KafkaStreamsJobServerDriver.KafkaStreamsServerConfiguration config = + new KafkaStreamsJobServerDriver.KafkaStreamsServerConfiguration(); + + assertThat(config.getHost(), is("localhost")); + assertThat(config.getPort(), is(8099)); + assertThat(config.getArtifactPort(), is(8098)); + assertThat(config.getExpansionPort(), is(8097)); + assertThat(config.isCleanArtifactsPerJob(), is(true)); + + KafkaStreamsJobServerDriver driver = KafkaStreamsJobServerDriver.fromConfig(config); + assertThat(driver, is(not(nullValue()))); + } + + @Test + public void testConfigurationFromArgs() { + KafkaStreamsJobServerDriver.KafkaStreamsServerConfiguration config = + KafkaStreamsJobServerDriver.parseArgs( + new String[] { + "--job-host=test-host", + "--job-port", + "42", + "--artifact-port", + "43", + "--expansion-port", + "44", + "--clean-artifacts-per-job=false", + }); + + assertThat(config.getHost(), is("test-host")); + assertThat(config.getPort(), is(42)); + assertThat(config.getArtifactPort(), is(43)); + assertThat(config.getExpansionPort(), is(44)); + assertThat(config.isCleanArtifactsPerJob(), is(false)); + } + + @Test(expected = IllegalArgumentException.class) + public void testInvalidArgsRejected() { + KafkaStreamsJobServerDriver.parseArgs(new String[] {"--unknown-flag=value"}); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptionsTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptionsTest.java new file mode 100644 index 000000000000..45f96637acfe --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptionsTest.java @@ -0,0 +1,74 @@ +/* + * 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; + +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.MatcherAssert.assertThat; + +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.junit.Test; + +/** Tests for {@link KafkaStreamsPipelineOptions}. */ +public class KafkaStreamsPipelineOptionsTest { + + @Test + public void testDefaultValues() { + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + + assertThat(options.getBootstrapServers(), is("localhost:9092")); + assertThat(options.getApplicationId(), is("beam-kafka-streams-runner")); + assertThat(options.getMaxBundleSize(), is(1000)); + assertThat(options.getMaxBundleTimeMs(), is(1000)); + assertThat(options.getStateDir(), is(notNullValue())); + assertThat(options.getStateDir(), containsString("beam-kafka-streams-state")); + assertThat(options.getStateDir(), containsString(options.getJobName())); + } + + @Test + public void testOverrides() { + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.fromArgs( + "--bootstrapServers=broker-1:9092,broker-2:9092", + "--applicationId=custom-app", + "--maxBundleSize=500", + "--maxBundleTimeMs=250", + "--stateDir=/var/data/beam-ks") + .as(KafkaStreamsPipelineOptions.class); + + assertThat(options.getBootstrapServers(), is("broker-1:9092,broker-2:9092")); + assertThat(options.getApplicationId(), is("custom-app")); + assertThat(options.getMaxBundleSize(), is(500)); + assertThat(options.getMaxBundleTimeMs(), is(250)); + assertThat(options.getStateDir(), is("/var/data/beam-ks")); + } + + @Test + public void testStateDirIsolatesByJobName() { + KafkaStreamsPipelineOptions optionsA = + PipelineOptionsFactory.fromArgs("--jobName=job-a").as(KafkaStreamsPipelineOptions.class); + KafkaStreamsPipelineOptions optionsB = + PipelineOptionsFactory.fromArgs("--jobName=job-b").as(KafkaStreamsPipelineOptions.class); + + assertThat(optionsA.getStateDir(), containsString("job-a")); + assertThat(optionsB.getStateDir(), containsString("job-b")); + assertThat(optionsA.getStateDir().equals(optionsB.getStateDir()), is(false)); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java new file mode 100644 index 000000000000..86e9cf878497 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java @@ -0,0 +1,113 @@ +/* + * 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.containsString; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertThrows; + +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.util.construction.PTransformTranslation; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.junit.Test; + +/** + * Tests for {@link KafkaStreamsPipelineTranslator}. + * + *

The skeleton translator does not yet handle any transforms; these tests pin the current + * "fail-fast with a clear URN" contract so that follow-up sub-issues can replace the assertions as + * real translators are added. + */ +public class KafkaStreamsPipelineTranslatorTest { + + private static final String JOB_ID = "kafka-streams-test-job"; + + @Test + public void translateRejectsUnknownTransformWithUrnInMessage() { + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + KafkaStreamsTranslationContext context = newContext(); + + RunnerApi.Pipeline pipeline = + RunnerApi.Pipeline.newBuilder() + .setComponents( + RunnerApi.Components.newBuilder() + .putTransforms( + "impulse", + RunnerApi.PTransform.newBuilder() + .setUniqueName("Impulse") + .setSpec( + RunnerApi.FunctionSpec.newBuilder() + .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) + .build())) + .build(); + + UnsupportedOperationException ex = + assertThrows( + UnsupportedOperationException.class, + () -> translator.translate(context, translator.prepareForTranslation(pipeline))); + + assertThat(ex.getMessage(), containsString("No translator registered for URN")); + assertThat(ex.getMessage(), containsString(PTransformTranslation.IMPULSE_TRANSFORM_URN)); + assertThat(ex.getMessage(), containsString(JOB_ID)); + } + + @Test + public void translateRejectsEmptyPipeline() { + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + KafkaStreamsTranslationContext context = newContext(); + + RunnerApi.Pipeline pipeline = + RunnerApi.Pipeline.newBuilder() + .setComponents(RunnerApi.Components.newBuilder().build()) + .build(); + + UnsupportedOperationException ex = + assertThrows( + UnsupportedOperationException.class, + () -> translator.translate(context, translator.prepareForTranslation(pipeline))); + + assertThat(ex.getMessage(), containsString("No translator registered")); + } + + @Test + public void createTranslationContextExposesJobInfoAndOptions() { + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + JobInfo jobInfo = + JobInfo.create( + JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); + + KafkaStreamsTranslationContext context = translator.createTranslationContext(jobInfo, options); + + assertThat(context.getJobInfo().jobId(), containsString(JOB_ID)); + assertThat(context.getPipelineOptions().getBootstrapServers(), containsString("localhost")); + } + + private static KafkaStreamsTranslationContext newContext() { + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + JobInfo jobInfo = + JobInfo.create( + JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); + return KafkaStreamsTranslationContext.create(jobInfo, options); + } +} From 0d445f79e0c61d8f2c69c6611abe6078ac608fd8 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Fri, 22 May 2026 19:09:38 +0500 Subject: [PATCH 04/37] Drop setRunner(null) suppression; make applicationId required --- .../kafka/streams/KafkaStreamsJobInvoker.java | 13 ---------- .../streams/KafkaStreamsPipelineOptions.java | 6 +++-- .../KafkaStreamsPipelineOptionsTest.java | 24 +++++++++++++++---- 3 files changed, 24 insertions(+), 19 deletions(-) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java index ad12e17da5af..a32c7c487737 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsJobInvoker.java @@ -23,7 +23,6 @@ import org.apache.beam.runners.jobsubmission.JobInvocation; import org.apache.beam.runners.jobsubmission.JobInvoker; import org.apache.beam.runners.jobsubmission.PortablePipelineRunner; -import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.Struct; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; @@ -68,7 +67,6 @@ protected JobInvocation invokeWithExecutor( String.format("%s_%s", kafkaStreamsOptions.getJobName(), UUID.randomUUID().toString()); PortablePipelineRunner pipelineRunner = new KafkaStreamsPipelineRunner(kafkaStreamsOptions); - clearRunner(kafkaStreamsOptions); LOG.info("Invoking job {} with pipeline runner {}", invocationId, pipelineRunner); return createJobInvocation( @@ -95,15 +93,4 @@ protected JobInvocation createJobInvocation( PipelineOptionsTranslation.toProto(kafkaStreamsOptions)); return new JobInvocation(jobInfo, executorService, pipeline, pipelineRunner); } - - /** - * Clears the runner class on the pipeline options before serialization. The Beam {@link - * PipelineOptions#setRunner} API accepts {@code null} to clear the runner — this mirrors the same - * pattern used by Flink/Spark portable runners so the runner class is not required on the SDK - * harness classpath. - */ - @SuppressWarnings("nullness") - private static void clearRunner(PipelineOptions options) { - options.setRunner(null); - } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java index f46e0a024ed0..2fa992e66e7e 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java @@ -23,6 +23,7 @@ import org.apache.beam.sdk.options.Description; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.options.Validation; /** Pipeline options for the Kafka Streams runner. */ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { @@ -35,8 +36,9 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { @Description( "Kafka Streams application.id (must be unique for each distinct topology using the same " - + "input topics in a Kafka cluster).") - @Default.String("beam-kafka-streams-runner") + + "input topics in a Kafka cluster). Must be specified explicitly: a shared default " + + "would let concurrent jobs collide on the same Kafka Streams consumer group.") + @Validation.Required String getApplicationId(); void setApplicationId(String applicationId); diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptionsTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptionsTest.java index 45f96637acfe..27e2a63cc1f5 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptionsTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptionsTest.java @@ -21,8 +21,10 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.notNullValue; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertThrows; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.PipelineOptionsValidator; import org.junit.Test; /** Tests for {@link KafkaStreamsPipelineOptions}. */ @@ -31,10 +33,10 @@ public class KafkaStreamsPipelineOptionsTest { @Test public void testDefaultValues() { KafkaStreamsPipelineOptions options = - PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + PipelineOptionsFactory.fromArgs("--applicationId=test-app") + .as(KafkaStreamsPipelineOptions.class); assertThat(options.getBootstrapServers(), is("localhost:9092")); - assertThat(options.getApplicationId(), is("beam-kafka-streams-runner")); assertThat(options.getMaxBundleSize(), is(1000)); assertThat(options.getMaxBundleTimeMs(), is(1000)); assertThat(options.getStateDir(), is(notNullValue())); @@ -42,6 +44,18 @@ public void testDefaultValues() { assertThat(options.getStateDir(), containsString(options.getJobName())); } + @Test + public void testApplicationIdIsRequired() { + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, + () -> PipelineOptionsValidator.validate(KafkaStreamsPipelineOptions.class, options)); + assertThat(ex.getMessage(), containsString("getApplicationId")); + } + @Test public void testOverrides() { KafkaStreamsPipelineOptions options = @@ -63,9 +77,11 @@ public void testOverrides() { @Test public void testStateDirIsolatesByJobName() { KafkaStreamsPipelineOptions optionsA = - PipelineOptionsFactory.fromArgs("--jobName=job-a").as(KafkaStreamsPipelineOptions.class); + PipelineOptionsFactory.fromArgs("--jobName=job-a", "--applicationId=app-a") + .as(KafkaStreamsPipelineOptions.class); KafkaStreamsPipelineOptions optionsB = - PipelineOptionsFactory.fromArgs("--jobName=job-b").as(KafkaStreamsPipelineOptions.class); + PipelineOptionsFactory.fromArgs("--jobName=job-b", "--applicationId=app-b") + .as(KafkaStreamsPipelineOptions.class); assertThat(optionsA.getStateDir(), containsString("job-a")); assertThat(optionsB.getStateDir(), containsString("job-b")); From d4c7dae98b8141b7aec80b95c75877995c5ee6b6 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Sat, 23 May 2026 19:58:11 +0500 Subject: [PATCH 05/37] Catch Exception in KafkaStreamsRunner.run() to avoid job-server leak --- .../beam/runners/kafka/streams/KafkaStreamsRunner.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java index a6f3621d1ffd..6a8f105bb2ee 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java @@ -17,7 +17,6 @@ */ package org.apache.beam.runners.kafka.streams; -import java.io.IOException; import java.util.ArrayList; import java.util.List; import org.apache.beam.runners.portability.PortableRunner; @@ -70,10 +69,13 @@ public PipelineResult run(Pipeline pipeline) { return new KafkaStreamsPipelineResult(result, driverForStop::stop); } return result; - } catch (IOException e) { + } catch (Exception e) { if (jobServerDriver != null) { jobServerDriver.stop(); } + if (e instanceof RuntimeException) { + throw (RuntimeException) e; + } throw new RuntimeException(e); } } From faefc95e9f05115b19d16a974c1310c594d4250d Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Thu, 28 May 2026 10:39:43 +0500 Subject: [PATCH 06/37] =?UTF-8?q?[GSoC=202026]=20Kafka=20Streams=20runner?= =?UTF-8?q?=20=E2=80=94=20translation=20framework=20+=20Impulse=20translat?= =?UTF-8?q?or=20(#38689)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add Impulse translator and URN-dispatch framework - KafkaStreamsPipelineTranslator now walks the pipeline in topological order via QueryablePipeline and dispatches each transform to a PTransformTranslator keyed by URN. Unknown URNs still fail fast with a clear "No translator registered for URN ..." message. - ImpulseTranslator implements beam:transform:impulse:v1 per design doc §4.1: a per-application bootstrap topic source (__beam_impulse_) satisfies Kafka Streams' real-source requirement, ImpulseProcessor emits exactly one WindowedValue in the GlobalWindow via a one-shot wall-clock punctuator scheduled on init, and a persistent state store records a "fired" flag so task restarts do not duplicate. Bootstrap-topic auto-creation is deferred to a follow-up sub-issue (design doc §12.1); the topic is expected to pre-exist in production. - KafkaStreamsTranslationContext now holds the Topology being built and a PCollection-id -> processor-name map so downstream translators can wire to their parent nodes. - KafkaStreamsPipelineRunner.run now translates and starts the KafkaStreams application, returning a KafkaStreamsPortablePipelineResult that maps KafkaStreams.State to Beam's PipelineResult.State. Forces processing.guarantee=exactly_once_v2. - Tests: * KafkaStreamsPipelineTranslatorTest also covers the Impulse success path; the unsupported-URN check now uses GroupByKey. * ImpulseTranslatorTest exercises the topology via TopologyTestDriver: exactly one empty byte[] in GlobalWindow is emitted, and a second wall-clock advance does not re-emit. * Address review feedback on translation framework + Impulse PR - KafkaStreamsPortablePipelineResult: close the race where KafkaStreams could transition to a terminal state before the state listener was registered, leaving waitUntilFinish() to block forever. Also add a volatile cancelled flag so that getState() returns State.CANCELLED after a user cancel(), instead of mapping NOT_RUNNING to State.DONE. - ImpulseProcessor: capture the Cancellable returned by context.schedule and cancel the wall-clock punctuator once the impulse has fired (or if the state store already records a prior emission), so the processor stops doing periodic state-store lookups for the lifetime of the task. - KafkaStreamsPipelineRunner.run: invoke PipelineOptionsValidator.validate on the pipeline options at the start of run() so a missing required option (e.g. applicationId) fails with a clear IllegalArgumentException rather than a raw NullPointerException on Properties.put further down. - ImpulseTranslatorTest: wrap CapturingProcessor.received in Collections.synchronizedList for best-practice thread-safety even though TopologyTestDriver runs single-threaded. * Address review feedback on Impulse translator --- runners/kafka-streams/build.gradle | 3 +- .../streams/KafkaStreamsPipelineRunner.java | 38 ++++- .../KafkaStreamsPortablePipelineResult.java | 135 +++++++++++++++ .../streams/translation/ImpulseProcessor.java | 140 +++++++++++++++ .../translation/ImpulseTranslator.java | 90 ++++++++++ .../streams/translation/KStreamsPayload.java | 126 ++++++++++++++ .../KafkaStreamsPipelineTranslator.java | 65 ++++--- .../KafkaStreamsTranslationContext.java | 62 ++++++- .../translation/PTransformTranslator.java | 41 +++++ .../kafka/streams/KafkaStreamsRunnerTest.java | 161 ++++++++++++++++++ .../translation/ImpulseTranslatorTest.java | 143 ++++++++++++++++ .../KafkaStreamsPipelineTranslatorTest.java | 73 +++++--- 12 files changed, 1020 insertions(+), 57 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/PTransformTranslator.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 9204fef4e768..3f34a3ca76b6 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -44,6 +44,7 @@ dependencies { implementation project(path: ":sdks:java:core", configuration: "shadow") implementation project(path: ":model:pipeline", configuration: "shadow") + implementation project(path: ":model:job-management", configuration: "shadow") implementation project(":runners:core-java") permitUnusedDeclared project(":runners:core-java") implementation project(":runners:java-fn-execution") @@ -58,10 +59,10 @@ dependencies { implementation "org.apache.kafka:kafka-clients:$kafka_version" implementation "org.apache.kafka:kafka-streams:$kafka_version" permitUnusedDeclared "org.apache.kafka:kafka-clients:$kafka_version" - permitUnusedDeclared "org.apache.kafka:kafka-streams:$kafka_version" testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") testImplementation library.java.hamcrest testImplementation library.java.junit testImplementation library.java.mockito_core + testImplementation "org.apache.kafka:kafka-streams-test-utils:$kafka_version" } 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 7c78d3a751b4..3e97638695e1 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 @@ -17,16 +17,25 @@ */ package org.apache.beam.runners.kafka.streams; +import java.util.Properties; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.jobsubmission.PortablePipelineResult; import org.apache.beam.runners.jobsubmission.PortablePipelineRunner; import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsPipelineTranslator; import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsTranslationContext; +import org.apache.beam.sdk.options.PipelineOptionsValidator; +import org.apache.kafka.streams.KafkaStreams; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.Topology; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; -/** Executes a portable pipeline by translating it to Kafka Streams. */ +/** Executes a portable pipeline by translating it to a Kafka Streams {@link Topology}. */ public class KafkaStreamsPipelineRunner implements PortablePipelineRunner { + private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsPipelineRunner.class); + private final KafkaStreamsPipelineOptions pipelineOptions; public KafkaStreamsPipelineRunner(KafkaStreamsPipelineOptions pipelineOptions) { @@ -34,12 +43,35 @@ public KafkaStreamsPipelineRunner(KafkaStreamsPipelineOptions pipelineOptions) { } @Override - public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) throws Exception { + public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) { + // Surface a clear error if a required option (e.g. applicationId) is missing instead of + // letting Properties.put fail with a raw NullPointerException further down. + PipelineOptionsValidator.validate(KafkaStreamsPipelineOptions.class, pipelineOptions); + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); KafkaStreamsTranslationContext context = translator.createTranslationContext(jobInfo, pipelineOptions); RunnerApi.Pipeline prepared = translator.prepareForTranslation(pipeline); translator.translate(context, prepared); - throw new IllegalStateException("Translation unexpectedly completed without an executor"); + + Topology topology = context.getTopology(); + LOG.info( + "Translated pipeline {} into Kafka Streams topology:\n{}", + jobInfo.jobId(), + topology.describe()); + + KafkaStreams kafkaStreams = new KafkaStreams(topology, streamsConfig(jobInfo)); + kafkaStreams.start(); + return new KafkaStreamsPortablePipelineResult(kafkaStreams); + } + + private Properties streamsConfig(JobInfo jobInfo) { + Properties props = new Properties(); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, pipelineOptions.getBootstrapServers()); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, pipelineOptions.getApplicationId()); + props.put(StreamsConfig.STATE_DIR_CONFIG, pipelineOptions.getStateDir()); + props.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2); + props.put(StreamsConfig.CLIENT_ID_CONFIG, jobInfo.jobId()); + return props; } } 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 new file mode 100644 index 000000000000..817746bf002f --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java @@ -0,0 +1,135 @@ +/* + * 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; + +import java.io.IOException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.apache.beam.model.jobmanagement.v1.JobApi; +import org.apache.beam.runners.jobsubmission.PortablePipelineResult; +import org.apache.beam.sdk.metrics.MetricResults; +import org.apache.kafka.streams.KafkaStreams; +import org.joda.time.Duration; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Result of executing a portable pipeline as a {@link KafkaStreams} application. + * + *

Translates the underlying {@link KafkaStreams.State} into Beam's {@link + * org.apache.beam.sdk.PipelineResult.State} and forwards {@link #cancel()} / {@link + * #waitUntilFinish()} to the {@code KafkaStreams} instance. + */ +class KafkaStreamsPortablePipelineResult implements PortablePipelineResult { + + private static final Logger LOG = + LoggerFactory.getLogger(KafkaStreamsPortablePipelineResult.class); + + private final KafkaStreams kafkaStreams; + private final CountDownLatch terminated = new CountDownLatch(1); + private volatile boolean cancelled = false; + + KafkaStreamsPortablePipelineResult(KafkaStreams kafkaStreams) { + this.kafkaStreams = kafkaStreams; + kafkaStreams.setStateListener( + (newState, oldState) -> { + if (newState == KafkaStreams.State.NOT_RUNNING || newState == KafkaStreams.State.ERROR) { + terminated.countDown(); + } + }); + // Guard against the race where the KafkaStreams instance transitions to a terminal state + // (e.g. immediate startup failure) before the state listener is registered above. Without + // this check, the latch would never be counted down and waitUntilFinish() would block forever. + KafkaStreams.State current = kafkaStreams.state(); + if (current == KafkaStreams.State.NOT_RUNNING || current == KafkaStreams.State.ERROR) { + terminated.countDown(); + } + } + + @Override + public State getState() { + if (cancelled) { + return State.CANCELLED; + } + return mapState(kafkaStreams.state()); + } + + @Override + public State cancel() throws IOException { + cancelled = true; + kafkaStreams.close(); + terminated.countDown(); + return getState(); + } + + @Override + public State waitUntilFinish(Duration duration) { + try { + boolean reachedTerminal = terminated.await(duration.getMillis(), TimeUnit.MILLISECONDS); + if (!reachedTerminal) { + return getState(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return State.UNKNOWN; + } + return getState(); + } + + @Override + public State waitUntilFinish() { + try { + terminated.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return State.UNKNOWN; + } + return getState(); + } + + @Override + public MetricResults metrics() { + throw new UnsupportedOperationException( + "Metrics are not yet implemented in the Kafka Streams runner."); + } + + @Override + public JobApi.MetricResults portableMetrics() throws UnsupportedOperationException { + LOG.debug("portableMetrics() not yet implemented in the Kafka Streams runner"); + return JobApi.MetricResults.newBuilder().build(); + } + + private static State mapState(KafkaStreams.State state) { + switch (state) { + case CREATED: + case REBALANCING: + return State.RUNNING; + case RUNNING: + return State.RUNNING; + case PENDING_SHUTDOWN: + return State.CANCELLED; + case PENDING_ERROR: + case ERROR: + return State.FAILED; + case NOT_RUNNING: + return State.DONE; + default: + return State.UNKNOWN; + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java new file mode 100644 index 000000000000..e9fc8ddb36aa --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java @@ -0,0 +1,140 @@ +/* + * 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 java.time.Duration; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.kafka.streams.processor.Cancellable; +import org.apache.kafka.streams.processor.PunctuationType; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Kafka Streams {@link Processor} implementing Beam's {@code Impulse} transform. + * + *

For each task instance, emits exactly two {@link KStreamsPayload}s downstream: + * + *

    + *
  1. A {@link KStreamsPayload#data data} payload wrapping a {@link WindowedValue} of an empty + * {@code byte[]} in the {@link org.apache.beam.sdk.transforms.windowing.GlobalWindow}, with + * event-time {@link BoundedWindow#TIMESTAMP_MIN_VALUE}. + *
  2. A {@link KStreamsPayload#watermark watermark} payload at {@link + * BoundedWindow#TIMESTAMP_MAX_VALUE} that tells downstream transforms the source is done. + *
+ * + *

A persistent state store records whether the data element has already been emitted so that + * task restarts do not duplicate the data. The terminal watermark, on the other hand, is re-emitted + * on every restart so downstream watermark holds release correctly after recovery (per Jan's review + * on PR #38689). + * + *

The trigger comes from a wall-clock punctuator scheduled on {@link #init} — this lets the + * processor fire even when the dedicated bootstrap source topic is empty, which is the expected + * production state. + * + *

Kafka Streams disallows negative record timestamps, so the forwarded {@link Record} carries + * the Unix epoch ({@code 0L}). The Beam event-time lives inside the {@link KStreamsPayload} + * variant: inside the {@link WindowedValue} for data, or as the explicit watermark millis. + */ +class ImpulseProcessor implements Processor> { + + private static final Logger LOG = LoggerFactory.getLogger(ImpulseProcessor.class); + + /** Sole entry in the state store; the value tracks whether this processor has already emitted. */ + static final String FIRED_KEY = "fired"; + + /** How soon after {@link #init} the punctuator first fires. */ + private static final Duration PUNCTUATION_DELAY = Duration.ofMillis(50); + + private final String stateStoreName; + private final String transformId; + + private @Nullable ProcessorContext> context; + private @Nullable KeyValueStore firedStore; + private @Nullable Cancellable scheduledPunctuator; + + ImpulseProcessor(String stateStoreName, String transformId) { + this.stateStoreName = stateStoreName; + this.transformId = transformId; + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + this.firedStore = context.getStateStore(stateStoreName); + this.scheduledPunctuator = + context.schedule(PUNCTUATION_DELAY, PunctuationType.WALL_CLOCK_TIME, ts -> maybeFire()); + } + + @Override + public void process(Record record) { + // Records that happen to land on the bootstrap topic are not actual data; they just provide an + // extra opportunity to fire the impulse on restart. The state store still gates the emit. + maybeFire(); + } + + private void maybeFire() { + ProcessorContext> ctx = context; + KeyValueStore store = firedStore; + if (ctx == null || store == null) { + return; + } + if (Boolean.TRUE.equals(store.get(FIRED_KEY))) { + // Data was already emitted in a previous task lifetime, but downstream watermark holds may + // still need to be released after the restart — re-emit the terminal watermark and stop the + // punctuator. + forwardWatermarkMax(ctx); + cancelPunctuator(); + return; + } + WindowedValue impulse = WindowedValues.valueInGlobalWindow(new byte[0]); + // The output PCollection is not keyed (PCollection); use an empty byte[] as a + // placeholder key so downstream processors that adopt the byte[]-key convention see a + // consistent shape. + ctx.forward( + new Record>( + new byte[0], KStreamsPayload.data(impulse), 0L)); + forwardWatermarkMax(ctx); + store.put(FIRED_KEY, Boolean.TRUE); + cancelPunctuator(); + LOG.debug("Impulse {} emitted single element and terminal watermark", transformId); + } + + /** Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. */ + private static void forwardWatermarkMax(ProcessorContext> ctx) { + long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); + ctx.forward( + new Record>( + new byte[0], KStreamsPayload.watermark(maxMillis), 0L)); + } + + /** Cancels the wall-clock punctuator after the impulse has fired to stop periodic wakeups. */ + private void cancelPunctuator() { + Cancellable handle = scheduledPunctuator; + if (handle != null) { + handle.cancel(); + scheduledPunctuator = null; + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java new file mode 100644 index 000000000000..a90987ba6383 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java @@ -0,0 +1,90 @@ +/* + * 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 org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.state.KeyValueBytesStoreSupplier; +import org.apache.kafka.streams.state.Stores; + +/** + * Translates the {@code beam:transform:impulse:v1} URN. + * + *

Adds three nodes to the Kafka Streams {@link Topology}: + * + *

    + *
  • A {@code byte[]} source bound to a dedicated per-application bootstrap topic (see {@link + * KafkaStreamsTranslationContext#getImpulseBootstrapTopic()}). Kafka Streams refuses to start + * a topology that has no real source topic, so the bootstrap topic exists purely to satisfy + * that requirement — records published to it are ignored by {@link ImpulseProcessor}. + *
  • The {@link ImpulseProcessor} itself, which schedules a one-shot wall-clock punctuator on + * {@code init} and emits a single empty data {@link KStreamsPayload} followed by a terminal + * watermark payload at {@link + * org.apache.beam.sdk.transforms.windowing.BoundedWindow#TIMESTAMP_MAX_VALUE}. + *
  • A per-processor {@link KeyValueBytesStoreSupplier persistent state store} that records + * whether the impulse has already fired so task restarts do not duplicate it. + *
+ * + *

The processor's output PCollection is registered with the translation context so subsequent + * translators can wire themselves to this node by id. + * + *

Bootstrap topic lifecycle: this translator does not auto-create the bootstrap + * topic. The topic is expected to exist on the broker before the job starts; otherwise Kafka + * Streams raises {@code MissingSourceTopicException} on startup. The auto-create-vs-pre-create + * decision (design doc §12.1) is deferred to a follow-up sub-issue along with the {@code + * AdminClient} wiring; pre-creation is sufficient for the {@code TopologyTestDriver}-based unit + * tests in this PR. + */ +class ImpulseTranslator implements PTransformTranslator { + + static final String SOURCE_SUFFIX = "-source"; + static final String STATE_STORE_SUFFIX = "-state"; + + @Override + public void translate( + String transformId, RunnerApi.Pipeline pipeline, KafkaStreamsTranslationContext context) { + RunnerApi.PTransform transform = pipeline.getComponents().getTransformsOrThrow(transformId); + // Impulse produces exactly one output PCollection. This is the produced-outputs map on the + // transform, not the consumer count — downstream transforms that consume this PCollection are + // modeled as separate PTransforms whose `inputs` reference the same PCollection id, and they + // are wired up by their own translators. Iterables.getOnlyElement throws a clear + // IllegalArgumentException if the proto is malformed. + String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + + Topology topology = context.getTopology(); + String sourceNodeName = transformId + SOURCE_SUFFIX; + String stateStoreName = transformId + STATE_STORE_SUFFIX; + String bootstrapTopic = context.getImpulseBootstrapTopic(); + + topology.addSource( + sourceNodeName, + Serdes.ByteArray().deserializer(), + Serdes.ByteArray().deserializer(), + bootstrapTopic); + topology.addProcessor( + transformId, () -> new ImpulseProcessor(stateStoreName, transformId), sourceNodeName); + topology.addStateStore( + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(stateStoreName), Serdes.String(), Serdes.Boolean()), + transformId); + + context.registerPCollectionProducer(outputPCollectionId, transformId); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java new file mode 100644 index 000000000000..47c94eea6eff --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java @@ -0,0 +1,126 @@ +/* + * 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 java.util.Objects; +import org.apache.beam.sdk.values.WindowedValue; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Sum-type envelope flowing between Kafka Streams processors in the Beam Kafka Streams runner. + * + *

Every record value emitted by a runner-introduced processor is one of: + * + *

    + *
  • A {@link #isData() data} element wrapping a {@link WindowedValue}, or + *
  • A {@link #isWatermark() watermark} signal carrying an event-time milliseconds value. + *
+ * + *

The envelope lets a single Kafka Streams output channel carry both Beam data and the watermark + * / synchronization primitives that Kafka Streams does not natively support. Future control + * messages (e.g. the {@code (epoch, assigned_partitions)} propagation from design doc §5) can be + * added here as additional variants. + * + *

This class is intentionally in-JVM only for now; serialization across topic boundaries + * (repartition or sink topics) will be introduced when the first translator that emits to a topic + * lands, at which point a corresponding Kafka {@link org.apache.kafka.common.serialization.Serde} + * will be added. + * + * @param element type carried by data variants + */ +public final class KStreamsPayload { + + private enum Kind { + DATA, + WATERMARK + } + + private final Kind kind; + private final @Nullable WindowedValue data; + private final long watermarkMillis; + + private KStreamsPayload(Kind kind, @Nullable WindowedValue data, long watermarkMillis) { + this.kind = kind; + this.data = data; + this.watermarkMillis = watermarkMillis; + } + + /** Returns a data payload wrapping the given {@link WindowedValue}. */ + public static KStreamsPayload data(WindowedValue value) { + return new KStreamsPayload<>(Kind.DATA, value, 0L); + } + + /** Returns a watermark payload carrying the given event-time milliseconds. */ + public static KStreamsPayload watermark(long watermarkMillis) { + return new KStreamsPayload<>(Kind.WATERMARK, null, watermarkMillis); + } + + public boolean isData() { + return kind == Kind.DATA; + } + + public boolean isWatermark() { + return kind == Kind.WATERMARK; + } + + /** + * Returns the wrapped data element. Caller must check {@link #isData()} first; calling this on a + * watermark payload throws. + */ + public WindowedValue getData() { + if (kind != Kind.DATA || data == null) { + throw new IllegalStateException("Payload is not a data element: kind=" + kind); + } + return data; + } + + /** + * Returns the watermark event-time milliseconds. Caller must check {@link #isWatermark()} first; + * calling this on a data payload throws. + */ + public long getWatermarkMillis() { + if (kind != Kind.WATERMARK) { + throw new IllegalStateException("Payload is not a watermark: kind=" + kind); + } + return watermarkMillis; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof KStreamsPayload)) { + return false; + } + KStreamsPayload that = (KStreamsPayload) o; + return kind == that.kind + && watermarkMillis == that.watermarkMillis + && Objects.equals(data, that.data); + } + + @Override + public int hashCode() { + return Objects.hash(kind, data, watermarkMillis); + } + + @Override + public String toString() { + return kind == Kind.DATA ? "Data{" + data + "}" : "Watermark{" + watermarkMillis + "}"; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index cc915d604b68..eb8567146143 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -18,19 +18,38 @@ package org.apache.beam.runners.kafka.streams.translation; import java.util.Map; -import java.util.TreeMap; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.util.construction.PTransformTranslation; +import org.apache.beam.sdk.util.construction.graph.PipelineNode; +import org.apache.beam.sdk.util.construction.graph.QueryablePipeline; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; /** - * Translates a portable Beam pipeline into a Kafka Streams {@code Topology}. + * Translates a portable Beam pipeline into a Kafka Streams {@link + * org.apache.kafka.streams.Topology}. * - *

The initial implementation only validates the graph and fails fast with an explicit message - * for transforms that are not yet supported. + *

Walks the pipeline in topological order via {@link QueryablePipeline} and dispatches each + * transform to a {@link PTransformTranslator} keyed by URN. Transforms whose URN has no registered + * translator fail fast with a clear {@link UnsupportedOperationException} so the failure points at + * the exact transform that is not yet supported. */ public class KafkaStreamsPipelineTranslator { + private final Map urnToTranslator; + + public KafkaStreamsPipelineTranslator() { + this( + ImmutableMap.builder() + .put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator()) + .build()); + } + + KafkaStreamsPipelineTranslator(Map urnToTranslator) { + this.urnToTranslator = urnToTranslator; + } + public KafkaStreamsTranslationContext createTranslationContext( JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { return KafkaStreamsTranslationContext.create(jobInfo, pipelineOptions); @@ -42,29 +61,27 @@ public RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) { } /** - * Translates the pipeline. Throws {@link UnsupportedOperationException} with a clear URN message - * for the first unsupported primitive encountered. + * Walks the pipeline in topological order and translates each transform whose URN is supported. + * Throws {@link UnsupportedOperationException} on the first unsupported URN. */ public void translate(KafkaStreamsTranslationContext context, RunnerApi.Pipeline pipeline) { - Map transforms = pipeline.getComponents().getTransformsMap(); - TreeMap ordered = new TreeMap<>(transforms); - for (Map.Entry entry : ordered.entrySet()) { - RunnerApi.PTransform transform = entry.getValue(); - if (!transform.hasSpec()) { - continue; - } - String urn = transform.getSpec().getUrn(); - if (urn.isEmpty()) { - continue; + QueryablePipeline queryable = + QueryablePipeline.forTransforms( + pipeline.getRootTransformIdsList(), pipeline.getComponents()); + for (PipelineNode.PTransformNode node : queryable.getTopologicallyOrderedTransforms()) { + String urn = node.getTransform().getSpec().getUrn(); + PTransformTranslator translator = urnToTranslator.get(urn); + if (translator == null) { + throw new UnsupportedOperationException( + "No translator registered for URN " + + urn + + " (transformId=" + + node.getId() + + ", jobId=" + + context.getJobInfo().jobId() + + ")"); } - throw new UnsupportedOperationException( - "No translator registered for URN " - + urn - + " (jobId=" - + context.getJobInfo().jobId() - + ")"); + translator.translate(node.getId(), pipeline, context); } - throw new UnsupportedOperationException( - "No translator registered for pipeline (no transform URNs found)"); } } 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 7c6d3d079159..3d95eabafed6 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 @@ -17,24 +17,44 @@ */ package org.apache.beam.runners.kafka.streams.translation; +import java.util.HashMap; +import java.util.Map; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.kafka.streams.Topology; -/** Mutable state shared while translating a portable pipeline into a Kafka Streams topology. */ +/** + * Mutable state shared while translating a portable pipeline into a Kafka Streams {@link Topology}. + * + *

Holds the topology being built and a {@code PCollection-id → processor-node-name} map so that + * downstream transforms can wire themselves to the right parent node. + */ public class KafkaStreamsTranslationContext { + /** Prefix for the per-job bootstrap topic Impulse reads from. */ + private static final String IMPULSE_BOOTSTRAP_TOPIC_PREFIX = "__beam_impulse_"; + private final JobInfo jobInfo; private final KafkaStreamsPipelineOptions pipelineOptions; + private final Topology topology; + private final Map pCollectionIdToProcessorName; public static KafkaStreamsTranslationContext create( JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { - return new KafkaStreamsTranslationContext(jobInfo, pipelineOptions); + return new KafkaStreamsTranslationContext(jobInfo, pipelineOptions, new Topology()); + } + + static KafkaStreamsTranslationContext createWithTopology( + JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions, Topology topology) { + return new KafkaStreamsTranslationContext(jobInfo, pipelineOptions, topology); } private KafkaStreamsTranslationContext( - JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { + JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions, Topology topology) { this.jobInfo = jobInfo; this.pipelineOptions = pipelineOptions; + this.topology = topology; + this.pCollectionIdToProcessorName = new HashMap<>(); } public JobInfo getJobInfo() { @@ -44,4 +64,40 @@ public JobInfo getJobInfo() { public KafkaStreamsPipelineOptions getPipelineOptions() { return pipelineOptions; } + + /** Returns the {@link Topology} being built by the translation. */ + public Topology getTopology() { + return topology; + } + + /** + * Registers the processor node that produces the given Beam PCollection. Downstream translators + * resolve their parent processor names by looking up the input PCollection id. + */ + public void registerPCollectionProducer(String pCollectionId, String processorName) { + String existing = pCollectionIdToProcessorName.putIfAbsent(pCollectionId, processorName); + if (existing != null && !existing.equals(processorName)) { + throw new IllegalStateException( + "PCollection " + + pCollectionId + + " already produced by processor " + + existing + + "; cannot reassign to " + + processorName); + } + } + + /** Returns the processor node name producing the given PCollection. */ + public String getProcessorNameForPCollection(String pCollectionId) { + String name = pCollectionIdToProcessorName.get(pCollectionId); + if (name == null) { + throw new IllegalStateException("No processor registered for PCollection " + pCollectionId); + } + return name; + } + + /** Returns the dedicated bootstrap topic name used by Impulse for this application. */ + public String getImpulseBootstrapTopic() { + return IMPULSE_BOOTSTRAP_TOPIC_PREFIX + pipelineOptions.getApplicationId(); + } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/PTransformTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/PTransformTranslator.java new file mode 100644 index 000000000000..c6bd44b8049c --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/PTransformTranslator.java @@ -0,0 +1,41 @@ +/* + * 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 org.apache.beam.model.pipeline.v1.RunnerApi; + +/** + * Translates a single Beam {@link RunnerApi.PTransform} into one or more nodes in the Kafka Streams + * {@link org.apache.kafka.streams.Topology} held by the {@link KafkaStreamsTranslationContext}. + */ +@FunctionalInterface +public interface PTransformTranslator { + + /** + * Translates the transform identified by {@code transformId} into nodes on the topology held by + * {@code context}. + * + * @param transformId the id of the transform to translate, as keyed in {@code + * pipeline.getComponents().getTransformsMap()} + * @param pipeline the full pipeline proto, in case the translator needs to walk subcomponents + * (coders, windowing strategies, etc.) + * @param context the shared translation context (topology under construction, producer map, etc.) + */ + void translate( + String transformId, RunnerApi.Pipeline pipeline, KafkaStreamsTranslationContext context); +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java new file mode 100644 index 000000000000..30000d77d864 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java @@ -0,0 +1,161 @@ +/* + * 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; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Properties; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.translation.KStreamsPayload; +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.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.TopologyTestDriver; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.ProcessorSupplier; +import org.apache.kafka.streams.processor.api.Record; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Test; + +/** + * Pipeline-level integration tests that build a Beam {@link Pipeline} via the high-level Java SDK + * ({@code Pipeline.create().apply(Impulse.create())}), translate it via the runner, and execute the + * resulting Kafka Streams topology under {@link TopologyTestDriver}. + * + *

This is the test layer Jan requested on PR #38689: rather than building hand-rolled {@link + * RunnerApi.Pipeline} protos, drive translation from the same surface a user would write. The tests + * stop short of calling {@code pipeline.run()} because that would require a real Kafka broker — + * {@code TopologyTestDriver} replaces the broker for unit-test purposes. + */ +public class KafkaStreamsRunnerTest { + + private static final String JOB_ID = "kafka-streams-runner-test"; + private static final String APPLICATION_ID = "ks-runner-test"; + + @Test + public void impulseOnlyPipelineEmitsDataAndTerminalWatermark() { + Pipeline pipeline = Pipeline.create(pipelineOptions()); + pipeline.apply("impulse", Impulse.create()); + + RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); + + KafkaStreamsPipelineOptions options = + pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + JobInfo jobInfo = + JobInfo.create( + JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); + KafkaStreamsTranslationContext context = translator.createTranslationContext(jobInfo, options); + translator.translate(context, translator.prepareForTranslation(pipelineProto)); + + CapturingProcessor capture = new CapturingProcessor(); + Topology topology = context.getTopology(); + // Wire a downstream test sink to every translated transform node so we can capture emissions. + // Impulse is the only transform here, so we attach to "impulse" (the processor name registered + // by ImpulseTranslator). + topology.addProcessor("capture", capture, expectedImpulseProcessorName(pipelineProto)); + + try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { + driver.advanceWallClockTime(Duration.ofSeconds(1)); + driver.advanceWallClockTime(Duration.ofSeconds(1)); + } + + assertThat(capture.received.size(), is(2)); + assertThat(capture.received.get(0).isData(), is(true)); + assertThat(capture.received.get(0).getData().getValue().length, is(0)); + assertThat(capture.received.get(1).isWatermark(), is(true)); + assertThat( + capture.received.get(1).getWatermarkMillis(), + is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); + } + + /** + * Finds the transform id that {@link Impulse} got assigned by the SDK so the test can attach a + * capturing processor to the matching Kafka Streams processor node (the translator names the + * processor after the transform id). + */ + private static String expectedImpulseProcessorName(RunnerApi.Pipeline pipelineProto) { + for (java.util.Map.Entry entry : + pipelineProto.getComponents().getTransformsMap().entrySet()) { + if ("beam:transform:impulse:v1".equals(entry.getValue().getSpec().getUrn())) { + return entry.getKey(); + } + } + throw new IllegalStateException("Impulse transform not found in pipeline proto"); + } + + private static PipelineOptions pipelineOptions() { + PipelineOptions options = + PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); + // Pipeline.create() requires a runner; CrashingRunner is the conventional "this pipeline is + // not going to be run() directly" choice used by other portable-runner tests. + options.setRunner(CrashingRunner.class); + options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); + return options; + } + + private static Properties streamsConfig() { + Properties props = new Properties(); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put( + StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + props.put( + StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + return props; + } + + private static class CapturingProcessor + implements ProcessorSupplier< + byte[], KStreamsPayload, byte[], KStreamsPayload> { + + final List> received = Collections.synchronizedList(new ArrayList<>()); + + @Override + public Processor, byte[], KStreamsPayload> get() { + return new Processor, byte[], KStreamsPayload>() { + @Override + public void init(@Nullable ProcessorContext> context) { + // no-op + } + + @Override + public void process(Record> record) { + received.add(record.value()); + } + }; + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java new file mode 100644 index 000000000000..a092b10e02c4 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java @@ -0,0 +1,143 @@ +/* + * 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.CoreMatchers.notNullValue; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Properties; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.TopologyTestDriver; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.ProcessorSupplier; +import org.apache.kafka.streams.processor.api.Record; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Test; + +/** + * Behavioural tests for {@link ImpulseTranslator} using {@link TopologyTestDriver}. + * + *

The translator builds a topology with a real source + processor pair. The tests sit a {@link + * CapturingProcessor} downstream so emitted {@link KStreamsPayload} elements can be inspected + * directly without going through a Kafka sink topic (the runner does not produce one because no + * downstream PCollections exist yet). + */ +public class ImpulseTranslatorTest { + + @Test + public void impulseEmitsDataElementFollowedByTerminalWatermark() { + KafkaStreamsTranslationContext context = KafkaStreamsPipelineTranslatorTest.newContext(); + new KafkaStreamsPipelineTranslator() + .translate(context, KafkaStreamsPipelineTranslatorTest.singleImpulsePipeline()); + + CapturingProcessor capture = new CapturingProcessor(); + Topology topology = context.getTopology(); + topology.addProcessor("capture", capture, "impulse"); + + try (TopologyTestDriver driver = new TopologyTestDriver(topology, baseProps())) { + driver.advanceWallClockTime(Duration.ofSeconds(1)); + driver.advanceWallClockTime(Duration.ofSeconds(1)); + } + + assertThat(capture.received.size(), is(2)); + + KStreamsPayload dataPayload = capture.received.get(0); + assertThat(dataPayload, is(notNullValue())); + assertThat(dataPayload.isData(), is(true)); + WindowedValue data = dataPayload.getData(); + assertThat(data.getValue().length, is(0)); + assertThat(data.getWindows().size(), is(1)); + assertThat(data.getTimestamp().getMillis(), is(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis())); + + KStreamsPayload watermarkPayload = capture.received.get(1); + assertThat(watermarkPayload.isWatermark(), is(true)); + assertThat( + watermarkPayload.getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); + } + + @Test + public void impulseDoesNotReEmitDataOnRepeatedPunctuation() { + KafkaStreamsTranslationContext context = KafkaStreamsPipelineTranslatorTest.newContext(); + new KafkaStreamsPipelineTranslator() + .translate(context, KafkaStreamsPipelineTranslatorTest.singleImpulsePipeline()); + + CapturingProcessor capture = new CapturingProcessor(); + Topology topology = context.getTopology(); + topology.addProcessor("capture", capture, "impulse"); + + try (TopologyTestDriver driver = new TopologyTestDriver(topology, baseProps())) { + driver.advanceWallClockTime(Duration.ofSeconds(1)); + // Trigger again — the data element is gated by the state store; the punctuator should also + // have been cancelled after the first emission, so no further events should be captured. + driver.advanceWallClockTime(Duration.ofSeconds(5)); + } + + assertThat(capture.received.size(), is(2)); + assertThat(capture.received.get(0).isData(), is(true)); + assertThat(capture.received.get(1).isWatermark(), is(true)); + } + + private static Properties baseProps() { + Properties props = new Properties(); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, "ks-translator-test"); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put( + StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + props.put( + StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + return props; + } + + /** + * Captures {@link KStreamsPayload} records forwarded by {@link ImpulseProcessor}. The supplier + * returns a fresh forwarder each call (required by Kafka Streams) but all forwarders write into + * the shared {@link #received} list so the test can read the captured elements after the topology + * is closed. + */ + private static class CapturingProcessor + implements ProcessorSupplier< + byte[], KStreamsPayload, byte[], KStreamsPayload> { + + final List> received = Collections.synchronizedList(new ArrayList<>()); + + @Override + public Processor, byte[], KStreamsPayload> get() { + return new Processor, byte[], KStreamsPayload>() { + @Override + public void init(@Nullable ProcessorContext> context) { + // no-op + } + + @Override + public void process(Record> record) { + received.add(record.value()); + } + }; + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java index 86e9cf878497..44cc00bcebbd 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java @@ -18,6 +18,7 @@ package org.apache.beam.runners.kafka.streams.translation; import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertThrows; @@ -27,18 +28,14 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.util.construction.PTransformTranslation; import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.kafka.streams.TopologyDescription; import org.junit.Test; -/** - * Tests for {@link KafkaStreamsPipelineTranslator}. - * - *

The skeleton translator does not yet handle any transforms; these tests pin the current - * "fail-fast with a clear URN" contract so that follow-up sub-issues can replace the assertions as - * real translators are added. - */ +/** Tests for {@link KafkaStreamsPipelineTranslator}. */ public class KafkaStreamsPipelineTranslatorTest { private static final String JOB_ID = "kafka-streams-test-job"; + private static final String OUTPUT_PCOLLECTION_ID = "impulse.out"; @Test public void translateRejectsUnknownTransformWithUrnInMessage() { @@ -47,15 +44,16 @@ public void translateRejectsUnknownTransformWithUrnInMessage() { RunnerApi.Pipeline pipeline = RunnerApi.Pipeline.newBuilder() + .addRootTransformIds("gbk") .setComponents( RunnerApi.Components.newBuilder() .putTransforms( - "impulse", + "gbk", RunnerApi.PTransform.newBuilder() - .setUniqueName("Impulse") + .setUniqueName("GroupByKey") .setSpec( RunnerApi.FunctionSpec.newBuilder() - .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) + .setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)) .build())) .build(); @@ -65,33 +63,31 @@ public void translateRejectsUnknownTransformWithUrnInMessage() { () -> translator.translate(context, translator.prepareForTranslation(pipeline))); assertThat(ex.getMessage(), containsString("No translator registered for URN")); - assertThat(ex.getMessage(), containsString(PTransformTranslation.IMPULSE_TRANSFORM_URN)); + assertThat(ex.getMessage(), containsString(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)); + assertThat(ex.getMessage(), containsString("gbk")); assertThat(ex.getMessage(), containsString(JOB_ID)); } @Test - public void translateRejectsEmptyPipeline() { + public void translateImpulsePipelineAddsSourceAndProcessorNodes() { KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); KafkaStreamsTranslationContext context = newContext(); - RunnerApi.Pipeline pipeline = - RunnerApi.Pipeline.newBuilder() - .setComponents(RunnerApi.Components.newBuilder().build()) - .build(); + RunnerApi.Pipeline pipeline = singleImpulsePipeline(); + translator.translate(context, translator.prepareForTranslation(pipeline)); - UnsupportedOperationException ex = - assertThrows( - UnsupportedOperationException.class, - () -> translator.translate(context, translator.prepareForTranslation(pipeline))); + TopologyDescription description = context.getTopology().describe(); + String describeText = description.toString(); - assertThat(ex.getMessage(), containsString("No translator registered")); + assertThat(describeText, containsString("impulse-source")); + assertThat(describeText, containsString("impulse")); + assertThat(context.getProcessorNameForPCollection(OUTPUT_PCOLLECTION_ID), is("impulse")); } @Test public void createTranslationContextExposesJobInfoAndOptions() { KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); - KafkaStreamsPipelineOptions options = - PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + KafkaStreamsPipelineOptions options = testOptions(); JobInfo jobInfo = JobInfo.create( JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); @@ -102,12 +98,37 @@ public void createTranslationContextExposesJobInfoAndOptions() { assertThat(context.getPipelineOptions().getBootstrapServers(), containsString("localhost")); } - private static KafkaStreamsTranslationContext newContext() { - KafkaStreamsPipelineOptions options = - PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + static RunnerApi.Pipeline singleImpulsePipeline() { + return RunnerApi.Pipeline.newBuilder() + .addRootTransformIds("impulse") + .setComponents( + RunnerApi.Components.newBuilder() + .putTransforms( + "impulse", + RunnerApi.PTransform.newBuilder() + .setUniqueName("Impulse") + .putOutputs("out", OUTPUT_PCOLLECTION_ID) + .setSpec( + RunnerApi.FunctionSpec.newBuilder() + .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) + .build()) + .putPcollections( + OUTPUT_PCOLLECTION_ID, + RunnerApi.PCollection.newBuilder().setUniqueName(OUTPUT_PCOLLECTION_ID).build()) + .build()) + .build(); + } + + static KafkaStreamsTranslationContext newContext() { + KafkaStreamsPipelineOptions options = testOptions(); JobInfo jobInfo = JobInfo.create( JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); return KafkaStreamsTranslationContext.create(jobInfo, options); } + + static KafkaStreamsPipelineOptions testOptions() { + return PipelineOptionsFactory.fromArgs("--applicationId=ks-translator-test") + .as(KafkaStreamsPipelineOptions.class); + } } From 0f875a071627e824b124d9538a631cf150da1095 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Thu, 28 May 2026 16:27:59 +0500 Subject: [PATCH 07/37] Add temporary feature-branch CI for Kafka Streams runner (#38725) Beam's standard PreCommit workflows gate on self-hosted runners and a branches: [master, release-*] filter, so they never trigger for PRs targeting the feat/18479-* integration branch. This lightweight github-hosted workflow bypasses that: checkout -> JDK 11 -> ./gradlew :runners:kafka-streams:build. To be removed when the runner work merges to master and the standard PreCommit takes over. Refs #18479 --- .../beam_KafkaStreamsRunner_FeatureBranch.yml | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 .github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml diff --git a/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml b/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml new file mode 100644 index 000000000000..4ec7b0778cc2 --- /dev/null +++ b/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml @@ -0,0 +1,72 @@ +# 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. + +# Temporary, feature-branch-only build for the Kafka Streams runner GSoC work. +# +# Beam's standard PreCommit workflows are gated on self-hosted runners and a +# `branches: ['master', 'release-*']` filter, so they never trigger for PRs +# targeting the `feat/18479-*` integration branch. This lightweight workflow +# bypasses that machinery: it runs on a GitHub-hosted runner and simply builds +# and tests the `runners/kafka-streams` module. +# +# REMOVE THIS WORKFLOW when the Kafka Streams runner work merges to master and +# the standard PreCommit (beam_PreCommit_Java_Kafka_Streams_Runner.yml) takes +# over. + +name: KafkaStreams Runner Feature Branch Build + +on: + pull_request: + branches: + - 'feat/18479-*' + paths: + - 'runners/kafka-streams/**' + - '.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml' + push: + branches: + - 'feat/18479-*' + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: '${{ github.workflow }} @ ${{ github.event.pull_request.number || github.ref }}' + cancel-in-progress: true + +jobs: + build: + name: Kafka Streams runner build + runs-on: ubuntu-latest + timeout-minutes: 60 + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Set up JDK 11 + uses: actions/setup-java@v4 + with: + distribution: temurin + java-version: '11' + - name: Set up Gradle cache + uses: actions/cache@v4 + with: + path: | + ~/.gradle/caches + ~/.gradle/wrapper + key: ${{ runner.os }}-gradle-kafka-streams-${{ hashFiles('**/*.gradle', '**/*.gradle.kts', 'gradle/wrapper/gradle-wrapper.properties') }} + restore-keys: | + ${{ runner.os }}-gradle-kafka-streams- + - name: Build and test Kafka Streams runner + run: ./gradlew :runners:kafka-streams:build --no-daemon --stacktrace From ff554f7438d658aa887f59ca78942340473a4ebd Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Sun, 7 Jun 2026 17:13:55 +0500 Subject: [PATCH 08/37] =?UTF-8?q?[GSoC=202026]=20Kafka=20Streams=20runner?= =?UTF-8?q?=20=E2=80=94=20ExecutableStage=20(stateless=20ParDo)=20translat?= =?UTF-8?q?or=20(#38764)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add ExecutableStage (stateless ParDo) translator with SDK-harness bridge --- runners/kafka-streams/build.gradle | 1 + .../translation/ExecutableStageProcessor.java | 217 ++++++++++++++++++ .../ExecutableStageTranslator.java | 96 ++++++++ .../streams/translation/KStreamsPayload.java | 9 +- ...aStreamsExecutableStageContextFactory.java | 66 ++++++ .../KafkaStreamsPipelineTranslator.java | 19 +- .../ExecutableStageTranslatorTest.java | 140 +++++++++++ .../KafkaStreamsPipelineTranslatorTest.java | 28 ++- .../translation/SharedTestCollector.java | 92 ++++++++ 9 files changed, 657 insertions(+), 11 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsExecutableStageContextFactory.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslatorTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/SharedTestCollector.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 3f34a3ca76b6..52b320dd70a8 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -61,6 +61,7 @@ dependencies { permitUnusedDeclared "org.apache.kafka:kafka-clients:$kafka_version" testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") + testImplementation project(":sdks:java:harness") testImplementation library.java.hamcrest testImplementation library.java.junit testImplementation library.java.mockito_core 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 new file mode 100644 index 000000000000..7417088bb672 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java @@ -0,0 +1,217 @@ +/* + * 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 java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.control.BundleProgressHandler; +import org.apache.beam.runners.fnexecution.control.ExecutableStageContext; +import org.apache.beam.runners.fnexecution.control.OutputReceiverFactory; +import org.apache.beam.runners.fnexecution.control.RemoteBundle; +import org.apache.beam.runners.fnexecution.control.StageBundleFactory; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.fnexecution.state.StateRequestHandler; +import org.apache.beam.sdk.fn.data.FnDataReceiver; +import org.apache.beam.sdk.util.construction.graph.ExecutableStage; +import org.apache.beam.sdk.values.WindowedValue; +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; +import org.apache.kafka.streams.processor.api.Record; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Kafka Streams {@link Processor} that executes a fused {@link ExecutableStage} (stateless user + * code such as ParDo) in the Beam SDK harness over the Fn API. + * + *

For each {@link KStreamsPayload#isData() data} payload it unwraps the {@link WindowedValue} + * and feeds it to the harness through the stage's main input {@link FnDataReceiver}. Harness + * outputs are collected on the harness threads into {@link #pendingOutputs} and then flushed + * downstream on the Kafka Streams processing thread when the bundle closes — Kafka Streams' {@link + * ProcessorContext#forward} must only be called from the processing thread, so outputs are never + * forwarded directly from a harness callback. + * + *

A {@link KStreamsPayload#isWatermark() watermark} payload marks a bundle boundary: the open + * bundle (if any) is closed (flushing outputs), and the watermark is then forwarded downstream so + * that subsequent stages observe it after all data of the bundle. + * + *

This is the Kafka Streams analogue of Flink's {@code ExecutableStageDoFnOperator} and Spark's + * {@code SparkExecutableStageFunction}. State, timers, and side inputs are out of scope for this + * first version: the stage is executed with {@link StateRequestHandler#unsupported()} and no timer + * receivers. + */ +class ExecutableStageProcessor + implements Processor, byte[], KStreamsPayload> { + + private static final Logger LOG = LoggerFactory.getLogger(ExecutableStageProcessor.class); + + private final RunnerApi.ExecutableStagePayload stagePayload; + private final JobInfo jobInfo; + + // 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. + private final Queue> pendingOutputs = new ConcurrentLinkedQueue<>(); + + private @Nullable ProcessorContext> context; + private @Nullable ExecutableStageContext stageContext; + private @Nullable StageBundleFactory stageBundleFactory; + private @Nullable RemoteBundle currentBundle; + + ExecutableStageProcessor(RunnerApi.ExecutableStagePayload stagePayload, JobInfo jobInfo) { + this.stagePayload = stagePayload; + this.jobInfo = jobInfo; + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload); + this.stageContext = KafkaStreamsExecutableStageContextFactory.getInstance().get(jobInfo); + this.stageBundleFactory = stageContext.getStageBundleFactory(executableStage); + } + + @Override + public void process(Record> record) { + KStreamsPayload payload = record.value(); + if (payload.isWatermark()) { + // NOTE: flushing the bundle on every received watermark is provisional. Once the + // WatermarkManager lands, a stage will receive watermarks from multiple parent instances and + // the output watermark becomes min() across them — the bundle should flush / the output + // watermark advance only when that minimum actually moves forward, not on every received + // watermark. Tracked in #38743. + closeBundleAndFlush(record); + forwardWatermark(record, payload.getWatermarkMillis()); + return; + } + try { + ensureBundleOpen(); + mainInputReceiver().accept(payload.getData()); + } catch (Exception e) { + throw new RuntimeException("Failed to process element through SDK harness", e); + } + } + + private void ensureBundleOpen() throws Exception { + if (currentBundle != null) { + return; + } + StageBundleFactory factory = checkInitialized(stageBundleFactory); + OutputReceiverFactory outputReceiverFactory = + new OutputReceiverFactory() { + @Override + public FnDataReceiver create(String pCollectionId) { + // Outputs are queued here on harness threads and drained on the processing thread + // after the bundle closes. + return receivedElement -> { + if (receivedElement != null) { + pendingOutputs.add((WindowedValue) receivedElement); + } + }; + } + }; + currentBundle = + factory.getBundle( + outputReceiverFactory, + StateRequestHandler.unsupported(), + BundleProgressHandler.ignored()); + } + + private FnDataReceiver> mainInputReceiver() { + RemoteBundle bundle = checkInitialized(currentBundle); + @SuppressWarnings("unchecked") + FnDataReceiver> receiver = + (FnDataReceiver>) + (FnDataReceiver) Iterables.getOnlyElement(bundle.getInputReceivers().values()); + return receiver; + } + + private void closeBundleAndFlush(Record> record) { + RemoteBundle bundle = currentBundle; + if (bundle == null) { + return; + } + try { + // close() blocks until the harness finishes the bundle and all outputs have been delivered + // to the output receiver (and hence enqueued in pendingOutputs). + bundle.close(); + } catch (Exception e) { + throw new RuntimeException("Failed to close SDK harness bundle", e); + } finally { + currentBundle = null; + } + ProcessorContext> ctx = checkInitialized(context); + // The harness has finished the bundle (close() returned) so no further enqueues happen. + // ConcurrentLinkedQueue's weakly-consistent iterator is therefore safe to drain via forEach. + pendingOutputs.forEach( + output -> + ctx.forward( + new Record>( + record.key(), KStreamsPayload.data(output), record.timestamp()))); + pendingOutputs.clear(); + } + + private void forwardWatermark( + Record> record, long watermarkMillis) { + // TODO(#38743 / WatermarkManager): a watermark must reach every parallel instance of every + // downstream processor, but ctx.forward routes to one downstream partition per Kafka Streams' + // partitioning. The simplest correct approach is to fan the watermark out to all downstream + // partitions; that wiring lands with the WatermarkManager sub-issue (per Jan on PR #38764). + ProcessorContext> ctx = checkInitialized(context); + ctx.forward( + new Record>( + record.key(), KStreamsPayload.watermark(watermarkMillis), record.timestamp())); + } + + @Override + public void close() { + try { + if (currentBundle != null) { + currentBundle.close(); + currentBundle = null; + } + } catch (Exception e) { + LOG.warn("Error closing in-flight SDK harness bundle", e); + } + try { + if (stageBundleFactory != null) { + stageBundleFactory.close(); + stageBundleFactory = null; + } + } catch (Exception e) { + LOG.warn("Error closing stage bundle factory", e); + } + try { + if (stageContext != null) { + stageContext.close(); + stageContext = null; + } + } catch (Exception e) { + LOG.warn("Error closing executable stage context", e); + } + } + + private static T checkInitialized(@Nullable T value) { + if (value == null) { + throw new IllegalStateException("ExecutableStageProcessor used before init()"); + } + return value; + } +} 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 new file mode 100644 index 000000000000..dc56d57f57c2 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java @@ -0,0 +1,96 @@ +/* + * 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 java.io.IOException; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.kafka.streams.Topology; + +/** + * Translates the {@code beam:runner:executable_stage:v1} URN. + * + *

Adds an {@link ExecutableStageProcessor} node to the topology, wired to the processor that + * produces the stage's input PCollection (resolved through {@link + * KafkaStreamsTranslationContext#getProcessorNameForPCollection}). The processor runs the fused + * user code in the SDK harness; its single output PCollection is registered so downstream + * translators can attach to this node. + * + *

Multi-output stages (additional outputs / side inputs / state / timers) are out of scope for + * this first version and are rejected so the limitation fails fast rather than silently dropping + * outputs. + */ +class ExecutableStageTranslator implements PTransformTranslator { + + @Override + public void translate( + String transformId, RunnerApi.Pipeline pipeline, KafkaStreamsTranslationContext context) { + RunnerApi.PTransform transform = pipeline.getComponents().getTransformsOrThrow(transformId); + + RunnerApi.ExecutableStagePayload stagePayload; + try { + stagePayload = RunnerApi.ExecutableStagePayload.parseFrom(transform.getSpec().getPayload()); + } catch (IOException e) { + throw new IllegalArgumentException( + "Failed to parse ExecutableStagePayload for transform " + transformId, e); + } + + // Fail fast on stage features that are not yet supported, so users get a clear message rather + // than a silent miss further down the harness/topology path. + if (stagePayload.getSideInputsCount() > 0) { + throw new UnsupportedOperationException( + "ExecutableStage " + + transformId + + " has side inputs; side inputs are not yet supported by the Kafka Streams runner."); + } + if (stagePayload.getUserStatesCount() > 0 || stagePayload.getTimersCount() > 0) { + throw new UnsupportedOperationException( + "ExecutableStage " + + transformId + + " uses user state or timers; stateful ParDo is not yet supported by the Kafka" + + " Streams runner."); + } + if (transform.getOutputsMap().size() > 1) { + // Multi-output stages (DoFns with side outputs, etc.) are a planned follow-up — they need + // an output-tag dispatch in the processor + per-output PCollection routing. The current + // rejection just fails loudly until that's wired in. + throw new UnsupportedOperationException( + "ExecutableStage " + + transformId + + " has " + + transform.getOutputsMap().size() + + " outputs; multi-output stages are not yet supported by the Kafka Streams runner."); + } + + // The payload distinguishes the main input from side inputs, so reading it from the payload + // is unambiguous even before we add side-input support. + String inputPCollectionId = stagePayload.getInput(); + String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); + + Topology topology = context.getTopology(); + topology.addProcessor( + transformId, + () -> new ExecutableStageProcessor(stagePayload, context.getJobInfo()), + parentProcessor); + + if (!transform.getOutputsMap().isEmpty()) { + String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + context.registerPCollectionProducer(outputPCollectionId, transformId); + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java index 47c94eea6eff..53e47b1216bb 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java @@ -19,6 +19,7 @@ import java.util.Objects; import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -121,6 +122,12 @@ public int hashCode() { @Override public String toString() { - return kind == Kind.DATA ? "Data{" + data + "}" : "Watermark{" + watermarkMillis + "}"; + MoreObjects.ToStringHelper helper = MoreObjects.toStringHelper(this).add("kind", kind); + if (kind == Kind.DATA) { + helper.add("data", data); + } else { + helper.add("watermarkMillis", watermarkMillis); + } + return helper.toString(); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsExecutableStageContextFactory.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsExecutableStageContextFactory.java new file mode 100644 index 000000000000..d376ea0ba455 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsExecutableStageContextFactory.java @@ -0,0 +1,66 @@ +/* + * 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 java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import org.apache.beam.runners.fnexecution.control.DefaultExecutableStageContext; +import org.apache.beam.runners.fnexecution.control.ExecutableStageContext; +import org.apache.beam.runners.fnexecution.control.ReferenceCountingExecutableStageContextFactory; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; + +/** + * Provides one {@link ExecutableStageContext.Factory} per job for the Kafka Streams runner. + * + *

Mirrors {@code FlinkExecutableStageContextFactory}: a singleton that hands out reference- + * counted {@link DefaultExecutableStageContext}s keyed by job id, so the SDK harness environment + * for a job is created once and shared across the {@link ImpulseProcessor}/executable-stage + * processors that run within the same JVM instance. + */ +public class KafkaStreamsExecutableStageContextFactory implements ExecutableStageContext.Factory { + + private static final KafkaStreamsExecutableStageContextFactory INSTANCE = + new KafkaStreamsExecutableStageContextFactory(); + + private final ConcurrentMap jobFactories = + new ConcurrentHashMap<>(); + + private KafkaStreamsExecutableStageContextFactory() {} + + public static KafkaStreamsExecutableStageContextFactory getInstance() { + return INSTANCE; + } + + @Override + public ExecutableStageContext get(JobInfo jobInfo) { + ExecutableStageContext.Factory jobFactory = + jobFactories.computeIfAbsent( + jobInfo.jobId(), + k -> + ReferenceCountingExecutableStageContextFactory.create( + DefaultExecutableStageContext::create, + // Release the context synchronously once its reference count drops to zero, + // and also drop the per-job factory entry so a long-lived JVM that runs many + // jobs does not accumulate one entry per finished job. + (caller) -> { + jobFactories.remove(k); + return true; + })); + return jobFactory.get(jobInfo); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index eb8567146143..5042f5426169 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -22,6 +22,8 @@ import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; import org.apache.beam.sdk.util.construction.PTransformTranslation; +import org.apache.beam.sdk.util.construction.graph.ExecutableStage; +import org.apache.beam.sdk.util.construction.graph.GreedyPipelineFuser; import org.apache.beam.sdk.util.construction.graph.PipelineNode; import org.apache.beam.sdk.util.construction.graph.QueryablePipeline; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; @@ -43,6 +45,7 @@ public KafkaStreamsPipelineTranslator() { this( ImmutableMap.builder() .put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator()) + .put(ExecutableStage.URN, new ExecutableStageTranslator()) .build()); } @@ -55,9 +58,21 @@ public KafkaStreamsTranslationContext createTranslationContext( return KafkaStreamsTranslationContext.create(jobInfo, pipelineOptions); } - /** Returns the pipeline to translate (placeholder for future fusion / expansion steps). */ + /** + * Fuses the pipeline so that stateless user code is grouped into {@code ExecutableStage} nodes. + * + *

Runner-executed primitives that have their own translator (e.g. Impulse) are left intact; + * everything else is fused. If the pipeline already contains {@code ExecutableStage} transforms + * it is returned unchanged. + */ public RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) { - return pipeline; + boolean alreadyFused = + pipeline.getComponents().getTransformsMap().values().stream() + .anyMatch(t -> ExecutableStage.URN.equals(t.getSpec().getUrn())); + if (alreadyFused) { + return pipeline; + } + return GreedyPipelineFuser.fuse(pipeline).toPipeline(); } /** diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslatorTest.java new file mode 100644 index 000000000000..98ec45ae2885 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslatorTest.java @@ -0,0 +1,140 @@ +/* + * 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 java.time.Duration; +import java.util.List; +import java.util.Properties; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.TopologyTestDriver; +import org.junit.Test; + +/** + * End-to-end test for {@link ExecutableStageTranslator}: builds an {@code Impulse -> ParDo} + * pipeline with the high-level Beam Java SDK, fuses + translates it, and runs the resulting Kafka + * Streams topology under {@link TopologyTestDriver}. The fused ParDo executes in an in-process + * (EMBEDDED) Java SDK harness, so the {@link DoFn}'s {@code @ProcessElement} body runs for real — + * no Docker, no broker. + * + *

Because the ParDo's output PCollection has no downstream consumer, it is not a stage output + * and is never forwarded out of the harness — that is the documented behaviour. The test verifies + * the bridge works by having the DoFn record into a {@link SharedTestCollector} as a side effect + * and asserting the recorded input from the test thread. + */ +public class ExecutableStageTranslatorTest { + + private static final String JOB_ID = "kafka-streams-executable-stage-test"; + private static final String APPLICATION_ID = "ks-executable-stage-test"; + + /** + * Records the length of every input element seen by the harness so the test can verify the DoFn + * ran. {@link SharedTestCollector} carries its identity via a UUID stored on the instance itself, + * so it survives any serialization the runner may perform on the DoFn. + */ + private static class RecordingFn extends DoFn { + private final SharedTestCollector collector; + + RecordingFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element byte[] input, OutputReceiver out) { + collector.record(input.length); + // Still emit something so the output codepath of the harness is exercised, even though no + // downstream consumer means the runner never observes the value. + out.output(new byte[] {1}); + } + } + + @Test + public void impulseThenParDoExecutesDoFnInHarnessOncePerImpulseElement() throws Exception { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(pipelineOptions()); + pipeline + .apply("impulse", Impulse.create()) + .apply("pardo", ParDo.of(new RecordingFn(collector))); + + RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); + + KafkaStreamsPipelineOptions options = + pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + JobInfo jobInfo = + JobInfo.create( + JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); + KafkaStreamsTranslationContext context = + translator.createTranslationContext(jobInfo, options); + + translator.translate(context, translator.prepareForTranslation(pipelineProto)); + + Topology topology = context.getTopology(); + try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { + driver.advanceWallClockTime(Duration.ofSeconds(1)); + driver.advanceWallClockTime(Duration.ofSeconds(1)); + } + + List recorded = collector.recorded(); + // Impulse emits exactly one empty byte[] in the GlobalWindow, so the DoFn must run exactly + // once and see a zero-length input. + assertThat(recorded.size(), is(1)); + assertThat(recorded.get(0), is(0)); + } + } + + private static PipelineOptions pipelineOptions() { + PipelineOptions options = + PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); + options.setRunner(CrashingRunner.class); + options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); + options + .as(PortablePipelineOptions.class) + .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); + return options; + } + + private static Properties streamsConfig() { + Properties props = new Properties(); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put( + StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + props.put( + StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + return props; + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java index 44cc00bcebbd..13baa551ebbf 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java @@ -18,16 +18,19 @@ package org.apache.beam.runners.kafka.streams.translation; import static org.hamcrest.CoreMatchers.containsString; -import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertThrows; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.transforms.Impulse; import org.apache.beam.sdk.util.construction.PTransformTranslation; import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; import org.apache.kafka.streams.TopologyDescription; import org.junit.Test; @@ -57,10 +60,11 @@ public void translateRejectsUnknownTransformWithUrnInMessage() { .build())) .build(); + // translate() directly — this test pins the URN-rejection contract on the dispatch loop + // itself, independent of the fuser/validator that prepareForTranslation runs. UnsupportedOperationException ex = assertThrows( - UnsupportedOperationException.class, - () -> translator.translate(context, translator.prepareForTranslation(pipeline))); + UnsupportedOperationException.class, () -> translator.translate(context, pipeline)); assertThat(ex.getMessage(), containsString("No translator registered for URN")); assertThat(ex.getMessage(), containsString(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)); @@ -73,15 +77,23 @@ public void translateImpulsePipelineAddsSourceAndProcessorNodes() { KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); KafkaStreamsTranslationContext context = newContext(); - RunnerApi.Pipeline pipeline = singleImpulsePipeline(); + // Build the pipeline through the SDK so the resulting RunnerApi.Pipeline carries the coders + // and windowing strategies that PipelineValidator requires (run inside the fuser). + Pipeline sdkPipeline = + Pipeline.create( + PipelineOptionsFactory.fromArgs( + "--applicationId=ks-translator-test", + "--runner=" + CrashingRunner.class.getName()) + .create()); + sdkPipeline.apply("impulse", Impulse.create()); + RunnerApi.Pipeline pipeline = PipelineTranslation.toProto(sdkPipeline); + translator.translate(context, translator.prepareForTranslation(pipeline)); TopologyDescription description = context.getTopology().describe(); String describeText = description.toString(); - - assertThat(describeText, containsString("impulse-source")); - assertThat(describeText, containsString("impulse")); - assertThat(context.getProcessorNameForPCollection(OUTPUT_PCOLLECTION_ID), is("impulse")); + assertThat(describeText, containsString("Source:")); + assertThat(describeText, containsString("Processor:")); } @Test diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/SharedTestCollector.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/SharedTestCollector.java new file mode 100644 index 000000000000..0986b0f2dc19 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/SharedTestCollector.java @@ -0,0 +1,92 @@ +/* + * 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 java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Test-only side-effect sink that survives Beam serialization without losing collected elements. + * + *

An ExecutableStage that contains a user {@link org.apache.beam.sdk.transforms.DoFn} runs the + * DoFn in the SDK harness even when its output PCollection has no downstream consumer — the work is + * still performed for its side effects. The natural unit test for that is to have the DoFn record + * into a side-effect container and assert the container's contents from the test thread. + * + *

A plain static {@code AtomicReference} / {@code List} works only as long as the runner does + * not serialize the {@code DoFn} (and therefore the container instance it holds). The EMBEDDED + * environment may already, and could in the future, serialize the user code, in which case a cloned + * container would silently drop its writes. + * + *

This class works around that by keying the actual storage on a {@link UUID} held by an + * otherwise-empty instance. The instance itself is cheaply {@link Serializable}; clones still carry + * the same {@code UUID} and therefore see the same backing list in the static {@link #REGISTRY}. + * + *

Implements {@link AutoCloseable} so tests can use try-with-resources to drop the per-UUID + * entry from {@link #REGISTRY} once they finish reading the recorded elements — without {@code + * close}, a long-lived JVM running many tests would accumulate one orphan entry per test. + * + * @param element type + */ +final class SharedTestCollector implements Serializable, AutoCloseable { + + private static final long serialVersionUID = 1L; + + /** Per-UUID storage, populated lazily on the first {@code record} for each instance. */ + private static final Map> REGISTRY = new ConcurrentHashMap<>(); + + private final UUID id = UUID.randomUUID(); + + /** Returns a fresh, empty collector instance with its own UUID. */ + static SharedTestCollector create() { + return new SharedTestCollector<>(); + } + + /** Records a single element. Safe to call from any thread. */ + void record(T element) { + REGISTRY.computeIfAbsent(id, k -> Collections.synchronizedList(new ArrayList<>())).add(element); + } + + /** Returns an immutable snapshot of all recorded elements, in order. */ + @SuppressWarnings("unchecked") + List recorded() { + List raw = REGISTRY.get(id); + if (raw == null) { + return Collections.emptyList(); + } + synchronized (raw) { + return Collections.unmodifiableList(new ArrayList<>((List) (List) raw)); + } + } + + /** + * Removes the per-UUID entry from the static registry. After {@code close}, any subsequent {@link + * #recorded()} call returns an empty list and {@link #record(Object)} will repopulate the entry + * for any new writes — but the typical use is to call {@code close} once at the end of the test + * (via try-with-resources) to keep {@link #REGISTRY} from accumulating orphan entries. + */ + @Override + public void close() { + REGISTRY.remove(id); + } +} From 4c02a7789a233aaf84c06caf19bd7e4d86a1308c Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Wed, 10 Jun 2026 14:24:27 +0500 Subject: [PATCH 09/37] =?UTF-8?q?[GSoC=202026]=20Kafka=20Streams=20runner?= =?UTF-8?q?=20=E2=80=94=20Redistribute=20translator=20+=20ExecutableStage?= =?UTF-8?q?=20type-agnostic=20edge=20(#38843)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add Redistribute (arbitrarily) translator + type-agnostic ExecutableStage edge * Drain pendingOutputs via poll() loop instead of forEach + clear * Assert ChainedExecutableStageTest pipeline actually has two ExecutableStages --- .../translation/ExecutableStageProcessor.java | 43 ++--- .../KafkaStreamsPipelineTranslator.java | 53 +++++- .../translation/RedistributeTranslator.java | 54 ++++++ .../ChainedExecutableStageTest.java | 163 ++++++++++++++++++ 4 files changed, 288 insertions(+), 25 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ChainedExecutableStageTest.java 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 7417088bb672..b8eb5413b27e 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 @@ -59,7 +59,7 @@ * receivers. */ class ExecutableStageProcessor - implements Processor, byte[], KStreamsPayload> { + implements Processor, byte[], KStreamsPayload> { private static final Logger LOG = LoggerFactory.getLogger(ExecutableStageProcessor.class); @@ -68,9 +68,13 @@ class ExecutableStageProcessor // 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. - private final Queue> pendingOutputs = new ConcurrentLinkedQueue<>(); + // The element type is intentionally wildcarded: the runner does not need to know the runtime + // value type — the bundle factory handles all coder application at the Fn-API boundary using + // the PCollection coders from the ExecutableStagePayload. Pretending the type was byte[] was + // only safe because the Impulse output coder happens to be ByteArrayCoder. + private final Queue> pendingOutputs = new ConcurrentLinkedQueue<>(); - private @Nullable ProcessorContext> context; + private @Nullable ProcessorContext> context; private @Nullable ExecutableStageContext stageContext; private @Nullable StageBundleFactory stageBundleFactory; private @Nullable RemoteBundle currentBundle; @@ -81,7 +85,7 @@ class ExecutableStageProcessor } @Override - public void init(ProcessorContext> context) { + public void init(ProcessorContext> context) { this.context = context; ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload); this.stageContext = KafkaStreamsExecutableStageContextFactory.getInstance().get(jobInfo); @@ -89,8 +93,8 @@ public void init(ProcessorContext> context) { } @Override - public void process(Record> record) { - KStreamsPayload payload = record.value(); + public void process(Record> record) { + KStreamsPayload payload = record.value(); if (payload.isWatermark()) { // NOTE: flushing the bundle on every received watermark is provisional. Once the // WatermarkManager lands, a stage will receive watermarks from multiple parent instances and @@ -122,7 +126,7 @@ public FnDataReceiver create(String pCollectionId) { // after the bundle closes. return receivedElement -> { if (receivedElement != null) { - pendingOutputs.add((WindowedValue) receivedElement); + pendingOutputs.add((WindowedValue) receivedElement); } }; } @@ -143,7 +147,7 @@ private FnDataReceiver> mainInputReceiver() { return receiver; } - private void closeBundleAndFlush(Record> record) { + private void closeBundleAndFlush(Record> record) { RemoteBundle bundle = currentBundle; if (bundle == null) { return; @@ -157,26 +161,25 @@ private void closeBundleAndFlush(Record> record) } finally { currentBundle = null; } - ProcessorContext> ctx = checkInitialized(context); + ProcessorContext> ctx = checkInitialized(context); // The harness has finished the bundle (close() returned) so no further enqueues happen. - // ConcurrentLinkedQueue's weakly-consistent iterator is therefore safe to drain via forEach. - pendingOutputs.forEach( - output -> - ctx.forward( - new Record>( - record.key(), KStreamsPayload.data(output), record.timestamp()))); - pendingOutputs.clear(); + // Drain via poll() so each element is removed as it is forwarded. + WindowedValue output; + while ((output = pendingOutputs.poll()) != null) { + ctx.forward( + new Record>( + record.key(), KStreamsPayload.data(output), record.timestamp())); + } } - private void forwardWatermark( - Record> record, long watermarkMillis) { + private void forwardWatermark(Record> record, long watermarkMillis) { // TODO(#38743 / WatermarkManager): a watermark must reach every parallel instance of every // downstream processor, but ctx.forward routes to one downstream partition per Kafka Streams' // partitioning. The simplest correct approach is to fan the watermark out to all downstream // partitions; that wiring lands with the WatermarkManager sub-issue (per Jan on PR #38764). - ProcessorContext> ctx = checkInitialized(context); + ProcessorContext> ctx = checkInitialized(context); ctx.forward( - new Record>( + new Record>( record.key(), KStreamsPayload.watermark(watermarkMillis), record.timestamp())); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index 5042f5426169..4e227749e1a3 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -17,16 +17,21 @@ */ package org.apache.beam.runners.kafka.streams.translation; +import com.google.auto.service.AutoService; import java.util.Map; +import java.util.Set; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.util.construction.NativeTransforms; import org.apache.beam.sdk.util.construction.PTransformTranslation; import org.apache.beam.sdk.util.construction.graph.ExecutableStage; import org.apache.beam.sdk.util.construction.graph.GreedyPipelineFuser; import org.apache.beam.sdk.util.construction.graph.PipelineNode; import org.apache.beam.sdk.util.construction.graph.QueryablePipeline; +import org.apache.beam.sdk.util.construction.graph.TrivialNativeTransformExpander; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; /** * Translates a portable Beam pipeline into a Kafka Streams {@link @@ -45,6 +50,7 @@ public KafkaStreamsPipelineTranslator() { this( ImmutableMap.builder() .put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator()) + .put(PTransformTranslation.REDISTRIBUTE_ARBITRARILY_URN, new RedistributeTranslator()) .put(ExecutableStage.URN, new ExecutableStageTranslator()) .build()); } @@ -59,11 +65,28 @@ public KafkaStreamsTranslationContext createTranslationContext( } /** - * Fuses the pipeline so that stateless user code is grouped into {@code ExecutableStage} nodes. + * Returns the set of URNs this translator handles natively. {@link + * TrivialNativeTransformExpander} uses this set to strip the sub-transforms of runner-native + * composites (e.g. {@code Redistribute.arbitrarily}) before fusion, so they survive into + * translation as leaves instead of being expanded into primitives the runner does not implement + * yet (e.g. GroupByKey). + */ + public Set knownUrns() { + return urnToTranslator.keySet(); + } + + /** + * Prepares the pipeline for translation: + * + *
    + *
  1. Trim sub-transforms of runner-native composites listed in {@link #knownUrns()} so the + * fuser leaves them as primitives. + *
  2. Fuse remaining stateless user code into {@code ExecutableStage} nodes via {@link + * GreedyPipelineFuser}. + *
* - *

Runner-executed primitives that have their own translator (e.g. Impulse) are left intact; - * everything else is fused. If the pipeline already contains {@code ExecutableStage} transforms - * it is returned unchanged. + *

If the pipeline already contains {@code ExecutableStage} transforms it is returned + * unchanged. */ public RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) { boolean alreadyFused = @@ -72,7 +95,8 @@ public RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) { if (alreadyFused) { return pipeline; } - return GreedyPipelineFuser.fuse(pipeline).toPipeline(); + RunnerApi.Pipeline trimmed = TrivialNativeTransformExpander.forKnownUrns(pipeline, knownUrns()); + return GreedyPipelineFuser.fuse(trimmed).toPipeline(); } /** @@ -99,4 +123,23 @@ public void translate(KafkaStreamsTranslationContext context, RunnerApi.Pipeline translator.translate(node.getId(), pipeline, context); } } + + /** + * Tells the SDK that URNs handled directly by the Kafka Streams runner should be treated as + * primitives by {@link QueryablePipeline}. Mirrors Flink's {@code IsFlinkNativeTransform} + * pattern. Without this, {@link TrivialNativeTransformExpander} strips a composite's + * sub-transforms but {@link QueryablePipeline} still does not recognise the composite itself as a + * producer of its outputs, and pipeline validation fails with "consumed but never produced". + */ + @AutoService(NativeTransforms.IsNativeTransform.class) + public static class IsKafkaStreamsNativeTransform implements NativeTransforms.IsNativeTransform { + private static final Set URNS = + ImmutableSet.of(PTransformTranslation.REDISTRIBUTE_ARBITRARILY_URN); + + @Override + public boolean test(RunnerApi.PTransform pTransform) { + String urn = PTransformTranslation.urnForTransformOrNull(pTransform); + return urn != null && URNS.contains(urn); + } + } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java new file mode 100644 index 000000000000..72c47db8d4b1 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java @@ -0,0 +1,54 @@ +/* + * 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 org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; + +/** + * Runner-native translator for {@code beam:transform:redistribute_arbitrarily:v1}. + * + *

The default Beam expansion of {@code Redistribute.arbitrarily()} goes through {@code + * GroupByKey} for materialization. The Kafka Streams runner does not need to do any actual + * redistribution in the single-instance topology — Kafka Streams is already handling per-task + * processing — so we provide a passthrough: the output PCollection is mapped to the same processor + * that already produces the input PCollection. No topology node is added. + * + *

The translator is registered in {@link KafkaStreamsPipelineTranslator#knownUrns()} so {@link + * org.apache.beam.sdk.util.construction.graph.TrivialNativeTransformExpander} strips the + * sub-transforms of {@code Redistribute} before the fuser runs. That keeps the Redistribute + * boundary intact and lets the fuser split adjacent stages correctly. + * + *

The keyed variant ({@code redistribute_by_key:v1}) is intentionally not handled here: it + * implies a rehash/partition step that the runner can implement on top of GroupByKey when that + * lands, but cannot reasonably no-op the way the arbitrary variant can. + */ +class RedistributeTranslator implements PTransformTranslator { + + @Override + public void translate( + String transformId, RunnerApi.Pipeline pipeline, KafkaStreamsTranslationContext context) { + RunnerApi.PTransform transform = pipeline.getComponents().getTransformsOrThrow(transformId); + String inputPCollectionId = Iterables.getOnlyElement(transform.getInputsMap().values()); + String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); + String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + // Passthrough: downstream lookups for the output PCollection resolve to the producer of the + // input PCollection. No KS Processor / state store / source is added. + context.registerPCollectionProducer(outputPCollectionId, parentProcessor); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ChainedExecutableStageTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ChainedExecutableStageTest.java new file mode 100644 index 000000000000..8ac2175dc391 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ChainedExecutableStageTest.java @@ -0,0 +1,163 @@ +/* + * 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 java.time.Duration; +import java.util.List; +import java.util.Properties; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Redistribute; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.beam.sdk.util.construction.graph.ExecutableStage; +import org.apache.beam.sdk.values.TypeDescriptor; +import org.apache.beam.sdk.values.TypeDescriptors; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.TopologyTestDriver; +import org.junit.Test; + +/** + * End-to-end test that exercises non-{@code byte[]} value types flowing across an {@code + * ExecutableStage} boundary. + * + *

Builds a pipeline {@code Impulse -> MapElements -> Redistribute.arbitrarily() + * -> ParDo(record into SharedTestCollector)}. The Redistribute boundary terminates + * the upstream {@code ExecutableStage} (because the runner registers Redistribute as a + * runner-native transform in {@link KafkaStreamsPipelineTranslator#knownUrns()}), so the + * Integer-emitting stage and the recording stage live in separate stages and the Integer payload + * actually flows stage-to-stage through the runner — not just through user code inside a single + * fused stage. That is what proves the {@code byte[]}-erasure cast we used to do was a latent bug, + * and that the type-agnostic {@code KStreamsPayload} edge handles it correctly. + */ +public class ChainedExecutableStageTest { + + private static final String JOB_ID = "ks-chained-stage-test"; + private static final String APPLICATION_ID = "ks-chained-stage-test"; + private static final int EXPECTED_VALUE = 42; + + /** + * Records the integer payload it sees so the test can verify it survived the harness round-trip + * across the Redistribute boundary. + */ + private static class RecordingFn extends DoFn { + private final SharedTestCollector collector; + + RecordingFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element Integer input) { + collector.record(input); + } + } + + @Test + public void chainedStagesPropagateIntegerValueAcrossRedistributeBoundary() throws Exception { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(pipelineOptions()); + pipeline + .apply("impulse", Impulse.create()) + .apply( + "to-integer", + MapElements.into(TypeDescriptors.integers()).via((byte[] ignored) -> EXPECTED_VALUE)) + .setTypeDescriptor(TypeDescriptor.of(Integer.class)) + .apply("redistribute", Redistribute.arbitrarily()) + .apply("record", ParDo.of(new RecordingFn(collector))); + + RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); + + KafkaStreamsPipelineOptions options = + pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + JobInfo jobInfo = + JobInfo.create( + JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); + KafkaStreamsTranslationContext context = + translator.createTranslationContext(jobInfo, options); + + RunnerApi.Pipeline preparedPipeline = translator.prepareForTranslation(pipelineProto); + + // Verify the Redistribute boundary actually split the user code into two separate stages. + // Without this, the test could pass even if the fuser collapsed MapElements and the + // RecordingFn into one stage — in which case the Integer payload would flow through user + // code inside a single stage and never cross the runner-side ExecutableStage edge that + // this test exists to exercise. + long executableStageCount = + preparedPipeline.getComponents().getTransformsMap().values().stream() + .filter(t -> ExecutableStage.URN.equals(t.getSpec().getUrn())) + .count(); + assertThat( + "expected two ExecutableStages (upstream MapElements + downstream RecordingFn ParDo) " + + "separated by the Redistribute boundary", + executableStageCount, + is(2L)); + + translator.translate(context, preparedPipeline); + + Topology topology = context.getTopology(); + try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { + driver.advanceWallClockTime(Duration.ofSeconds(1)); + driver.advanceWallClockTime(Duration.ofSeconds(1)); + } + + List recorded = collector.recorded(); + assertThat(recorded.size(), is(1)); + assertThat(recorded.get(0), is(EXPECTED_VALUE)); + } + } + + private static PipelineOptions pipelineOptions() { + PipelineOptions options = + PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); + options.setRunner(CrashingRunner.class); + options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); + options + .as(PortablePipelineOptions.class) + .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); + return options; + } + + private static Properties streamsConfig() { + Properties props = new Properties(); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put( + StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + props.put( + StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + return props; + } +} From 4636b1f333eae8b0969b91d42bc436af59a66e47 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Mon, 15 Jun 2026 22:37:55 +0500 Subject: [PATCH 10/37] #38957: Add in-memory WatermarkManager core (per-source-partition tracking --- .../streams/translation/WatermarkManager.java | 157 ++++++++++++++++++ .../translation/WatermarkManagerTest.java | 155 +++++++++++++++++ 2 files changed, 312 insertions(+) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java new file mode 100644 index 000000000000..cd0fca3654ae --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java @@ -0,0 +1,157 @@ +/* + * 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 java.util.HashMap; +import java.util.Map; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.joda.time.Instant; + +/** + * In-memory tracker of a single fused stage's input watermark, computed from the committed + * watermarks reported by the upstream source partitions that feed it (the output / + * repartition-topic partitions of the parent stage). + * + *

This is the core of the Kafka Streams runner's watermark propagation, decoupled from the Kafka + * wiring so it can be unit-tested in isolation. The wiring that produces the reports (flushing + * {@code (sourcePartition, committedWatermark, totalSourcePartitions)} atomically with each offset + * commit and fanning it out to every downstream partition) and consumes them lands in a follow-up. + * + *

Why source partitions, not producer instances

+ * + *

The question a stage has to answer is "have I received the watermark from every upstream + * producer, so that {@code min()} across them is meaningful?". Counting producer instances + * is hard: an instance can be killed without notice, leaving stale state, and the number changes on + * every rebalance. Counting source partitions is robust instead, because the partition count + * is fixed and known: it travels in-band with every report ({@code totalSourcePartitions}), a + * partition is always owned by exactly one live instance, and when an instance dies its partitions + * are reassigned and the new owner keeps reporting. So the manager only ever reasons about + * partitions, never about instances. (Design agreed with the mentor; see the watermark + * coordination-channel PoC findings.) + * + *

Holding until ready

+ * + *

Until a committed watermark has been seen for every source partition, the stage's input + * watermark is undefined and {@link #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} — + * i.e. the stage emits no meaningful watermark downstream. A change in {@code + * totalSourcePartitions} (e.g. a repartition) clears the accumulated reports and re-opens this hold + * until the new full set has reported, which subsumes the "new epoch / revert" rule without an + * explicit epoch. + * + *

Monotonicity

+ * + *

Beam watermarks must be non-decreasing. Each source partition's watermark is held monotonic (a + * lower report is ignored), and the emitted stage watermark is additionally clamped so it never + * regresses below the previously emitted value — relevant if a newly appeared partition reports an + * older watermark after the stage had already advanced. + * + *

Not thread-safe; the caller (a single Kafka Streams processor thread) serializes access. + */ +public final class WatermarkManager { + + /** Total source partitions feeding this stage, learned in-band; -1 until the first report. */ + private int expectedSourcePartitionCount = -1; + + /** Latest committed watermark per source partition (kept monotonic non-decreasing). */ + private final Map committedWatermarkByPartition = new HashMap<>(); + + /** Last watermark {@link #advance()} emitted, to enforce a non-decreasing output. */ + private Instant lastEmitted = BoundedWindow.TIMESTAMP_MIN_VALUE; + + /** + * Record a committed watermark reported for one source partition, together with the total source + * partition count carried in-band with the report. + * + * @param sourcePartition the source partition the report is for, in {@code [0, + * totalSourcePartitions)} + * @param committedWatermark the committed watermark for that partition + * @param totalSourcePartitions the total number of upstream source partitions feeding this stage + */ + public void observe(int sourcePartition, Instant committedWatermark, int totalSourcePartitions) { + if (committedWatermark == null) { + throw new IllegalArgumentException("committedWatermark must not be null"); + } + if (totalSourcePartitions <= 0) { + throw new IllegalArgumentException( + "totalSourcePartitions must be positive: " + totalSourcePartitions); + } + if (sourcePartition < 0 || sourcePartition >= totalSourcePartitions) { + throw new IllegalArgumentException( + "sourcePartition " + + sourcePartition + + " out of range for totalSourcePartitions " + + totalSourcePartitions); + } + if (totalSourcePartitions != expectedSourcePartitionCount) { + // The source partition set changed (e.g. a repartition). The previous per-partition + // watermarks describe a different partitioning, so drop them entirely and re-open the hold + // until the new full set reports. The output watermark still cannot regress (lastEmitted is + // retained). + expectedSourcePartitionCount = totalSourcePartitions; + committedWatermarkByPartition.clear(); + } + // A source partition's watermark is monotonic non-decreasing; ignore an out-of-order lower + // report. + committedWatermarkByPartition.merge( + sourcePartition, committedWatermark, (oldW, newW) -> newW.isAfter(oldW) ? newW : oldW); + } + + /** True once a committed watermark has been seen for every current source partition. */ + public boolean isReady() { + return expectedSourcePartitionCount > 0 + && committedWatermarkByPartition.size() == expectedSourcePartitionCount; + } + + /** + * Advance and return the stage input watermark. + * + *

Returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} while the stage is still holding (not + * every source partition has reported) — the caller emits nothing meaningful downstream in that + * case. Once ready, returns {@code min()} over all source partitions, clamped to never regress + * below the previously emitted value. The sequence of values returned across calls is + * non-decreasing. + */ + public Instant advance() { + if (!isReady()) { + return BoundedWindow.TIMESTAMP_MIN_VALUE; + } + // isReady() guarantees the map is non-empty, so the seed is always replaced by a real value. + Instant min = BoundedWindow.TIMESTAMP_MAX_VALUE; + for (Instant w : committedWatermarkByPartition.values()) { + if (w.isBefore(min)) { + min = w; + } + } + Instant emit = min.isAfter(lastEmitted) ? min : lastEmitted; + lastEmitted = emit; + return emit; + } + + /** The total source partition count learned in-band, or -1 if nothing reported yet. */ + @VisibleForTesting + int expectedSourcePartitionCount() { + return expectedSourcePartitionCount; + } + + /** How many distinct source partitions have reported so far. */ + @VisibleForTesting + int reportedPartitionCount() { + return committedWatermarkByPartition.size(); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java new file mode 100644 index 000000000000..1d642e5ea1a7 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java @@ -0,0 +1,155 @@ +/* + * 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 static org.junit.Assert.assertThrows; + +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.joda.time.Instant; +import org.junit.Test; + +/** Tests for {@link WatermarkManager}. */ +public class WatermarkManagerTest { + + private static Instant ts(long millis) { + return new Instant(millis); + } + + @Test + public void holdsBeforeAnyReport() { + WatermarkManager manager = new WatermarkManager(); + assertThat(manager.isReady(), is(false)); + assertThat(manager.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + } + + @Test + public void holdsUntilEverySourcePartitionReports() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, ts(100L), 4); + manager.observe(1, ts(100L), 4); + manager.observe(2, ts(100L), 4); + // Three of four partitions reported — still holding. + assertThat(manager.isReady(), is(false)); + assertThat(manager.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + manager.observe(3, ts(100L), 4); + assertThat(manager.isReady(), is(true)); + assertThat(manager.advance(), is(ts(100L))); + } + + @Test + public void emitsMinAcrossPartitionsOnceReady() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, ts(300L), 4); + manager.observe(1, ts(100L), 4); + manager.observe(2, ts(500L), 4); + manager.observe(3, ts(200L), 4); + // min(300, 100, 500, 200) = 100 + assertThat(manager.advance(), is(ts(100L))); + } + + @Test + public void perPartitionWatermarkIsMonotonic() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, ts(100L), 1); + assertThat(manager.advance(), is(ts(100L))); + // A lower out-of-order report for the same partition is ignored. + manager.observe(0, ts(50L), 1); + assertThat(manager.advance(), is(ts(100L))); + } + + @Test + public void partitionCountChangeClearsReportsAndReopensHold() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, ts(100L), 2); + manager.observe(1, ts(100L), 2); + assertThat(manager.advance(), is(ts(100L))); + + // Source set grows to 4. The previous reports are dropped, so the stage holds again until all + // four of the new partition set have reported. + manager.observe(0, ts(200L), 4); + assertThat(manager.reportedPartitionCount(), is(1)); + assertThat(manager.isReady(), is(false)); + assertThat(manager.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + manager.observe(1, ts(200L), 4); + manager.observe(2, ts(200L), 4); + manager.observe(3, ts(200L), 4); + assertThat(manager.isReady(), is(true)); + assertThat(manager.advance(), is(ts(200L))); + } + + @Test + public void emittedWatermarkDoesNotRegressAfterRepartition() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, ts(100L), 2); + manager.observe(1, ts(100L), 2); + assertThat(manager.advance(), is(ts(100L))); + + // After a repartition to 3, the new partitions report older watermarks. Once ready again the + // min is 50, but the emitted stage watermark must not go backwards below 100. + manager.observe(0, ts(50L), 3); + manager.observe(1, ts(50L), 3); + manager.observe(2, ts(50L), 3); + assertThat(manager.isReady(), is(true)); + assertThat(manager.advance(), is(ts(100L))); + } + + @Test + public void partitionCountDecreaseClearsAndRecomputes() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, ts(100L), 4); + manager.observe(1, ts(100L), 4); + manager.observe(2, ts(100L), 4); + manager.observe(3, ts(50L), 4); + assertThat(manager.advance(), is(ts(50L))); + + // Source set shrinks to 2. Reports are cleared; the stage holds until {0, 1} report again. + manager.observe(0, ts(100L), 2); + assertThat(manager.expectedSourcePartitionCount(), is(2)); + assertThat(manager.reportedPartitionCount(), is(1)); + assertThat(manager.isReady(), is(false)); + + manager.observe(1, ts(100L), 2); + assertThat(manager.isReady(), is(true)); + // min is 100; the non-regression clamp keeps it >= the previously emitted 50. + assertThat(manager.advance(), is(ts(100L))); + } + + @Test + public void rejectsNullWatermark() { + WatermarkManager manager = new WatermarkManager(); + assertThrows(IllegalArgumentException.class, () -> manager.observe(0, null, 4)); + } + + @Test + public void rejectsNonPositiveTotalSourcePartitions() { + WatermarkManager manager = new WatermarkManager(); + assertThrows(IllegalArgumentException.class, () -> manager.observe(0, ts(100L), 0)); + assertThrows(IllegalArgumentException.class, () -> manager.observe(0, ts(100L), -1)); + } + + @Test + public void rejectsOutOfRangeSourcePartition() { + WatermarkManager manager = new WatermarkManager(); + assertThrows(IllegalArgumentException.class, () -> manager.observe(-1, ts(100L), 4)); + assertThrows(IllegalArgumentException.class, () -> manager.observe(4, ts(100L), 4)); + } +} From 1058e94027e5925a6324c0cbd9f3013a5a870218 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Thu, 18 Jun 2026 16:30:35 +0500 Subject: [PATCH 11/37] [GSoC 2026] Kafka Streams runner #38987: Wire WatermarkManager into ExecutableStageProcessor --- .../translation/ExecutableStageProcessor.java | 62 +++++-- .../streams/translation/ImpulseProcessor.java | 10 +- .../streams/translation/KStreamsPayload.java | 77 +++++++-- .../streams/translation/WatermarkPayload.java | 41 +++++ .../kafka/streams/KafkaStreamsRunnerTest.java | 2 +- ...ExecutableStageProcessorWatermarkTest.java | 124 ++++++++++++++ .../translation/ImpulseTranslatorTest.java | 3 +- .../translation/WatermarkPropagationTest.java | 155 ++++++++++++++++++ 8 files changed, 440 insertions(+), 34 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java 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 b8eb5413b27e..00f85032d041 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 @@ -28,6 +28,7 @@ import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.fnexecution.state.StateRequestHandler; import org.apache.beam.sdk.fn.data.FnDataReceiver; +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.guava.v32_1_2_jre.com.google.common.collect.Iterables; @@ -35,6 +36,7 @@ import org.apache.kafka.streams.processor.api.ProcessorContext; import org.apache.kafka.streams.processor.api.Record; import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -49,9 +51,12 @@ * ProcessorContext#forward} must only be called from the processing thread, so outputs are never * forwarded directly from a harness callback. * - *

A {@link KStreamsPayload#isWatermark() watermark} payload marks a bundle boundary: the open - * bundle (if any) is closed (flushing outputs), and the watermark is then forwarded downstream so - * that subsequent stages observe it after all data of the bundle. + *

A {@link KStreamsPayload#isWatermark() watermark} payload is a per-source-partition report and + * marks a bundle boundary: the open bundle (if any) is closed (flushing outputs), the report is fed + * to the {@link WatermarkManager}, and the stage's output watermark is forwarded downstream only + * when the {@code min()} across its source partitions actually advances. Until every source + * partition has reported, the watermark is held and nothing is forwarded — but data is still + * processed in the meantime. * *

This is the Kafka Streams analogue of Flink's {@code ExecutableStageDoFnOperator} and Spark's * {@code SparkExecutableStageFunction}. State, timers, and side inputs are out of scope for this @@ -74,6 +79,12 @@ class ExecutableStageProcessor // only safe because the Impulse output coder happens to be ByteArrayCoder. private final Queue> pendingOutputs = new ConcurrentLinkedQueue<>(); + // Computes this stage's output watermark as min() over its source partitions' reported + // watermarks, holding until every source partition has reported (see WatermarkManager). + private final WatermarkManager watermarkManager = new WatermarkManager(); + // The last watermark actually forwarded downstream, so we only forward when it advances. + private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + private @Nullable ProcessorContext> context; private @Nullable ExecutableStageContext stageContext; private @Nullable StageBundleFactory stageBundleFactory; @@ -87,22 +98,40 @@ class ExecutableStageProcessor @Override public void init(ProcessorContext> context) { this.context = context; + // The SDK harness (stage context + bundle factory) is created lazily on the first data + // element, so a stage that only forwards watermarks never spins one up. This mirrors Spark's + // SparkExecutableStageFunction, which likewise does not build a bundle factory when there are + // no inputs to process. + } + + private void ensureStageBundleFactory() { + if (stageBundleFactory != null) { + return; + } ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload); - this.stageContext = KafkaStreamsExecutableStageContextFactory.getInstance().get(jobInfo); - this.stageBundleFactory = stageContext.getStageBundleFactory(executableStage); + stageContext = KafkaStreamsExecutableStageContextFactory.getInstance().get(jobInfo); + stageBundleFactory = stageContext.getStageBundleFactory(executableStage); } @Override public void process(Record> record) { KStreamsPayload payload = record.value(); if (payload.isWatermark()) { - // NOTE: flushing the bundle on every received watermark is provisional. Once the - // WatermarkManager lands, a stage will receive watermarks from multiple parent instances and - // the output watermark becomes min() across them — the bundle should flush / the output - // watermark advance only when that minimum actually moves forward, not on every received - // watermark. Tracked in #38743. + // Emit any buffered outputs before the watermark. Data is processed regardless of watermark + // readiness; only the watermark itself is held until every source partition has reported. closeBundleAndFlush(record); - forwardWatermark(record, payload.getWatermarkMillis()); + // Feed the report into the WatermarkManager and forward the stage's output watermark only + // when min() across the source partitions actually advances, not on every received watermark. + WatermarkPayload report = payload.asWatermark(); + watermarkManager.observe( + report.getSourcePartition(), + new Instant(report.getWatermarkMillis()), + report.getTotalSourcePartitions()); + Instant advanced = watermarkManager.advance(); + if (advanced.isAfter(lastForwardedWatermark)) { + lastForwardedWatermark = advanced; + forwardWatermark(record, advanced.getMillis()); + } return; } try { @@ -117,6 +146,7 @@ private void ensureBundleOpen() throws Exception { if (currentBundle != null) { return; } + ensureStageBundleFactory(); StageBundleFactory factory = checkInitialized(stageBundleFactory); OutputReceiverFactory outputReceiverFactory = new OutputReceiverFactory() { @@ -173,14 +203,14 @@ private void closeBundleAndFlush(Record> record) { } private void forwardWatermark(Record> record, long watermarkMillis) { - // TODO(#38743 / WatermarkManager): a watermark must reach every parallel instance of every - // downstream processor, but ctx.forward routes to one downstream partition per Kafka Streams' - // partitioning. The simplest correct approach is to fan the watermark out to all downstream - // partitions; that wiring lands with the WatermarkManager sub-issue (per Jan on PR #38764). + // This stage is a single instance for now, so it forwards its watermark as the only source + // partition (0 of 1). Fanning the watermark out to every downstream partition — and producing + // it atomically with the offset commit so it is durable — lands with the topic-based shuffle + // work, when there are real source partitions to track (#18479). ProcessorContext> ctx = checkInitialized(context); ctx.forward( new Record>( - record.key(), KStreamsPayload.watermark(watermarkMillis), record.timestamp())); + record.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1), record.timestamp())); } @Override diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java index e9fc8ddb36aa..675bee4d8591 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java @@ -121,12 +121,18 @@ private void maybeFire() { LOG.debug("Impulse {} emitted single element and terminal watermark", transformId); } - /** Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. */ + /** + * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. + * + *

Impulse is a single-instance source, so the report is stamped as the only source partition: + * {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real per-partition identities + * arrive once the topology gains topic-based shuffle. + */ private static void forwardWatermarkMax(ProcessorContext> ctx) { long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); ctx.forward( new Record>( - new byte[0], KStreamsPayload.watermark(maxMillis), 0L)); + new byte[0], KStreamsPayload.watermark(maxMillis, 0, 1), 0L)); } /** Cancels the wall-clock punctuator after the impulse has fired to stop periodic wakeups. */ diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java index 53e47b1216bb..93b40346761a 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java @@ -20,6 +20,7 @@ import java.util.Objects; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -29,7 +30,9 @@ * *

    *
  • A {@link #isData() data} element wrapping a {@link WindowedValue}, or - *
  • A {@link #isWatermark() watermark} signal carrying an event-time milliseconds value. + *
  • A {@link #isWatermark() watermark} report carrying an event-time milliseconds value plus + * the in-band coordination fields (source partition and total source partition count) the + * downstream {@link WatermarkManager} needs. *
* *

The envelope lets a single Kafka Streams output channel carry both Beam data and the watermark @@ -54,21 +57,45 @@ private enum Kind { private final Kind kind; private final @Nullable WindowedValue data; private final long watermarkMillis; + private final int sourcePartition; + private final int totalSourcePartitions; - private KStreamsPayload(Kind kind, @Nullable WindowedValue data, long watermarkMillis) { + private KStreamsPayload( + Kind kind, + @Nullable WindowedValue data, + long watermarkMillis, + int sourcePartition, + int totalSourcePartitions) { this.kind = kind; this.data = data; this.watermarkMillis = watermarkMillis; + this.sourcePartition = sourcePartition; + this.totalSourcePartitions = totalSourcePartitions; } /** Returns a data payload wrapping the given {@link WindowedValue}. */ public static KStreamsPayload data(WindowedValue value) { - return new KStreamsPayload<>(Kind.DATA, value, 0L); + return new KStreamsPayload<>(Kind.DATA, value, 0L, 0, 0); } - /** Returns a watermark payload carrying the given event-time milliseconds. */ - public static KStreamsPayload watermark(long watermarkMillis) { - return new KStreamsPayload<>(Kind.WATERMARK, null, watermarkMillis); + /** + * Returns a watermark report payload: the event-time milliseconds together with the in-band + * coordination fields the downstream stage's {@link WatermarkManager} needs — which source + * partition this report is for and how many source partitions feed the stage in total. + */ + public static KStreamsPayload watermark( + long watermarkMillis, int sourcePartition, int totalSourcePartitions) { + Preconditions.checkArgument( + totalSourcePartitions > 0, + "totalSourcePartitions must be positive: %s", + totalSourcePartitions); + Preconditions.checkArgument( + sourcePartition >= 0 && sourcePartition < totalSourcePartitions, + "sourcePartition %s out of range for totalSourcePartitions %s", + sourcePartition, + totalSourcePartitions); + return new KStreamsPayload<>( + Kind.WATERMARK, null, watermarkMillis, sourcePartition, totalSourcePartitions); } public boolean isData() { @@ -91,14 +118,31 @@ public WindowedValue getData() { } /** - * Returns the watermark event-time milliseconds. Caller must check {@link #isWatermark()} first; - * calling this on a data payload throws. + * Narrows this payload to its {@link WatermarkPayload} view, through which the watermark report + * fields are read. Caller must check {@link #isWatermark()} first; calling this on a data payload + * throws. */ - public long getWatermarkMillis() { - if (kind != Kind.WATERMARK) { - throw new IllegalStateException("Payload is not a watermark: kind=" + kind); + public WatermarkPayload asWatermark() { + Preconditions.checkState(isWatermark(), "Payload is not a watermark: kind=%s", kind); + return new WatermarkView(); + } + + /** {@link WatermarkPayload} view backed by this payload's fields. */ + private final class WatermarkView implements WatermarkPayload { + @Override + public long getWatermarkMillis() { + return watermarkMillis; + } + + @Override + public int getSourcePartition() { + return sourcePartition; + } + + @Override + public int getTotalSourcePartitions() { + return totalSourcePartitions; } - return watermarkMillis; } @Override @@ -112,12 +156,14 @@ public boolean equals(@Nullable Object o) { KStreamsPayload that = (KStreamsPayload) o; return kind == that.kind && watermarkMillis == that.watermarkMillis + && sourcePartition == that.sourcePartition + && totalSourcePartitions == that.totalSourcePartitions && Objects.equals(data, that.data); } @Override public int hashCode() { - return Objects.hash(kind, data, watermarkMillis); + return Objects.hash(kind, data, watermarkMillis, sourcePartition, totalSourcePartitions); } @Override @@ -126,7 +172,10 @@ public String toString() { if (kind == Kind.DATA) { helper.add("data", data); } else { - helper.add("watermarkMillis", watermarkMillis); + helper + .add("watermarkMillis", watermarkMillis) + .add("sourcePartition", sourcePartition) + .add("totalSourcePartitions", totalSourcePartitions); } return helper.toString(); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java new file mode 100644 index 000000000000..194be701ecec --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java @@ -0,0 +1,41 @@ +/* + * 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; + +/** + * The watermark-only view of a {@link KStreamsPayload}, obtained via {@link + * KStreamsPayload#asWatermark()}. Keeping the watermark accessors on this interface — rather than + * on {@link KStreamsPayload} itself — means they are only reachable after the caller has checked + * {@link KStreamsPayload#isWatermark()} and narrowed the payload, so there is no kind check to do + * on each accessor. + * + *

A watermark report is the in-band coordination message a downstream stage's {@link + * WatermarkManager} consumes: the watermark value plus which source partition reported it and how + * many source partitions feed the stage in total. + */ +public interface WatermarkPayload { + + /** The reported watermark, in event-time milliseconds. */ + long getWatermarkMillis(); + + /** The source partition this report is for, in {@code [0, getTotalSourcePartitions())}. */ + int getSourcePartition(); + + /** The total number of source partitions feeding the downstream stage. */ + int getTotalSourcePartitions(); +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java index 30000d77d864..0c576d27122f 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java @@ -97,7 +97,7 @@ public void impulseOnlyPipelineEmitsDataAndTerminalWatermark() { assertThat(capture.received.get(0).getData().getValue().length, is(0)); assertThat(capture.received.get(1).isWatermark(), is(true)); assertThat( - capture.received.get(1).getWatermarkMillis(), + capture.received.get(1).asWatermark().getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); } 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 new file mode 100644 index 000000000000..fc5797a12c26 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java @@ -0,0 +1,124 @@ +/* + * 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.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.kafka.streams.processor.api.MockProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.junit.Test; + +/** + * Tests the watermark wiring of {@link ExecutableStageProcessor}: how it feeds incoming watermark + * reports to the {@link WatermarkManager} and forwards the stage's output watermark. + * + *

Only the watermark path is exercised, so the SDK harness is never started (it is created + * lazily on the first data element). A {@link MockProcessorContext} captures what the processor + * forwards downstream. + */ +public class ExecutableStageProcessorWatermarkTest { + + private static ExecutableStageProcessor newProcessor() { + JobInfo jobInfo = + JobInfo.create( + "job-id", + "job-name", + "", + PipelineOptionsTranslation.toProto(PipelineOptionsFactory.create())); + return new ExecutableStageProcessor( + RunnerApi.ExecutableStagePayload.getDefaultInstance(), jobInfo); + } + + private static Record> watermark( + long millis, int sourcePartition, int totalSourcePartitions) { + KStreamsPayload payload = + KStreamsPayload.watermark(millis, sourcePartition, totalSourcePartitions); + return new Record<>(new byte[0], payload, 0L); + } + + private static KStreamsPayload onlyForwarded( + MockProcessorContext> ctx) { + assertThat(ctx.forwarded().size(), is(1)); + return ctx.forwarded().get(0).record().value(); + } + + @Test + public void singleSourcePartitionForwardsImmediatelyStampedAsItsOwnSource() { + MockProcessorContext> ctx = new MockProcessorContext<>(); + ExecutableStageProcessor processor = newProcessor(); + processor.init(ctx); + + processor.process(watermark(100L, 0, 1)); + + KStreamsPayload out = onlyForwarded(ctx); + assertThat(out.isWatermark(), is(true)); + WatermarkPayload report = out.asWatermark(); + assertThat(report.getWatermarkMillis(), is(100L)); + // The stage forwards as its own single source (0 of 1), not the upstream's identity. + assertThat(report.getSourcePartition(), is(0)); + assertThat(report.getTotalSourcePartitions(), is(1)); + } + + @Test + public void holdsUntilAllSourcePartitionsReportThenForwardsMin() { + MockProcessorContext> ctx = new MockProcessorContext<>(); + ExecutableStageProcessor processor = newProcessor(); + processor.init(ctx); + + processor.process(watermark(300L, 0, 3)); + processor.process(watermark(100L, 1, 3)); + // Two of three source partitions reported — still holding, nothing forwarded. + assertThat(ctx.forwarded().isEmpty(), is(true)); + + processor.process(watermark(500L, 2, 3)); + // All three reported — forward min(300, 100, 500) = 100. + assertThat(onlyForwarded(ctx).asWatermark().getWatermarkMillis(), is(100L)); + } + + @Test + public void doesNotReforwardWhenWatermarkDoesNotAdvance() { + MockProcessorContext> ctx = new MockProcessorContext<>(); + ExecutableStageProcessor processor = newProcessor(); + processor.init(ctx); + + processor.process(watermark(100L, 0, 1)); + assertThat(ctx.forwarded().size(), is(1)); + + // A repeated, non-advancing watermark must not be forwarded again. + processor.process(watermark(100L, 0, 1)); + assertThat(ctx.forwarded().size(), is(1)); + } + + @Test + public void forwardsTerminalMaxWatermark() { + long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); + MockProcessorContext> ctx = new MockProcessorContext<>(); + ExecutableStageProcessor processor = newProcessor(); + processor.init(ctx); + + processor.process(watermark(maxMillis, 0, 1)); + + assertThat(onlyForwarded(ctx).asWatermark().getWatermarkMillis(), is(maxMillis)); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java index a092b10e02c4..f15e818f43ed 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java @@ -77,7 +77,8 @@ public void impulseEmitsDataElementFollowedByTerminalWatermark() { KStreamsPayload watermarkPayload = capture.received.get(1); assertThat(watermarkPayload.isWatermark(), is(true)); assertThat( - watermarkPayload.getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); + watermarkPayload.asWatermark().getWatermarkMillis(), + is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); } @Test diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java new file mode 100644 index 000000000000..e4bdbe4d9733 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java @@ -0,0 +1,155 @@ +/* + * 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 java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Properties; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.TopologyDescription; +import org.apache.kafka.streams.TopologyTestDriver; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.Record; +import org.junit.Test; + +/** + * End-to-end test that a watermark propagates through the topology: {@code Impulse -> + * ExecutableStage -> a recording sink}. The Impulse source emits a terminal {@code + * TIMESTAMP_MAX_VALUE} watermark report; the ExecutableStage routes it through its {@link + * WatermarkManager} and forwards it on, stamped as its own single source partition. A sink + * processor attached to the leaf captures the forwarded watermark and the test asserts on it. + */ +public class WatermarkPropagationTest { + + private static final String APPLICATION_ID = "ks-watermark-propagation-test"; + + /** Identity DoFn so the pipeline contains a fused ExecutableStage. */ + private static class IdentityFn extends DoFn { + @ProcessElement + public void processElement(@Element byte[] input, OutputReceiver out) { + out.output(input); + } + } + + /** Sink processor that records the watermark payloads it is forwarded. */ + private static final class WatermarkCapture + implements Processor, Void, Void> { + private final List> watermarks; + + WatermarkCapture(List> watermarks) { + this.watermarks = watermarks; + } + + @Override + public void process(Record> record) { + if (record.value().isWatermark()) { + watermarks.add(record.value()); + } + } + } + + @Test + public void terminalWatermarkPropagatesToDownstreamStampedAsSingleSource() throws Exception { + Pipeline pipeline = Pipeline.create(pipelineOptions()); + pipeline.apply("impulse", Impulse.create()).apply("identity", ParDo.of(new IdentityFn())); + + RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); + KafkaStreamsPipelineOptions options = + pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + JobInfo jobInfo = + JobInfo.create( + APPLICATION_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); + KafkaStreamsTranslationContext context = translator.createTranslationContext(jobInfo, options); + translator.translate(context, translator.prepareForTranslation(pipelineProto)); + + // Attach a sink to the leaf ExecutableStage processor to capture the watermark it forwards. + Topology topology = context.getTopology(); + List> captured = new ArrayList<>(); + topology.addProcessor( + "watermark-capture", () -> new WatermarkCapture(captured), findLeafProcessor(topology)); + + try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { + driver.advanceWallClockTime(Duration.ofSeconds(1)); + driver.advanceWallClockTime(Duration.ofSeconds(1)); + } + + assertThat("a watermark reached the downstream sink", captured.isEmpty(), is(false)); + WatermarkPayload terminal = captured.get(captured.size() - 1).asWatermark(); + assertThat(terminal.getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); + assertThat(terminal.getSourcePartition(), is(0)); + assertThat(terminal.getTotalSourcePartitions(), is(1)); + } + + /** + * Returns the name of the single processor node with no successors (the leaf of the topology). + */ + private static String findLeafProcessor(Topology topology) { + for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { + for (TopologyDescription.Node node : subtopology.nodes()) { + if (node instanceof TopologyDescription.Processor && node.successors().isEmpty()) { + return node.name(); + } + } + } + throw new IllegalStateException("no leaf processor found in topology"); + } + + private static PipelineOptions pipelineOptions() { + PipelineOptions options = + PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); + options.setRunner(CrashingRunner.class); + options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); + options + .as(PortablePipelineOptions.class) + .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); + return options; + } + + private static Properties streamsConfig() { + Properties props = new Properties(); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put( + StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + props.put( + StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + return props; + } +} From 3e963a683433ddbbfd555ab05ee982ba9158d40b Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Sat, 27 Jun 2026 12:18:44 +0500 Subject: [PATCH 12/37] [GSoC 2026] Kafka Streams runner #39051: Add KStreamsPayload Serde for crossing topic boundaries --- runners/kafka-streams/build.gradle | 1 + runners/kafka-streams/proto/build.gradle | 34 +++++ .../main/proto/kafka_streams_payload.proto | 52 ++++++++ .../translation/KStreamsPayloadSerde.java | 121 ++++++++++++++++++ .../translation/KStreamsPayloadSerdeTest.java | 96 ++++++++++++++ settings.gradle.kts | 1 + 6 files changed, 305 insertions(+) create mode 100644 runners/kafka-streams/proto/build.gradle create mode 100644 runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 52b320dd70a8..2ba55e618308 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -43,6 +43,7 @@ dependencies { permitUnusedDeclared project(":sdks:java:build-tools") 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:job-management", configuration: "shadow") implementation project(":runners:core-java") diff --git a/runners/kafka-streams/proto/build.gradle b/runners/kafka-streams/proto/build.gradle new file mode 100644 index 000000000000..d3e42aeb617f --- /dev/null +++ b/runners/kafka-streams/proto/build.gradle @@ -0,0 +1,34 @@ +/* + * 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. + */ + +plugins { id 'org.apache.beam.module' } + +// Portability nature compiles the .proto against the vendored gRPC/protobuf (relocated to +// org.apache.beam.vendor.grpc...), so consumers use the vendored protobuf runtime — no raw +// com.google.protobuf on their classpath. Same pattern as the dataflow windmill proto module. +applyPortabilityNature( + publish: false, + shadowJarValidationExcludes: ["org/apache/beam/runners/kafka/streams/v1/**"], + archivesBaseName: 'beam-runners-kafka-streams-proto', + generatedClassPatterns: [ + /^org\.apache\.beam\.runners\.kafka\.streams\.v1.*/ + ] +) + +description = "Apache Beam :: Runners :: Kafka Streams :: Proto" +ext.summary = "Kafka Streams runner control-message protos" diff --git a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto new file mode 100644 index 000000000000..87e3de44397b --- /dev/null +++ b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto @@ -0,0 +1,52 @@ +/* + * 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. + */ + +syntax = "proto3"; + +package org.apache.beam.runners.kafka.streams.v1; + +option java_package = "org.apache.beam.runners.kafka.streams.v1"; +option java_outer_classname = "KafkaStreamsPayloadProtos"; + +// On-wire form of the in-JVM KStreamsPayload envelope, used when the payload must cross a Kafka +// topic boundary (e.g. the GroupByKey repartition topic and the watermark fan-out). Protobuf is +// used for compatible schema evolution and compact varint encoding. +message KafkaStreamsPayload { + // A watermark report: the watermark plus the in-band coordination fields the downstream + // WatermarkManager needs. + message WatermarkPayload { + // Event-time watermark in milliseconds. Signed (sint64, zigzag-encoded) because Beam event + // times can be negative, e.g. BoundedWindow.TIMESTAMP_MIN_VALUE. + sint64 millis = 1; + // The source partition this report is for. + uint32 source_partition = 2; + // The total number of source partitions feeding the downstream stage. + uint32 total_partitions = 3; + } + + // A data element: the Beam WindowedValue encoded with the PCollection's windowed-value coder. + message DataPayload { + bytes value = 1; + } + + // Exactly one variant is set; the oneof case discriminates data vs watermark. + oneof payload { + WatermarkPayload watermark = 1; + DataPayload data = 2; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java new file mode 100644 index 000000000000..912eb5fb6047 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java @@ -0,0 +1,121 @@ +/* + * 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 java.io.ByteArrayOutputStream; +import java.io.IOException; +import org.apache.beam.runners.kafka.streams.v1.KafkaStreamsPayloadProtos.KafkaStreamsPayload; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException; +import org.apache.kafka.common.errors.SerializationException; +import org.apache.kafka.common.serialization.Deserializer; +import org.apache.kafka.common.serialization.Serde; +import org.apache.kafka.common.serialization.Serializer; + +/** + * Kafka {@link Serde} for {@link KStreamsPayload}, enabling the envelope to cross topic boundaries + * (e.g. the repartition topic a {@code GroupByKey} introduces). Until now {@link KStreamsPayload} + * only flowed in-JVM via {@code ProcessorContext#forward}, so no serialization was needed. + * + *

The wire form is the {@link KafkaStreamsPayload} protobuf message — protobuf gives compatible + * schema evolution and compact varint encoding. The data variant carries the {@link WindowedValue} + * encoded with the {@link Coder} supplied for the topic's PCollection; the watermark variant + * carries the coder-independent watermark report. A {@link KStreamsPayloadSerde} is therefore + * parameterized by the data {@link Coder} (different topics carry different element types). + * + *

The serde assumes non-null payloads: the topics it is used on (repartition and watermark + * fan-out) are not log-compacted, so no tombstone (null-valued) records occur. + * + * @param the data element type carried by data payloads on this topic + */ +public final class KStreamsPayloadSerde implements Serde> { + + private final Coder> dataCoder; + + public KStreamsPayloadSerde(Coder> dataCoder) { + this.dataCoder = dataCoder; + } + + @Override + public Serializer> serializer() { + return new PayloadSerializer(); + } + + @Override + public Deserializer> deserializer() { + return new PayloadDeserializer(); + } + + private final class PayloadSerializer implements Serializer> { + @Override + public byte[] serialize(String topic, KStreamsPayload payload) { + KafkaStreamsPayload.Builder proto = KafkaStreamsPayload.newBuilder(); + if (payload.isData()) { + ByteArrayOutputStream encoded = new ByteArrayOutputStream(); + try { + dataCoder.encode(payload.getData(), encoded); + } catch (IOException e) { + throw new SerializationException("Failed to encode KStreamsPayload data element", e); + } + proto.setData( + KafkaStreamsPayload.DataPayload.newBuilder() + .setValue(ByteString.copyFrom(encoded.toByteArray()))); + } else { + WatermarkPayload watermark = payload.asWatermark(); + proto.setWatermark( + KafkaStreamsPayload.WatermarkPayload.newBuilder() + .setMillis(watermark.getWatermarkMillis()) + .setSourcePartition(watermark.getSourcePartition()) + .setTotalPartitions(watermark.getTotalSourcePartitions())); + } + return proto.build().toByteArray(); + } + } + + private final class PayloadDeserializer implements Deserializer> { + @Override + public KStreamsPayload deserialize(String topic, byte[] bytes) { + KafkaStreamsPayload proto; + try { + proto = KafkaStreamsPayload.parseFrom(bytes); + } catch (InvalidProtocolBufferException e) { + throw new SerializationException("Failed to parse KStreamsPayload", e); + } + switch (proto.getPayloadCase()) { + case DATA: + try { + return KStreamsPayload.data(dataCoder.decode(proto.getData().getValue().newInput())); + } catch (IOException e) { + throw new SerializationException("Failed to decode KStreamsPayload data element", e); + } + case WATERMARK: + KafkaStreamsPayload.WatermarkPayload watermark = proto.getWatermark(); + return KStreamsPayload.watermark( + watermark.getMillis(), + watermark.getSourcePartition(), + watermark.getTotalPartitions()); + case PAYLOAD_NOT_SET: + default: + throw new SerializationException( + "KStreamsPayload has no payload variant set: " + proto.getPayloadCase()); + } + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java new file mode 100644 index 000000000000..19346f54763e --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java @@ -0,0 +1,96 @@ +/* + * 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 static org.junit.Assert.assertThrows; + +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.kafka.common.errors.SerializationException; +import org.apache.kafka.common.serialization.Deserializer; +import org.apache.kafka.common.serialization.Serializer; +import org.junit.Test; + +/** Tests for {@link KStreamsPayloadSerde}. */ +public class KStreamsPayloadSerdeTest { + + private static final String TOPIC = "ks-payload-serde-test"; + + private final Coder> dataCoder = + WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE); + private final KStreamsPayloadSerde serde = new KStreamsPayloadSerde<>(dataCoder); + + private KStreamsPayload roundTrip(KStreamsPayload payload) { + Serializer> serializer = serde.serializer(); + Deserializer> deserializer = serde.deserializer(); + return deserializer.deserialize(TOPIC, serializer.serialize(TOPIC, payload)); + } + + @Test + public void roundTripsDataPayload() { + KStreamsPayload payload = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(42)); + KStreamsPayload out = roundTrip(payload); + assertThat(out.isData(), is(true)); + assertThat(out.getData().getValue(), is(42)); + assertThat(out, is(payload)); + } + + @Test + public void roundTripsWatermarkPayload() { + KStreamsPayload payload = KStreamsPayload.watermark(12345L, 2, 4); + KStreamsPayload out = roundTrip(payload); + assertThat(out.isWatermark(), is(true)); + assertThat(out.asWatermark().getWatermarkMillis(), is(12345L)); + assertThat(out.asWatermark().getSourcePartition(), is(2)); + assertThat(out.asWatermark().getTotalSourcePartitions(), is(4)); + assertThat(out, is(payload)); + } + + @Test + public void roundTripsTerminalMaxWatermark() { + KStreamsPayload payload = + KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(), 0, 1); + assertThat( + roundTrip(payload).asWatermark().getWatermarkMillis(), + is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); + } + + @Test + public void roundTripsNegativeWatermark() { + // Beam event times can be negative; sint64 must round-trip them losslessly. + KStreamsPayload payload = + KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(), 0, 1); + assertThat( + roundTrip(payload).asWatermark().getWatermarkMillis(), + is(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis())); + } + + @Test + public void malformedBytesThrow() { + // 0x7f encodes field 15 with the invalid wire type 7, so protobuf parsing fails. + byte[] bogus = new byte[] {(byte) 0x7f}; + assertThrows( + SerializationException.class, () -> serde.deserializer().deserialize(TOPIC, bogus)); + } +} diff --git a/settings.gradle.kts b/settings.gradle.kts index b92b254981fe..9a9cb1dd1acb 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -144,6 +144,7 @@ include(":runners:java-fn-execution") include(":runners:java-job-service") include(":runners:jet") include(":runners:kafka-streams") +include(":runners:kafka-streams:proto") include(":runners:local-java") include(":runners:portability:java") include(":runners:prism") From 5e65d4772fb72265ec1d6ebe3235f6777a30bec6 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Tue, 30 Jun 2026 14:21:02 +0500 Subject: [PATCH 13/37] [GSoC 2026] Kafka Streams runner #39141: Add GroupByKey (GlobalWindow, fire at watermark) --- .../GroupByKeyBroadcastPartitioner.java | 62 ++++++ .../translation/GroupByKeyProcessor.java | 196 ++++++++++++++++++ .../translation/GroupByKeyTranslator.java | 130 ++++++++++++ .../KafkaStreamsPipelineTranslator.java | 1 + .../translation/ShuffleByKeyProcessor.java | 86 ++++++++ .../streams/translation/GroupByKeyTest.java | 184 ++++++++++++++++ .../KafkaStreamsPipelineTranslatorTest.java | 16 +- 7 files changed, 667 insertions(+), 8 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java new file mode 100644 index 000000000000..b5ddcf2536a2 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java @@ -0,0 +1,62 @@ +/* + * 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 java.util.Collections; +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; +import org.apache.kafka.common.utils.Utils; +import org.apache.kafka.streams.processor.StreamPartitioner; + +/** + * Partitions records on the GroupByKey repartition topic. + * + *

    + *
  • data records go to the single partition selected by hashing the (already encoded + * Beam key) Kafka record key — the same scheme Kafka's default partitioner uses — so every + * value of a key lands together; + *
  • watermark reports are broadcast to every partition, so each downstream + * GroupByKey task observes the terminal watermark and fires its keys. + *
+ * + * @param the data element type carried by data payloads + */ +class GroupByKeyBroadcastPartitioner implements StreamPartitioner> { + + @Override + public Integer partition(String topic, byte[] key, KStreamsPayload value, int numPartitions) { + // Required by the interface but unused: Kafka Streams calls partitions() (overridden below) + // when it is present. Kept consistent with the data-hash path for safety. + return key == null ? 0 : Utils.toPositive(Utils.murmur2(key)) % numPartitions; + } + + @Override + public Optional> partitions( + String topic, byte[] key, KStreamsPayload value, int numPartitions) { + if (value.isWatermark()) { + Set all = new HashSet<>(); + for (int partition = 0; partition < numPartitions; partition++) { + all.add(partition); + } + return Optional.of(all); + } + int partition = Utils.toPositive(Utils.murmur2(key)) % numPartitions; + return Optional.of(Collections.singleton(partition)); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java new file mode 100644 index 000000000000..a6c0cb8c4061 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java @@ -0,0 +1,196 @@ +/* + * 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 java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.IterableCoder; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * Executes a {@code GroupByKey} (GlobalWindow, default trigger, no allowed lateness). + * + *

Records arrive on the repartition topic keyed by the encoded Beam key, so every value of a key + * is co-located here. Each value is appended to a per-key buffer in a Kafka Streams state store. + * Watermark reports are fed to a {@link WatermarkManager}; when the input watermark reaches {@link + * BoundedWindow#TIMESTAMP_MAX_VALUE} (the end of the global window) every buffered key is emitted + * once as {@code KV>} and the buffer cleared, then the watermark is forwarded + * downstream. + * + *

Buffering whole value lists and re-encoding on each append is O(n^2) per key; fine for this + * first GroupByKey, and replaced when this moves to runner-core {@code GroupAlsoByWindow}. + */ +class GroupByKeyProcessor + implements Processor, byte[], KStreamsPayload> { + + private final String stateStoreName; + private final Coder keyCoder; + private final IterableCoder<@Nullable Object> bufferCoder; + + private final WatermarkManager watermarkManager = new WatermarkManager(); + private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + // The global window fires exactly once, when the watermark first reaches its end. Later watermark + // reports (e.g. the same terminal watermark broadcast across repartition partitions) must not + // re-fire. This flag is in-memory only; restart correctness comes from the state store plus + // exactly-once-v2: the buffered values and consumer offsets are committed atomically, and the + // store is empty once a key has fired, so a restart cannot double-emit. Persisting watermark + // holds is part of the separate WatermarkManager persistence work, not this initial GroupByKey. + private boolean fired = false; + + private @Nullable ProcessorContext> context; + private @Nullable KeyValueStore store; + + GroupByKeyProcessor( + String stateStoreName, Coder keyCoder, Coder<@Nullable Object> valueCoder) { + this.stateStoreName = stateStoreName; + this.keyCoder = keyCoder; + this.bufferCoder = IterableCoder.of(valueCoder); + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + this.store = context.getStateStore(stateStoreName); + } + + @Override + public void process(Record> record) { + KStreamsPayload payload = record.value(); + if (payload.isData()) { + byte[] encodedKey = record.key(); + Object element = payload.getData().getValue(); + if (encodedKey == null || element == null) { + throw new IllegalStateException("GroupByKey data record is missing its key or value"); + } + appendValue(encodedKey, element); + return; + } + WatermarkPayload report = payload.asWatermark(); + watermarkManager.observe( + report.getSourcePartition(), + new Instant(report.getWatermarkMillis()), + report.getTotalSourcePartitions()); + Instant advanced = watermarkManager.advance(); + if (!fired && !advanced.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) { + fireAll(record); + fired = true; + } + if (advanced.isAfter(lastForwardedWatermark)) { + lastForwardedWatermark = advanced; + forwardWatermark(record, advanced.getMillis()); + } + } + + private void appendValue(byte[] encodedKey, Object kvObject) { + KV kv = (KV) kvObject; + KeyValueStore kvStore = checkInitialized(store); + byte[] existing = kvStore.get(encodedKey); + List<@Nullable Object> values = existing == null ? new ArrayList<>() : decodeBuffer(existing); + values.add(kv.getValue()); + kvStore.put(encodedKey, encodeBuffer(values)); + } + + private void fireAll(Record> trigger) { + // NOTE: this emits every buffered key in a single watermark turn. For a very large key space + // that risks memory pressure and exceeding the poll / transaction timeout. Acceptable for this + // initial GlobalWindow GroupByKey (fire once at end of input); incremental, timer-driven output + // via runner-core GroupAlsoByWindow lands with the windowing/timers work. + ProcessorContext> ctx = checkInitialized(context); + KeyValueStore kvStore = checkInitialized(store); + List firedKeys = new ArrayList<>(); + try (KeyValueIterator it = kvStore.all()) { + while (it.hasNext()) { + org.apache.kafka.streams.KeyValue entry = it.next(); + Object key = decodeKey(entry.key); + List<@Nullable Object> values = decodeBuffer(entry.value); + // The pane fires at the end of the global window, so the grouped element carries the + // window's max timestamp (END_OF_GLOBAL_WINDOW). Emitting at TIMESTAMP_MIN_VALUE (the + // default of valueInGlobalWindow) would make the output appear arbitrarily late and be + // dropped downstream once the watermark has advanced. + WindowedValue>> output = + WindowedValues.timestampedValueInGlobalWindow( + KV.of(key, (Iterable<@Nullable Object>) values), + GlobalWindow.INSTANCE.maxTimestamp()); + ctx.forward( + new Record>( + entry.key, KStreamsPayload.data(output), trigger.timestamp())); + firedKeys.add(entry.key); + } + } + for (byte[] key : firedKeys) { + kvStore.delete(key); + } + } + + private void forwardWatermark(Record> trigger, long watermarkMillis) { + ProcessorContext> ctx = checkInitialized(context); + // GroupByKey is a single logical source for the next stage; report it as partition 0 of 1. + ctx.forward( + new Record>( + trigger.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1), trigger.timestamp())); + } + + private byte[] encodeBuffer(List<@Nullable Object> values) { + try { + return CoderUtils.encodeToByteArray(bufferCoder, values); + } catch (CoderException e) { + throw new RuntimeException("Failed to encode GroupByKey value buffer", e); + } + } + + private List<@Nullable Object> decodeBuffer(byte[] bytes) { + try { + List<@Nullable Object> values = new ArrayList<>(); + for (@Nullable Object value : CoderUtils.decodeFromByteArray(bufferCoder, bytes)) { + values.add(value); + } + return values; + } catch (CoderException e) { + throw new RuntimeException("Failed to decode GroupByKey value buffer", e); + } + } + + private Object decodeKey(byte[] bytes) { + try { + return CoderUtils.decodeFromByteArray(keyCoder, bytes); + } catch (CoderException e) { + throw new RuntimeException("Failed to decode GroupByKey key", e); + } + } + + private static T checkInitialized(@Nullable T value) { + if (value == null) { + throw new IllegalStateException("GroupByKeyProcessor used before init()"); + } + return value; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java new file mode 100644 index 000000000000..d7c4a309d9c3 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -0,0 +1,130 @@ +/* + * 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.apache.beam.runners.fnexecution.translation.PipelineTranslatorUtils.instantiateCoder; + +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.state.Stores; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Translates the {@code beam:transform:group_by_key:v1} URN — the runner's first stateful, + * shuffle-bearing transform. + * + *

This is the simplest GroupByKey: GlobalWindow, default trigger, no allowed lateness (per the + * plan agreed with the mentor). Each key's values are buffered in a Kafka Streams state store and + * emitted once as {@code KV>} when the watermark reaches {@link + * org.apache.beam.sdk.transforms.windowing.BoundedWindow#TIMESTAMP_MAX_VALUE}. + * + *

Topology added (the Beam key becomes the Kafka record key so Kafka Streams shuffles by it): + * + *

    + *
  • a {@link ShuffleByKeyProcessor} wired to the input's producer, which sets the Kafka record + * key to the encoded Beam key for data records and passes watermark reports through; + *
  • a {@link Topology#addSink sink} to an internal repartition topic, with the payload encoded + * via {@link KStreamsPayloadSerde} and a {@link GroupByKeyBroadcastPartitioner} that hashes + * data by key and fans watermark reports out to every partition; + *
  • a {@link Topology#addSource source} reading the repartition topic back; + *
  • the {@link GroupByKeyProcessor} plus a persistent state store, wired to the source. + *
+ * + *

The repartition topic is expected to exist on the broker before the job starts (same + * pre-create assumption as the Impulse bootstrap topic); auto-creation lands with the AdminClient + * wiring in a follow-up. + */ +class GroupByKeyTranslator implements PTransformTranslator { + + static final String SHUFFLE_SUFFIX = "-shuffle-by-key"; + static final String SINK_SUFFIX = "-repartition-sink"; + static final String SOURCE_SUFFIX = "-repartition-source"; + static final String STATE_STORE_SUFFIX = "-state"; + static final String REPARTITION_TOPIC_PREFIX = "__beam_gbk_"; + + @Override + public void translate( + String transformId, RunnerApi.Pipeline pipeline, KafkaStreamsTranslationContext context) { + RunnerApi.PTransform transform = pipeline.getComponents().getTransformsOrThrow(transformId); + String inputPCollectionId = Iterables.getOnlyElement(transform.getInputsMap().values()); + String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + + @SuppressWarnings({"unchecked", "rawtypes"}) + WindowedValues.WindowedValueCoder> inputCoder = + (WindowedValues.WindowedValueCoder) + instantiateCoder(inputPCollectionId, pipeline.getComponents()); + KvCoder kvCoder = (KvCoder) inputCoder.getValueCoder(); + Coder keyCoder = kvCoder.getKeyCoder(); + // User values may be null; the checker tracks that through to the buffered iterables. + @SuppressWarnings("unchecked") + Coder<@Nullable Object> valueCoder = + (Coder<@Nullable Object>) (Coder) kvCoder.getValueCoder(); + + String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); + + String shuffleName = transformId + SHUFFLE_SUFFIX; + String sinkName = transformId + SINK_SUFFIX; + String sourceName = transformId + SOURCE_SUFFIX; + String stateStoreName = transformId + STATE_STORE_SUFFIX; + String repartitionTopic = repartitionTopic(transformId); + + KStreamsPayloadSerde> payloadSerde = new KStreamsPayloadSerde<>(inputCoder); + + Topology topology = context.getTopology(); + + // Re-key data records by the encoded Beam key; pass watermark reports through. + topology.addProcessor(shuffleName, () -> new ShuffleByKeyProcessor(keyCoder), parentProcessor); + + // Shuffle through the repartition topic: data partitioned by key, watermark broadcast. + topology.addSink( + sinkName, + repartitionTopic, + Serdes.ByteArray().serializer(), + payloadSerde.serializer(), + new GroupByKeyBroadcastPartitioner<>(), + shuffleName); + topology.addSource( + sourceName, + Serdes.ByteArray().deserializer(), + payloadSerde.deserializer(), + repartitionTopic); + + // Buffer values per key and fire KV> at the terminal watermark. + topology.addProcessor( + transformId, + () -> new GroupByKeyProcessor(stateStoreName, keyCoder, valueCoder), + sourceName); + topology.addStateStore( + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(stateStoreName), Serdes.ByteArray(), Serdes.ByteArray()), + transformId); + + context.registerPCollectionProducer(outputPCollectionId, transformId); + } + + /** The internal repartition topic name for a GroupByKey transform. */ + static String repartitionTopic(String transformId) { + return REPARTITION_TOPIC_PREFIX + transformId.replaceAll("[^a-zA-Z0-9._-]", "_"); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index 4e227749e1a3..189482ed5f9d 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -51,6 +51,7 @@ public KafkaStreamsPipelineTranslator() { ImmutableMap.builder() .put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator()) .put(PTransformTranslation.REDISTRIBUTE_ARBITRARILY_URN, new RedistributeTranslator()) + .put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new GroupByKeyTranslator()) .put(ExecutableStage.URN, new ExecutableStageTranslator()) .build()); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java new file mode 100644 index 000000000000..3184d838e09d --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java @@ -0,0 +1,86 @@ +/* + * 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 org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.KV; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Re-keys a {@code KV}-valued stream by the Beam key so Kafka Streams shuffles by it. + * + *

This is not GroupByKey-specific: any transform that needs the values of a key co-located on + * one partition uses it — GroupByKey today, and stateful ParDo later. For a data record it sets the + * Kafka record key to the encoded Beam key (taken from the {@code KV}), so the downstream + * repartition sink co-locates every value of a key. Watermark reports are forwarded unchanged — the + * {@link GroupByKeyBroadcastPartitioner} fans them out to all partitions so every downstream task + * can fire. + */ +class ShuffleByKeyProcessor + implements Processor, byte[], KStreamsPayload> { + + private final Coder keyCoder; + private @Nullable ProcessorContext> context; + + ShuffleByKeyProcessor(Coder keyCoder) { + this.keyCoder = keyCoder; + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + } + + @Override + public void process(Record> record) { + ProcessorContext> ctx = checkInitialized(context); + KStreamsPayload payload = record.value(); + if (payload.isData()) { + Object element = payload.getData().getValue(); + if (element == null) { + throw new IllegalStateException("shuffle data element must not be null"); + } + Object key = ((KV) element).getKey(); + if (key == null) { + throw new IllegalStateException("shuffle key must not be null"); + } + byte[] encodedKey; + try { + encodedKey = CoderUtils.encodeToByteArray(keyCoder, key); + } catch (CoderException e) { + throw new RuntimeException("Failed to encode shuffle key", e); + } + ctx.forward(record.withKey(encodedKey)); + } else { + // Watermark report: forward as-is; the sink's partitioner broadcasts it to all partitions. + ctx.forward(record); + } + } + + private static T checkInitialized(@Nullable T value) { + if (value == null) { + throw new IllegalStateException("ShuffleByKeyProcessor used before init()"); + } + return value; + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTest.java new file mode 100644 index 000000000000..0aed80fa8585 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTest.java @@ -0,0 +1,184 @@ +/* + * 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.hasItems; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Properties; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.PTransformTranslation; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.beam.sdk.values.KV; +import org.apache.kafka.common.serialization.ByteArrayDeserializer; +import org.apache.kafka.common.serialization.ByteArraySerializer; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.TestInputTopic; +import org.apache.kafka.streams.TestOutputTopic; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.TopologyTestDriver; +import org.apache.kafka.streams.test.TestRecord; +import org.junit.Test; + +/** + * End-to-end test of GroupByKey: {@code Impulse -> emit KVs -> GroupByKey -> record groups}. + * + *

GroupByKey shuffles through an internal repartition topic. {@link TopologyTestDriver} does not + * loop a low-level sink topic back into its source, so the test drives the upstream, drains the + * repartition topic, and pipes those records back into it — standing in for the broker round-trip. + * The downstream {@code RecordGroupFn} records each emitted group into a {@link + * SharedTestCollector}. + */ +public class GroupByKeyTest { + + private static final String JOB_ID = "ks-gbk-test"; + private static final String APPLICATION_ID = "ks-gbk-test"; + + /** Emits a few KVs from the single impulse element so there is something to group. */ + private static class EmitKvsFn extends DoFn> { + @ProcessElement + public void processElement(OutputReceiver> out) { + out.output(KV.of("a", 1)); + out.output(KV.of("a", 2)); + out.output(KV.of("b", 3)); + } + } + + /** Records each grouped result as {@code "key=[sorted values]"}. */ + private static class RecordGroupFn extends DoFn>, Void> { + private final SharedTestCollector collector; + + RecordGroupFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element KV> group) { + List values = new ArrayList<>(); + group.getValue().forEach(values::add); + Collections.sort(values); + collector.record(group.getKey() + "=" + values); + } + } + + @Test + public void groupsValuesByKeyAndFiresAtWatermark() throws Exception { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(pipelineOptions()); + pipeline + .apply("impulse", Impulse.create()) + .apply("emit", ParDo.of(new EmitKvsFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())) + .apply("gbk", GroupByKey.create()) + .apply("record", ParDo.of(new RecordGroupFn(collector))); + + RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); + KafkaStreamsPipelineOptions options = + pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + JobInfo jobInfo = + JobInfo.create( + JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); + KafkaStreamsTranslationContext context = + translator.createTranslationContext(jobInfo, options); + + RunnerApi.Pipeline prepared = translator.prepareForTranslation(pipelineProto); + translator.translate(context, prepared); + String repartitionTopic = GroupByKeyTranslator.repartitionTopic(findGroupByKeyId(prepared)); + + Topology topology = context.getTopology(); + try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { + // Fire the impulse; the upstream stage emits the KVs and the terminal watermark, which the + // re-key processor sends to the repartition sink. + driver.advanceWallClockTime(Duration.ofSeconds(1)); + driver.advanceWallClockTime(Duration.ofSeconds(1)); + + // Round-trip the repartition topic: drain what the sink wrote and feed it back to the + // source so the GroupByKey processor buffers the values and fires at the watermark. + TestOutputTopic repartitionOut = + driver.createOutputTopic( + repartitionTopic, new ByteArrayDeserializer(), new ByteArrayDeserializer()); + TestInputTopic repartitionIn = + driver.createInputTopic( + repartitionTopic, new ByteArraySerializer(), new ByteArraySerializer()); + for (TestRecord record : repartitionOut.readRecordsToList()) { + repartitionIn.pipeInput(record); + } + } + + List groups = collector.recorded(); + assertThat(groups.size(), is(2)); + assertThat(groups, hasItems("a=[1, 2]", "b=[3]")); + } + } + + private static String findGroupByKeyId(RunnerApi.Pipeline pipeline) { + return pipeline.getComponents().getTransformsMap().entrySet().stream() + .filter( + e -> + PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN.equals( + e.getValue().getSpec().getUrn())) + .map(java.util.Map.Entry::getKey) + .findFirst() + .orElseThrow(() -> new AssertionError("no GroupByKey transform in the pipeline")); + } + + private static PipelineOptions pipelineOptions() { + PipelineOptions options = + PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); + options.setRunner(CrashingRunner.class); + options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); + options + .as(PortablePipelineOptions.class) + .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); + return options; + } + + private static Properties streamsConfig() { + Properties props = new Properties(); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put( + StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + props.put( + StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + return props; + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java index 13baa551ebbf..8b56915fe26f 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslatorTest.java @@ -45,18 +45,18 @@ public void translateRejectsUnknownTransformWithUrnInMessage() { KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); KafkaStreamsTranslationContext context = newContext(); + // A URN the runner does not register a translator for. + String unsupportedUrn = "beam:transform:kafka_streams_unsupported_test:v1"; RunnerApi.Pipeline pipeline = RunnerApi.Pipeline.newBuilder() - .addRootTransformIds("gbk") + .addRootTransformIds("unsupported") .setComponents( RunnerApi.Components.newBuilder() .putTransforms( - "gbk", + "unsupported", RunnerApi.PTransform.newBuilder() - .setUniqueName("GroupByKey") - .setSpec( - RunnerApi.FunctionSpec.newBuilder() - .setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)) + .setUniqueName("Unsupported") + .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(unsupportedUrn)) .build())) .build(); @@ -67,8 +67,8 @@ public void translateRejectsUnknownTransformWithUrnInMessage() { UnsupportedOperationException.class, () -> translator.translate(context, pipeline)); assertThat(ex.getMessage(), containsString("No translator registered for URN")); - assertThat(ex.getMessage(), containsString(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)); - assertThat(ex.getMessage(), containsString("gbk")); + assertThat(ex.getMessage(), containsString(unsupportedUrn)); + assertThat(ex.getMessage(), containsString("unsupported")); assertThat(ex.getMessage(), containsString(JOB_ID)); } From 4d5847f6bfa3202b2f94bb956cd47d567a3b6875 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Tue, 7 Jul 2026 15:22:28 +0500 Subject: [PATCH 14/37] [GSoC 2026] Kafka Streams runner #39211: Add KafkaStreamsTestRunner test harness --- .../kafka/streams/KafkaStreamsRunnerTest.java | 79 +------ .../kafka/streams/KafkaStreamsTestRunner.java | 220 ++++++++++++++++++ .../ExecutableStageTranslatorTest.java | 72 +----- .../streams/translation/GroupByKeyTest.java | 105 +-------- .../translation/WatermarkPropagationTest.java | 74 +----- 5 files changed, 251 insertions(+), 299 deletions(-) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java index 0c576d27122f..02bc1a887c73 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java @@ -24,22 +24,11 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.Properties; import org.apache.beam.model.pipeline.v1.RunnerApi; -import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.kafka.streams.translation.KStreamsPayload; -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.options.PipelineOptions; -import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.apache.beam.sdk.testing.CrashingRunner; import org.apache.beam.sdk.transforms.Impulse; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; -import org.apache.beam.sdk.util.construction.PipelineTranslation; -import org.apache.kafka.common.serialization.Serdes; -import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.Topology; import org.apache.kafka.streams.TopologyTestDriver; import org.apache.kafka.streams.processor.api.Processor; @@ -51,8 +40,8 @@ /** * Pipeline-level integration tests that build a Beam {@link Pipeline} via the high-level Java SDK - * ({@code Pipeline.create().apply(Impulse.create())}), translate it via the runner, and execute the - * resulting Kafka Streams topology under {@link TopologyTestDriver}. + * ({@code Pipeline.create().apply(Impulse.create())}) and run it through {@link + * KafkaStreamsTestRunner}. * *

This is the test layer Jan requested on PR #38689: rather than building hand-rolled {@link * RunnerApi.Pipeline} protos, drive translation from the same surface a user would write. The tests @@ -61,33 +50,19 @@ */ public class KafkaStreamsRunnerTest { - private static final String JOB_ID = "kafka-streams-runner-test"; - private static final String APPLICATION_ID = "ks-runner-test"; - @Test public void impulseOnlyPipelineEmitsDataAndTerminalWatermark() { - Pipeline pipeline = Pipeline.create(pipelineOptions()); + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); pipeline.apply("impulse", Impulse.create()); - RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); - - KafkaStreamsPipelineOptions options = - pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); - KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); - JobInfo jobInfo = - JobInfo.create( - JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); - KafkaStreamsTranslationContext context = translator.createTranslationContext(jobInfo, options); - translator.translate(context, translator.prepareForTranslation(pipelineProto)); - CapturingProcessor capture = new CapturingProcessor(); - Topology topology = context.getTopology(); - // Wire a downstream test sink to every translated transform node so we can capture emissions. - // Impulse is the only transform here, so we attach to "impulse" (the processor name registered - // by ImpulseTranslator). - topology.addProcessor("capture", capture, expectedImpulseProcessorName(pipelineProto)); + Topology topology = KafkaStreamsTestRunner.translate(pipeline).getTopology(); + // Impulse is the only transform, so it is the topology leaf; capture what it forwards. + topology.addProcessor( + "capture", capture, KafkaStreamsTestRunner.findAnyLeafProcessorName(topology)); - try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { + try (TopologyTestDriver driver = + new TopologyTestDriver(topology, KafkaStreamsTestRunner.streamsConfig(pipeline))) { driver.advanceWallClockTime(Duration.ofSeconds(1)); driver.advanceWallClockTime(Duration.ofSeconds(1)); } @@ -101,42 +76,6 @@ public void impulseOnlyPipelineEmitsDataAndTerminalWatermark() { is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); } - /** - * Finds the transform id that {@link Impulse} got assigned by the SDK so the test can attach a - * capturing processor to the matching Kafka Streams processor node (the translator names the - * processor after the transform id). - */ - private static String expectedImpulseProcessorName(RunnerApi.Pipeline pipelineProto) { - for (java.util.Map.Entry entry : - pipelineProto.getComponents().getTransformsMap().entrySet()) { - if ("beam:transform:impulse:v1".equals(entry.getValue().getSpec().getUrn())) { - return entry.getKey(); - } - } - throw new IllegalStateException("Impulse transform not found in pipeline proto"); - } - - private static PipelineOptions pipelineOptions() { - PipelineOptions options = - PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); - // Pipeline.create() requires a runner; CrashingRunner is the conventional "this pipeline is - // not going to be run() directly" choice used by other portable-runner tests. - options.setRunner(CrashingRunner.class); - options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); - return options; - } - - private static Properties streamsConfig() { - Properties props = new Properties(); - props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); - props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); - props.put( - StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); - props.put( - StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); - return props; - } - private static class CapturingProcessor implements ProcessorSupplier< byte[], KStreamsPayload, byte[], KStreamsPayload> { 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 new file mode 100644 index 000000000000..53b7714548a0 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java @@ -0,0 +1,220 @@ +/* + * 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; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Properties; +import java.util.Set; +import java.util.UUID; +import org.apache.beam.model.pipeline.v1.RunnerApi; +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.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.kafka.common.serialization.ByteArrayDeserializer; +import org.apache.kafka.common.serialization.ByteArraySerializer; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.TestInputTopic; +import org.apache.kafka.streams.TestOutputTopic; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.TopologyDescription; +import org.apache.kafka.streams.TopologyTestDriver; +import org.apache.kafka.streams.test.TestRecord; + +/** + * Test harness that runs a Beam {@link Pipeline} through the Kafka Streams runner's translation and + * a {@link TopologyTestDriver}, so tests do not repeat the translate + drive boilerplate. + * + *

Usage: build a pipeline with {@link #testOptions()}, then call {@link #run(Pipeline)}. Side + * effects (e.g. a {@code SharedTestCollector} written by a recording DoFn) have completed when it + * returns. + * + *

{@link TopologyTestDriver} does not loop a low-level sink topic back into its source, so an + * internal repartition topic (one that is both a sink and a source in the topology — e.g. the one + * GroupByKey introduces) would otherwise dead-end. {@link #run(Pipeline)} discovers those topics + * from the {@link TopologyDescription} and round-trips them until no more records flow, standing in + * for the broker. + */ +public final class KafkaStreamsTestRunner { + + private static final int MAX_ROUND_TRIPS = 100; + + private KafkaStreamsTestRunner() {} + + /** Pipeline options for a Kafka Streams runner test: the EMBEDDED harness and a unique app id. */ + public static PipelineOptions testOptions() { + String applicationId = "ks-test-" + UUID.randomUUID(); + PipelineOptions options = + PipelineOptionsFactory.fromArgs("--applicationId=" + applicationId).create(); + options.setRunner(CrashingRunner.class); + options.as(KafkaStreamsPipelineOptions.class).setApplicationId(applicationId); + options + .as(PortablePipelineOptions.class) + .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); + return options; + } + + /** + * Translates the pipeline into a Kafka Streams {@link KafkaStreamsTranslationContext}. Tests that + * need the {@link Topology} (e.g. to attach a capture processor before driving) use this and + * build their own {@link TopologyTestDriver}; simpler tests use {@link #run(Pipeline)}. + */ + public static KafkaStreamsTranslationContext translate(Pipeline pipeline) { + KafkaStreamsPipelineOptions options = + pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + JobInfo jobInfo = + JobInfo.create( + options.getApplicationId(), + options.getJobName(), + "", + PipelineOptionsTranslation.toProto(options)); + KafkaStreamsTranslationContext context = translator.createTranslationContext(jobInfo, options); + translator.translate(context, translator.prepareForTranslation(pipelineProto)); + return context; + } + + /** Translates and drives the pipeline to quiescence through a {@link TopologyTestDriver}. */ + public static void run(Pipeline pipeline) { + KafkaStreamsTranslationContext context = translate(pipeline); + Topology topology = context.getTopology(); + try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig(pipeline))) { + // Fire the Impulse wall-clock punctuator and let the initial records flow. + driver.advanceWallClockTime(Duration.ofSeconds(1)); + driver.advanceWallClockTime(Duration.ofSeconds(1)); + roundTripInternalTopics(driver, internalTopics(topology)); + } + } + + /** + * The name of some processor node with no successors (a topology leaf). Tests with a single leaf + * attach a capture processor here to observe what the last stage forwards; if a topology has more + * than one leaf, which one is returned is unspecified. + */ + public static String findAnyLeafProcessorName(Topology topology) { + for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { + for (TopologyDescription.Node node : subtopology.nodes()) { + if (node instanceof TopologyDescription.Processor && node.successors().isEmpty()) { + return node.name(); + } + } + } + throw new IllegalStateException("no leaf processor found in topology"); + } + + /** Repartition/internal topics are the ones that appear as both a sink and a source. */ + private static Set internalTopics(Topology topology) { + Set sinkTopics = new HashSet<>(); + Set sourceTopics = new HashSet<>(); + for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { + for (TopologyDescription.Node node : subtopology.nodes()) { + if (node instanceof TopologyDescription.Sink) { + String topic = ((TopologyDescription.Sink) node).topic(); + if (topic != null) { + sinkTopics.add(topic); + } + } else if (node instanceof TopologyDescription.Source) { + sourceTopics.addAll(((TopologyDescription.Source) node).topicSet()); + } + } + } + sinkTopics.retainAll(sourceTopics); + return sinkTopics; + } + + /** + * Simulates the broker for internal repartition topics. + * + *

The runner shuffles data (and the watermark) through internal topics that a processor both + * writes to (a sink) and reads back from (a source) — e.g. the topic GroupByKey introduces to + * partition by key. On a real broker those records make the round trip automatically, but {@link + * TopologyTestDriver} does not connect a sink back to a source, so the downstream half of the + * topology would never see them. This drains what each internal topic's sink wrote and pipes it + * into that topic's source, repeating until nothing new flows (a fixpoint), which stands in for + * the broker and lets the pipeline run to completion. + */ + private static void roundTripInternalTopics(TopologyTestDriver driver, Set topics) { + // Create the sink-output and source-input handles once and reuse them across rounds; a single + // TestOutputTopic keeps returning newly produced records on each read. + List roundTrips = new ArrayList<>(); + for (String topic : topics) { + roundTrips.add( + new TopicRoundTrip( + driver.createOutputTopic( + topic, new ByteArrayDeserializer(), new ByteArrayDeserializer()), + driver.createInputTopic( + topic, new ByteArraySerializer(), new ByteArraySerializer()))); + } + + for (int round = 0; round < MAX_ROUND_TRIPS; round++) { + boolean progressed = false; + for (TopicRoundTrip roundTrip : roundTrips) { + List> records = roundTrip.output.readRecordsToList(); + if (records.isEmpty()) { + continue; + } + progressed = true; + for (TestRecord record : records) { + roundTrip.input.pipeInput(record); + } + } + if (!progressed) { + return; + } + } + throw new IllegalStateException( + "Internal topics did not reach quiescence after " + MAX_ROUND_TRIPS + " round trips"); + } + + /** The reusable sink-output and source-input handles for one internal topic. */ + private static final class TopicRoundTrip { + final TestOutputTopic output; + final TestInputTopic input; + + TopicRoundTrip(TestOutputTopic output, TestInputTopic input) { + this.output = output; + this.input = input; + } + } + + /** Kafka Streams config for a {@link TopologyTestDriver} built from the pipeline's app id. */ + public static Properties streamsConfig(Pipeline pipeline) { + String applicationId = + pipeline.getOptions().as(KafkaStreamsPipelineOptions.class).getApplicationId(); + Properties props = new Properties(); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, applicationId); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put( + StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + props.put( + StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + return props; + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslatorTest.java index 98ec45ae2885..4750718c00f6 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslatorTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslatorTest.java @@ -20,35 +20,19 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; -import java.time.Duration; import java.util.List; -import java.util.Properties; -import org.apache.beam.model.pipeline.v1.RunnerApi; -import org.apache.beam.runners.fnexecution.provisioning.JobInfo; -import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; import org.apache.beam.sdk.Pipeline; -import org.apache.beam.sdk.options.PipelineOptions; -import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.apache.beam.sdk.options.PortablePipelineOptions; -import org.apache.beam.sdk.testing.CrashingRunner; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.Impulse; import org.apache.beam.sdk.transforms.ParDo; -import org.apache.beam.sdk.util.construction.Environments; -import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; -import org.apache.beam.sdk.util.construction.PipelineTranslation; -import org.apache.kafka.common.serialization.Serdes; -import org.apache.kafka.streams.StreamsConfig; -import org.apache.kafka.streams.Topology; -import org.apache.kafka.streams.TopologyTestDriver; import org.junit.Test; /** * End-to-end test for {@link ExecutableStageTranslator}: builds an {@code Impulse -> ParDo} - * pipeline with the high-level Beam Java SDK, fuses + translates it, and runs the resulting Kafka - * Streams topology under {@link TopologyTestDriver}. The fused ParDo executes in an in-process - * (EMBEDDED) Java SDK harness, so the {@link DoFn}'s {@code @ProcessElement} body runs for real — - * no Docker, no broker. + * pipeline with the high-level Beam Java SDK and runs it through {@link KafkaStreamsTestRunner}. + * The fused ParDo executes in an in-process (EMBEDDED) Java SDK harness, so the {@link DoFn}'s + * {@code @ProcessElement} body runs for real — no Docker, no broker. * *

Because the ParDo's output PCollection has no downstream consumer, it is not a stage output * and is never forwarded out of the harness — that is the documented behaviour. The test verifies @@ -57,9 +41,6 @@ */ public class ExecutableStageTranslatorTest { - private static final String JOB_ID = "kafka-streams-executable-stage-test"; - private static final String APPLICATION_ID = "ks-executable-stage-test"; - /** * Records the length of every input element seen by the harness so the test can verify the DoFn * ran. {@link SharedTestCollector} carries its identity via a UUID stored on the instance itself, @@ -82,31 +63,14 @@ public void processElement(@Element byte[] input, OutputReceiver out) { } @Test - public void impulseThenParDoExecutesDoFnInHarnessOncePerImpulseElement() throws Exception { + public void impulseThenParDoExecutesDoFnInHarnessOncePerImpulseElement() { try (SharedTestCollector collector = SharedTestCollector.create()) { - Pipeline pipeline = Pipeline.create(pipelineOptions()); + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); pipeline .apply("impulse", Impulse.create()) .apply("pardo", ParDo.of(new RecordingFn(collector))); - RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); - - KafkaStreamsPipelineOptions options = - pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); - KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); - JobInfo jobInfo = - JobInfo.create( - JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); - KafkaStreamsTranslationContext context = - translator.createTranslationContext(jobInfo, options); - - translator.translate(context, translator.prepareForTranslation(pipelineProto)); - - Topology topology = context.getTopology(); - try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { - driver.advanceWallClockTime(Duration.ofSeconds(1)); - driver.advanceWallClockTime(Duration.ofSeconds(1)); - } + KafkaStreamsTestRunner.run(pipeline); List recorded = collector.recorded(); // Impulse emits exactly one empty byte[] in the GlobalWindow, so the DoFn must run exactly @@ -115,26 +79,4 @@ public void impulseThenParDoExecutesDoFnInHarnessOncePerImpulseElement() throws assertThat(recorded.get(0), is(0)); } } - - private static PipelineOptions pipelineOptions() { - PipelineOptions options = - PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); - options.setRunner(CrashingRunner.class); - options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); - options - .as(PortablePipelineOptions.class) - .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); - return options; - } - - private static Properties streamsConfig() { - Properties props = new Properties(); - props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); - props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); - props.put( - StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); - props.put( - StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); - return props; - } } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTest.java index 0aed80fa8585..38f19ded356c 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTest.java @@ -21,56 +21,30 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; -import java.time.Duration; import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.Properties; -import org.apache.beam.model.pipeline.v1.RunnerApi; -import org.apache.beam.runners.fnexecution.provisioning.JobInfo; -import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.coders.VarIntCoder; -import org.apache.beam.sdk.options.PipelineOptions; -import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.apache.beam.sdk.options.PortablePipelineOptions; -import org.apache.beam.sdk.testing.CrashingRunner; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.GroupByKey; import org.apache.beam.sdk.transforms.Impulse; import org.apache.beam.sdk.transforms.ParDo; -import org.apache.beam.sdk.util.construction.Environments; -import org.apache.beam.sdk.util.construction.PTransformTranslation; -import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; -import org.apache.beam.sdk.util.construction.PipelineTranslation; import org.apache.beam.sdk.values.KV; -import org.apache.kafka.common.serialization.ByteArrayDeserializer; -import org.apache.kafka.common.serialization.ByteArraySerializer; -import org.apache.kafka.common.serialization.Serdes; -import org.apache.kafka.streams.StreamsConfig; -import org.apache.kafka.streams.TestInputTopic; -import org.apache.kafka.streams.TestOutputTopic; -import org.apache.kafka.streams.Topology; -import org.apache.kafka.streams.TopologyTestDriver; -import org.apache.kafka.streams.test.TestRecord; import org.junit.Test; /** * End-to-end test of GroupByKey: {@code Impulse -> emit KVs -> GroupByKey -> record groups}. * - *

GroupByKey shuffles through an internal repartition topic. {@link TopologyTestDriver} does not - * loop a low-level sink topic back into its source, so the test drives the upstream, drains the - * repartition topic, and pipes those records back into it — standing in for the broker round-trip. - * The downstream {@code RecordGroupFn} records each emitted group into a {@link - * SharedTestCollector}. + *

Driven through {@link KafkaStreamsTestRunner}, which round-trips the internal repartition + * topic that GroupByKey introduces. The downstream {@code RecordGroupFn} records each emitted group + * into a {@link SharedTestCollector}. */ public class GroupByKeyTest { - private static final String JOB_ID = "ks-gbk-test"; - private static final String APPLICATION_ID = "ks-gbk-test"; - /** Emits a few KVs from the single impulse element so there is something to group. */ private static class EmitKvsFn extends DoFn> { @ProcessElement @@ -99,9 +73,9 @@ public void processElement(@Element KV> group) { } @Test - public void groupsValuesByKeyAndFiresAtWatermark() throws Exception { + public void groupsValuesByKeyAndFiresAtWatermark() { try (SharedTestCollector collector = SharedTestCollector.create()) { - Pipeline pipeline = Pipeline.create(pipelineOptions()); + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); pipeline .apply("impulse", Impulse.create()) .apply("emit", ParDo.of(new EmitKvsFn())) @@ -109,76 +83,11 @@ public void groupsValuesByKeyAndFiresAtWatermark() throws Exception { .apply("gbk", GroupByKey.create()) .apply("record", ParDo.of(new RecordGroupFn(collector))); - RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); - KafkaStreamsPipelineOptions options = - pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); - KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); - JobInfo jobInfo = - JobInfo.create( - JOB_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); - KafkaStreamsTranslationContext context = - translator.createTranslationContext(jobInfo, options); - - RunnerApi.Pipeline prepared = translator.prepareForTranslation(pipelineProto); - translator.translate(context, prepared); - String repartitionTopic = GroupByKeyTranslator.repartitionTopic(findGroupByKeyId(prepared)); - - Topology topology = context.getTopology(); - try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { - // Fire the impulse; the upstream stage emits the KVs and the terminal watermark, which the - // re-key processor sends to the repartition sink. - driver.advanceWallClockTime(Duration.ofSeconds(1)); - driver.advanceWallClockTime(Duration.ofSeconds(1)); - - // Round-trip the repartition topic: drain what the sink wrote and feed it back to the - // source so the GroupByKey processor buffers the values and fires at the watermark. - TestOutputTopic repartitionOut = - driver.createOutputTopic( - repartitionTopic, new ByteArrayDeserializer(), new ByteArrayDeserializer()); - TestInputTopic repartitionIn = - driver.createInputTopic( - repartitionTopic, new ByteArraySerializer(), new ByteArraySerializer()); - for (TestRecord record : repartitionOut.readRecordsToList()) { - repartitionIn.pipeInput(record); - } - } + KafkaStreamsTestRunner.run(pipeline); List groups = collector.recorded(); assertThat(groups.size(), is(2)); assertThat(groups, hasItems("a=[1, 2]", "b=[3]")); } } - - private static String findGroupByKeyId(RunnerApi.Pipeline pipeline) { - return pipeline.getComponents().getTransformsMap().entrySet().stream() - .filter( - e -> - PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN.equals( - e.getValue().getSpec().getUrn())) - .map(java.util.Map.Entry::getKey) - .findFirst() - .orElseThrow(() -> new AssertionError("no GroupByKey transform in the pipeline")); - } - - private static PipelineOptions pipelineOptions() { - PipelineOptions options = - PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); - options.setRunner(CrashingRunner.class); - options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); - options - .as(PortablePipelineOptions.class) - .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); - return options; - } - - private static Properties streamsConfig() { - Properties props = new Properties(); - props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); - props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); - props.put( - StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); - props.put( - StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); - return props; - } } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java index e4bdbe4d9733..0de19d70f7d3 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java @@ -23,26 +23,13 @@ import java.time.Duration; import java.util.ArrayList; import java.util.List; -import java.util.Properties; -import org.apache.beam.model.pipeline.v1.RunnerApi; -import org.apache.beam.runners.fnexecution.provisioning.JobInfo; -import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; import org.apache.beam.sdk.Pipeline; -import org.apache.beam.sdk.options.PipelineOptions; -import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.apache.beam.sdk.options.PortablePipelineOptions; -import org.apache.beam.sdk.testing.CrashingRunner; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.Impulse; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.util.construction.Environments; -import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; -import org.apache.beam.sdk.util.construction.PipelineTranslation; -import org.apache.kafka.common.serialization.Serdes; -import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.Topology; -import org.apache.kafka.streams.TopologyDescription; import org.apache.kafka.streams.TopologyTestDriver; import org.apache.kafka.streams.processor.api.Processor; import org.apache.kafka.streams.processor.api.Record; @@ -57,8 +44,6 @@ */ public class WatermarkPropagationTest { - private static final String APPLICATION_ID = "ks-watermark-propagation-test"; - /** Identity DoFn so the pipeline contains a fused ExecutableStage. */ private static class IdentityFn extends DoFn { @ProcessElement @@ -86,26 +71,19 @@ public void process(Record> record) { @Test public void terminalWatermarkPropagatesToDownstreamStampedAsSingleSource() throws Exception { - Pipeline pipeline = Pipeline.create(pipelineOptions()); + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); pipeline.apply("impulse", Impulse.create()).apply("identity", ParDo.of(new IdentityFn())); - RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); - KafkaStreamsPipelineOptions options = - pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); - KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); - JobInfo jobInfo = - JobInfo.create( - APPLICATION_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); - KafkaStreamsTranslationContext context = translator.createTranslationContext(jobInfo, options); - translator.translate(context, translator.prepareForTranslation(pipelineProto)); - // Attach a sink to the leaf ExecutableStage processor to capture the watermark it forwards. - Topology topology = context.getTopology(); + Topology topology = KafkaStreamsTestRunner.translate(pipeline).getTopology(); List> captured = new ArrayList<>(); topology.addProcessor( - "watermark-capture", () -> new WatermarkCapture(captured), findLeafProcessor(topology)); + "watermark-capture", + () -> new WatermarkCapture(captured), + KafkaStreamsTestRunner.findAnyLeafProcessorName(topology)); - try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { + try (TopologyTestDriver driver = + new TopologyTestDriver(topology, KafkaStreamsTestRunner.streamsConfig(pipeline))) { driver.advanceWallClockTime(Duration.ofSeconds(1)); driver.advanceWallClockTime(Duration.ofSeconds(1)); } @@ -116,40 +94,4 @@ public void terminalWatermarkPropagatesToDownstreamStampedAsSingleSource() throw assertThat(terminal.getSourcePartition(), is(0)); assertThat(terminal.getTotalSourcePartitions(), is(1)); } - - /** - * Returns the name of the single processor node with no successors (the leaf of the topology). - */ - private static String findLeafProcessor(Topology topology) { - for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { - for (TopologyDescription.Node node : subtopology.nodes()) { - if (node instanceof TopologyDescription.Processor && node.successors().isEmpty()) { - return node.name(); - } - } - } - throw new IllegalStateException("no leaf processor found in topology"); - } - - private static PipelineOptions pipelineOptions() { - PipelineOptions options = - PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); - options.setRunner(CrashingRunner.class); - options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); - options - .as(PortablePipelineOptions.class) - .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); - return options; - } - - private static Properties streamsConfig() { - Properties props = new Properties(); - props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); - props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); - props.put( - StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); - props.put( - StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); - return props; - } } From 27c4522ae5e4de8aad086aa917be4ceae3b35b7b Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Thu, 9 Jul 2026 13:20:20 +0500 Subject: [PATCH 15/37] [GSoC 2026] Kafka Streams runner #39249: Support Create --- .../KafkaStreamsPipelineTranslator.java | 1 + .../KafkaStreamsTranslationContext.java | 20 ++ .../streams/translation/ReadProcessor.java | 207 ++++++++++++++++++ .../streams/translation/ReadTranslator.java | 173 +++++++++++++++ .../kafka/streams/KafkaStreamsTestRunner.java | 6 + .../kafka/streams/translation/CreateTest.java | 73 ++++++ .../kafka/streams/translation/ReadTest.java | 75 +++++++ 7 files changed, 555 insertions(+) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/CreateTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ReadTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index 189482ed5f9d..943673dfe98c 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -50,6 +50,7 @@ public KafkaStreamsPipelineTranslator() { this( ImmutableMap.builder() .put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator()) + .put(PTransformTranslation.READ_TRANSFORM_URN, new ReadTranslator()) .put(PTransformTranslation.REDISTRIBUTE_ARBITRARILY_URN, new RedistributeTranslator()) .put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new GroupByKeyTranslator()) .put(ExecutableStage.URN, new ExecutableStageTranslator()) 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 3d95eabafed6..0a045ff7395b 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 @@ -19,6 +19,7 @@ import java.util.HashMap; import java.util.Map; +import java.util.regex.Pattern; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; import org.apache.kafka.streams.Topology; @@ -34,6 +35,12 @@ public class KafkaStreamsTranslationContext { /** Prefix for the per-job bootstrap topic Impulse reads from. */ private static final String IMPULSE_BOOTSTRAP_TOPIC_PREFIX = "__beam_impulse_"; + /** Prefix for the per-transform bootstrap topic a primitive Read reads from. */ + private static final String READ_BOOTSTRAP_TOPIC_PREFIX = "__beam_read_"; + + /** Characters not legal in a Kafka topic name; a topic's legal set is {@code [a-zA-Z0-9._-]}. */ + private static final Pattern ILLEGAL_TOPIC_CHARS = Pattern.compile("[^a-zA-Z0-9._-]"); + private final JobInfo jobInfo; private final KafkaStreamsPipelineOptions pipelineOptions; private final Topology topology; @@ -100,4 +107,17 @@ public String getProcessorNameForPCollection(String pCollectionId) { public String getImpulseBootstrapTopic() { return IMPULSE_BOOTSTRAP_TOPIC_PREFIX + pipelineOptions.getApplicationId(); } + + /** + * Returns the dedicated bootstrap topic name a primitive Read reads from. Keyed by transform id + * (sanitized to Kafka's legal topic-name character set) so multiple Reads — and Impulse — never + * register the same topic on two source nodes, which Kafka Streams rejects. + */ + public String getReadBootstrapTopic(String transformId) { + String sanitizedTransformId = ILLEGAL_TOPIC_CHARS.matcher(transformId).replaceAll("_"); + return READ_BOOTSTRAP_TOPIC_PREFIX + + pipelineOptions.getApplicationId() + + "_" + + sanitizedTransformId; + } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java new file mode 100644 index 000000000000..604eca54014a --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java @@ -0,0 +1,207 @@ +/* + * 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 java.io.IOException; +import java.time.Duration; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.io.BoundedSource; +import org.apache.beam.sdk.io.BoundedSource.BoundedReader; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.kafka.streams.processor.Cancellable; +import org.apache.kafka.streams.processor.PunctuationType; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Kafka Streams {@link Processor} implementing Beam's deprecated primitive {@code Read} + * (beam:transform:read:v1) over a {@link BoundedSource}. + * + *

For each task instance, reads the whole {@link BoundedSource} once and emits, in order: + * + *

    + *
  1. One {@link KStreamsPayload#data data} payload per source element, each wrapping a {@link + * WindowedValue} in the {@link org.apache.beam.sdk.transforms.windowing.GlobalWindow} at the + * element's own event time (from {@link BoundedReader#getCurrentTimestamp()}). + *
  2. A {@link KStreamsPayload#watermark watermark} payload at {@link + * BoundedWindow#TIMESTAMP_MAX_VALUE} telling downstream transforms the source is done. + *
+ * + *

Wire form. Unlike Impulse (whose element is already an opaque {@code byte[]}), a Read + * produces decoded Java objects. Downstream {@link ExecutableStageProcessor} feeds + * whatever it receives straight into the SDK harness, whose main-input receiver expects each + * element in the runner-side wire form — a raw object for a model coder, but a length-prefixed + * {@code byte[]} for a coder the runner does not know (e.g. {@code VarIntCoder}). Stage-to-stage + * edges already carry that wire form because harness outputs are decoded with the runner-side wire + * coder; this processor reproduces it for the source edge by transcoding each element through the + * SDK-side wire coder (encode) and back through the runner-side wire coder (decode). The two are + * byte-compatible by construction, so the transcode yields exactly the object the receiver expects, + * nesting and all. + * + *

This mirrors {@link ImpulseProcessor}: a persistent state store records whether the elements + * have already been emitted so task restarts do not duplicate them, while the terminal watermark is + * re-emitted on every restart so downstream watermark holds still release after recovery. The + * trigger is a wall-clock punctuator scheduled on {@link #init} so the processor fires even though + * its bootstrap source topic is empty. + * + *

The source is read in a single instance with no splitting — parallelism across the source's + * splits arrives with the topic-based shuffle work (#18479). Kafka Streams disallows negative + * record timestamps, so each forwarded {@link Record} carries the Unix epoch ({@code 0L}); the Beam + * event time lives inside the {@link WindowedValue}. + */ +class ReadProcessor implements Processor> { + + private static final Logger LOG = LoggerFactory.getLogger(ReadProcessor.class); + + /** Sole entry in the state store; the value tracks whether this processor has already emitted. */ + static final String FIRED_KEY = "fired"; + + /** How soon after {@link #init} the punctuator first fires. */ + private static final Duration PUNCTUATION_DELAY = Duration.ofMillis(50); + + private final BoundedSource source; + // Held as SerializablePipelineOptions (Beam's idiom for a reader holding options) so the + // processor's captured state is uniformly serializable alongside the BoundedSource and coders. + private final SerializablePipelineOptions options; + // Encodes a raw WindowedValue as the SDK harness would on the wire (length-prefixing the + // element coder if the runner does not know it); the runner-side coder then decodes it into the + // wire form the downstream stage's input receiver expects. See the class javadoc. + private final Coder> sdkWireCoder; + private final Coder> runnerWireCoder; + private final String stateStoreName; + private final String transformId; + + private @Nullable ProcessorContext> context; + private @Nullable KeyValueStore firedStore; + private @Nullable Cancellable scheduledPunctuator; + + ReadProcessor( + BoundedSource source, + SerializablePipelineOptions options, + Coder> sdkWireCoder, + Coder> runnerWireCoder, + String stateStoreName, + String transformId) { + this.source = source; + this.options = options; + this.sdkWireCoder = sdkWireCoder; + this.runnerWireCoder = runnerWireCoder; + this.stateStoreName = stateStoreName; + this.transformId = transformId; + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + this.firedStore = context.getStateStore(stateStoreName); + this.scheduledPunctuator = + context.schedule(PUNCTUATION_DELAY, PunctuationType.WALL_CLOCK_TIME, ts -> maybeFire()); + } + + @Override + public void process(Record record) { + // Records that happen to land on the bootstrap topic are not actual data; they just provide an + // extra opportunity to fire the read on restart. The state store still gates the emit. + maybeFire(); + } + + private void maybeFire() { + ProcessorContext> ctx = context; + KeyValueStore store = firedStore; + if (ctx == null || store == null) { + return; + } + if (Boolean.TRUE.equals(store.get(FIRED_KEY))) { + // Elements were already emitted in a previous task lifetime, but downstream watermark holds + // may still need to release after the restart — re-emit the terminal watermark and stop. + forwardWatermarkMax(ctx); + cancelPunctuator(); + return; + } + int count = readAndForward(ctx); + forwardWatermarkMax(ctx); + store.put(FIRED_KEY, Boolean.TRUE); + cancelPunctuator(); + LOG.debug("Read {} emitted {} elements and terminal watermark", transformId, count); + } + + /** Reads the whole bounded source and forwards each element, in wire form, as a data payload. */ + private int readAndForward(ProcessorContext> ctx) { + int count = 0; + try (BoundedReader reader = source.createReader(options.get())) { + for (boolean hasElement = reader.start(); hasElement; hasElement = reader.advance()) { + WindowedValue element = + WindowedValues.timestampedValueInGlobalWindow( + reader.getCurrent(), reader.getCurrentTimestamp()); + // The Read output PCollection is not keyed; use an empty byte[] as a placeholder key so + // downstream processors that adopt the byte[]-key convention see a consistent shape. + ctx.forward( + new Record>( + new byte[0], KStreamsPayload.data(toRunnerWire(element)), 0L)); + count++; + } + } catch (IOException e) { + throw new RuntimeException("Failed to read bounded source for transform " + transformId, e); + } + return count; + } + + /** Transcodes a raw element into the runner-side wire form the SDK harness input expects. */ + private WindowedValue toRunnerWire(WindowedValue element) { + try { + byte[] wireBytes = CoderUtils.encodeToByteArray(sdkWireCoder, element); + return CoderUtils.decodeFromByteArray(runnerWireCoder, wireBytes); + } catch (CoderException e) { + throw new RuntimeException( + "Failed to transcode a read element to wire form for transform " + transformId, e); + } + } + + /** + * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. + * + *

Read is a single-instance source, so the report is stamped as the only source partition: + * {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real per-partition identities + * arrive once the topology gains topic-based shuffle. + */ + private static void forwardWatermarkMax(ProcessorContext> ctx) { + long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); + ctx.forward( + new Record>( + new byte[0], KStreamsPayload.watermark(maxMillis, 0, 1), 0L)); + } + + /** Cancels the wall-clock punctuator after the read has fired to stop periodic wakeups. */ + private void cancelPunctuator() { + Cancellable handle = scheduledPunctuator; + if (handle != null) { + handle.cancel(); + scheduledPunctuator = null; + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java new file mode 100644 index 000000000000..403499ca616c --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java @@ -0,0 +1,173 @@ +/* + * 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 java.io.IOException; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.WireCoderSetting; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.runners.fnexecution.wire.WireCoders; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.BoundedSource; +import org.apache.beam.sdk.util.construction.ReadTranslation; +import org.apache.beam.sdk.util.construction.RehydratedComponents; +import org.apache.beam.sdk.util.construction.graph.PipelineNode; +import org.apache.beam.sdk.util.construction.graph.PipelineNode.PCollectionNode; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.state.KeyValueBytesStoreSupplier; +import org.apache.kafka.streams.state.Stores; + +/** + * Translates the deprecated primitive {@code Read} URN ({@code beam:transform:read:v1}) over a + * {@link BoundedSource}. + * + *

The runner forces every {@code Read.Bounded} (including the one {@code Create} of two or more + * elements expands to) into this primitive read before translation — see {@code + * KafkaStreamsTestRunner.translate}, which applies {@code + * SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads}. This deliberately avoids the + * default {@code BoundedSourceAsSDFWrapperFn} splittable-DoFn expansion, which the runner cannot + * execute yet (no SDF restriction protocol), as agreed with the mentor. + * + *

Adds the same three-node shape as {@link ImpulseTranslator}: + * + *

    + *
  • A {@code byte[]} source bound to a dedicated per-transform bootstrap topic (see {@link + * KafkaStreamsTranslationContext#getReadBootstrapTopic(String)}). Kafka Streams refuses to + * start a topology with no real source topic; records published to it are ignored by {@link + * ReadProcessor}. + *
  • The {@link ReadProcessor}, which reads the {@link BoundedSource} on a one-shot wall-clock + * punctuator and emits one data payload per element followed by a terminal watermark. + *
  • A per-processor persistent state store recording whether the read has already fired so task + * restarts do not duplicate elements. + *
+ * + *

The processor emits elements in the runner-side wire form the downstream stage's SDK harness + * expects, so it is handed the SDK-side and runner-side wire coders for the read's output + * PCollection (see {@link ReadProcessor} for why). Only {@link + * org.apache.beam.model.pipeline.v1.RunnerApi.IsBounded.Enum#BOUNDED bounded} sources are + * supported; {@link ReadTranslation#boundedSourceFromProto} rejects an unbounded payload. + */ +class ReadTranslator implements PTransformTranslator { + + static final String SOURCE_SUFFIX = "-source"; + static final String STATE_STORE_SUFFIX = "-state"; + + @Override + public void translate( + String transformId, RunnerApi.Pipeline pipeline, KafkaStreamsTranslationContext context) { + RunnerApi.PTransform transform = pipeline.getComponents().getTransformsOrThrow(transformId); + // Read produces exactly one output PCollection; downstream consumers are separate PTransforms + // whose inputs reference this PCollection id and are wired by their own translators. + String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + addReadNodes( + transformId, + boundedSource(transform), + pipeline.getComponents(), + outputPCollectionId, + context); + } + + private static BoundedSource boundedSource(RunnerApi.PTransform transform) { + try { + RunnerApi.ReadPayload payload = + RunnerApi.ReadPayload.parseFrom(transform.getSpec().getPayload()); + return ReadTranslation.boundedSourceFromProto(payload); + } catch (IOException e) { + throw new RuntimeException( + "Failed to read the BoundedSource from transform " + transform.getUniqueName(), e); + } + } + + /** + * Adds the source, {@link ReadProcessor}, and state store for the read. The type variable {@code + * T} captures the {@link BoundedSource}'s element type so the processor and its wire coders are + * built consistently. + */ + private void addReadNodes( + String transformId, + BoundedSource source, + RunnerApi.Components components, + String outputPCollectionId, + KafkaStreamsTranslationContext context) { + PCollectionNode outputNode = + PipelineNode.pCollection( + outputPCollectionId, components.getPcollectionsOrThrow(outputPCollectionId)); + Coder> sdkWireCoder = sdkWireCoder(outputNode, components); + Coder> runnerWireCoder = runnerWireCoder(outputNode, components); + + Topology topology = context.getTopology(); + String sourceNodeName = transformId + SOURCE_SUFFIX; + String stateStoreName = transformId + STATE_STORE_SUFFIX; + String bootstrapTopic = context.getReadBootstrapTopic(transformId); + SerializablePipelineOptions options = + new SerializablePipelineOptions(context.getPipelineOptions()); + + topology.addSource( + sourceNodeName, + Serdes.ByteArray().deserializer(), + Serdes.ByteArray().deserializer(), + bootstrapTopic); + topology.addProcessor( + transformId, + () -> + new ReadProcessor<>( + source, options, sdkWireCoder, runnerWireCoder, stateStoreName, transformId), + sourceNodeName); + KeyValueBytesStoreSupplier storeSupplier = Stores.persistentKeyValueStore(stateStoreName); + topology.addStateStore( + Stores.keyValueStoreBuilder(storeSupplier, Serdes.String(), Serdes.Boolean()), transformId); + + context.registerPCollectionProducer(outputPCollectionId, transformId); + } + + /** The coder the SDK harness would use on the wire, keeping unknown element coders intact. */ + private static Coder> sdkWireCoder( + PCollectionNode outputNode, RunnerApi.Components components) { + try { + RunnerApi.Components.Builder builder = components.toBuilder(); + String coderId = + WireCoders.addSdkWireCoder(outputNode, builder, WireCoderSetting.getDefaultInstance()); + @SuppressWarnings("unchecked") + Coder> coder = + (Coder>) + RehydratedComponents.forComponents(builder.build()).getCoder(coderId); + return coder; + } catch (IOException e) { + throw new RuntimeException( + "Failed to build the SDK wire coder for PCollection " + outputNode.getId(), e); + } + } + + /** The coder the runner uses on the wire, replacing unknown element coders with byte arrays. */ + private static Coder> runnerWireCoder( + PCollectionNode outputNode, RunnerApi.Components components) { + try { + @SuppressWarnings("unchecked") + Coder> coder = + (Coder>) + (Coder) WireCoders.instantiateRunnerWireCoder(outputNode, components); + return coder; + } catch (IOException e) { + throw new RuntimeException( + "Failed to build the runner wire coder for PCollection " + outputNode.getId(), e); + } + } +} 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 53b7714548a0..4d73eee9581a 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 @@ -36,6 +36,7 @@ import org.apache.beam.sdk.util.construction.Environments; import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.beam.sdk.util.construction.SplittableParDo; import org.apache.kafka.common.serialization.ByteArrayDeserializer; import org.apache.kafka.common.serialization.ByteArraySerializer; import org.apache.kafka.common.serialization.Serdes; @@ -88,6 +89,11 @@ public static PipelineOptions testOptions() { public static KafkaStreamsTranslationContext translate(Pipeline pipeline) { KafkaStreamsPipelineOptions options = pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + // Force every Read.Bounded (including the one Create of 2+ elements expands to) into the + // deprecated primitive Read the runner translates, instead of the default + // BoundedSourceAsSDFWrapperFn splittable-DoFn expansion the runner cannot execute yet. This is + // unconditional — it does not depend on the use_deprecated_read experiment (mentor's steer). + SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(pipeline); RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); JobInfo jobInfo = diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/CreateTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/CreateTest.java new file mode 100644 index 000000000000..e2fab243d4bf --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/CreateTest.java @@ -0,0 +1,73 @@ +/* + * 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.hasItems; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.util.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.Test; + +/** + * End-to-end test that {@code Create} of two or more elements runs on the Kafka Streams runner. + * + *

{@code Create.of(1, 2, 3)} expands to {@code Read.from(CreateSource)}, a bounded {@code + * Read.Bounded}. {@code KafkaStreamsTestRunner.translate} forces it into the deprecated primitive + * read the runner translates (see {@link ReadTranslator}), so the elements flow through {@link + * ReadProcessor} into a recording ParDo in the EMBEDDED harness. This is the multi-element Create + * case that previously failed because the default expansion is an unsupported splittable DoFn. + */ +public class CreateTest { + + /** Records every element the harness feeds it so the test can assert Create produced them. */ + private static class RecordingFn extends DoFn { + private final SharedTestCollector collector; + + RecordingFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver out) { + collector.record(input); + out.output(input); + } + } + + @Test + public void createOfMultipleElementsRunsThroughPrimitiveRead() { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + pipeline + .apply("create", Create.of(1, 2, 3)) + .apply("record", ParDo.of(new RecordingFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List recorded = collector.recorded(); + assertThat(recorded.size(), is(3)); + assertThat(recorded, hasItems(1, 2, 3)); + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ReadTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ReadTest.java new file mode 100644 index 000000000000..5b50aa272145 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ReadTest.java @@ -0,0 +1,75 @@ +/* + * 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.hasItems; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.util.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.io.CountingSource; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.Test; + +/** + * End-to-end test for {@link ReadTranslator}: a primitive bounded {@code Read} over a {@link + * CountingSource} feeds a recording ParDo executed in the in-process (EMBEDDED) Java SDK harness. + * Verifies every source element reaches the DoFn. + * + *

{@code KafkaStreamsTestRunner.translate} forces the {@code Read.Bounded} into the deprecated + * primitive read the runner translates, rather than the default splittable-DoFn expansion — so this + * exercises the {@link ReadProcessor} + {@link ReadTranslator} path end to end. + */ +public class ReadTest { + + /** Records every element the harness feeds it so the test can assert the read produced them. */ + private static class RecordingFn extends DoFn { + private final SharedTestCollector collector; + + RecordingFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element Long input, OutputReceiver out) { + collector.record(input); + out.output(input); + } + } + + @Test + public void boundedReadEmitsEverySourceElement() { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + pipeline + .apply("read", Read.from(CountingSource.upTo(5))) + .apply("record", ParDo.of(new RecordingFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List recorded = collector.recorded(); + // CountingSource.upTo(5) yields 0..4 exactly once each. + assertThat(recorded.size(), is(5)); + assertThat(recorded, hasItems(0L, 1L, 2L, 3L, 4L)); + } + } +} From ff323ebd6dbcb47fa1dd11c300aad721527bd0b8 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Wed, 15 Jul 2026 15:19:13 +0500 Subject: [PATCH 16/37] [GSoC 2026] Kafka Streams runner #39273: Kafka Streams runner: Flatten support --- .../main/proto/kafka_streams_payload.proto | 14 +- .../translation/ExecutableStageProcessor.java | 66 ++++-- .../ExecutableStageTranslator.java | 8 +- .../streams/translation/FlattenProcessor.java | 119 ++++++++++ .../translation/FlattenTranslator.java | 82 +++++++ .../translation/GroupByKeyProcessor.java | 56 +++-- .../translation/GroupByKeyTranslator.java | 13 +- .../streams/translation/ImpulseProcessor.java | 13 +- .../streams/translation/KStreamsPayload.java | 27 ++- .../translation/KStreamsPayloadSerde.java | 2 + .../KafkaStreamsPipelineTranslator.java | 1 + .../streams/translation/ReadProcessor.java | 13 +- .../translation/ShuffleByKeyProcessor.java | 10 + .../translation/WatermarkAggregator.java | 107 +++++++++ .../streams/translation/WatermarkPayload.java | 18 +- ...ExecutableStageProcessorWatermarkTest.java | 21 +- .../streams/translation/FlattenTest.java | 212 ++++++++++++++++++ .../translation/KStreamsPayloadSerdeTest.java | 7 +- .../translation/WatermarkAggregatorTest.java | 160 +++++++++++++ 19 files changed, 873 insertions(+), 76 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java diff --git a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto index 87e3de44397b..4dadf2d14b70 100644 --- a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto +++ b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto @@ -27,16 +27,22 @@ option java_outer_classname = "KafkaStreamsPayloadProtos"; // topic boundary (e.g. the GroupByKey repartition topic and the watermark fan-out). Protobuf is // used for compatible schema evolution and compact varint encoding. message KafkaStreamsPayload { - // A watermark report: the watermark plus the in-band coordination fields the downstream - // WatermarkManager needs. + // A watermark report: the watermark plus the in-band coordination fields a downstream + // watermark aggregator needs to reconstruct its input watermark. message WatermarkPayload { // Event-time watermark in milliseconds. Signed (sint64, zigzag-encoded) because Beam event // times can be negative, e.g. BoundedWindow.TIMESTAMP_MIN_VALUE. sint64 millis = 1; - // The source partition this report is for. + // Which partition (physical instance) of the producing transform this report is for, in + // [0, total_partitions). uint32 source_partition = 2; - // The total number of source partitions feeding the downstream stage. + // How many partitions (physical instances) the producing transform has in total. uint32 total_partitions = 3; + // Globally unique id of the transform that produced this report. A producer stamps its own id + // without regard to who consumes the report; a consumer with several upstream transforms + // (e.g. Flatten) aggregates per producing transform, holding its output watermark until every + // partition of every upstream transform it expects has reported. + string transform_id = 4; } // A data element: the Beam WindowedValue encoded with the PCollection's windowed-value coder. 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 00f85032d041..38d3601e814f 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 @@ -18,6 +18,7 @@ package org.apache.beam.runners.kafka.streams.translation; import java.util.Queue; +import java.util.Set; import java.util.concurrent.ConcurrentLinkedQueue; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.fnexecution.control.BundleProgressHandler; @@ -51,12 +52,12 @@ * ProcessorContext#forward} must only be called from the processing thread, so outputs are never * forwarded directly from a harness callback. * - *

A {@link KStreamsPayload#isWatermark() watermark} payload is a per-source-partition report and - * marks a bundle boundary: the open bundle (if any) is closed (flushing outputs), the report is fed - * to the {@link WatermarkManager}, and the stage's output watermark is forwarded downstream only - * when the {@code min()} across its source partitions actually advances. Until every source - * partition has reported, the watermark is held and nothing is forwarded — but data is still - * processed in the meantime. + *

A {@link KStreamsPayload#isWatermark() watermark} payload is a report from one partition of + * one upstream transform and marks a bundle boundary: the open bundle (if any) is closed (flushing + * outputs), the report is fed to the {@link WatermarkAggregator}, and the stage's output watermark + * is forwarded downstream — stamped with this stage's own transform id — only when the aggregate + * across the upstream transform's partitions actually advances. Until every partition has reported, + * the watermark is held and nothing is forwarded — but data is still processed in the meantime. * *

This is the Kafka Streams analogue of Flink's {@code ExecutableStageDoFnOperator} and Spark's * {@code SparkExecutableStageFunction}. State, timers, and side inputs are out of scope for this @@ -70,6 +71,9 @@ class ExecutableStageProcessor private final RunnerApi.ExecutableStagePayload stagePayload; private final JobInfo jobInfo; + // 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; // 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. @@ -79,9 +83,9 @@ class ExecutableStageProcessor // only safe because the Impulse output coder happens to be ByteArrayCoder. private final Queue> pendingOutputs = new ConcurrentLinkedQueue<>(); - // Computes this stage's output watermark as min() over its source partitions' reported - // watermarks, holding until every source partition has reported (see WatermarkManager). - private final WatermarkManager watermarkManager = new WatermarkManager(); + // Computes this stage's input watermark from its upstream transform's reports, holding until + // every partition of the upstream transform has reported (see WatermarkAggregator). + private final WatermarkAggregator watermarkAggregator; // The last watermark actually forwarded downstream, so we only forward when it advances. private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; @@ -90,9 +94,20 @@ class ExecutableStageProcessor private @Nullable StageBundleFactory stageBundleFactory; private @Nullable RemoteBundle currentBundle; - ExecutableStageProcessor(RunnerApi.ExecutableStagePayload stagePayload, JobInfo jobInfo) { + /** + * @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 + */ + ExecutableStageProcessor( + RunnerApi.ExecutableStagePayload stagePayload, + JobInfo jobInfo, + String transformId, + Set upstreamTransformIds) { this.stagePayload = stagePayload; this.jobInfo = jobInfo; + this.transformId = transformId; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); } @Override @@ -116,18 +131,21 @@ private void ensureStageBundleFactory() { @Override public void process(Record> record) { KStreamsPayload payload = record.value(); + if (payload == null) { + // A topic feeding the runner can always be written to from outside (or carry a tombstone), + // so recover from the obvious error instead of crashing the task: warn and drop. + LOG.warn( + "Stage {} dropping record with null payload (external write or tombstone)", transformId); + return; + } if (payload.isWatermark()) { // Emit any buffered outputs before the watermark. Data is processed regardless of watermark // readiness; only the watermark itself is held until every source partition has reported. closeBundleAndFlush(record); - // Feed the report into the WatermarkManager and forward the stage's output watermark only - // when min() across the source partitions actually advances, not on every received watermark. - WatermarkPayload report = payload.asWatermark(); - watermarkManager.observe( - report.getSourcePartition(), - new Instant(report.getWatermarkMillis()), - report.getTotalSourcePartitions()); - Instant advanced = watermarkManager.advance(); + // Feed the report into the aggregator and forward the stage's output watermark only when the + // aggregate across the upstream transform's partitions actually advances. + watermarkAggregator.observe(payload.asWatermark()); + Instant advanced = watermarkAggregator.advance(); if (advanced.isAfter(lastForwardedWatermark)) { lastForwardedWatermark = advanced; forwardWatermark(record, advanced.getMillis()); @@ -203,14 +221,16 @@ private void closeBundleAndFlush(Record> record) { } private void forwardWatermark(Record> record, long watermarkMillis) { - // This stage is a single instance for now, so it forwards its watermark as the only source - // partition (0 of 1). Fanning the watermark out to every downstream partition — and producing - // it atomically with the offset commit so it is durable — lands with the topic-based shuffle - // work, when there are real source partitions to track (#18479). + // Stamped with this stage's own transform id; this stage is a single instance for now, so the + // report is for its only partition (0 of 1). Fanning the watermark out to every downstream + // partition — and producing it atomically with the offset commit so it is durable — lands with + // the topic-based shuffle work (#18479). ProcessorContext> ctx = checkInitialized(context); ctx.forward( new Record>( - record.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1), record.timestamp())); + record.key(), + KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), + record.timestamp())); } @Override 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 dc56d57f57c2..a2e6ed837c7d 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 @@ -19,6 +19,7 @@ import java.io.IOException; import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.kafka.streams.Topology; @@ -83,9 +84,14 @@ public void translate( String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); 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). topology.addProcessor( transformId, - () -> new ExecutableStageProcessor(stagePayload, context.getJobInfo()), + () -> + new ExecutableStageProcessor( + stagePayload, context.getJobInfo(), transformId, ImmutableSet.of(parentProcessor)), parentProcessor); if (!transform.getOutputsMap().isEmpty()) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java new file mode 100644 index 000000000000..d09fed8185a1 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java @@ -0,0 +1,119 @@ +/* + * 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 java.util.Set; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Kafka Streams {@link Processor} implementing Beam's {@code Flatten} primitive ({@code + * beam:transform:flatten:v1}): the union of N input PCollections into one output PCollection. + * + *

Data records are forwarded straight through unchanged — the merge of the N parents' + * data streams is the flatten. + * + *

Watermark reports are where Flatten does real work, and it owns its output watermark + * the same way GroupByKey does: it runs a {@link WatermarkAggregator} over its inputs, forwards its + * own watermark only when the {@code min()} across them advances, and stamps that as a single + * source ({@code 0 of 1}) to its downstream. This holds the output watermark back until + * every input branch has reported, so a downstream GroupByKey does not fire before all + * flattened branches are drained. + * + *

The {@link WatermarkAggregator} tells the input branches apart by the transform id each + * branch's producer stamps on its watermark (Kafka Streams does not tell a processor which parent + * forwarded a record). Each producer stamps its own identity regardless of who consumes it, so a + * PCollection feeding several Flattens reports one identity and every Flatten still waits only for + * the upstream transforms it expects — the set handed to it at construction from the pipeline + * graph. + */ +class FlattenProcessor + implements Processor, byte[], KStreamsPayload> { + + private static final Logger LOG = LoggerFactory.getLogger(FlattenProcessor.class); + + // This transform's own id, stamped on every watermark it forwards downstream. + private final String transformId; + // Computes the output watermark as min() over the upstream transforms' reports, holding until + // every partition of every expected upstream transform has reported (see WatermarkAggregator). + private final WatermarkAggregator watermarkAggregator; + // The last watermark actually forwarded downstream, so we only forward when it advances. + private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + + private @Nullable ProcessorContext> context; + + /** + * @param transformId this Flatten's own transform id, stamped on the watermarks it emits + * @param upstreamTransformIds the producers of this Flatten's input PCollections (known from the + * pipeline graph), whose reports the {@link WatermarkAggregator} waits for + */ + FlattenProcessor(String transformId, Set upstreamTransformIds) { + this.transformId = transformId; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + } + + @Override + public void process(Record> record) { + KStreamsPayload payload = record.value(); + if (payload == null) { + // A topic feeding the runner can always be written to from outside (or carry a tombstone), + // so recover from the obvious error instead of crashing the task: warn and drop. + LOG.warn( + "Flatten {} dropping record with null payload (external write or tombstone)", + transformId); + return; + } + ProcessorContext> ctx = checkInitialized(context); + if (!payload.isWatermark()) { + // Data: the union of the parents' data streams is the flatten — forward unchanged. + ctx.forward(record); + return; + } + watermarkAggregator.observe(payload.asWatermark()); + Instant advanced = watermarkAggregator.advance(); + if (advanced.isAfter(lastForwardedWatermark)) { + lastForwardedWatermark = advanced; + // Stamped with this Flatten's own transform id; Flatten is a single instance for now, so the + // report is for its only partition (0 of 1). + ctx.forward( + new Record>( + record.key(), + KStreamsPayload.watermark(advanced.getMillis(), transformId, 0, 1), + record.timestamp())); + } + } + + private static ProcessorContext> checkInitialized( + @Nullable ProcessorContext> context) { + if (context == null) { + throw new IllegalStateException("FlattenProcessor used before init()"); + } + return context; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java new file mode 100644 index 000000000000..8f6ce2b546dd --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java @@ -0,0 +1,82 @@ +/* + * 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 java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.kafka.streams.Topology; + +/** + * Translates Beam's {@code Flatten} primitive ({@code beam:transform:flatten:v1}): the union of N + * input PCollections into one output PCollection. + * + *

Wires a single {@link FlattenProcessor} node to the producer of every input PCollection (Kafka + * Streams lets a processor have many parents), so the parents' data streams merge into it, and + * registers it as the producer of the flattened output so downstream translators wire to it. The + * processor forwards data through and owns its output watermark via a {@link WatermarkAggregator}, + * which is handed the producers of the input PCollections — the upstream transform ids whose + * watermark reports the Flatten must hear from. Producers stamp their own transform id on the + * reports they emit, without regard to who consumes them, so an input shared with another Flatten + * needs no special handling. + * + *

A user-written self-flatten never reaches this translator: the fuser folds the Flatten into + * the consuming SDK-harness stage, which performs the duplication itself. The duplicate-input check + * below is defensive — if a runner-executed Flatten ever did receive the same PCollection twice, + * Kafka Streams could not wire the same parent to a child twice and the duplicate copy would be + * silently dropped, so failing fast is safer. + */ +class FlattenTranslator implements PTransformTranslator { + + @Override + public void translate( + String transformId, RunnerApi.Pipeline pipeline, KafkaStreamsTranslationContext context) { + RunnerApi.PTransform transform = pipeline.getComponents().getTransformsOrThrow(transformId); + // Flatten produces exactly one output PCollection, fed by all of its input PCollections. + String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + + Set seenInputs = new HashSet<>(); + List parentProcessors = new ArrayList<>(); + Set upstreamTransformIds = new HashSet<>(); + for (String inputPCollectionId : transform.getInputsMap().values()) { + if (!seenInputs.add(inputPCollectionId)) { + throw new UnsupportedOperationException( + "Flatten " + + transform.getUniqueName() + + " has PCollection " + + inputPCollectionId + + " as an input more than once; a self-flatten is not yet supported by the Kafka" + + " Streams runner."); + } + String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); + parentProcessors.add(parentProcessor); + upstreamTransformIds.add(parentProcessor); + } + + Topology topology = context.getTopology(); + topology.addProcessor( + transformId, + () -> new FlattenProcessor(transformId, upstreamTransformIds), + parentProcessors.toArray(new String[0])); + + context.registerPCollectionProducer(outputPCollectionId, transformId); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java index a6c0cb8c4061..3e82935b807d 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java @@ -19,6 +19,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Set; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; import org.apache.beam.sdk.coders.IterableCoder; @@ -35,16 +36,18 @@ import org.apache.kafka.streams.state.KeyValueStore; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Executes a {@code GroupByKey} (GlobalWindow, default trigger, no allowed lateness). * *

Records arrive on the repartition topic keyed by the encoded Beam key, so every value of a key * is co-located here. Each value is appended to a per-key buffer in a Kafka Streams state store. - * Watermark reports are fed to a {@link WatermarkManager}; when the input watermark reaches {@link - * BoundedWindow#TIMESTAMP_MAX_VALUE} (the end of the global window) every buffered key is emitted - * once as {@code KV>} and the buffer cleared, then the watermark is forwarded - * downstream. + * Watermark reports are fed to a {@link WatermarkAggregator}; when the input watermark reaches + * {@link BoundedWindow#TIMESTAMP_MAX_VALUE} (the end of the global window) every buffered key is + * emitted once as {@code KV>} and the buffer cleared, then the watermark is + * forwarded downstream. * *

Buffering whole value lists and re-encoding on each append is O(n^2) per key; fine for this * first GroupByKey, and replaced when this moves to runner-core {@code GroupAlsoByWindow}. @@ -52,11 +55,18 @@ class GroupByKeyProcessor implements Processor, byte[], KStreamsPayload> { + private static final Logger LOG = LoggerFactory.getLogger(GroupByKeyProcessor.class); + private final String stateStoreName; + // This transform's own id, stamped on every watermark it forwards downstream. + private final String transformId; private final Coder keyCoder; private final IterableCoder<@Nullable Object> bufferCoder; - private final WatermarkManager watermarkManager = new WatermarkManager(); + // Aggregates the input watermark from the upstream transform's reports, which arrive through the + // repartition topic with the upstream producer's transform id intact (the shuffle forwards + // watermark payloads unchanged). + private final WatermarkAggregator watermarkAggregator; private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; // The global window fires exactly once, when the watermark first reaches its end. Later watermark // reports (e.g. the same terminal watermark broadcast across repartition partitions) must not @@ -69,9 +79,20 @@ class GroupByKeyProcessor private @Nullable ProcessorContext> context; private @Nullable KeyValueStore store; + /** + * @param transformId this transform's own id, stamped on the watermarks it emits + * @param upstreamTransformIds the transform ids feeding this GroupByKey (known from the pipeline + * graph), whose reports the {@link WatermarkAggregator} waits for + */ GroupByKeyProcessor( - String stateStoreName, Coder keyCoder, Coder<@Nullable Object> valueCoder) { + String stateStoreName, + String transformId, + Set upstreamTransformIds, + Coder keyCoder, + Coder<@Nullable Object> valueCoder) { this.stateStoreName = stateStoreName; + this.transformId = transformId; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); this.keyCoder = keyCoder; this.bufferCoder = IterableCoder.of(valueCoder); } @@ -85,6 +106,14 @@ public void init(ProcessorContext> context) { @Override public void process(Record> record) { KStreamsPayload payload = record.value(); + if (payload == null) { + // The repartition topic can be written to from outside the runner (or carry a tombstone), + // so recover from the obvious error instead of crashing the task: warn and drop. + LOG.warn( + "GroupByKey {} dropping record with null payload (external write or tombstone)", + transformId); + return; + } if (payload.isData()) { byte[] encodedKey = record.key(); Object element = payload.getData().getValue(); @@ -94,12 +123,8 @@ public void process(Record> record) { appendValue(encodedKey, element); return; } - WatermarkPayload report = payload.asWatermark(); - watermarkManager.observe( - report.getSourcePartition(), - new Instant(report.getWatermarkMillis()), - report.getTotalSourcePartitions()); - Instant advanced = watermarkManager.advance(); + watermarkAggregator.observe(payload.asWatermark()); + Instant advanced = watermarkAggregator.advance(); if (!fired && !advanced.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) { fireAll(record); fired = true; @@ -153,10 +178,13 @@ private void fireAll(Record> trigger) { private void forwardWatermark(Record> trigger, long watermarkMillis) { ProcessorContext> ctx = checkInitialized(context); - // GroupByKey is a single logical source for the next stage; report it as partition 0 of 1. + // Stamped with this transform's own id; GroupByKey is a single instance for now, so the report + // is for its only partition (0 of 1). ctx.forward( new Record>( - trigger.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1), trigger.timestamp())); + trigger.key(), + KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), + trigger.timestamp())); } private byte[] encodeBuffer(List<@Nullable Object> values) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index d7c4a309d9c3..9e23dbb5cfb0 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.streams.Topology; @@ -110,10 +111,18 @@ public void translate( payloadSerde.deserializer(), repartitionTopic); - // Buffer values per key and fire KV> at the terminal watermark. + // Buffer values per key and fire KV> at the terminal watermark. Watermark + // reports cross the repartition topic unchanged, so they still carry the id of the transform + // that produced this GroupByKey's input — the parent the shuffle is attached to. topology.addProcessor( transformId, - () -> new GroupByKeyProcessor(stateStoreName, keyCoder, valueCoder), + () -> + new GroupByKeyProcessor( + stateStoreName, + transformId, + ImmutableSet.of(parentProcessor), + keyCoder, + valueCoder), sourceName); topology.addStateStore( Stores.keyValueStoreBuilder( diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java index 675bee4d8591..bac91978a298 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java @@ -122,17 +122,16 @@ private void maybeFire() { } /** - * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. - * - *

Impulse is a single-instance source, so the report is stamped as the only source partition: - * {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real per-partition identities - * arrive once the topology gains topic-based shuffle. + * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors, + * stamped with this transform's id. Impulse is a single-instance source, so the report is for its + * only partition: {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real + * per-partition identities arrive once the topology gains topic-based shuffle. */ - private static void forwardWatermarkMax(ProcessorContext> ctx) { + private void forwardWatermarkMax(ProcessorContext> ctx) { long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); ctx.forward( new Record>( - new byte[0], KStreamsPayload.watermark(maxMillis, 0, 1), 0L)); + new byte[0], KStreamsPayload.watermark(maxMillis, transformId, 0, 1), 0L)); } /** Cancels the wall-clock punctuator after the impulse has fired to stop periodic wakeups. */ diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java index 93b40346761a..c165f0e875d4 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java @@ -57,6 +57,7 @@ private enum Kind { private final Kind kind; private final @Nullable WindowedValue data; private final long watermarkMillis; + private final String transformId; private final int sourcePartition; private final int totalSourcePartitions; @@ -64,27 +65,33 @@ private KStreamsPayload( Kind kind, @Nullable WindowedValue data, long watermarkMillis, + String transformId, int sourcePartition, int totalSourcePartitions) { this.kind = kind; this.data = data; this.watermarkMillis = watermarkMillis; + this.transformId = transformId; this.sourcePartition = sourcePartition; this.totalSourcePartitions = totalSourcePartitions; } /** Returns a data payload wrapping the given {@link WindowedValue}. */ public static KStreamsPayload data(WindowedValue value) { - return new KStreamsPayload<>(Kind.DATA, value, 0L, 0, 0); + return new KStreamsPayload<>(Kind.DATA, value, 0L, "", 0, 0); } /** * Returns a watermark report payload: the event-time milliseconds together with the in-band - * coordination fields the downstream stage's {@link WatermarkManager} needs — which source - * partition this report is for and how many source partitions feed the stage in total. + * coordination fields a downstream watermark aggregator needs — which transform produced the + * report ({@code transformId}, stamped by the producer without regard to who consumes it), which + * of that transform's partitions this report is for, and how many partitions that transform has + * in total. */ public static KStreamsPayload watermark( - long watermarkMillis, int sourcePartition, int totalSourcePartitions) { + long watermarkMillis, String transformId, int sourcePartition, int totalSourcePartitions) { + Preconditions.checkArgument( + transformId != null && !transformId.isEmpty(), "transformId must be non-empty"); Preconditions.checkArgument( totalSourcePartitions > 0, "totalSourcePartitions must be positive: %s", @@ -95,7 +102,7 @@ public static KStreamsPayload watermark( sourcePartition, totalSourcePartitions); return new KStreamsPayload<>( - Kind.WATERMARK, null, watermarkMillis, sourcePartition, totalSourcePartitions); + Kind.WATERMARK, null, watermarkMillis, transformId, sourcePartition, totalSourcePartitions); } public boolean isData() { @@ -134,6 +141,11 @@ public long getWatermarkMillis() { return watermarkMillis; } + @Override + public String getTransformId() { + return transformId; + } + @Override public int getSourcePartition() { return sourcePartition; @@ -156,6 +168,7 @@ public boolean equals(@Nullable Object o) { KStreamsPayload that = (KStreamsPayload) o; return kind == that.kind && watermarkMillis == that.watermarkMillis + && transformId.equals(that.transformId) && sourcePartition == that.sourcePartition && totalSourcePartitions == that.totalSourcePartitions && Objects.equals(data, that.data); @@ -163,7 +176,8 @@ public boolean equals(@Nullable Object o) { @Override public int hashCode() { - return Objects.hash(kind, data, watermarkMillis, sourcePartition, totalSourcePartitions); + return Objects.hash( + kind, data, watermarkMillis, transformId, sourcePartition, totalSourcePartitions); } @Override @@ -174,6 +188,7 @@ public String toString() { } else { helper .add("watermarkMillis", watermarkMillis) + .add("transformId", transformId) .add("sourcePartition", sourcePartition) .add("totalSourcePartitions", totalSourcePartitions); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java index 912eb5fb6047..1363740d58bd 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java @@ -82,6 +82,7 @@ public byte[] serialize(String topic, KStreamsPayload payload) { proto.setWatermark( KafkaStreamsPayload.WatermarkPayload.newBuilder() .setMillis(watermark.getWatermarkMillis()) + .setTransformId(watermark.getTransformId()) .setSourcePartition(watermark.getSourcePartition()) .setTotalPartitions(watermark.getTotalSourcePartitions())); } @@ -109,6 +110,7 @@ public KStreamsPayload deserialize(String topic, byte[] bytes) { KafkaStreamsPayload.WatermarkPayload watermark = proto.getWatermark(); return KStreamsPayload.watermark( watermark.getMillis(), + watermark.getTransformId(), watermark.getSourcePartition(), watermark.getTotalPartitions()); case PAYLOAD_NOT_SET: diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index 943673dfe98c..a9e26ecd6412 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -52,6 +52,7 @@ public KafkaStreamsPipelineTranslator() { .put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator()) .put(PTransformTranslation.READ_TRANSFORM_URN, new ReadTranslator()) .put(PTransformTranslation.REDISTRIBUTE_ARBITRARILY_URN, new RedistributeTranslator()) + .put(PTransformTranslation.FLATTEN_TRANSFORM_URN, new FlattenTranslator()) .put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new GroupByKeyTranslator()) .put(ExecutableStage.URN, new ExecutableStageTranslator()) .build()); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java index 604eca54014a..eb20a8b83588 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java @@ -183,17 +183,16 @@ private WindowedValue toRunnerWire(WindowedValue element) { } /** - * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. - * - *

Read is a single-instance source, so the report is stamped as the only source partition: - * {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real per-partition identities - * arrive once the topology gains topic-based shuffle. + * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors, + * stamped with this transform's id. Read is a single-instance source, so the report is for its + * only partition: {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real + * per-partition identities arrive once the topology gains topic-based shuffle. */ - private static void forwardWatermarkMax(ProcessorContext> ctx) { + private void forwardWatermarkMax(ProcessorContext> ctx) { long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); ctx.forward( new Record>( - new byte[0], KStreamsPayload.watermark(maxMillis, 0, 1), 0L)); + new byte[0], KStreamsPayload.watermark(maxMillis, transformId, 0, 1), 0L)); } /** Cancels the wall-clock punctuator after the read has fired to stop periodic wakeups. */ diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java index 3184d838e09d..774d36db0f2e 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java @@ -25,6 +25,8 @@ import org.apache.kafka.streams.processor.api.ProcessorContext; import org.apache.kafka.streams.processor.api.Record; import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Re-keys a {@code KV}-valued stream by the Beam key so Kafka Streams shuffles by it. @@ -39,6 +41,8 @@ class ShuffleByKeyProcessor implements Processor, byte[], KStreamsPayload> { + private static final Logger LOG = LoggerFactory.getLogger(ShuffleByKeyProcessor.class); + private final Coder keyCoder; private @Nullable ProcessorContext> context; @@ -55,6 +59,12 @@ public void init(ProcessorContext> context) { public void process(Record> record) { ProcessorContext> ctx = checkInitialized(context); KStreamsPayload payload = record.value(); + if (payload == null) { + // A topic feeding the runner can always be written to from outside (or carry a tombstone), + // so recover from the obvious error instead of crashing the task: warn and drop. + LOG.warn("Shuffle dropping record with null payload (external write or tombstone)"); + return; + } if (payload.isData()) { Object element = payload.getData().getValue(); if (element == null) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java new file mode 100644 index 000000000000..c9af03df26a9 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java @@ -0,0 +1,107 @@ +/* + * 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 java.util.HashMap; +import java.util.Map; +import java.util.Set; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.joda.time.Instant; + +/** + * Computes a transform's input watermark from the watermark reports of its upstream transforms. + * + *

A watermark report carries three orthogonal pieces of information (see {@link + * WatermarkPayload}): which transform produced it, which partition (physical + * instance) of that transform it is for, and how many partitions that transform has. A + * producer stamps its own identity without regard to who consumes the report. This aggregator is + * the consuming side, used by every transform that aggregates a watermark — ExecutableStage, + * GroupByKey, Flatten (and CombinePerKey later): + * + *

    + *
  • It is constructed with the set of upstream transform ids the consumer expects, known from + * the pipeline graph at translation time (a single-input transform passes its one parent; a + * Flatten passes the producers of all of its input PCollections). + *
  • Per upstream transform it tracks partitions with a dedicated {@link WatermarkManager}, + * which holds until every partition of that transform has reported and keeps each partition + * monotonic. + *
  • The aggregate input watermark is the {@code min()} across the upstream transforms' + * watermarks, defined only once every expected upstream transform is ready; until + * then {@link #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} and the caller + * emits nothing. + *
+ * + *

Not thread-safe; the caller (a single Kafka Streams processor thread) serializes access. + */ +final class WatermarkAggregator { + + /** Upstream transform ids this consumer must hear from, fixed by the pipeline graph. */ + private final Set expectedUpstreamTransformIds; + + /** Per-upstream-transform partition tracking. */ + private final Map managerByTransformId = new HashMap<>(); + + WatermarkAggregator(Set expectedUpstreamTransformIds) { + Preconditions.checkArgument( + !expectedUpstreamTransformIds.isEmpty(), "expectedUpstreamTransformIds must not be empty"); + this.expectedUpstreamTransformIds = ImmutableSet.copyOf(expectedUpstreamTransformIds); + } + + /** + * Records one upstream watermark report. A report from a transform this consumer does not expect + * indicates a translation wiring bug and fails fast. + */ + void observe(WatermarkPayload report) { + String transformId = report.getTransformId(); + if (!expectedUpstreamTransformIds.contains(transformId)) { + throw new IllegalStateException( + "Received a watermark report from unexpected transform " + + transformId + + "; expected one of " + + expectedUpstreamTransformIds); + } + managerByTransformId + .computeIfAbsent(transformId, id -> new WatermarkManager()) + .observe( + report.getSourcePartition(), + new Instant(report.getWatermarkMillis()), + report.getTotalSourcePartitions()); + } + + /** + * Returns the aggregate input watermark: {@code min()} across all expected upstream transforms, + * or {@link BoundedWindow#TIMESTAMP_MIN_VALUE} while any upstream transform has not yet fully + * reported (the hold). + */ + Instant advance() { + if (managerByTransformId.size() < expectedUpstreamTransformIds.size()) { + return BoundedWindow.TIMESTAMP_MIN_VALUE; + } + Instant min = BoundedWindow.TIMESTAMP_MAX_VALUE; + for (WatermarkManager manager : managerByTransformId.values()) { + // A not-yet-ready manager advances to TIMESTAMP_MIN_VALUE, which correctly holds the min. + Instant watermark = manager.advance(); + if (watermark.isBefore(min)) { + min = watermark; + } + } + return min; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java index 194be701ecec..bac5314a6da7 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java @@ -24,18 +24,26 @@ * {@link KStreamsPayload#isWatermark()} and narrowed the payload, so there is no kind check to do * on each accessor. * - *

A watermark report is the in-band coordination message a downstream stage's {@link - * WatermarkManager} consumes: the watermark value plus which source partition reported it and how - * many source partitions feed the stage in total. + *

A watermark report is the in-band coordination message a downstream watermark aggregator + * consumes: the watermark value, which transform produced it, which of that transform's partitions + * reported it, and how many partitions that transform has in total. The producer stamps its own + * identity without regard to who consumes the report; a consumer with several upstream transforms + * (e.g. Flatten) aggregates per producing transform. */ public interface WatermarkPayload { /** The reported watermark, in event-time milliseconds. */ long getWatermarkMillis(); - /** The source partition this report is for, in {@code [0, getTotalSourcePartitions())}. */ + /** Globally unique id of the transform that produced this report. */ + String getTransformId(); + + /** + * Which partition (physical instance) of the producing transform this report is for, in {@code + * [0, getTotalSourcePartitions())}. + */ int getSourcePartition(); - /** The total number of source partitions feeding the downstream stage. */ + /** How many partitions (physical instances) the producing transform has in total. */ int getTotalSourcePartitions(); } 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 fc5797a12c26..74169acb81cb 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 @@ -25,13 +25,14 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.kafka.streams.processor.api.MockProcessorContext; import org.apache.kafka.streams.processor.api.Record; import org.junit.Test; /** * Tests the watermark wiring of {@link ExecutableStageProcessor}: how it feeds incoming watermark - * reports to the {@link WatermarkManager} and forwards the stage's output watermark. + * reports to the {@link WatermarkAggregator} and forwards the stage's output watermark. * *

Only the watermark path is exercised, so the SDK harness is never started (it is created * lazily on the first data element). A {@link MockProcessorContext} captures what the processor @@ -39,6 +40,12 @@ */ public class ExecutableStageProcessorWatermarkTest { + /** The stage's own transform id, expected on every watermark it forwards. */ + private static final String STAGE_ID = "stage"; + + /** The single upstream transform whose reports the stage aggregates. */ + private static final String UPSTREAM_ID = "upstream"; + private static ExecutableStageProcessor newProcessor() { JobInfo jobInfo = JobInfo.create( @@ -47,13 +54,17 @@ private static ExecutableStageProcessor newProcessor() { "", PipelineOptionsTranslation.toProto(PipelineOptionsFactory.create())); return new ExecutableStageProcessor( - RunnerApi.ExecutableStagePayload.getDefaultInstance(), jobInfo); + RunnerApi.ExecutableStagePayload.getDefaultInstance(), + jobInfo, + STAGE_ID, + ImmutableSet.of(UPSTREAM_ID)); } + /** A report from the upstream transform's given partition. */ private static Record> watermark( long millis, int sourcePartition, int totalSourcePartitions) { KStreamsPayload payload = - KStreamsPayload.watermark(millis, sourcePartition, totalSourcePartitions); + KStreamsPayload.watermark(millis, UPSTREAM_ID, sourcePartition, totalSourcePartitions); return new Record<>(new byte[0], payload, 0L); } @@ -75,7 +86,9 @@ public void singleSourcePartitionForwardsImmediatelyStampedAsItsOwnSource() { assertThat(out.isWatermark(), is(true)); WatermarkPayload report = out.asWatermark(); assertThat(report.getWatermarkMillis(), is(100L)); - // The stage forwards as its own single source (0 of 1), not the upstream's identity. + // The stage forwards under its own identity — its transform id and its own single partition + // (0 of 1) — not the upstream's. + assertThat(report.getTransformId(), is(STAGE_ID)); assertThat(report.getSourcePartition(), is(0)); assertThat(report.getTotalSourcePartitions(), is(1)); } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java new file mode 100644 index 000000000000..a5af0cdc9223 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java @@ -0,0 +1,212 @@ +/* + * 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.hasItems; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.junit.Test; + +/** + * End-to-end test for {@link FlattenTranslator}: two branches are flattened into one PCollection + * and a recording ParDo sees every element from both. + * + *

Each branch is a {@code Create -> identity ParDo}, so its producer feeding the Flatten is an + * {@link ExecutableStageProcessor} — the same shape PAssert's {@code GroupGlobally} produces. This + * exercises the per-producing-transform watermark aggregation: each branch's producer stamps its + * own transform id on its watermark, and the Flatten holds its output watermark until every + * upstream transform it expects has reported. Without that, the Flatten would release its watermark + * after the first branch drained and the downstream stage's bundle would close early, dropping the + * second branch's elements. + */ +public class FlattenTest { + + private static class IdentityFn extends DoFn { + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver out) { + out.output(input); + } + } + + /** Records every element the harness feeds it so the test can assert the flatten's union. */ + private static class RecordingFn extends DoFn { + private final SharedTestCollector collector; + + RecordingFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver out) { + collector.record(input); + out.output(input); + } + } + + @Test + public void flattenUnionsEveryBranch() { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection a = + pipeline.apply("createA", Create.of(1, 2)).apply("idA", ParDo.of(new IdentityFn())); + PCollection b = + pipeline.apply("createB", Create.of(3, 4)).apply("idB", ParDo.of(new IdentityFn())); + PCollectionList.of(a) + .and(b) + .apply("flatten", Flatten.pCollections()) + .apply("record", ParDo.of(new RecordingFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List recorded = collector.recorded(); + assertThat(recorded.size(), is(4)); + assertThat(recorded, hasItems(1, 2, 3, 4)); + } + } + + @Test + public void pCollectionFeedingTwoFlattensIsSupported() { + // input2 feeds both flattens, so its producer's watermark report is consumed by two different + // aggregators. The producer stamps its own transform id once, and each flatten holds until its + // own two branches drain. Verify both flattens produce the right union. + try (SharedTestCollector left = SharedTestCollector.create(); + SharedTestCollector right = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection input1 = + pipeline.apply("c1", Create.of(1, 2)).apply("id1", ParDo.of(new IdentityFn())); + PCollection input2 = + pipeline.apply("c2", Create.of(3, 4)).apply("id2", ParDo.of(new IdentityFn())); + PCollection input3 = + pipeline.apply("c3", Create.of(5, 6)).apply("id3", ParDo.of(new IdentityFn())); + PCollectionList.of(input1) + .and(input2) + .apply("l1", Flatten.pCollections()) + .apply("recordL1", ParDo.of(new RecordingFn(left))); + PCollectionList.of(input2) + .and(input3) + .apply("l2", Flatten.pCollections()) + .apply("recordL2", ParDo.of(new RecordingFn(right))); + + KafkaStreamsTestRunner.run(pipeline); + + assertThat(left.recorded().size(), is(4)); + assertThat(left.recorded(), hasItems(1, 2, 3, 4)); + assertThat(right.recorded().size(), is(4)); + assertThat(right.recorded(), hasItems(3, 4, 5, 6)); + } + } + + /** Maps each int to {@code KV("k", int)} so a downstream GroupByKey groups all branches. */ + private static class ToKvFn extends DoFn> { + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver> out) { + out.output(KV.of("k", input)); + } + } + + /** Records each grouped result as {@code "key=[sorted values]"}. */ + private static class RecordGroupFn extends DoFn>, Void> { + private final SharedTestCollector collector; + + RecordGroupFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element KV> group) { + List values = new ArrayList<>(); + group.getValue().forEach(values::add); + Collections.sort(values); + collector.record(group.getKey() + "=" + values); + } + } + + @Test + public void watermarkPropagatesThroughFlattenAndFiresDownstreamGroupByKey() { + // GroupByKey fires exactly once, when its input watermark reaches the end of the global + // window, and the Flatten forwards its watermark only after every branch has drained. Both + // branches share the key, so a single group holding the elements of both branches proves the + // watermark propagated through the Flatten at the right time — a premature release would fire + // a partial group instead. + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection> a = + pipeline + .apply("createA", Create.of(1, 2)) + .apply("kvA", ParDo.of(new ToKvFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); + PCollection> b = + pipeline + .apply("createB", Create.of(3, 4)) + .apply("kvB", ParDo.of(new ToKvFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); + PCollectionList.of(a) + .and(b) + .apply("flatten", Flatten.pCollections()) + .apply("gbk", GroupByKey.create()) + .apply("record", ParDo.of(new RecordGroupFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List groups = collector.recorded(); + assertThat(groups.size(), is(1)); + assertThat(groups, hasItems("k=[1, 2, 3, 4]")); + } + } + + @Test + public void flattenOfOneBranchTwiceDuplicatesEveryElement() { + // A self-flatten is a bag union with itself: every element must appear twice. The fuser folds + // the Flatten into the SDK-harness stage (it never reaches the runner's Flatten translator as a + // duplicate-input node), so the duplication happens in the harness. + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection branch = + pipeline.apply("create", Create.of(1, 2)).apply("id", ParDo.of(new IdentityFn())); + PCollectionList.of(branch) + .and(branch) + .apply("flatten", Flatten.pCollections()) + .apply("record", ParDo.of(new RecordingFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List recorded = collector.recorded(); + assertThat(recorded.size(), is(4)); + assertThat(recorded, hasItems(1, 2)); + assertThat(recorded.stream().filter(v -> v == 1).count(), is(2L)); + assertThat(recorded.stream().filter(v -> v == 2).count(), is(2L)); + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java index 19346f54763e..95ce70b88578 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java @@ -58,10 +58,11 @@ public void roundTripsDataPayload() { @Test public void roundTripsWatermarkPayload() { - KStreamsPayload payload = KStreamsPayload.watermark(12345L, 2, 4); + KStreamsPayload payload = KStreamsPayload.watermark(12345L, "transform-a", 2, 4); KStreamsPayload out = roundTrip(payload); assertThat(out.isWatermark(), is(true)); assertThat(out.asWatermark().getWatermarkMillis(), is(12345L)); + assertThat(out.asWatermark().getTransformId(), is("transform-a")); assertThat(out.asWatermark().getSourcePartition(), is(2)); assertThat(out.asWatermark().getTotalSourcePartitions(), is(4)); assertThat(out, is(payload)); @@ -70,7 +71,7 @@ public void roundTripsWatermarkPayload() { @Test public void roundTripsTerminalMaxWatermark() { KStreamsPayload payload = - KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(), 0, 1); + KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(), "t", 0, 1); assertThat( roundTrip(payload).asWatermark().getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); @@ -80,7 +81,7 @@ public void roundTripsTerminalMaxWatermark() { public void roundTripsNegativeWatermark() { // Beam event times can be negative; sint64 must round-trip them losslessly. KStreamsPayload payload = - KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(), 0, 1); + KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(), "t", 0, 1); assertThat( roundTrip(payload).asWatermark().getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis())); diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java new file mode 100644 index 000000000000..88d541bdc0c2 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java @@ -0,0 +1,160 @@ +/* + * 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.containsString; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertThrows; + +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.joda.time.Instant; +import org.junit.Test; + +/** + * Unit tests for {@link WatermarkAggregator}: aggregation of watermark reports across multiple + * upstream transforms, each with its own partition set. Per-partition behaviour within one upstream + * transform (monotonicity, repartition reset) is covered in depth by {@code WatermarkManagerTest}; + * here it is exercised through the aggregate. + */ +public class WatermarkAggregatorTest { + + private static Instant ts(long millis) { + return new Instant(millis); + } + + /** A report from the given transform's partition {@code sourcePartition} of {@code total}. */ + private static WatermarkPayload report( + String transformId, long millis, int sourcePartition, int total) { + return KStreamsPayload.watermark(millis, transformId, sourcePartition, total).asWatermark(); + } + + @Test + public void holdsBeforeAnyReport() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + } + + @Test + public void singleUpstreamSinglePartitionAdvances() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + aggregator.observe(report("a", 100L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void holdsUntilEveryExpectedUpstreamReports() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b")); + aggregator.observe(report("a", 100L, 0, 1)); + // Only one of the two expected upstream transforms has reported — still holding. + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + aggregator.observe(report("b", 200L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void holdsUntilEveryPartitionOfEachUpstreamReports() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b")); + aggregator.observe(report("a", 100L, 0, 1)); + aggregator.observe(report("b", 200L, 0, 2)); + // Upstream "b" has reported only one of its two partitions — still holding. + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + aggregator.observe(report("b", 300L, 1, 2)); + // Ready: min(100, min(200, 300)) = 100. + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void aggregateIsMinAcrossUpstreams() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b", "c")); + aggregator.observe(report("a", 300L, 0, 1)); + aggregator.observe(report("b", 100L, 0, 1)); + aggregator.observe(report("c", 500L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + + // The slowest upstream advances; the aggregate follows the new min. + aggregator.observe(report("b", 400L, 0, 1)); + assertThat(aggregator.advance(), is(ts(300L))); + } + + @Test + public void perUpstreamWatermarkIsMonotonic() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + aggregator.observe(report("a", 200L, 0, 1)); + assertThat(aggregator.advance(), is(ts(200L))); + + // A late lower report from the same upstream partition is ignored. + aggregator.observe(report("a", 100L, 0, 1)); + assertThat(aggregator.advance(), is(ts(200L))); + } + + @Test + public void duplicateReportDoesNotAdvance() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + aggregator.observe(report("a", 100L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + + // The same report again (e.g. a broadcast duplicate) leaves the aggregate unchanged. + aggregator.observe(report("a", 100L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void terminalMaxWatermarkAggregates() { + long max = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b")); + aggregator.observe(report("a", max, 0, 1)); + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + aggregator.observe(report("b", max, 0, 1)); + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MAX_VALUE)); + } + + @Test + public void repartitionOfOneUpstreamReopensTheHold() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a", "b")); + aggregator.observe(report("a", 100L, 0, 1)); + aggregator.observe(report("b", 200L, 0, 1)); + assertThat(aggregator.advance(), is(ts(100L))); + + // Upstream "b" changes its partition count (repartition): its per-partition state resets and + // the aggregate holds again until b's new full partition set has reported. + aggregator.observe(report("b", 250L, 0, 2)); + assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + aggregator.observe(report("b", 300L, 1, 2)); + assertThat(aggregator.advance(), is(ts(100L))); + } + + @Test + public void reportFromUnexpectedTransformFailsFast() { + WatermarkAggregator aggregator = new WatermarkAggregator(ImmutableSet.of("a")); + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, () -> aggregator.observe(report("intruder", 100L, 0, 1))); + assertThat(thrown.getMessage(), containsString("intruder")); + } + + @Test + public void emptyExpectedUpstreamsIsRejected() { + assertThrows(IllegalArgumentException.class, () -> new WatermarkAggregator(ImmutableSet.of())); + } +} From 75adf49ee8deb411bf8851a8e86e4c34fd420cbd Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:00:06 +0500 Subject: [PATCH 17/37] [GSoC 2026] Kafka Streams runner: surface SDK-harness metrics as MetricResults (#39341) --- runners/kafka-streams/build.gradle | 1 + .../streams/KafkaStreamsPipelineRunner.java | 3 +- .../KafkaStreamsPortablePipelineResult.java | 13 ++- .../translation/ExecutableStageProcessor.java | 46 ++++++++++- .../ExecutableStageTranslator.java | 9 ++- .../KafkaStreamsTranslationContext.java | 19 +++++ .../kafka/streams/KafkaStreamsTestRunner.java | 12 ++- ...ExecutableStageProcessorWatermarkTest.java | 4 +- .../streams/translation/MetricsTest.java | 79 +++++++++++++++++++ 9 files changed, 173 insertions(+), 13 deletions(-) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/MetricsTest.java 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)); + } +} From 87911f00973cb7ae13e78bd07ce71d9b2395e044 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Fri, 17 Jul 2026 12:40:09 +0500 Subject: [PATCH 18/37] [GSoC 2026] Kafka Streams runner: TestPipeline-dispatchable test runner (PAssert works) (#39362) --- .../kafka/streams/TestKafkaStreamsRunner.java | 126 ++++++++++++++++++ .../streams/TestKafkaStreamsRunnerTest.java | 82 ++++++++++++ 2 files changed, 208 insertions(+) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunner.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunnerTest.java diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunner.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunner.java new file mode 100644 index 000000000000..1fe4b310789f --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunner.java @@ -0,0 +1,126 @@ +/* + * 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; + +import java.io.IOException; +import java.util.UUID; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.PipelineRunner; +import org.apache.beam.sdk.metrics.MetricResults; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.util.construction.Environments; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; + +/** + * A {@link PipelineRunner} for tests, so {@code TestPipeline}-based suites — including Beam's + * {@code @ValidatesRunner} tests — can run against the Kafka Streams runner without a broker. + * + *

{@link #run} translates the pipeline and drives it to quiescence through a {@code + * TopologyTestDriver} (via {@link KafkaStreamsTestRunner}), with the SDK harness in-process + * (EMBEDDED environment). The returned {@link PipelineResult} is terminal ({@code DONE}) and + * exposes the metrics the harness reported — which is what {@code + * TestPipeline.verifyPAssertsSucceeded} uses to check that every {@code PAssert} in the pipeline + * actually ran and succeeded. + * + *

A {@code PAssert} failure inside a DoFn fails the pipeline run — over the Fn API the DoFn's + * error travels as an error string (carrying the assertion text), not as a Java {@link + * AssertionError} object. {@link #run} additionally unwraps and rethrows an {@link AssertionError} + * when one is present in the exception chain (e.g. assertions thrown runner-side). A {@code + * PAssert} that silently never runs is caught by {@code TestPipeline.verifyPAssertsSucceeded} + * comparing the success-counter metric against the number of assertions in the pipeline. + * + *

Select it with {@code --runner=org.apache.beam.runners.kafka.streams.TestKafkaStreamsRunner} + * in {@code beamTestPipelineOptions}. + */ +public final class TestKafkaStreamsRunner extends PipelineRunner { + + private TestKafkaStreamsRunner() {} + + /** Called reflectively by {@link PipelineRunner#fromOptions}. */ + public static TestKafkaStreamsRunner fromOptions(PipelineOptions options) { + return new TestKafkaStreamsRunner(); + } + + @Override + public PipelineResult run(Pipeline pipeline) { + // Tests supply generic options; fill in what the Kafka Streams translation needs. The + // pipeline's own options are the ones the translation reads. + KafkaStreamsPipelineOptions options = + pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + if (options.getApplicationId() == null || options.getApplicationId().isEmpty()) { + options.setApplicationId("ks-validates-runner-" + UUID.randomUUID()); + } + options + .as(PortablePipelineOptions.class) + .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); + + MetricResults metrics; + try { + metrics = KafkaStreamsTestRunner.run(pipeline); + } catch (Throwable t) { + // A PAssert failure throws an AssertionError inside the harness DoFn, which reaches us + // wrapped in the bundle-processing exception chain. Rethrow the innermost AssertionError so + // the test framework reports the actual assertion. + for (Throwable current = t; current != null; current = current.getCause()) { + if (current instanceof AssertionError) { + throw (AssertionError) current; + } + } + throw t; + } + return new TestKafkaStreamsPipelineResult(metrics); + } + + /** Terminal result of a test run: state {@code DONE} with the harness-reported metrics. */ + private static final class TestKafkaStreamsPipelineResult implements PipelineResult { + private final MetricResults metrics; + + TestKafkaStreamsPipelineResult(MetricResults metrics) { + this.metrics = metrics; + } + + @Override + public State getState() { + return State.DONE; + } + + @Override + public State cancel() throws IOException { + throw new UnsupportedOperationException( + "A TestKafkaStreamsRunner pipeline has already finished when its result is returned."); + } + + @Override + public State waitUntilFinish(@Nullable Duration duration) { + return State.DONE; + } + + @Override + public State waitUntilFinish() { + return State.DONE; + } + + @Override + public MetricResults metrics() { + return metrics; + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunnerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunnerTest.java new file mode 100644 index 000000000000..ba62a9d7f49e --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunnerTest.java @@ -0,0 +1,82 @@ +/* + * 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; + +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertThrows; + +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.junit.Rule; +import org.junit.Test; + +/** + * End-to-end tests for {@link TestKafkaStreamsRunner}: a {@link TestPipeline} with {@link PAssert} + * runs against the Kafka Streams runner, and {@code TestPipeline.verifyPAssertsSucceeded} confirms + * through the runner's metrics that every assertion actually executed. + * + *

This is the full {@code @ValidatesRunner} mechanism exercised in-module: PAssert's {@code + * containsInAnyOrder} materializes the actual PCollection through the {@code GBK + Flatten} + * bootstrap path (no side inputs), the assertion DoFn runs in the EMBEDDED harness, its success + * counter travels back through the runner's MetricResults, and a failing assertion surfaces as an + * {@link AssertionError}. + */ +public class TestKafkaStreamsRunnerTest { + + private static PipelineOptions options() { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(TestKafkaStreamsRunner.class); + return options; + } + + @Rule public final transient TestPipeline pipeline = TestPipeline.fromOptions(options()); + + @Test + public void pAssertContainsInAnyOrderSucceeds() { + PAssert.that(pipeline.apply("create", Create.of(1, 2, 3))).containsInAnyOrder(3, 2, 1); + pipeline.run(); + } + + @Test + public void pAssertOnGroupedResultSucceeds() { + PAssert.that( + pipeline.apply("create", Create.of("a", "b", "a")).apply("count", Count.perElement())) + .containsInAnyOrder(KV.of("a", 2L), KV.of("b", 1L)); + pipeline.run(); + } + + @Test + public void failingPAssertFailsTheRun() { + // Over the Fn API a DoFn failure travels as an error string, not a Java AssertionError object, + // so the guarantee for a portable runner is that a failing PAssert fails the pipeline run (and + // its message carries the assertion text). A dropped assertion — one that never runs — would + // instead be caught by TestPipeline.verifyPAssertsSucceeded via the runner's metrics. + // Throwable rather than RuntimeException: if the exception chain ever preserves the DoFn's + // AssertionError (an Error), the runner unwraps and rethrows it, which is also a pass here. + PAssert.that(pipeline.apply("create", Create.of(1, 2, 3))).containsInAnyOrder(1, 2, 4); + Throwable thrown = assertThrows(Throwable.class, pipeline::run); + assertThat(Throwables.getStackTraceAsString(thrown), containsString("AssertionError")); + } +} From d9cea589e5b9e2da28730620ca1ffb045607fdbc Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Mon, 20 Jul 2026 13:12:26 +0500 Subject: [PATCH 19/37] [GSoC 2026] Kafka Streams runner: validatesRunner task; Create and Flatten suites green (#39380) --- .../beam_KafkaStreamsRunner_FeatureBranch.yml | 2 + runners/kafka-streams/build.gradle | 76 ++++++++++++++++ .../translation/EmptyBoundedSource.java | 89 +++++++++++++++++++ .../translation/FlattenTranslator.java | 2 +- .../translation/ImpulseTranslator.java | 8 +- .../KafkaStreamsPipelineTranslator.java | 64 ++++++++++++- .../KafkaStreamsTranslationContext.java | 15 +++- 7 files changed, 246 insertions(+), 10 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/EmptyBoundedSource.java diff --git a/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml b/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml index 4ec7b0778cc2..e290c9beee1a 100644 --- a/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml +++ b/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml @@ -70,3 +70,5 @@ jobs: ${{ runner.os }}-gradle-kafka-streams- - name: Build and test Kafka Streams runner run: ./gradlew :runners:kafka-streams:build --no-daemon --stacktrace + - name: Run ValidatesRunner suite + run: ./gradlew :runners:kafka-streams:validatesRunner --no-daemon --stacktrace diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index a794299cb608..81099560ba79 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -16,6 +16,8 @@ * limitations under the License. */ +import groovy.json.JsonOutput + plugins { id 'org.apache.beam.module' } def kafka_version = '3.9.0' @@ -29,6 +31,10 @@ description = "Apache Beam :: Runners :: Kafka Streams" evaluationDependsOn(":sdks:java:core") evaluationDependsOn(":runners:core-java") +configurations { + validatesRunner +} + configurations.configureEach { resolutionStrategy.eachDependency { details -> if (details.requested.group == "org.apache.kafka") { @@ -68,4 +74,74 @@ dependencies { testImplementation library.java.junit testImplementation library.java.mockito_core testImplementation "org.apache.kafka:kafka-streams-test-utils:$kafka_version" + + // Beam's @ValidatesRunner suite: the test classes come from the SDK core test jar; the runner + // (TestKafkaStreamsRunner) and its TopologyTestDriver harness come from this module's test + // output and test runtime classpath. + validatesRunner project(path: ":sdks:java:core", configuration: "shadowTest") + validatesRunner project(project.path) + validatesRunner sourceSets.test.output + validatesRunner sourceSets.test.runtimeClasspath +} + + +// Known-failing @ValidatesRunner tests, excluded until the feature they need lands. +def sickbayTests = [ + // Needs multi-output executable stages (output-tag dispatch in ExecutableStageProcessor). + 'org.apache.beam.sdk.transforms.FlattenTest.testFlattenMultiplePCollectionsHavingMultipleConsumers', +] + +tasks.register("validatesRunner", Test) { + group = "Verification" + description = "Runs the subset of Beam's ValidatesRunner suite the Kafka Streams runner supports." + // Never consider up-to-date; the suite is the correctness gate. + outputs.upToDateWhen { false } + systemProperty "beamTestPipelineOptions", + JsonOutput.toJson(["--runner=org.apache.beam.runners.kafka.streams.TestKafkaStreamsRunner"]) + classpath = configurations.validatesRunner + testClassesDirs = files(project(":sdks:java:core").sourceSets.test.output.classesDirs) + maxParallelForks 2 + useJUnit { + includeCategories 'org.apache.beam.sdk.testing.ValidatesRunner' + // Environment / harness features that need a properly configured external environment. + excludeCategories 'org.apache.beam.sdk.testing.UsesExternalService' + excludeCategories 'org.apache.beam.sdk.testing.UsesSdkHarnessEnvironment' + excludeCategories 'org.apache.beam.sdk.testing.UsesBundleFinalizer' + excludeCategories 'org.apache.beam.sdk.testing.UsesJavaExpansionService' + excludeCategories 'org.apache.beam.sdk.testing.UsesPythonExpansionService' + // Features the runner does not support yet. + excludeCategories 'org.apache.beam.sdk.testing.UsesSideInputs' + excludeCategories 'org.apache.beam.sdk.testing.UsesStatefulParDo' + excludeCategories 'org.apache.beam.sdk.testing.UsesTimersInParDo' + excludeCategories 'org.apache.beam.sdk.testing.UsesTimerMap' + excludeCategories 'org.apache.beam.sdk.testing.UsesLoopingTimer' + excludeCategories 'org.apache.beam.sdk.testing.UsesStrictTimerOrdering' + excludeCategories 'org.apache.beam.sdk.testing.UsesProcessingTimeTimers' + excludeCategories 'org.apache.beam.sdk.testing.UsesOnWindowExpiration' + excludeCategories 'org.apache.beam.sdk.testing.UsesTestStream' + excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedPCollections' + excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedSplittableParDo' + excludeCategories 'org.apache.beam.sdk.testing.UsesBoundedSplittableParDo' + excludeCategories 'org.apache.beam.sdk.testing.UsesCustomWindowMerging' + excludeCategories 'org.apache.beam.sdk.testing.UsesMetricsPusher' + excludeCategories 'org.apache.beam.sdk.testing.UsesCommittedMetrics' + excludeCategories 'org.apache.beam.sdk.testing.UsesSystemMetrics' + excludeCategories 'org.apache.beam.sdk.testing.UsesOrderedListState' + excludeCategories 'org.apache.beam.sdk.testing.UsesMultimapState' + excludeCategories 'org.apache.beam.sdk.testing.UsesMapState' + excludeCategories 'org.apache.beam.sdk.testing.UsesSetState' + excludeCategories 'org.apache.beam.sdk.testing.FlattenWithHeterogeneousCoders' + excludeCategories 'org.apache.beam.sdk.testing.LargeKeys$Above100MB' + } + filter { + // The suites enabled so far, extended class by class as runner support grows. An explicit + // include list is needed because feature gaps like windowing beyond the global window have no + // JUnit category to exclude. + includeTestsMatching 'org.apache.beam.sdk.transforms.CreateTest' + includeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest' + for (String test : sickbayTests) { + excludeTestsMatching test + } + failOnNoMatchingTests = false + } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/EmptyBoundedSource.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/EmptyBoundedSource.java new file mode 100644 index 000000000000..14bee0c3cfd0 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/EmptyBoundedSource.java @@ -0,0 +1,89 @@ +/* + * 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 java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.NoSuchElementException; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.BoundedSource; +import org.apache.beam.sdk.options.PipelineOptions; + +/** + * A {@link BoundedSource} with no elements. The runner substitutes it for a {@code Flatten} of zero + * PCollections (see {@link KafkaStreamsPipelineTranslator}): reading it produces no data and a + * terminal watermark — exactly the semantics of an empty PCollection. The element type is never + * observed since no element is ever produced. + */ +class EmptyBoundedSource extends BoundedSource { + + @Override + public List> split( + long desiredBundleSizeBytes, PipelineOptions options) { + return Collections.singletonList(this); + } + + @Override + public long getEstimatedSizeBytes(PipelineOptions options) { + return 0; + } + + @Override + public BoundedReader createReader(PipelineOptions options) { + return new EmptyReader(this); + } + + @Override + public Coder getOutputCoder() { + return ByteArrayCoder.of(); + } + + /** A reader that is exhausted from the start. */ + private static final class EmptyReader extends BoundedReader { + private final EmptyBoundedSource source; + + EmptyReader(EmptyBoundedSource source) { + this.source = source; + } + + @Override + public boolean start() { + return false; + } + + @Override + public boolean advance() { + return false; + } + + @Override + public byte[] getCurrent() throws NoSuchElementException { + throw new NoSuchElementException("EmptyBoundedSource has no elements"); + } + + @Override + public void close() throws IOException {} + + @Override + public BoundedSource getCurrentSource() { + return source; + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java index 8f6ce2b546dd..3ac4c1304daa 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java @@ -53,6 +53,7 @@ public void translate( // Flatten produces exactly one output PCollection, fed by all of its input PCollections. String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + Topology topology = context.getTopology(); Set seenInputs = new HashSet<>(); List parentProcessors = new ArrayList<>(); Set upstreamTransformIds = new HashSet<>(); @@ -71,7 +72,6 @@ public void translate( upstreamTransformIds.add(parentProcessor); } - Topology topology = context.getTopology(); topology.addProcessor( transformId, () -> new FlattenProcessor(transformId, upstreamTransformIds), diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java index a90987ba6383..79d8c3cf577f 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java @@ -30,9 +30,9 @@ *

Adds three nodes to the Kafka Streams {@link Topology}: * *

    - *
  • A {@code byte[]} source bound to a dedicated per-application bootstrap topic (see {@link - * KafkaStreamsTranslationContext#getImpulseBootstrapTopic()}). Kafka Streams refuses to start - * a topology that has no real source topic, so the bootstrap topic exists purely to satisfy + *
  • A {@code byte[]} source bound to a dedicated per-transform bootstrap topic (see {@link + * KafkaStreamsTranslationContext#getImpulseBootstrapTopic}). Kafka Streams refuses to start a + * topology that has no real source topic, so the bootstrap topic exists purely to satisfy * that requirement — records published to it are ignored by {@link ImpulseProcessor}. *
  • The {@link ImpulseProcessor} itself, which schedules a one-shot wall-clock punctuator on * {@code init} and emits a single empty data {@link KStreamsPayload} followed by a terminal @@ -71,7 +71,7 @@ public void translate( Topology topology = context.getTopology(); String sourceNodeName = transformId + SOURCE_SUFFIX; String stateStoreName = transformId + STATE_STORE_SUFFIX; - String bootstrapTopic = context.getImpulseBootstrapTopic(); + String bootstrapTopic = context.getImpulseBootstrapTopic(transformId); topology.addSource( sourceNodeName, diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java index a9e26ecd6412..a71434bfd8d0 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java @@ -23,6 +23,7 @@ import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.util.SerializableUtils; import org.apache.beam.sdk.util.construction.NativeTransforms; import org.apache.beam.sdk.util.construction.PTransformTranslation; import org.apache.beam.sdk.util.construction.graph.ExecutableStage; @@ -30,8 +31,10 @@ import org.apache.beam.sdk.util.construction.graph.PipelineNode; import org.apache.beam.sdk.util.construction.graph.QueryablePipeline; import org.apache.beam.sdk.util.construction.graph.TrivialNativeTransformExpander; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; /** * Translates a portable Beam pipeline into a Kafka Streams {@link @@ -75,7 +78,14 @@ public KafkaStreamsTranslationContext createTranslationContext( * yet (e.g. GroupByKey). */ public Set knownUrns() { - return urnToTranslator.keySet(); + // Flatten is deliberately not exposed: the fuser has its own handling for Flatten (unzipping + // it into producer stages and re-introducing a runner Flatten via the OutputDeduplicator), and + // declaring it native makes TrivialNativeTransformExpander interact badly with multi-branch + // flattens whose inputs come from composite expansions (their producing stages get dropped + // from the fused pipeline, failing "consumed but never produced" validation). Mirrors the + // Flink runner, which likewise keeps Read out of its known URNs for expander reasons. + return Sets.difference( + urnToTranslator.keySet(), ImmutableSet.of(PTransformTranslation.FLATTEN_TRANSFORM_URN)); } /** @@ -98,10 +108,60 @@ public RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) { if (alreadyFused) { return pipeline; } - RunnerApi.Pipeline trimmed = TrivialNativeTransformExpander.forKnownUrns(pipeline, knownUrns()); + RunnerApi.Pipeline withoutEmptyFlattens = replaceEmptyFlattensWithEmptyReads(pipeline); + RunnerApi.Pipeline trimmed = + TrivialNativeTransformExpander.forKnownUrns(withoutEmptyFlattens, knownUrns()); return GreedyPipelineFuser.fuse(trimmed).toPipeline(); } + /** + * Rewrites every {@code Flatten} of zero PCollections into a primitive Read of an {@link + * EmptyBoundedSource}. A zero-input Flatten produces an empty PCollection, but it is also a root + * of the pipeline graph, and {@link GreedyPipelineFuser} only accepts Impulse or Read roots. The + * Read of an empty source has exactly the right semantics — no elements, then the terminal + * watermark — and reuses the existing {@link ReadTranslator} path. + */ + private static RunnerApi.Pipeline replaceEmptyFlattensWithEmptyReads( + RunnerApi.Pipeline pipeline) { + RunnerApi.Pipeline.Builder pipelineBuilder = null; + for (Map.Entry entry : + pipeline.getComponents().getTransformsMap().entrySet()) { + RunnerApi.PTransform transform = entry.getValue(); + if (!PTransformTranslation.FLATTEN_TRANSFORM_URN.equals(transform.getSpec().getUrn()) + || !transform.getInputsMap().isEmpty()) { + continue; + } + if (pipelineBuilder == null) { + pipelineBuilder = pipeline.toBuilder(); + } + RunnerApi.ReadPayload emptyReadPayload = + RunnerApi.ReadPayload.newBuilder() + .setIsBounded(RunnerApi.IsBounded.Enum.BOUNDED) + .setSource( + RunnerApi.FunctionSpec.newBuilder() + .setUrn(JAVA_SERIALIZED_BOUNDED_SOURCE_URN) + .setPayload( + ByteString.copyFrom( + SerializableUtils.serializeToByteArray(new EmptyBoundedSource())))) + .build(); + pipelineBuilder + .getComponentsBuilder() + .putTransforms( + entry.getKey(), + transform + .toBuilder() + .setSpec( + RunnerApi.FunctionSpec.newBuilder() + .setUrn(PTransformTranslation.READ_TRANSFORM_URN) + .setPayload(emptyReadPayload.toByteString())) + .build()); + } + return pipelineBuilder == null ? pipeline : pipelineBuilder.build(); + } + + /** The URN {@code ReadTranslation} uses for a Java-serialized {@code BoundedSource} payload. */ + private static final String JAVA_SERIALIZED_BOUNDED_SOURCE_URN = "beam:java:boundedsource:v1"; + /** * Walks the pipeline in topological order and translates each transform whose URN is supported. * Throws {@link UnsupportedOperationException} on the first unsupported URN. 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 89a2d9825cbc..ec1b3f26aded 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 @@ -122,9 +122,18 @@ public String getProcessorNameForPCollection(String pCollectionId) { return name; } - /** Returns the dedicated bootstrap topic name used by Impulse for this application. */ - public String getImpulseBootstrapTopic() { - return IMPULSE_BOOTSTRAP_TOPIC_PREFIX + pipelineOptions.getApplicationId(); + /** + * Returns the dedicated bootstrap topic name for one Impulse transform. Keyed by transform id + * (sanitized to Kafka's legal topic-name character set) because a pipeline can contain several + * Impulses (e.g. an empty {@code Create} plus the dummy branch {@code PAssert} adds), and Kafka + * Streams rejects registering the same topic on two source nodes. + */ + public String getImpulseBootstrapTopic(String transformId) { + String sanitizedTransformId = ILLEGAL_TOPIC_CHARS.matcher(transformId).replaceAll("_"); + return IMPULSE_BOOTSTRAP_TOPIC_PREFIX + + pipelineOptions.getApplicationId() + + "_" + + sanitizedTransformId; } /** From fca14356614faa904e69633463d710bf860b63bd Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Wed, 22 Jul 2026 16:28:18 +0500 Subject: [PATCH 20/37] [GSoC 2026] Kafka Streams runner: multi-output executable stages (#39410) --- runners/kafka-streams/build.gradle | 8 +- .../translation/ExecutableStageProcessor.java | 53 ++++++--- .../ExecutableStageTranslator.java | 52 ++++++--- .../translation/StageOutputProcessor.java | 93 +++++++++++++++ .../kafka/streams/MultiOutputStageTest.java | 87 ++++++++++++++ ...ExecutableStageProcessorWatermarkTest.java | 5 +- .../translation/StageOutputProcessorTest.java | 108 ++++++++++++++++++ 7 files changed, 373 insertions(+), 33 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/MultiOutputStageTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 81099560ba79..52535c6654e6 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -87,8 +87,11 @@ dependencies { // Known-failing @ValidatesRunner tests, excluded until the feature they need lands. def sickbayTests = [ - // Needs multi-output executable stages (output-tag dispatch in ExecutableStageProcessor). - 'org.apache.beam.sdk.transforms.FlattenTest.testFlattenMultiplePCollectionsHavingMultipleConsumers', + // Non-global windowing (FixedWindows, merging windows, timestamp combiners) is not supported + // yet; these apply a window and assert on window-derived output, hitting a GlobalWindow cast. + 'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests', + 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerLatest', + 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerEarliest', ] tasks.register("validatesRunner", Test) { @@ -139,6 +142,7 @@ tasks.register("validatesRunner", Test) { // JUnit category to exclude. includeTestsMatching 'org.apache.beam.sdk.transforms.CreateTest' includeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest' + includeTestsMatching 'org.apache.beam.sdk.transforms.GroupByKeyTest*' for (String test : sickbayTests) { excludeTestsMatching test } 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 fb26e9e71f6b..15265073dd21 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 @@ -17,6 +17,7 @@ */ package org.apache.beam.runners.kafka.streams.translation; +import java.util.Map; import java.util.Queue; import java.util.Set; import java.util.concurrent.ConcurrentLinkedQueue; @@ -36,6 +37,7 @@ 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.ImmutableMap; 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; @@ -85,11 +87,14 @@ class ExecutableStageProcessor // 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. - // The element type is intentionally wildcarded: the runner does not need to know the runtime - // value type — the bundle factory handles all coder application at the Fn-API boundary using - // the PCollection coders from the ExecutableStagePayload. Pretending the type was byte[] was - // only safe because the Impulse output coder happens to be ByteArrayCoder. - private final Queue> pendingOutputs = new ConcurrentLinkedQueue<>(); + // Each entry carries the output PCollection id so it can be routed to that output's downstream on + // flush. The element type is intentionally wildcarded: the runner does not need to know the + // runtime value type — the bundle factory handles all coder application at the Fn-API boundary + // using the PCollection coders from the ExecutableStagePayload. + private final Queue pendingOutputs = new ConcurrentLinkedQueue<>(); + // Output PCollection id -> the child node (a StageOutputProcessor relay) to forward that output + // to. Empty for a single-output stage, which forwards to its one downstream directly. + private final Map outputChildByPCollectionId; // Computes this stage's input watermark from its upstream transform's reports, holding until // every partition of the upstream transform has reported (see WatermarkAggregator). @@ -114,12 +119,25 @@ class ExecutableStageProcessor JobInfo jobInfo, String transformId, Set upstreamTransformIds, - MetricsContainerImpl metricsContainer) { + MetricsContainerImpl metricsContainer, + Map outputChildByPCollectionId) { this.stagePayload = stagePayload; this.jobInfo = jobInfo; this.transformId = transformId; this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); this.metricsContainer = metricsContainer; + this.outputChildByPCollectionId = ImmutableMap.copyOf(outputChildByPCollectionId); + } + + /** A harness output element together with the id of the output PCollection it belongs to. */ + private static final class PendingOutput { + final String pCollectionId; + final WindowedValue value; + + PendingOutput(String pCollectionId, WindowedValue value) { + this.pCollectionId = pCollectionId; + this.value = value; + } } @Override @@ -182,11 +200,12 @@ private void ensureBundleOpen() throws Exception { new OutputReceiverFactory() { @Override public FnDataReceiver create(String pCollectionId) { - // Outputs are queued here on harness threads and drained on the processing thread - // after the bundle closes. + // Outputs are queued here on harness threads, tagged with their output PCollection id, + // and drained on the processing thread after the bundle closes. return receivedElement -> { if (receivedElement != null) { - pendingOutputs.add((WindowedValue) receivedElement); + pendingOutputs.add( + new PendingOutput(pCollectionId, (WindowedValue) receivedElement)); } }; } @@ -249,12 +268,20 @@ private void closeBundleAndFlush(Record> record) { } ProcessorContext> ctx = checkInitialized(context); // The harness has finished the bundle (close() returned) so no further enqueues happen. - // Drain via poll() so each element is removed as it is forwarded. - WindowedValue output; + // Drain via poll() so each element is removed as it is forwarded. Each output is routed to its + // own output's relay child for a multi-output stage; a single-output stage forwards directly to + // its one downstream (empty routing map). + PendingOutput output; while ((output = pendingOutputs.poll()) != null) { - ctx.forward( + Record> outputRecord = new Record>( - record.key(), KStreamsPayload.data(output), record.timestamp())); + record.key(), KStreamsPayload.data(output.value), record.timestamp()); + String childNode = outputChildByPCollectionId.get(output.pCollectionId); + if (childNode == null) { + ctx.forward(outputRecord); + } else { + ctx.forward(outputRecord, childNode); + } } } 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 59c233a78bb1..71e24f32c1f5 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 @@ -18,9 +18,13 @@ package org.apache.beam.runners.kafka.streams.translation; import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.kafka.streams.Topology; /** @@ -66,23 +70,28 @@ public void translate( + " uses user state or timers; stateful ParDo is not yet supported by the Kafka" + " Streams runner."); } - if (transform.getOutputsMap().size() > 1) { - // Multi-output stages (DoFns with side outputs, etc.) are a planned follow-up — they need - // an output-tag dispatch in the processor + per-output PCollection routing. The current - // rejection just fails loudly until that's wired in. - throw new UnsupportedOperationException( - "ExecutableStage " - + transformId - + " has " - + transform.getOutputsMap().size() - + " outputs; multi-output stages are not yet supported by the Kafka Streams runner."); - } - // The payload distinguishes the main input from side inputs, so reading it from the payload // is unambiguous even before we add side-input support. String inputPCollectionId = stagePayload.getInput(); String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); + // A multi-output stage (a DoFn with side outputs, or a Read whose SDF wrapper produces several + // outputs) needs each output routed to the right downstream. Since downstream transforms are + // wired to a producer node by PCollection id, and that node must exist when the stage is + // translated, give each output its own relay node (StageOutputProcessor) and route to it by + // name. A single-output stage needs none of this — it forwards to its one downstream directly + // and registers itself as that output's producer. Outputs are sorted so the routing is + // deterministic across topology builds. + List outputPCollectionIds = new ArrayList<>(transform.getOutputsMap().values()); + Collections.sort(outputPCollectionIds); + boolean multiOutput = outputPCollectionIds.size() > 1; + Map outputChildByPCollectionId = new LinkedHashMap<>(); + if (multiOutput) { + for (int i = 0; i < outputPCollectionIds.size(); i++) { + outputChildByPCollectionId.put(outputPCollectionIds.get(i), transformId + "-output-" + i); + } + } + 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 @@ -96,12 +105,21 @@ public void translate( context.getJobInfo(), transformId, ImmutableSet.of(parentProcessor), - context.getMetricsContainerStepMap().getContainer(transformId)), + context.getMetricsContainerStepMap().getContainer(transformId), + outputChildByPCollectionId), parentProcessor); - if (!transform.getOutputsMap().isEmpty()) { - String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); - context.registerPCollectionProducer(outputPCollectionId, transformId); + if (multiOutput) { + // One relay per output; downstream transforms wire to the relay, which re-stamps the stage's + // watermark with the relay's own id so their watermark aggregation stays consistent. + outputChildByPCollectionId.forEach( + (outputPCollectionId, relayName) -> { + topology.addProcessor( + relayName, () -> new StageOutputProcessor(relayName), transformId); + context.registerPCollectionProducer(outputPCollectionId, relayName); + }); + } else if (!outputPCollectionIds.isEmpty()) { + context.registerPCollectionProducer(outputPCollectionIds.get(0), transformId); } } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java new file mode 100644 index 000000000000..eec3f2bae08a --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java @@ -0,0 +1,93 @@ +/* + * 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 org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * One output port of a multi-output {@link ExecutableStageProcessor}: a relay node that stands in + * as the producer of one of the stage's output PCollections. + * + *

    A Kafka Streams processor forwards a record either to all of its children or to one named + * child, but downstream transforms are wired to a producer node by PCollection id, and + * that node's identity has to be known when the stage is translated — before the downstream + * transforms are. So each output PCollection of a multi-output stage gets its own relay: the stage + * routes each harness output to the matching relay by name, and downstream transforms wire to the + * relay. (A single-output stage needs none of this and forwards directly.) + * + *

    The relay forwards data records unchanged, and on the watermark it relabels only the + * transform id — keeping the stage instance's source partition and total partition count intact. + * This is a 1:1 pass-through of one stage task's watermark (relay task {@code i} sees stage task + * {@code i}), so preserving the partition identity is what lets a downstream aggregator both know + * the report comes from this output (the relay's id) and still infer how many parallel instances of + * the stage produced it. The relay does no aggregation of its own. + */ +class StageOutputProcessor + implements Processor, byte[], KStreamsPayload> { + + private static final Logger LOG = LoggerFactory.getLogger(StageOutputProcessor.class); + + private final String transformId; + private @Nullable ProcessorContext> context; + + StageOutputProcessor(String transformId) { + this.transformId = transformId; + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + } + + @Override + public void process(Record> record) { + KStreamsPayload payload = record.value(); + ProcessorContext> ctx = context; + if (ctx == null) { + throw new IllegalStateException("StageOutputProcessor used before init()"); + } + if (payload == null) { + LOG.warn( + "Stage output {} dropping record with null payload (external write or tombstone)", + transformId); + return; + } + if (!payload.isWatermark()) { + // Data for this output: forward unchanged. + ctx.forward(record); + return; + } + // Relabel the transform id to this output port's, but keep the reporting stage instance's + // partition and partition count so downstream can still infer the stage's parallelism. + WatermarkPayload report = payload.asWatermark(); + ctx.forward( + new Record>( + record.key(), + KStreamsPayload.watermark( + report.getWatermarkMillis(), + transformId, + report.getSourcePartition(), + report.getTotalSourcePartitions()), + record.timestamp())); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/MultiOutputStageTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/MultiOutputStageTest.java new file mode 100644 index 000000000000..7fe85d732f2b --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/MultiOutputStageTest.java @@ -0,0 +1,87 @@ +/* + * 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; + +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Count; +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.sdk.values.KV; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.junit.Rule; +import org.junit.Test; + +/** + * Verifies that a fused executable stage with more than one output routes each output to its own + * downstream — including the case the mentor asked about, where each output is followed by a + * separate GroupByKey (so each output ends up in its own repartition topic). + * + *

    A DoFn splits its input to a main and an additional (side) output. One output is grouped with + * {@code Count.perElement} (a GroupByKey under the hood, so that output lands in its own + * repartition topic); the other is asserted directly (a plain in-process output). A {@link PAssert} + * on each confirms the two outputs carried the right elements to the right downstream. + */ +public class MultiOutputStageTest { + + private static final TupleTag EVENS = new TupleTag() {}; + private static final TupleTag ODDS = new TupleTag() {}; + + /** Routes even elements to the main output and odd elements to the side output. */ + private static class SplitByParityFn extends DoFn { + @ProcessElement + public void processElement(@Element Integer input, MultiOutputReceiver out) { + if (input % 2 == 0) { + out.get(EVENS).output(input); + } else { + out.get(ODDS).output(input); + } + } + } + + private static PipelineOptions options() { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(TestKafkaStreamsRunner.class); + return options; + } + + @Rule public final transient TestPipeline pipeline = TestPipeline.fromOptions(options()); + + @Test + public void eachStageOutputFeedsItsOwnGroupByKey() { + PCollectionTuple split = + pipeline + .apply("create", Create.of(1, 2, 3, 4, 5, 6)) + .apply( + "split", + ParDo.of(new SplitByParityFn()).withOutputTags(EVENS, TupleTagList.of(ODDS))); + + // The evens branch goes through a GroupByKey (Count.perElement), so that output lands in its + // own repartition topic; the odds branch is asserted directly, exercising a plain output. + PAssert.that(split.get(EVENS).apply("countEvens", Count.perElement())) + .containsInAnyOrder(KV.of(2, 1L), KV.of(4, 1L), KV.of(6, 1L)); + PAssert.that(split.get(ODDS)).containsInAnyOrder(1, 3, 5); + + pipeline.run(); + } +} 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 de3a23911adc..010d98d52e88 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 @@ -26,6 +26,7 @@ import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.kafka.streams.processor.api.MockProcessorContext; import org.apache.kafka.streams.processor.api.Record; @@ -59,7 +60,9 @@ private static ExecutableStageProcessor newProcessor() { jobInfo, STAGE_ID, ImmutableSet.of(UPSTREAM_ID), - new MetricsContainerImpl(STAGE_ID)); + new MetricsContainerImpl(STAGE_ID), + // Single-output: no per-output routing (this test drives the watermark path directly). + ImmutableMap.of()); } /** 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/StageOutputProcessorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java new file mode 100644 index 000000000000..83653f13c891 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java @@ -0,0 +1,108 @@ +/* + * 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 java.util.List; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.kafka.streams.processor.api.MockProcessorContext; +import org.apache.kafka.streams.processor.api.MockProcessorContext.CapturedForward; +import org.apache.kafka.streams.processor.api.Record; +import org.junit.Test; + +/** + * Unit tests for {@link StageOutputProcessor}, the per-output relay of a multi-output stage: it + * forwards data unchanged and relabels the watermark's transform id while preserving the reporting + * stage instance's partition and partition count. + * + *

    The partition preservation is what makes the watermark propagate correctly once the multi- + * output stage runs in several parallel instances (per je-ik's review): each stage task's report + * must reach downstream carrying its own {@code (partition, totalPartitions)}, so the downstream + * aggregator can tell how many parallel instances produced the output. A {@link + * MockProcessorContext} lets the reports from different stage partitions be fed directly, which a + * single-instance {@code TopologyTestDriver} cannot do. + */ +public class StageOutputProcessorTest { + + private static final String RELAY_ID = "stage-output-0"; + private static final String STAGE_ID = "stage"; + + private static Record> watermark( + long millis, int sourcePartition, int totalPartitions) { + return new Record<>( + new byte[0], + KStreamsPayload.watermark(millis, STAGE_ID, sourcePartition, totalPartitions), + 0L); + } + + @Test + public void watermarkKeepsPartitionIdentityAndRelabelsTransformId() { + MockProcessorContext> ctx = new MockProcessorContext<>(); + StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID); + processor.init(ctx); + + // A report from partition 1 of a 3-instance stage. + processor.process(watermark(500L, 1, 3)); + + List>> forwarded = + ctx.forwarded(); + assertThat(forwarded.size(), is(1)); + WatermarkPayload out = forwarded.get(0).record().value().asWatermark(); + assertThat(out.getWatermarkMillis(), is(500L)); + // Relabeled to the relay's id, but the stage instance's partition and count are preserved. + assertThat(out.getTransformId(), is(RELAY_ID)); + assertThat(out.getSourcePartition(), is(1)); + assertThat(out.getTotalSourcePartitions(), is(3)); + } + + @Test + public void distinctStagePartitionsStayDistinctDownstream() { + MockProcessorContext> ctx = new MockProcessorContext<>(); + StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID); + processor.init(ctx); + + processor.process(watermark(100L, 0, 3)); + processor.process(watermark(200L, 2, 3)); + + List>> forwarded = + ctx.forwarded(); + assertThat(forwarded.size(), is(2)); + assertThat(forwarded.get(0).record().value().asWatermark().getSourcePartition(), is(0)); + assertThat(forwarded.get(1).record().value().asWatermark().getSourcePartition(), is(2)); + assertThat(forwarded.get(0).record().value().asWatermark().getTotalSourcePartitions(), is(3)); + } + + @Test + public void dataIsForwardedUnchanged() { + MockProcessorContext> ctx = new MockProcessorContext<>(); + StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID); + processor.init(ctx); + + WindowedValue element = WindowedValues.valueInGlobalWindow(new byte[] {7}); + KStreamsPayload data = KStreamsPayload.data(element); + processor.process(new Record<>(new byte[0], data, 0L)); + + List>> forwarded = + ctx.forwarded(); + assertThat(forwarded.size(), is(1)); + assertThat(forwarded.get(0).record().value(), is(data)); + } +} From b615ae85aedfce799a94e4dc76152c13adc7f392 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Thu, 23 Jul 2026 13:36:14 +0500 Subject: [PATCH 21/37] [GSoC 2026] Kafka Streams runner: enable ParDoTest in the ValidatesRunner suite (#39451) --- runners/kafka-streams/build.gradle | 10 +++++++ .../kafka/streams/TestKafkaStreamsRunner.java | 27 ++++++++++++++++++- 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 52535c6654e6..2d0209eafbf5 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -92,6 +92,15 @@ def sickbayTests = [ 'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests', 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerLatest', 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerEarliest', + // A DoFn whose @StartBundle throws never gets to report its error: SdkHarnessClient.newBundle + // sends the ProcessBundleRequest and then blocks in GrpcDataService.createOutboundAggregator + // waiting for the SDK harness to open its data stream, which a bundle that failed during setup + // never does — so the run blocks for the data service's three-minute timeout instead of + // surfacing the user's exception. This is shared java-fn-execution behaviour rather than + // anything specific to this runner; the Flink runner sickbays all of LifecycleTests and the + // Prism runner sickbays each of its three error tests. The @ProcessElement and @FinishBundle + // variants do pass here, because by then the data stream is established. + 'org.apache.beam.sdk.transforms.ParDoTest$LifecycleTests.testParDoWithErrorInStartBatch', ] tasks.register("validatesRunner", Test) { @@ -143,6 +152,7 @@ tasks.register("validatesRunner", Test) { includeTestsMatching 'org.apache.beam.sdk.transforms.CreateTest' includeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest' includeTestsMatching 'org.apache.beam.sdk.transforms.GroupByKeyTest*' + includeTestsMatching 'org.apache.beam.sdk.transforms.ParDoTest*' for (String test : sickbayTests) { excludeTestsMatching test } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunner.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunner.java index 1fe4b310789f..eb1ebf707746 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunner.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/TestKafkaStreamsRunner.java @@ -20,6 +20,7 @@ import java.io.IOException; import java.util.UUID; import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.Pipeline.PipelineExecutionException; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.PipelineRunner; import org.apache.beam.sdk.metrics.MetricResults; @@ -47,6 +48,10 @@ * PAssert} that silently never runs is caught by {@code TestPipeline.verifyPAssertsSucceeded} * comparing the success-counter metric against the number of assertions in the pipeline. * + *

    Any other failed run is restated as a {@link PipelineExecutionException} carrying the + * exception the user's code threw, which is the failure Beam's contract for {@link #run} — and the + * tests that assert on a DoFn throwing — expect to see. + * *

    Select it with {@code --runner=org.apache.beam.runners.kafka.streams.TestKafkaStreamsRunner} * in {@code beamTestPipelineOptions}. */ @@ -84,11 +89,31 @@ public PipelineResult run(Pipeline pipeline) { throw (AssertionError) current; } } - throw t; + throw pipelineExecutionExceptionFor(t); } return new TestKafkaStreamsPipelineResult(metrics); } + /** + * Restates a failed run as the {@link PipelineExecutionException} Beam's contract for {@link + * #run} expects, carrying the exception the user's code threw. + * + *

    What surfaces from the run is the outermost layer of runner plumbing — a Kafka Streams + * {@code StreamsException} naming the processor node that failed — wrapping several layers of + * bundle- and Fn-API-level wrappers, with the user's exception at the bottom. Callers assert on + * the user's failure, so the root cause is what {@link PipelineExecutionException} is given: + * because {@code RuntimeException(Throwable)} derives its message from the cause, the user's + * message ends up on the exception that {@link #run} throws. Other runners restate failures the + * same way (e.g. {@code SparkPipelineResult}, {@code DirectRunner}). + */ + private static PipelineExecutionException pipelineExecutionExceptionFor(Throwable t) { + Throwable rootCause = t; + while (rootCause.getCause() != null && rootCause.getCause() != rootCause) { + rootCause = rootCause.getCause(); + } + return new PipelineExecutionException(rootCause); + } + /** Terminal result of a test run: state {@code DONE} with the harness-reported metrics. */ private static final class TestKafkaStreamsPipelineResult implements PipelineResult { private final MetricResults metrics; From 47614a943a7046b77056e9fd087b0c0482f24ac6 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:12:59 +0500 Subject: [PATCH 22/37] [GSoC 2026] Kafka Streams runner: windowed GroupByKey via ReduceFnRunner (#39494) * [GSoC 2026] Kafka Streams runner: windowed GroupByKey via ReduceFnRunner Replaces the global-window-only GroupByKey with a windowed one that drives Beam's ReduceFnRunner on the runner side, the way the Flink and Spark portable runners do, backed by Kafka Streams state and timers. WindowedGroupByKeyProcessor builds a ReduceFnRunner per key (like GroupAlsoByWindowViaWindowSetNewDoFn) over two new backends: KafkaStreamsStateInternals, which stores each Beam state cell as one entry in a KeyValueStore under a composite key of key + namespace + tag (modelled on SparkStateInternals), and KafkaStreamsTimerInternals, which persists timers keyed by identity and is fired by the processor scanning for due event-time timers on each input-watermark advance. GroupByKeyTranslator hydrates the input windowing strategy from the pipeline proto and wires the state and timer stores. Windowing, the default trigger, panes, allowed lateness and timestamp combiners all come from ReduceFnRunner. --- runners/kafka-streams/build.gradle | 11 +- .../translation/GroupByKeyProcessor.java | 224 --------- .../translation/GroupByKeyTranslator.java | 78 ++- .../KafkaStreamsStateInternals.java | 463 ++++++++++++++++++ .../KafkaStreamsTimerInternals.java | 265 ++++++++++ .../kafka/streams/translation/StoreKeys.java | 105 ++++ .../WindowedGroupByKeyProcessor.java | 323 ++++++++++++ .../kafka/streams/KafkaStreamsTestRunner.java | 101 +--- .../FixedWindowGroupByKeyTest.java | 106 ++++ .../KafkaStreamsTimerInternalsTest.java | 208 ++++++++ 10 files changed, 1552 insertions(+), 332 deletions(-) delete mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StoreKeys.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FixedWindowGroupByKeyTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternalsTest.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 2d0209eafbf5..bdf7e3be0585 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -87,11 +87,11 @@ dependencies { // Known-failing @ValidatesRunner tests, excluded until the feature they need lands. def sickbayTests = [ - // Non-global windowing (FixedWindows, merging windows, timestamp combiners) is not supported - // yet; these apply a window and assert on window-derived output, hitting a GlobalWindow cast. - 'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests', - 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerLatest', - 'org.apache.beam.sdk.transforms.GroupByKeyTest$BasicTests.testTimestampCombinerEarliest', + // Merging (session) windows are not supported yet: ReduceFnRunner drives them through a merging + // window set that moves per-window state as windows merge, which this first windowing pass does + // not implement. Non-merging windows (fixed, sliding), the default trigger and timestamp + // combiners do work. Lands with the follow-up windowing PR. + 'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests.testGroupByKeyMergingWindows', // A DoFn whose @StartBundle throws never gets to report its error: SdkHarnessClient.newBundle // sends the ProcessBundleRequest and then blocks in GrpcDataService.createOutboundAggregator // waiting for the SDK harness to open its data stream, which a bundle that failed during setup @@ -100,6 +100,7 @@ def sickbayTests = [ // anything specific to this runner; the Flink runner sickbays all of LifecycleTests and the // Prism runner sickbays each of its three error tests. The @ProcessElement and @FinishBundle // variants do pass here, because by then the data stream is established. + // Tracked by https://github.com/apache/beam/issues/39452. 'org.apache.beam.sdk.transforms.ParDoTest$LifecycleTests.testParDoWithErrorInStartBatch', ] diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java deleted file mode 100644 index 3e82935b807d..000000000000 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java +++ /dev/null @@ -1,224 +0,0 @@ -/* - * 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 java.util.ArrayList; -import java.util.List; -import java.util.Set; -import org.apache.beam.sdk.coders.Coder; -import org.apache.beam.sdk.coders.CoderException; -import org.apache.beam.sdk.coders.IterableCoder; -import org.apache.beam.sdk.transforms.windowing.BoundedWindow; -import org.apache.beam.sdk.transforms.windowing.GlobalWindow; -import org.apache.beam.sdk.util.CoderUtils; -import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.WindowedValue; -import org.apache.beam.sdk.values.WindowedValues; -import org.apache.kafka.streams.processor.api.Processor; -import org.apache.kafka.streams.processor.api.ProcessorContext; -import org.apache.kafka.streams.processor.api.Record; -import org.apache.kafka.streams.state.KeyValueIterator; -import org.apache.kafka.streams.state.KeyValueStore; -import org.checkerframework.checker.nullness.qual.Nullable; -import org.joda.time.Instant; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Executes a {@code GroupByKey} (GlobalWindow, default trigger, no allowed lateness). - * - *

    Records arrive on the repartition topic keyed by the encoded Beam key, so every value of a key - * is co-located here. Each value is appended to a per-key buffer in a Kafka Streams state store. - * Watermark reports are fed to a {@link WatermarkAggregator}; when the input watermark reaches - * {@link BoundedWindow#TIMESTAMP_MAX_VALUE} (the end of the global window) every buffered key is - * emitted once as {@code KV>} and the buffer cleared, then the watermark is - * forwarded downstream. - * - *

    Buffering whole value lists and re-encoding on each append is O(n^2) per key; fine for this - * first GroupByKey, and replaced when this moves to runner-core {@code GroupAlsoByWindow}. - */ -class GroupByKeyProcessor - implements Processor, byte[], KStreamsPayload> { - - private static final Logger LOG = LoggerFactory.getLogger(GroupByKeyProcessor.class); - - private final String stateStoreName; - // This transform's own id, stamped on every watermark it forwards downstream. - private final String transformId; - private final Coder keyCoder; - private final IterableCoder<@Nullable Object> bufferCoder; - - // Aggregates the input watermark from the upstream transform's reports, which arrive through the - // repartition topic with the upstream producer's transform id intact (the shuffle forwards - // watermark payloads unchanged). - private final WatermarkAggregator watermarkAggregator; - private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; - // The global window fires exactly once, when the watermark first reaches its end. Later watermark - // reports (e.g. the same terminal watermark broadcast across repartition partitions) must not - // re-fire. This flag is in-memory only; restart correctness comes from the state store plus - // exactly-once-v2: the buffered values and consumer offsets are committed atomically, and the - // store is empty once a key has fired, so a restart cannot double-emit. Persisting watermark - // holds is part of the separate WatermarkManager persistence work, not this initial GroupByKey. - private boolean fired = false; - - private @Nullable ProcessorContext> context; - private @Nullable KeyValueStore store; - - /** - * @param transformId this transform's own id, stamped on the watermarks it emits - * @param upstreamTransformIds the transform ids feeding this GroupByKey (known from the pipeline - * graph), whose reports the {@link WatermarkAggregator} waits for - */ - GroupByKeyProcessor( - String stateStoreName, - String transformId, - Set upstreamTransformIds, - Coder keyCoder, - Coder<@Nullable Object> valueCoder) { - this.stateStoreName = stateStoreName; - this.transformId = transformId; - this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); - this.keyCoder = keyCoder; - this.bufferCoder = IterableCoder.of(valueCoder); - } - - @Override - public void init(ProcessorContext> context) { - this.context = context; - this.store = context.getStateStore(stateStoreName); - } - - @Override - public void process(Record> record) { - KStreamsPayload payload = record.value(); - if (payload == null) { - // The repartition topic can be written to from outside the runner (or carry a tombstone), - // so recover from the obvious error instead of crashing the task: warn and drop. - LOG.warn( - "GroupByKey {} dropping record with null payload (external write or tombstone)", - transformId); - return; - } - if (payload.isData()) { - byte[] encodedKey = record.key(); - Object element = payload.getData().getValue(); - if (encodedKey == null || element == null) { - throw new IllegalStateException("GroupByKey data record is missing its key or value"); - } - appendValue(encodedKey, element); - return; - } - watermarkAggregator.observe(payload.asWatermark()); - Instant advanced = watermarkAggregator.advance(); - if (!fired && !advanced.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) { - fireAll(record); - fired = true; - } - if (advanced.isAfter(lastForwardedWatermark)) { - lastForwardedWatermark = advanced; - forwardWatermark(record, advanced.getMillis()); - } - } - - private void appendValue(byte[] encodedKey, Object kvObject) { - KV kv = (KV) kvObject; - KeyValueStore kvStore = checkInitialized(store); - byte[] existing = kvStore.get(encodedKey); - List<@Nullable Object> values = existing == null ? new ArrayList<>() : decodeBuffer(existing); - values.add(kv.getValue()); - kvStore.put(encodedKey, encodeBuffer(values)); - } - - private void fireAll(Record> trigger) { - // NOTE: this emits every buffered key in a single watermark turn. For a very large key space - // that risks memory pressure and exceeding the poll / transaction timeout. Acceptable for this - // initial GlobalWindow GroupByKey (fire once at end of input); incremental, timer-driven output - // via runner-core GroupAlsoByWindow lands with the windowing/timers work. - ProcessorContext> ctx = checkInitialized(context); - KeyValueStore kvStore = checkInitialized(store); - List firedKeys = new ArrayList<>(); - try (KeyValueIterator it = kvStore.all()) { - while (it.hasNext()) { - org.apache.kafka.streams.KeyValue entry = it.next(); - Object key = decodeKey(entry.key); - List<@Nullable Object> values = decodeBuffer(entry.value); - // The pane fires at the end of the global window, so the grouped element carries the - // window's max timestamp (END_OF_GLOBAL_WINDOW). Emitting at TIMESTAMP_MIN_VALUE (the - // default of valueInGlobalWindow) would make the output appear arbitrarily late and be - // dropped downstream once the watermark has advanced. - WindowedValue>> output = - WindowedValues.timestampedValueInGlobalWindow( - KV.of(key, (Iterable<@Nullable Object>) values), - GlobalWindow.INSTANCE.maxTimestamp()); - ctx.forward( - new Record>( - entry.key, KStreamsPayload.data(output), trigger.timestamp())); - firedKeys.add(entry.key); - } - } - for (byte[] key : firedKeys) { - kvStore.delete(key); - } - } - - private void forwardWatermark(Record> trigger, long watermarkMillis) { - ProcessorContext> ctx = checkInitialized(context); - // Stamped with this transform's own id; GroupByKey is a single instance for now, so the report - // is for its only partition (0 of 1). - ctx.forward( - new Record>( - trigger.key(), - KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), - trigger.timestamp())); - } - - private byte[] encodeBuffer(List<@Nullable Object> values) { - try { - return CoderUtils.encodeToByteArray(bufferCoder, values); - } catch (CoderException e) { - throw new RuntimeException("Failed to encode GroupByKey value buffer", e); - } - } - - private List<@Nullable Object> decodeBuffer(byte[] bytes) { - try { - List<@Nullable Object> values = new ArrayList<>(); - for (@Nullable Object value : CoderUtils.decodeFromByteArray(bufferCoder, bytes)) { - values.add(value); - } - return values; - } catch (CoderException e) { - throw new RuntimeException("Failed to decode GroupByKey value buffer", e); - } - } - - private Object decodeKey(byte[] bytes) { - try { - return CoderUtils.decodeFromByteArray(keyCoder, bytes); - } catch (CoderException e) { - throw new RuntimeException("Failed to decode GroupByKey key", e); - } - } - - private static T checkInitialized(@Nullable T value) { - if (value == null) { - throw new IllegalStateException("GroupByKeyProcessor used before init()"); - } - return value; - } -} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index 9e23dbb5cfb0..c5327e28e069 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -22,8 +22,12 @@ import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.RehydratedComponents; +import org.apache.beam.sdk.util.construction.WindowingStrategyTranslation; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.kafka.common.serialization.Serdes; @@ -35,10 +39,11 @@ * Translates the {@code beam:transform:group_by_key:v1} URN — the runner's first stateful, * shuffle-bearing transform. * - *

    This is the simplest GroupByKey: GlobalWindow, default trigger, no allowed lateness (per the - * plan agreed with the mentor). Each key's values are buffered in a Kafka Streams state store and - * emitted once as {@code KV>} when the watermark reaches {@link - * org.apache.beam.sdk.transforms.windowing.BoundedWindow#TIMESTAMP_MAX_VALUE}. + *

    Windowing and triggering are executed by Beam's {@link + * org.apache.beam.runners.core.ReduceFnRunner} inside {@link WindowedGroupByKeyProcessor}, the same + * way the Flink and Spark portable runners do it — so fixed/sliding windows, the default trigger, + * allowed lateness and timestamp combiners all work. The input PCollection's windowing strategy is + * hydrated from the pipeline proto and handed to the processor. * *

    Topology added (the Beam key becomes the Kafka record key so Kafka Streams shuffles by it): * @@ -49,7 +54,8 @@ * via {@link KStreamsPayloadSerde} and a {@link GroupByKeyBroadcastPartitioner} that hashes * data by key and fans watermark reports out to every partition; *

  • a {@link Topology#addSource source} reading the repartition topic back; - *
  • the {@link GroupByKeyProcessor} plus a persistent state store, wired to the source. + *
  • the {@link WindowedGroupByKeyProcessor} plus persistent state and timer stores, wired to + * the source. * * *

    The repartition topic is expected to exist on the broker before the job starts (same @@ -62,6 +68,9 @@ class GroupByKeyTranslator implements PTransformTranslator { static final String SINK_SUFFIX = "-repartition-sink"; static final String SOURCE_SUFFIX = "-repartition-source"; static final String STATE_STORE_SUFFIX = "-state"; + static final String HOLDS_INDEX_STORE_SUFFIX = "-holds-index"; + static final String TIMER_STORE_SUFFIX = "-timers"; + static final String TIMER_INDEX_STORE_SUFFIX = "-timers-index"; static final String REPARTITION_TOPIC_PREFIX = "__beam_gbk_"; @Override @@ -82,12 +91,18 @@ public void translate( Coder<@Nullable Object> valueCoder = (Coder<@Nullable Object>) (Coder) kvCoder.getValueCoder(); + WindowingStrategy windowingStrategy = + hydrateWindowingStrategy(pipeline, inputPCollectionId); + String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); String shuffleName = transformId + SHUFFLE_SUFFIX; String sinkName = transformId + SINK_SUFFIX; String sourceName = transformId + SOURCE_SUFFIX; String stateStoreName = transformId + STATE_STORE_SUFFIX; + String holdsIndexStoreName = transformId + HOLDS_INDEX_STORE_SUFFIX; + String timerStoreName = transformId + TIMER_STORE_SUFFIX; + String timerIndexStoreName = transformId + TIMER_INDEX_STORE_SUFFIX; String repartitionTopic = repartitionTopic(transformId); KStreamsPayloadSerde> payloadSerde = new KStreamsPayloadSerde<>(inputCoder); @@ -111,27 +126,70 @@ public void translate( payloadSerde.deserializer(), repartitionTopic); - // Buffer values per key and fire KV> at the terminal watermark. Watermark - // reports cross the repartition topic unchanged, so they still carry the id of the transform - // that produced this GroupByKey's input — the parent the shuffle is attached to. + // Group by key and window through Beam's ReduceFnRunner, backed by the state and timer stores. + // Watermark reports cross the repartition topic unchanged, so they still carry the id of the + // transform that produced this GroupByKey's input — the parent the shuffle is attached to. topology.addProcessor( transformId, () -> - new GroupByKeyProcessor( + new WindowedGroupByKeyProcessor( stateStoreName, + holdsIndexStoreName, + timerStoreName, + timerIndexStoreName, transformId, ImmutableSet.of(parentProcessor), keyCoder, - valueCoder), + valueCoder, + windowingStrategy, + context.getPipelineOptions()), sourceName); topology.addStateStore( Stores.keyValueStoreBuilder( Stores.persistentKeyValueStore(stateStoreName), Serdes.ByteArray(), Serdes.ByteArray()), transformId); + topology.addStateStore( + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(timerStoreName), Serdes.ByteArray(), Serdes.ByteArray()), + transformId); + // Indexes ordered by timestamp, so due timers and the minimum watermark hold are range scans + // rather than scans of every timer or every held window. + topology.addStateStore( + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(timerIndexStoreName), + Serdes.ByteArray(), + Serdes.ByteArray()), + transformId); + topology.addStateStore( + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(holdsIndexStoreName), + Serdes.ByteArray(), + Serdes.ByteArray()), + transformId); context.registerPCollectionProducer(outputPCollectionId, transformId); } + /** Hydrates the input PCollection's windowing strategy from the pipeline proto. */ + private static WindowingStrategy hydrateWindowingStrategy( + RunnerApi.Pipeline pipeline, String inputPCollectionId) { + RunnerApi.Components components = pipeline.getComponents(); + String windowingStrategyId = + components.getPcollectionsOrThrow(inputPCollectionId).getWindowingStrategyId(); + try { + @SuppressWarnings("unchecked") + WindowingStrategy strategy = + (WindowingStrategy) + WindowingStrategyTranslation.fromProto( + components.getWindowingStrategiesOrThrow(windowingStrategyId), + RehydratedComponents.forComponents(components)); + return strategy; + } catch (Exception e) { + throw new IllegalStateException( + "Failed to hydrate GroupByKey windowing strategy " + windowingStrategyId, e); + } + } + /** The internal repartition topic name for a GroupByKey transform. */ static String repartitionTopic(String transformId) { return REPARTITION_TOPIC_PREFIX + transformId.replaceAll("[^a-zA-Z0-9._-]", "_"); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java new file mode 100644 index 000000000000..c17d7c45ad4d --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsStateInternals.java @@ -0,0 +1,463 @@ +/* + * 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 java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateTag; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.InstantCoder; +import org.apache.beam.sdk.coders.ListCoder; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.CombiningState; +import org.apache.beam.sdk.state.MapState; +import org.apache.beam.sdk.state.MultimapState; +import org.apache.beam.sdk.state.OrderedListState; +import org.apache.beam.sdk.state.ReadableState; +import org.apache.beam.sdk.state.SetState; +import org.apache.beam.sdk.state.State; +import org.apache.beam.sdk.state.StateBinder; +import org.apache.beam.sdk.state.StateContext; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.state.WatermarkHoldState; +import org.apache.beam.sdk.transforms.Combine.CombineFn; +import org.apache.beam.sdk.transforms.CombineWithContext; +import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.CombineFnUtil; +import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * A {@link StateInternals} for one key, backed by a Kafka Streams {@link KeyValueStore}. + * + *

    Beam addresses a state cell by {@code (key, StateNamespace, StateTag)}; a windowed pipeline + * puts each window's state in its own namespace. Every cell is stored as one entry in the shared + * per-transform store under a composite byte key {@code len(key)|key | len(ns)|ns | len(tag)|tag}, + * so all cells for one Beam key share a prefix and a whole key's state can be range-scanned. The + * value is the cell's contents encoded with its Beam {@link Coder}. Writing straight to the store + * (rather than buffering and flushing) keeps this restart-safe for free: the store is changelogged + * and, under exactly-once, its writes commit atomically with the input offsets. + * + *

    Modeled on the Spark runner's {@code SparkStateInternals}; the difference is that each cell + * reads and writes its own store entry instead of an in-memory table, so there is no separate + * persist step. + */ +class KafkaStreamsStateInternals implements StateInternals { + + /** The holds index is a set; only its keys carry information. */ + private static final byte[] EMPTY_VALUE = new byte[0]; + + /** + * Reads the minimum watermark hold held by any key and window, or {@code null} if none is held. + * The index is ordered by hold time, so this is the first entry rather than a scan. + */ + static @Nullable Instant minWatermarkHold(KeyValueStore holdsIndexStore) { + try (KeyValueIterator it = holdsIndexStore.all()) { + if (!it.hasNext()) { + return null; + } + return new Instant(StoreKeys.readTimestamp(it.next().key, 0)); + } + } + + private final @NonNull K key; + private final byte[] encodedKey; + private final KeyValueStore store; + private final KeyValueStore holdsIndexStore; + + /** + * The last namespace a composite key was built for, and the {@code key | namespace} prefix that + * was built for it. One turn of the windowing runner touches several tags in the same namespace + * back to back (read the buffer, read the hold, write both), so caching the prefix removes most + * of the per-access encoding work. + */ + private @Nullable StateNamespace cachedNamespace; + + private byte @Nullable [] cachedPrefix; + + KafkaStreamsStateInternals( + @NonNull K key, + byte[] encodedKey, + KeyValueStore store, + KeyValueStore holdsIndexStore) { + this.key = key; + this.encodedKey = encodedKey; + this.store = store; + this.holdsIndexStore = holdsIndexStore; + } + + @Override + public Object getKey() { + return key; + } + + @Override + public T state( + StateNamespace namespace, StateTag address, StateContext c) { + return address.getSpec().bind(address.getId(), new KafkaStreamsStateBinder(namespace, c)); + } + + /** + * The composite store key for one cell: {@code len|key len|namespace len|tagId}. + * + *

    Built into one exactly-sized array, reusing the cached {@code key | namespace} prefix. This + * runs on every state access, so it avoids the repeated growth and final copy a stream would do. + */ + private byte[] compositeKey(StateNamespace namespace, String id) { + byte[] prefix = prefixFor(namespace); + byte[] idBytes = id.getBytes(StandardCharsets.UTF_8); + byte[] compositeKey = new byte[prefix.length + StoreKeys.segmentLength(idBytes)]; + System.arraycopy(prefix, 0, compositeKey, 0, prefix.length); + StoreKeys.writeSegment(compositeKey, prefix.length, idBytes); + return compositeKey; + } + + /** The {@code key | namespace} prefix every cell in {@code namespace} starts with. */ + private byte[] prefixFor(StateNamespace namespace) { + byte[] cached = cachedPrefix; + if (cached != null && namespace.equals(cachedNamespace)) { + return cached; + } + byte[] namespaceBytes = namespace.stringKey().getBytes(StandardCharsets.UTF_8); + byte[] prefix = + new byte[StoreKeys.segmentLength(encodedKey) + StoreKeys.segmentLength(namespaceBytes)]; + int offset = StoreKeys.writeSegment(prefix, 0, encodedKey); + StoreKeys.writeSegment(prefix, offset, namespaceBytes); + cachedNamespace = namespace; + cachedPrefix = prefix; + return prefix; + } + + private class KafkaStreamsStateBinder implements StateBinder { + private final StateNamespace namespace; + private final StateContext stateContext; + + private KafkaStreamsStateBinder(StateNamespace namespace, StateContext stateContext) { + this.namespace = namespace; + this.stateContext = stateContext; + } + + @Override + public ValueState bindValue(String id, StateSpec> spec, Coder coder) { + return new KafkaStreamsValueState<>(namespace, id, coder); + } + + @Override + public BagState bindBag(String id, StateSpec> spec, Coder elemCoder) { + return new KafkaStreamsBagState<>(namespace, id, elemCoder); + } + + @Override + public SetState bindSet(String id, StateSpec> spec, Coder elemCoder) { + throw new UnsupportedOperationException( + SetState.class.getSimpleName() + " is not supported by the Kafka Streams runner yet"); + } + + @Override + public MapState bindMap( + String id, + StateSpec> spec, + Coder mapKeyCoder, + Coder mapValueCoder) { + throw new UnsupportedOperationException( + MapState.class.getSimpleName() + " is not supported by the Kafka Streams runner yet"); + } + + @Override + public MultimapState bindMultimap( + String id, + StateSpec> spec, + Coder keyCoder, + Coder valueCoder) { + throw new UnsupportedOperationException( + MultimapState.class.getSimpleName() + + " is not supported by the Kafka Streams runner yet"); + } + + @Override + public OrderedListState bindOrderedList( + String id, StateSpec> spec, Coder elemCoder) { + throw new UnsupportedOperationException( + OrderedListState.class.getSimpleName() + + " is not supported by the Kafka Streams runner yet"); + } + + @Override + public CombiningState bindCombining( + String id, + StateSpec> spec, + Coder accumCoder, + CombineFn combineFn) { + return new KafkaStreamsCombiningState<>(namespace, id, accumCoder, combineFn); + } + + @Override + public + CombiningState bindCombiningWithContext( + String id, + StateSpec> spec, + Coder accumCoder, + CombineWithContext.CombineFnWithContext combineFn) { + return new KafkaStreamsCombiningState<>( + namespace, id, accumCoder, CombineFnUtil.bindContext(combineFn, stateContext)); + } + + @Override + public WatermarkHoldState bindWatermark( + String id, StateSpec spec, TimestampCombiner timestampCombiner) { + return new KafkaStreamsWatermarkHoldState(namespace, id, timestampCombiner); + } + } + + /** Common read/write/clear against the backing store for one cell. */ + private abstract class AbstractState { + final StateNamespace namespace; + final String id; + final Coder coder; + + AbstractState(StateNamespace namespace, String id, Coder coder) { + this.namespace = namespace; + this.id = id; + this.coder = coder; + } + + @Nullable + T readValue() { + byte[] bytes = store.get(compositeKey(namespace, id)); + if (bytes == null) { + return null; + } + try { + return CoderUtils.decodeFromByteArray(coder, bytes); + } catch (CoderException e) { + throw new RuntimeException("Failed to decode state " + id, e); + } + } + + void writeValue(T input) { + try { + store.put(compositeKey(namespace, id), CoderUtils.encodeToByteArray(coder, input)); + } catch (CoderException e) { + throw new RuntimeException("Failed to encode state " + id, e); + } + } + + public void clear() { + store.delete(compositeKey(namespace, id)); + } + + ReadableState isEmptyState() { + return new ReadableState() { + @Override + public Boolean read() { + return store.get(compositeKey(namespace, id)) == null; + } + + @Override + public ReadableState readLater() { + return this; + } + }; + } + } + + private class KafkaStreamsValueState extends AbstractState implements ValueState { + KafkaStreamsValueState(StateNamespace namespace, String id, Coder coder) { + super(namespace, id, coder); + } + + @Override + public KafkaStreamsValueState readLater() { + return this; + } + + @Override + public @Nullable T read() { + return readValue(); + } + + @Override + public void write(T input) { + writeValue(input); + } + } + + private class KafkaStreamsBagState extends AbstractState> implements BagState { + KafkaStreamsBagState(StateNamespace namespace, String id, Coder elemCoder) { + super(namespace, id, ListCoder.of(elemCoder)); + } + + @Override + public KafkaStreamsBagState readLater() { + return this; + } + + @Override + public Iterable read() { + List value = readValue(); + return value == null ? new ArrayList<>() : value; + } + + @Override + public void add(T input) { + List value = readValue(); + if (value == null) { + value = new ArrayList<>(); + } + value.add(input); + writeValue(value); + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + } + + private class KafkaStreamsWatermarkHoldState extends AbstractState + implements WatermarkHoldState { + private final TimestampCombiner timestampCombiner; + + KafkaStreamsWatermarkHoldState( + StateNamespace namespace, String id, TimestampCombiner timestampCombiner) { + super(namespace, id, InstantCoder.of()); + this.timestampCombiner = timestampCombiner; + } + + @Override + public KafkaStreamsWatermarkHoldState readLater() { + return this; + } + + // GroupingState.read() is typed non-null, but an empty hold reads back null. Beam's state + // interfaces are under-annotated here (https://github.com/apache/beam/issues/20497), which is + // why the Spark and Flink StateInternals suppress nullness for the whole class; this runner + // narrows the suppression to just this method. + @Override + @SuppressWarnings("nullness") + public Instant read() { + return readValue(); + } + + @Override + public void add(Instant outputTime) { + Instant current = readValue(); + Instant combined = + current == null ? outputTime : timestampCombiner.combine(current, outputTime); + writeValue(combined); + // Mirror the hold into the index so the processor can find the minimum hold across every key + // and window with one lookup instead of reading all of them. + if (current != null) { + holdsIndexStore.delete(holdIndexKey(current)); + } + holdsIndexStore.put(holdIndexKey(combined), EMPTY_VALUE); + } + + @Override + public void clear() { + Instant current = readValue(); + if (current != null) { + holdsIndexStore.delete(holdIndexKey(current)); + } + super.clear(); + } + + /** {@code holdTimestamp | cell}, so the index is ordered by hold time. */ + private byte[] holdIndexKey(Instant hold) { + byte[] cellKey = compositeKey(namespace, id); + byte[] indexKey = new byte[StoreKeys.TIMESTAMP_BYTES + cellKey.length]; + int offset = StoreKeys.writeTimestamp(indexKey, 0, hold.getMillis()); + System.arraycopy(cellKey, 0, indexKey, offset, cellKey.length); + return indexKey; + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + + @Override + public TimestampCombiner getTimestampCombiner() { + return timestampCombiner; + } + } + + @SuppressWarnings("TypeParameterShadowing") + private class KafkaStreamsCombiningState extends AbstractState + implements CombiningState { + private final CombineFn combineFn; + + KafkaStreamsCombiningState( + StateNamespace namespace, + String id, + Coder accumCoder, + CombineFn combineFn) { + super(namespace, id, accumCoder); + this.combineFn = combineFn; + } + + @Override + public KafkaStreamsCombiningState readLater() { + return this; + } + + // GroupingState.read() is typed non-null but a CombineFn may extract a null output; the same + // under-annotation as WatermarkHoldState.read() (https://github.com/apache/beam/issues/20497). + @Override + @SuppressWarnings("nullness") + public OutputT read() { + return combineFn.extractOutput(getAccum()); + } + + @Override + public void add(InputT input) { + writeValue(combineFn.addInput(getAccum(), input)); + } + + @Override + public AccumT getAccum() { + AccumT accum = readValue(); + return accum == null ? combineFn.createAccumulator() : accum; + } + + @Override + public void addAccum(AccumT accum) { + writeValue(combineFn.mergeAccumulators(Arrays.asList(getAccum(), accum))); + } + + @Override + public AccumT mergeAccumulators(Iterable accumulators) { + return combineFn.mergeAccumulators(accumulators); + } + + @Override + public ReadableState isEmpty() { + return isEmptyState(); + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java new file mode 100644 index 000000000000..dddc28eb4381 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java @@ -0,0 +1,265 @@ +/* + * 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 java.nio.charset.StandardCharsets; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * A {@link TimerInternals} for one key, backed by two Kafka Streams stores shared by a GroupByKey. + * + *

    Kafka Streams has no per-key timer service, so timers are persisted like any other state, in + * two stores that serve the two ways a timer is looked up: + * + *

      + *
    • the identity store, keyed by {@code key | domain | timerFamily | timerId | + * namespace}, is how {@link #setTimer} overwrites and {@link #deleteTimer} removes exactly + * one timer, as {@link TimerInternals}' contract requires. Its value is the index key below, + * so a timer that is overwritten or deleted can have its index entry removed without knowing + * what time it had been set for. + *
    • the index store, keyed by {@code domain | fireTimestamp | identity}, is how due + * timers are found. Because the timestamp is written in the sortable form described on {@link + * StoreKeys}, all event-time timers due at a watermark are one range scan — {@link + * #dueEventTimeRangeStart} to {@link #dueEventTimeRangeEnd} — rather than a scan of every + * timer of every key. Its value is the {@link TimerData}, so firing needs no second lookup. + *
    + * + *

    Firing is driven by {@link WindowedGroupByKeyProcessor}: on a watermark advance it range-scans + * the index for event-time timers that are due and replays them through {@link + * org.apache.beam.runners.core.ReduceFnRunner#onTimers}. + * + *

    This instance reports the times it was constructed with; it never fires timers itself. + */ +class KafkaStreamsTimerInternals implements TimerInternals { + + private final byte[] encodedKey; + private final KeyValueStore identityStore; + private final KeyValueStore indexStore; + private final TimerInternals.TimerDataCoderV2 timerCoder; + private final Instant inputWatermarkTime; + private final Instant outputWatermarkTime; + private final Instant processingTime; + + KafkaStreamsTimerInternals( + byte[] encodedKey, + KeyValueStore identityStore, + KeyValueStore indexStore, + Coder windowCoder, + Instant inputWatermarkTime, + Instant outputWatermarkTime, + Instant processingTime) { + this.encodedKey = encodedKey; + this.identityStore = identityStore; + this.indexStore = indexStore; + this.timerCoder = TimerInternals.TimerDataCoderV2.of(windowCoder); + this.inputWatermarkTime = inputWatermarkTime; + this.outputWatermarkTime = outputWatermarkTime; + this.processingTime = processingTime; + } + + @Override + public void setTimer( + StateNamespace namespace, + String timerId, + String timerFamilyId, + Instant target, + Instant outputTimestamp, + TimeDomain timeDomain) { + setTimer(TimerData.of(timerId, timerFamilyId, namespace, target, outputTimestamp, timeDomain)); + } + + @Override + public void setTimer(TimerData timerData) { + byte[] identityKey = identityKey(encodedKey, timerData); + // Setting a timer that already exists replaces it, so drop the old index entry first — + // otherwise the timer would still be due at the time it was originally set for. + byte[] previousIndexKey = identityStore.get(identityKey); + if (previousIndexKey != null) { + indexStore.delete(previousIndexKey); + } + byte[] indexKey = + indexKey(timerData.getDomain(), timerData.getTimestamp().getMillis(), identityKey); + identityStore.put(identityKey, indexKey); + indexStore.put(indexKey, encodeTimer(timerData)); + } + + @Override + public void deleteTimer( + StateNamespace namespace, String timerId, String timerFamilyId, TimeDomain timeDomain) { + deleteByIdentity(identityKey(encodedKey, timerId, timerFamilyId, timeDomain, namespace)); + } + + @Override + public void deleteTimer(StateNamespace namespace, String timerId, String timerFamilyId) { + throw new UnsupportedOperationException( + "Deleting a timer without a time domain is not supported; the domain is part of a timer's" + + " store identity."); + } + + @Override + public void deleteTimer(TimerData timerKey) { + deleteByIdentity(identityKey(encodedKey, timerKey)); + } + + private void deleteByIdentity(byte[] identityKey) { + byte[] indexKey = identityStore.get(identityKey); + if (indexKey != null) { + indexStore.delete(indexKey); + } + identityStore.delete(identityKey); + } + + @Override + public Instant currentProcessingTime() { + return processingTime; + } + + /** + * Returns {@code null}: a synchronized processing time is the slowest processing time across the + * job's workers, which needs the cross-instance coordination that the runner's watermark reports + * only carry for event time. {@link TimerInternals} allows null here, and nothing on the paths + * this runner supports today reads it — it is consulted for processing-time triggers, which land + * with the processing-time timer support in a follow-up (the same work that would supply it). + */ + @Override + public @Nullable Instant currentSynchronizedProcessingTime() { + return null; + } + + @Override + public Instant currentInputWatermarkTime() { + return inputWatermarkTime; + } + + /** + * The watermark this GroupByKey has last forwarded downstream, which trails {@link + * #currentInputWatermarkTime} by the pending watermark holds. + */ + @Override + public Instant currentOutputWatermarkTime() { + return outputWatermarkTime; + } + + private byte[] encodeTimer(TimerData timerData) { + try { + return CoderUtils.encodeToByteArray(timerCoder, timerData); + } catch (CoderException e) { + throw new RuntimeException("Failed to encode timer " + timerData, e); + } + } + + /** Decodes an index store value back into its timer. */ + static TimerData decodeTimer(Coder windowCoder, byte[] bytes) { + try { + return CoderUtils.decodeFromByteArray(TimerInternals.TimerDataCoderV2.of(windowCoder), bytes); + } catch (CoderException e) { + throw new RuntimeException("Failed to decode timer", e); + } + } + + /** + * The identity store key for a timer: {@code key | domain | timerFamily | timerId | namespace}. + */ + static byte[] identityKey(byte[] encodedKey, TimerData timerData) { + return identityKey( + encodedKey, + timerData.getTimerId(), + timerData.getTimerFamilyId(), + timerData.getDomain(), + timerData.getNamespace()); + } + + static byte[] identityKey( + byte[] encodedKey, + String timerId, + String timerFamilyId, + TimeDomain domain, + StateNamespace namespace) { + byte[] domainBytes = {(byte) domain.ordinal()}; + byte[] familyBytes = timerFamilyId.getBytes(StandardCharsets.UTF_8); + byte[] idBytes = timerId.getBytes(StandardCharsets.UTF_8); + byte[] namespaceBytes = namespace.stringKey().getBytes(StandardCharsets.UTF_8); + byte[] key = + new byte + [StoreKeys.segmentLength(encodedKey) + + StoreKeys.segmentLength(domainBytes) + + StoreKeys.segmentLength(familyBytes) + + StoreKeys.segmentLength(idBytes) + + StoreKeys.segmentLength(namespaceBytes)]; + int offset = StoreKeys.writeSegment(key, 0, encodedKey); + offset = StoreKeys.writeSegment(key, offset, domainBytes); + offset = StoreKeys.writeSegment(key, offset, familyBytes); + offset = StoreKeys.writeSegment(key, offset, idBytes); + StoreKeys.writeSegment(key, offset, namespaceBytes); + return key; + } + + /** The index store key for a timer: {@code domain | fireTimestamp | identity}. */ + static byte[] indexKey(TimeDomain domain, long fireMillis, byte[] identityKey) { + byte[] key = new byte[1 + StoreKeys.TIMESTAMP_BYTES + identityKey.length]; + key[0] = (byte) domain.ordinal(); + int offset = StoreKeys.writeTimestamp(key, 1, fireMillis); + System.arraycopy(identityKey, 0, key, offset, identityKey.length); + return key; + } + + /** Inclusive lower bound of the range scan for due event-time timers. */ + static byte[] dueEventTimeRangeStart() { + byte[] bound = new byte[1 + StoreKeys.TIMESTAMP_BYTES]; + bound[0] = (byte) TimeDomain.EVENT_TIME.ordinal(); + StoreKeys.writeTimestamp(bound, 1, Long.MIN_VALUE); + return bound; + } + + /** + * Inclusive upper bound of the range scan for event-time timers due at {@code watermarkMillis}. + * + *

    Every index key carries a non-empty identity after its timestamp, so no key is equal to the + * bare {@code domain | watermark + 1} prefix returned here: an inclusive scan up to it yields + * exactly the timers whose fire time is at or before the watermark. Beam's maximum timestamp is + * far below {@link Long#MAX_VALUE}, so the increment cannot overflow. + */ + static byte[] dueEventTimeRangeEnd(long watermarkMillis) { + byte[] bound = new byte[1 + StoreKeys.TIMESTAMP_BYTES]; + bound[0] = (byte) TimeDomain.EVENT_TIME.ordinal(); + StoreKeys.writeTimestamp(bound, 1, watermarkMillis + 1); + return bound; + } + + /** Reads the identity key back out of an index key. */ + static byte[] identityKeyOf(byte[] indexKey) { + int offset = 1 + StoreKeys.TIMESTAMP_BYTES; + byte[] identityKey = new byte[indexKey.length - offset]; + System.arraycopy(indexKey, offset, identityKey, 0, identityKey.length); + return identityKey; + } + + /** Reads the encoded Beam key (the first segment) back out of an identity key. */ + static byte[] encodedKeyOf(byte[] identityKey) { + return StoreKeys.readSegment(identityKey, 0); + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StoreKeys.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StoreKeys.java new file mode 100644 index 000000000000..5cb00ca8b6cc --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StoreKeys.java @@ -0,0 +1,105 @@ +/* + * 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; + +/** + * Byte helpers for the composite keys the runner's state and timer stores are addressed by. + * + *

    Keys are built into a single pre-sized array rather than through a stream, because they are on + * the hot path: one is built for every state cell read or written, several times per element. + * + *

    Variable-length parts are written length-prefixed. A separator byte would be shorter, but both + * an encoded Beam key (arbitrary user coder output) and a {@link + * org.apache.beam.runners.core.StateNamespace#stringKey} can contain any byte value, so no + * separator is safe from collisions — {@code key="a/b", ns="c"} and {@code key="a", ns="b/c"} would + * produce the same bytes. Length prefixes also let the encoded key be read back out of a timer key, + * which the timer scan needs. + * + *

    Timestamps are written sign-flipped big-endian so that the unsigned lexicographic order Kafka + * Streams compares keys by is the same as numeric order. That is what makes a range scan over a + * timestamp-prefixed store return exactly the entries up to a point in time, which is how due + * timers and the minimum watermark hold are found without scanning everything. + */ +final class StoreKeys { + + /** Bytes taken by a length prefix. */ + static final int LENGTH_BYTES = 4; + + /** Bytes taken by a sortable timestamp. */ + static final int TIMESTAMP_BYTES = 8; + + private StoreKeys() {} + + /** Bytes a length-prefixed segment occupies. */ + static int segmentLength(byte[] segment) { + return LENGTH_BYTES + segment.length; + } + + /** Writes {@code segment} length-prefixed at {@code offset}, returning the offset after it. */ + static int writeSegment(byte[] target, int offset, byte[] segment) { + int next = writeLength(target, offset, segment.length); + System.arraycopy(segment, 0, target, next, segment.length); + return next + segment.length; + } + + private static int writeLength(byte[] target, int offset, int length) { + target[offset] = (byte) ((length >>> 24) & 0xff); + target[offset + 1] = (byte) ((length >>> 16) & 0xff); + target[offset + 2] = (byte) ((length >>> 8) & 0xff); + target[offset + 3] = (byte) (length & 0xff); + return offset + LENGTH_BYTES; + } + + /** Reads the length prefix at {@code offset}. */ + static int readLength(byte[] source, int offset) { + return ((source[offset] & 0xff) << 24) + | ((source[offset + 1] & 0xff) << 16) + | ((source[offset + 2] & 0xff) << 8) + | (source[offset + 3] & 0xff); + } + + /** Reads the length-prefixed segment starting at {@code offset}. */ + static byte[] readSegment(byte[] source, int offset) { + int length = readLength(source, offset); + byte[] segment = new byte[length]; + System.arraycopy(source, offset + LENGTH_BYTES, segment, 0, length); + return segment; + } + + /** + * Writes a timestamp so that unsigned byte order matches numeric order: flipping the sign bit + * maps {@link Long#MIN_VALUE}..{@link Long#MAX_VALUE} onto 0x00.. 0xff.. big-endian, so + * negative timestamps (valid in Beam) sort before positive ones. + */ + static int writeTimestamp(byte[] target, int offset, long millis) { + long sortable = millis ^ Long.MIN_VALUE; + for (int i = 0; i < TIMESTAMP_BYTES; i++) { + target[offset + i] = (byte) ((sortable >>> (8 * (TIMESTAMP_BYTES - 1 - i))) & 0xff); + } + return offset + TIMESTAMP_BYTES; + } + + /** Reads a timestamp written by {@link #writeTimestamp}. */ + static long readTimestamp(byte[] source, int offset) { + long sortable = 0; + for (int i = 0; i < TIMESTAMP_BYTES; i++) { + sortable = (sortable << 8) | (source[offset + i] & 0xffL); + } + return sortable ^ Long.MIN_VALUE; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java new file mode 100644 index 000000000000..1b0ffae23fde --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java @@ -0,0 +1,323 @@ +/* + * 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 java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.core.NullSideInputReader; +import org.apache.beam.runners.core.ReduceFnRunner; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.SystemReduceFn; +import org.apache.beam.runners.core.TimerInternals.TimerData; +import org.apache.beam.runners.core.triggers.ExecutableTriggerStateMachine; +import org.apache.beam.runners.core.triggers.TriggerStateMachines; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.util.construction.TriggerTranslation; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.BaseEncoding; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Executes a windowed {@code GroupByKey} by driving Beam's {@link ReduceFnRunner} — the same + * windowing + triggering state machine the Flink and Spark portable runners use — with Kafka + * Streams state and timers behind it. + * + *

    Records arrive on the repartition topic keyed by the encoded Beam key, so every value of a key + * is co-located here. For each data record this builds a {@link ReduceFnRunner} for that key over a + * {@link KafkaStreamsStateInternals} and {@link KafkaStreamsTimerInternals} (both backed by + * persistent stores) and feeds the element in; the runner assigns it to windows, updates the + * trigger state, and sets any timers it needs. When the aggregated input watermark advances, every + * event-time timer whose fire time has passed is replayed through {@code onTimers}, which is what + * makes windows emit their panes. The runner is stateless between records — all durable state lives + * in the stores — so a fresh one per record is correct, mirroring {@code + * GroupAlsoByWindowViaWindowSetNewDoFn}. + * + *

    This first version supports the default trigger and non-merging windows well; richer triggers, + * processing-time timers and session (merging) windows build on the same machinery in a follow-up. + */ +class WindowedGroupByKeyProcessor + implements Processor, byte[], KStreamsPayload> { + + private static final Logger LOG = LoggerFactory.getLogger(WindowedGroupByKeyProcessor.class); + + private final String stateStoreName; + private final String holdsIndexStoreName; + private final String timerStoreName; + private final String timerIndexStoreName; + private final String transformId; + private final Coder keyCoder; + private final WindowingStrategy windowingStrategy; + private final Coder windowCoder; + private final RunnerApi.Trigger triggerProto; + private final SystemReduceFn, Iterable, W> reduceFn; + private final PipelineOptions options; + + private final WatermarkAggregator watermarkAggregator; + private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + private Instant inputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + + private @Nullable ProcessorContext> context; + private @Nullable KeyValueStore stateStore; + private @Nullable KeyValueStore holdsIndexStore; + private @Nullable KeyValueStore timerStore; + private @Nullable KeyValueStore timerIndexStore; + + WindowedGroupByKeyProcessor( + String stateStoreName, + String holdsIndexStoreName, + String timerStoreName, + String timerIndexStoreName, + String transformId, + Set upstreamTransformIds, + Coder keyCoder, + Coder valueCoder, + WindowingStrategy windowingStrategy, + PipelineOptions options) { + this.stateStoreName = stateStoreName; + this.holdsIndexStoreName = holdsIndexStoreName; + this.timerStoreName = timerStoreName; + this.timerIndexStoreName = timerIndexStoreName; + this.transformId = transformId; + this.keyCoder = keyCoder; + this.windowingStrategy = windowingStrategy; + this.windowCoder = windowingStrategy.getWindowFn().windowCoder(); + this.triggerProto = TriggerTranslation.toProto(windowingStrategy.getTrigger()); + this.reduceFn = SystemReduceFn.buffering(valueCoder); + this.options = options; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + this.stateStore = context.getStateStore(stateStoreName); + this.holdsIndexStore = context.getStateStore(holdsIndexStoreName); + this.timerStore = context.getStateStore(timerStoreName); + this.timerIndexStore = context.getStateStore(timerIndexStoreName); + } + + @Override + public void process(Record> record) { + KStreamsPayload payload = record.value(); + if (payload == null) { + LOG.warn( + "GroupByKey {} dropping record with null payload (external write or tombstone)", + transformId); + return; + } + if (payload.isData()) { + processData(record, payload); + return; + } + watermarkAggregator.observe(payload.asWatermark()); + Instant advanced = watermarkAggregator.advance(); + if (advanced.isAfter(inputWatermark)) { + inputWatermark = advanced; + fireDueEventTimeTimers(record, advanced); + } + // Firing may have emitted panes and released their holds, so the output watermark is computed + // after it. + Instant output = outputWatermark(); + if (output.isAfter(lastForwardedWatermark)) { + lastForwardedWatermark = output; + forwardWatermark(record, output.getMillis()); + } + } + + /** + * The watermark to publish downstream: the input watermark, held back by the earliest watermark + * hold any pending pane has taken. + * + *

    {@link ReduceFnRunner} takes a hold for buffered elements that have not been emitted yet, at + * the timestamp their pane will carry. Forwarding the raw input watermark would tell downstream + * that nothing earlier is coming while those panes are still buffered, and the elements would + * then arrive late against the watermark we had already published. + */ + private Instant outputWatermark() { + Instant minHold = + KafkaStreamsStateInternals.minWatermarkHold(checkInitialized(holdsIndexStore)); + return minHold == null || inputWatermark.isBefore(minHold) ? inputWatermark : minHold; + } + + private void processData(Record> record, KStreamsPayload payload) { + byte[] encodedKey = record.key(); + if (encodedKey == null) { + throw new IllegalStateException("GroupByKey data record is missing its key"); + } + @SuppressWarnings("unchecked") + WindowedValue> element = (WindowedValue>) payload.getData(); + K key = decodeKey(encodedKey); + WindowedValue valueElement = element.withValue(element.getValue().getValue()); + runReduceFn( + record, encodedKey, key, Collections.singletonList(valueElement), Collections.emptyList()); + } + + /** + * Fires every event-time timer whose fire time is at or before the new input watermark. + * + *

    The timers to fire are found by range-scanning the fire-time-ordered index over exactly the + * window {@code (-inf, watermark]}, so the cost is proportional to the number of timers that are + * actually due rather than to the number of keys that hold a timer. + */ + private void fireDueEventTimeTimers( + Record> record, Instant watermark) { + KeyValueStore identityStore = checkInitialized(timerStore); + KeyValueStore indexStore = checkInitialized(timerIndexStore); + // Group the due timers by the Beam key they belong to; a key's timers fire together in one + // ReduceFnRunner turn. + Map dueByKey = new LinkedHashMap<>(); + List firedIndexKeys = new ArrayList<>(); + try (KeyValueIterator it = + indexStore.range( + KafkaStreamsTimerInternals.dueEventTimeRangeStart(), + KafkaStreamsTimerInternals.dueEventTimeRangeEnd(watermark.getMillis()))) { + while (it.hasNext()) { + org.apache.kafka.streams.KeyValue entry = it.next(); + TimerData timer = KafkaStreamsTimerInternals.decodeTimer(windowCoder, entry.value); + byte[] encodedKey = + KafkaStreamsTimerInternals.encodedKeyOf( + KafkaStreamsTimerInternals.identityKeyOf(entry.key)); + dueByKey + .computeIfAbsent( + BaseEncoding.base16().encode(encodedKey), k -> new DueTimers(encodedKey)) + .timers + .add(timer); + firedIndexKeys.add(entry.key); + } + } + // Clear the fired timers from both stores before replaying them, since onTimers may + // legitimately + // set new ones — including at the same identity. + for (byte[] indexKey : firedIndexKeys) { + indexStore.delete(indexKey); + identityStore.delete(KafkaStreamsTimerInternals.identityKeyOf(indexKey)); + } + for (DueTimers due : dueByKey.values()) { + runReduceFn( + record, due.encodedKey, decodeKey(due.encodedKey), Collections.emptyList(), due.timers); + } + } + + private void runReduceFn( + Record> record, + byte[] encodedKey, + @NonNull K key, + List> elements, + List timers) { + StateInternals stateInternals = + new KafkaStreamsStateInternals<>( + key, encodedKey, checkInitialized(stateStore), checkInitialized(holdsIndexStore)); + KafkaStreamsTimerInternals timerInternals = + new KafkaStreamsTimerInternals( + encodedKey, + checkInitialized(timerStore), + checkInitialized(timerIndexStore), + windowCoder, + inputWatermark, + lastForwardedWatermark, + Instant.now()); + ReduceFnRunner, W> runner = + new ReduceFnRunner<>( + key, + windowingStrategy, + ExecutableTriggerStateMachine.create( + TriggerStateMachines.stateMachineForTrigger(triggerProto)), + stateInternals, + timerInternals, + output -> forwardData(record, encodedKey, output), + NullSideInputReader.empty(), + reduceFn, + options); + try { + runner.processElements(elements); + runner.onTimers(timers); + runner.persist(); + } catch (Exception e) { + throw new RuntimeException("GroupByKey " + transformId + " failed to run windowing", e); + } + } + + private void forwardData( + Record> trigger, + byte[] encodedKey, + WindowedValue>> output) { + ProcessorContext> ctx = checkInitialized(context); + ctx.forward( + new Record>( + encodedKey, KStreamsPayload.data(output), trigger.timestamp())); + } + + private void forwardWatermark(Record> trigger, long watermarkMillis) { + ProcessorContext> ctx = checkInitialized(context); + // Stamped with this transform's own id; GroupByKey is a single instance for now (0 of 1). + ctx.forward( + new Record>( + trigger.key(), + KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), + trigger.timestamp())); + } + + private @NonNull K decodeKey(byte[] bytes) { + try { + K key = CoderUtils.decodeFromByteArray(keyCoder, bytes); + if (key == null) { + throw new IllegalStateException("GroupByKey key decoded to null"); + } + return key; + } catch (CoderException e) { + throw new RuntimeException("Failed to decode GroupByKey key", e); + } + } + + private static T checkInitialized(@Nullable T value) { + if (value == null) { + throw new IllegalStateException("WindowedGroupByKeyProcessor used before init()"); + } + return value; + } + + /** The event-time timers due for one Beam key, plus that key's encoded bytes. */ + private static final class DueTimers { + final byte[] encodedKey; + final List timers = new ArrayList<>(); + + DueTimers(byte[] encodedKey) { + this.encodedKey = encodedKey; + } + } +} 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 8ab29182c281..bdbde1db36dd 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 @@ -18,11 +18,7 @@ package org.apache.beam.runners.kafka.streams; import java.time.Duration; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; import java.util.Properties; -import java.util.Set; import java.util.UUID; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.core.metrics.MetricsContainerStepMap; @@ -39,16 +35,11 @@ import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; import org.apache.beam.sdk.util.construction.PipelineTranslation; import org.apache.beam.sdk.util.construction.SplittableParDo; -import org.apache.kafka.common.serialization.ByteArrayDeserializer; -import org.apache.kafka.common.serialization.ByteArraySerializer; import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.streams.StreamsConfig; -import org.apache.kafka.streams.TestInputTopic; -import org.apache.kafka.streams.TestOutputTopic; import org.apache.kafka.streams.Topology; import org.apache.kafka.streams.TopologyDescription; import org.apache.kafka.streams.TopologyTestDriver; -import org.apache.kafka.streams.test.TestRecord; /** * Test harness that runs a Beam {@link Pipeline} through the Kafka Streams runner's translation and @@ -58,16 +49,13 @@ * effects (e.g. a {@code SharedTestCollector} written by a recording DoFn) have completed when it * returns. * - *

    {@link TopologyTestDriver} does not loop a low-level sink topic back into its source, so an - * internal repartition topic (one that is both a sink and a source in the topology — e.g. the one - * GroupByKey introduces) would otherwise dead-end. {@link #run(Pipeline)} discovers those topics - * from the {@link TopologyDescription} and round-trips them until no more records flow, standing in - * for the broker. + *

    {@link TopologyTestDriver} loops each internal repartition topic (one that is both a sink and + * a source in the topology — e.g. the one GroupByKey introduces) from its sink back to its source + * within a single driver step, so advancing the wall clock is enough to drive the whole pipeline to + * completion; no manual broker simulation is needed. */ public final class KafkaStreamsTestRunner { - private static final int MAX_ROUND_TRIPS = 100; - private KafkaStreamsTestRunner() {} /** Pipeline options for a Kafka Streams runner test: the EMBEDDED harness and a unique app id. */ @@ -118,10 +106,12 @@ public static MetricResults run(Pipeline pipeline) { KafkaStreamsTranslationContext context = translate(pipeline); Topology topology = context.getTopology(); try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig(pipeline))) { - // Fire the Impulse wall-clock punctuator and let the initial records flow. + // Fire the Impulse wall-clock punctuator; TopologyTestDriver then flows the records through + // the whole topology, including looping each internal repartition topic (a sink that is also + // a source, e.g. the one GroupByKey introduces) back to its source, standing in for the + // broker. A second advance covers punctuators that need a later tick. driver.advanceWallClockTime(Duration.ofSeconds(1)); driver.advanceWallClockTime(Duration.ofSeconds(1)); - roundTripInternalTopics(driver, internalTopics(topology)); } return MetricsContainerStepMap.asAttemptedOnlyMetricResults( context.getMetricsContainerStepMap()); @@ -143,81 +133,6 @@ public static String findAnyLeafProcessorName(Topology topology) { throw new IllegalStateException("no leaf processor found in topology"); } - /** Repartition/internal topics are the ones that appear as both a sink and a source. */ - private static Set internalTopics(Topology topology) { - Set sinkTopics = new HashSet<>(); - Set sourceTopics = new HashSet<>(); - for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { - for (TopologyDescription.Node node : subtopology.nodes()) { - if (node instanceof TopologyDescription.Sink) { - String topic = ((TopologyDescription.Sink) node).topic(); - if (topic != null) { - sinkTopics.add(topic); - } - } else if (node instanceof TopologyDescription.Source) { - sourceTopics.addAll(((TopologyDescription.Source) node).topicSet()); - } - } - } - sinkTopics.retainAll(sourceTopics); - return sinkTopics; - } - - /** - * Simulates the broker for internal repartition topics. - * - *

    The runner shuffles data (and the watermark) through internal topics that a processor both - * writes to (a sink) and reads back from (a source) — e.g. the topic GroupByKey introduces to - * partition by key. On a real broker those records make the round trip automatically, but {@link - * TopologyTestDriver} does not connect a sink back to a source, so the downstream half of the - * topology would never see them. This drains what each internal topic's sink wrote and pipes it - * into that topic's source, repeating until nothing new flows (a fixpoint), which stands in for - * the broker and lets the pipeline run to completion. - */ - private static void roundTripInternalTopics(TopologyTestDriver driver, Set topics) { - // Create the sink-output and source-input handles once and reuse them across rounds; a single - // TestOutputTopic keeps returning newly produced records on each read. - List roundTrips = new ArrayList<>(); - for (String topic : topics) { - roundTrips.add( - new TopicRoundTrip( - driver.createOutputTopic( - topic, new ByteArrayDeserializer(), new ByteArrayDeserializer()), - driver.createInputTopic( - topic, new ByteArraySerializer(), new ByteArraySerializer()))); - } - - for (int round = 0; round < MAX_ROUND_TRIPS; round++) { - boolean progressed = false; - for (TopicRoundTrip roundTrip : roundTrips) { - List> records = roundTrip.output.readRecordsToList(); - if (records.isEmpty()) { - continue; - } - progressed = true; - for (TestRecord record : records) { - roundTrip.input.pipeInput(record); - } - } - if (!progressed) { - return; - } - } - throw new IllegalStateException( - "Internal topics did not reach quiescence after " + MAX_ROUND_TRIPS + " round trips"); - } - - /** The reusable sink-output and source-input handles for one internal topic. */ - private static final class TopicRoundTrip { - final TestOutputTopic output; - final TestInputTopic input; - - TopicRoundTrip(TestOutputTopic output, TestInputTopic input) { - this.output = output; - this.input = input; - } - } - /** Kafka Streams config for a {@link TopologyTestDriver} built from the pipeline's app id. */ public static Properties streamsConfig(Pipeline pipeline) { String applicationId = diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FixedWindowGroupByKeyTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FixedWindowGroupByKeyTest.java new file mode 100644 index 000000000000..73f4dac601eb --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FixedWindowGroupByKeyTest.java @@ -0,0 +1,106 @@ +/* + * 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.hasItems; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.values.KV; +import org.joda.time.Duration; +import org.joda.time.Instant; +import org.junit.Test; + +/** + * End-to-end test that GroupByKey groups per fixed window, not just per key: {@code Impulse -> emit + * timestamped KVs -> Window.into(FixedWindows) -> GroupByKey -> record groups}. + * + *

    The same key "a" has values in two different windows, so a correct windowed GroupByKey emits + * two groups for it (one per window) rather than one combined group. This exercises the {@link + * WindowedGroupByKeyProcessor} path (ReduceFnRunner over the Kafka Streams state and timer stores) + * that the earlier global-window GroupByKey did not. + */ +public class FixedWindowGroupByKeyTest { + + private static final Duration WINDOW_SIZE = Duration.millis(10); + + /** Emits KVs whose timestamps fall into two adjacent fixed windows. */ + private static class EmitTimestampedKvsFn extends DoFn> { + @ProcessElement + public void processElement(OutputReceiver> out) { + // Window [0, 10): a=1, a=2, b=5. + out.outputWithTimestamp(KV.of("a", 1), new Instant(1)); + out.outputWithTimestamp(KV.of("a", 2), new Instant(2)); + out.outputWithTimestamp(KV.of("b", 5), new Instant(3)); + // Window [10, 20): a=3. + out.outputWithTimestamp(KV.of("a", 3), new Instant(15)); + } + } + + /** Records each grouped result as {@code "key=[sorted values]"}. */ + private static class RecordGroupFn extends DoFn>, Void> { + private final SharedTestCollector collector; + + RecordGroupFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element KV> group) { + List values = new ArrayList<>(); + group.getValue().forEach(values::add); + Collections.sort(values); + collector.record(group.getKey() + "=" + values); + } + } + + @Test + public void groupsValuesPerFixedWindow() { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + pipeline + .apply("impulse", Impulse.create()) + .apply("emit", ParDo.of(new EmitTimestampedKvsFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())) + .apply("window", Window.into(FixedWindows.of(WINDOW_SIZE))) + .apply("gbk", GroupByKey.create()) + .apply("record", ParDo.of(new RecordGroupFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List groups = collector.recorded(); + // a splits across two windows -> two groups; b has one; three groups total. + assertThat(groups.size(), is(3)); + assertThat(groups, hasItems("a=[1, 2]", "a=[3]", "b=[5]")); + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternalsTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternalsTest.java new file mode 100644 index 000000000000..df5b0eda1b26 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternalsTest.java @@ -0,0 +1,208 @@ +/* + * 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 java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.core.StateNamespace; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.TimerInternals.TimerData; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.processor.api.MockProcessorContext; +import org.apache.kafka.streams.state.KeyValueIterator; +import org.apache.kafka.streams.state.KeyValueStore; +import org.apache.kafka.streams.state.Stores; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.Test; + +/** + * Tests the timer store and its fire-time index: that a timer can be set, replaced and deleted by + * identity, and that the timers due at a watermark are found by a range scan over the index rather + * than by inspecting every timer. + */ +public class KafkaStreamsTimerInternalsTest { + + private static final StateNamespace NAMESPACE = + StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE); + + private KeyValueStore identityStore; + private KeyValueStore indexStore; + + @Before + public void setUp() { + MockProcessorContext context = new MockProcessorContext<>(); + identityStore = newStore("timers", context); + indexStore = newStore("timers-index", context); + } + + private static KeyValueStore newStore( + String name, MockProcessorContext context) { + KeyValueStore store = + Stores.keyValueStoreBuilder( + Stores.inMemoryKeyValueStore(name), Serdes.ByteArray(), Serdes.ByteArray()) + .withLoggingDisabled() + .build(); + store.init(context.getStateStoreContext(), store); + return store; + } + + private KafkaStreamsTimerInternals timersFor(String key) { + return new KafkaStreamsTimerInternals( + encode(key), + identityStore, + indexStore, + GlobalWindow.Coder.INSTANCE, + BoundedWindow.TIMESTAMP_MIN_VALUE, + BoundedWindow.TIMESTAMP_MIN_VALUE, + new Instant(0)); + } + + private static byte[] encode(String key) { + try { + return CoderUtils.encodeToByteArray(StringUtf8Coder.of(), key); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private static TimerData eventTimer(String id, long millis) { + return TimerData.of( + id, "", NAMESPACE, new Instant(millis), new Instant(millis), TimeDomain.EVENT_TIME); + } + + /** The timers the processor would fire at {@code watermarkMillis}, in fire-time order. */ + private List dueAt(long watermarkMillis) { + List due = new ArrayList<>(); + try (KeyValueIterator it = + indexStore.range( + KafkaStreamsTimerInternals.dueEventTimeRangeStart(), + KafkaStreamsTimerInternals.dueEventTimeRangeEnd(watermarkMillis))) { + while (it.hasNext()) { + due.add( + KafkaStreamsTimerInternals.decodeTimer(GlobalWindow.Coder.INSTANCE, it.next().value)); + } + } + return due; + } + + private static int storeSize(KeyValueStore store) { + int size = 0; + try (KeyValueIterator it = store.all()) { + while (it.hasNext()) { + it.next(); + size++; + } + } + return size; + } + + @Test + public void dueScanReturnsOnlyTimersAtOrBeforeTheWatermark() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer(eventTimer("early", 100L)); + timers.setTimer(eventTimer("onWatermark", 200L)); + timers.setTimer(eventTimer("late", 300L)); + + List due = dueAt(200L); + + // Ordered by fire time, and the timer set exactly at the watermark is included. + assertThat(due.size(), is(2)); + assertThat(due.get(0).getTimerId(), is("early")); + assertThat(due.get(1).getTimerId(), is("onWatermark")); + } + + @Test + public void negativeTimestampsSortBeforePositiveOnes() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer(eventTimer("negative", -5000L)); + timers.setTimer(eventTimer("zero", 0L)); + timers.setTimer(eventTimer("positive", 5000L)); + + List due = dueAt(0L); + + assertThat(due.size(), is(2)); + assertThat(due.get(0).getTimerId(), is("negative")); + assertThat(due.get(1).getTimerId(), is("zero")); + } + + @Test + public void resettingATimerReplacesItsIndexEntry() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer(eventTimer("timer", 100L)); + // Re-setting the same timer identity for a later time must not leave the old entry behind, + // or the timer would still fire at the time it was first set for. + timers.setTimer(eventTimer("timer", 900L)); + + assertThat(dueAt(100L).isEmpty(), is(true)); + assertThat(dueAt(900L).size(), is(1)); + assertThat(storeSize(indexStore), is(1)); + assertThat(storeSize(identityStore), is(1)); + } + + @Test + public void deletingATimerRemovesItFromBothStores() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer(eventTimer("timer", 100L)); + timers.deleteTimer(NAMESPACE, "timer", "", TimeDomain.EVENT_TIME); + + assertThat(dueAt(1000L).isEmpty(), is(true)); + assertThat(storeSize(indexStore), is(0)); + assertThat(storeSize(identityStore), is(0)); + } + + @Test + public void timersOfDifferentKeysAreIndependentButShareTheIndex() { + timersFor("a").setTimer(eventTimer("timer", 100L)); + timersFor("b").setTimer(eventTimer("timer", 150L)); + + // Same timer id under two Beam keys are two distinct timers, and one scan finds both. + assertThat(storeSize(identityStore), is(2)); + assertThat(dueAt(200L).size(), is(2)); + + timersFor("a").deleteTimer(NAMESPACE, "timer", "", TimeDomain.EVENT_TIME); + assertThat(dueAt(200L).size(), is(1)); + } + + @Test + public void processingTimeTimersAreNotReturnedByTheEventTimeScan() { + KafkaStreamsTimerInternals timers = timersFor("key"); + timers.setTimer( + TimerData.of( + "processing", + "", + NAMESPACE, + new Instant(100L), + new Instant(100L), + TimeDomain.PROCESSING_TIME)); + timers.setTimer(eventTimer("event", 100L)); + + List due = dueAt(1000L); + + assertThat(due.size(), is(1)); + assertThat(due.get(0).getTimerId(), is("event")); + } +} From 27198768f89f5b177e65e9c00c32dbba2d94cfb1 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Sat, 1 Aug 2026 14:31:20 +0500 Subject: [PATCH 23/37] [GSoC 2026] Kafka Streams runner: run on a real broker, correctly across partitions (#39546) * [GSoC 2026] Kafka Streams runner: run against a real broker Adds what the runner needs to execute on an actual Kafka cluster, and an integration test that runs a pipeline through the production KafkaStreamsPipelineRunner against a Kafka container. Everything until now ran through TopologyTestDriver, which fakes the topics and never builds a KafkaStreams application, so the production path had not been executed. --- runners/kafka-streams/build.gradle | 23 ++ .../streams/KafkaStreamsPipelineOptions.java | 17 ++ .../streams/KafkaStreamsPipelineRunner.java | 38 ++- .../KafkaStreamsPortablePipelineResult.java | 4 + .../streams/KafkaStreamsTopicManager.java | 171 +++++++++++ .../translation/ExecutableStageProcessor.java | 9 +- .../ExecutableStageTranslator.java | 4 + .../streams/translation/FlattenProcessor.java | 3 +- .../translation/FlattenTranslator.java | 5 + .../translation/GroupByKeyTranslator.java | 12 +- .../KafkaStreamsTranslationContext.java | 30 ++ .../translation/RedistributeTranslator.java | 3 + .../translation/ShuffleByKeyProcessor.java | 37 ++- .../WindowedGroupByKeyProcessor.java | 3 +- .../streams/KafkaStreamsRunnerBrokerIT.java | 286 ++++++++++++++++++ .../ShuffleByKeyProcessorTest.java | 123 ++++++++ .../StandardWindowFnTranslationTest.java | 152 ++++++++++ 17 files changed, 900 insertions(+), 20 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTopicManager.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index bdf7e3be0585..d1326c079e5d 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -17,6 +17,7 @@ */ import groovy.json.JsonOutput +import java.time.Duration plugins { id 'org.apache.beam.module' } @@ -74,6 +75,7 @@ dependencies { testImplementation library.java.junit testImplementation library.java.mockito_core testImplementation "org.apache.kafka:kafka-streams-test-utils:$kafka_version" + testImplementation library.java.testcontainers_kafka // Beam's @ValidatesRunner suite: the test classes come from the SDK core test jar; the runner // (TestKafkaStreamsRunner) and its TopologyTestDriver harness come from this module's test @@ -85,6 +87,27 @@ dependencies { } +// The broker integration test drives the production runner against a real Kafka in Docker, so it +// is not part of the default build. Run it with :runners:kafka-streams:brokerIntegrationTest. +test { + filter { + excludeTestsMatching 'org.apache.beam.runners.kafka.streams.*IT' + } +} + +tasks.register("brokerIntegrationTest", Test) { + group = "Verification" + description = "Runs the Kafka Streams runner against a real broker (requires Docker)." + outputs.upToDateWhen { false } + testClassesDirs = sourceSets.test.output.classesDirs + classpath = sourceSets.test.runtimeClasspath + filter { + includeTestsMatching 'org.apache.beam.runners.kafka.streams.*IT' + } + // A container start plus a streaming run is well past the default per-test expectations. + timeout = Duration.ofMinutes(15) +} + // Known-failing @ValidatesRunner tests, excluded until the feature they need lands. def sickbayTests = [ // Merging (session) windows are not supported yet: ReduceFnRunner drives them through a merging diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java index 2fa992e66e7e..a5a8bb9328b1 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java @@ -55,6 +55,23 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { void setMaxBundleTimeMs(int maxBundleTimeMs); + @Description( + "How many partitions the runner gives the internal topics it creates to shuffle a pipeline" + + " through, which is the parallelism the shuffled parts of that pipeline can reach. A" + + " GroupByKey runs one task per partition of its repartition topic, so this is the" + + " number of instances its state and its downstream stages are spread over. Must be at" + + " least 1.") + @Default.Integer(1) + int getInternalParallelism(); + + void setInternalParallelism(int internalParallelism); + + @Description("Replication factor for the internal topics the runner creates for a pipeline.") + @Default.Short(1) + short getTopicReplicationFactor(); + + void setTopicReplicationFactor(short topicReplicationFactor); + @Description("Directory where Kafka Streams stores local state.") @Default.InstanceFactory(StateDirDefaultFactory.class) String getStateDir(); 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 cdb59f67cce9..96b4b2cd7f92 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 @@ -24,10 +24,10 @@ import org.apache.beam.runners.jobsubmission.PortablePipelineRunner; import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsPipelineTranslator; import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsTranslationContext; -import org.apache.beam.sdk.options.PipelineOptionsValidator; import org.apache.kafka.streams.KafkaStreams; import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.Topology; +import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -44,9 +44,22 @@ public KafkaStreamsPipelineRunner(KafkaStreamsPipelineOptions pipelineOptions) { @Override public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) { - // Surface a clear error if a required option (e.g. applicationId) is missing instead of - // letting Properties.put fail with a raw NullPointerException further down. - PipelineOptionsValidator.validate(KafkaStreamsPipelineOptions.class, pipelineOptions); + // Surface a clear error if an option this runner needs is missing, instead of letting + // Properties.put fail with a raw NullPointerException further down. Only the options that are + // meaningful here are checked, rather than validating the whole interface: this runs on the job + // server, executing a pipeline that has already been submitted, so the client-side options + // PortablePipelineOptions marks required — jobEndpoint above all — do not apply. Flink's + // equivalent PortablePipelineRunner does not validate here either. + checkRequiredOption("applicationId", pipelineOptions.getApplicationId()); + checkRequiredOption("bootstrapServers", pipelineOptions.getBootstrapServers()); + // A topic cannot have fewer than one partition, and the value is also the number of watermark + // reports a shuffle's consumer waits for, so a non-positive value would leave it waiting + // forever rather than failing. + if (pipelineOptions.getInternalParallelism() < 1) { + throw new IllegalArgumentException( + "--internalParallelism must be at least 1, but was " + + pipelineOptions.getInternalParallelism()); + } KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); KafkaStreamsTranslationContext context = @@ -55,15 +68,28 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) translator.translate(context, prepared); Topology topology = context.getTopology(); + // The runner names its own bootstrap and repartition topics, which Kafka Streams treats as + // user topics and will not create; it refuses to start if a source topic is missing. + KafkaStreamsTopicManager.createMissingTopics(topology, pipelineOptions); LOG.info( "Translated pipeline {} into Kafka Streams topology:\n{}", jobInfo.jobId(), topology.describe()); KafkaStreams kafkaStreams = new KafkaStreams(topology, streamsConfig(jobInfo)); + // Build the result before starting: it registers a state listener, and Kafka Streams only + // accepts one while the application is still in the CREATED state. + KafkaStreamsPortablePipelineResult result = + new KafkaStreamsPortablePipelineResult(kafkaStreams, context.getMetricsContainerStepMap()); kafkaStreams.start(); - return new KafkaStreamsPortablePipelineResult( - kafkaStreams, context.getMetricsContainerStepMap()); + return result; + } + + private static void checkRequiredOption(String name, @Nullable String value) { + if (value == null || value.isEmpty()) { + throw new IllegalArgumentException( + "Missing required pipeline option --" + name + " for the Kafka Streams runner"); + } } 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 972afa15c358..0d508b8189fd 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 @@ -48,6 +48,10 @@ class KafkaStreamsPortablePipelineResult implements PortablePipelineResult { private final CountDownLatch terminated = new CountDownLatch(1); private volatile boolean cancelled = false; + /** + * Must be constructed before {@link KafkaStreams#start()} is called: it registers a state + * listener, and Kafka Streams rejects one once the application has left the CREATED state. + */ KafkaStreamsPortablePipelineResult( KafkaStreams kafkaStreams, MetricsContainerStepMap metricsContainerStepMap) { this.kafkaStreams = kafkaStreams; diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTopicManager.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTopicManager.java new file mode 100644 index 000000000000..0e5f46f53eac --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTopicManager.java @@ -0,0 +1,171 @@ +/* + * 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; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashSet; +import java.util.List; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import org.apache.kafka.clients.admin.Admin; +import org.apache.kafka.clients.admin.AdminClientConfig; +import org.apache.kafka.clients.admin.NewTopic; +import org.apache.kafka.common.errors.TopicExistsException; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.TopologyDescription; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Creates the topics a translated pipeline needs before the Kafka Streams application starts. + * + *

    The runner shuffles data through topics it names itself: a bootstrap topic per Impulse and per + * primitive Read, and a repartition topic per GroupByKey. Kafka Streams does create the internal + * topics it manages on its own, but these are declared with explicit names through {@code + * addSource} and {@code addSink}, so to Kafka Streams they are ordinary user topics — it will not + * create them, and refuses to start with {@code MissingSourceTopicException} if a source topic is + * absent. Relying on the broker's {@code auto.create.topics.enable} is not an option either: it is + * off on many clusters, and a topic auto-created on first fetch gets the broker's default partition + * count rather than the pipeline's. + * + *

    Only topics carrying one of the runner's own prefixes are created. Any other topic in the + * topology belongs to the user (a source or sink they named), and creating those implicitly would + * hide a misconfiguration behind an empty topic. + */ +class KafkaStreamsTopicManager { + + private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsTopicManager.class); + + /** + * Prefixes of the bootstrap topics, which must have exactly one partition. + * + *

    An Impulse or a primitive Read emits its elements once per task, gated by a state store that + * is itself per task. Kafka Streams creates one task per partition of the source topic, so a + * bootstrap topic with several partitions would make the same Impulse fire once per partition and + * the same source be read once per partition. + */ + private static final List SINGLE_PARTITION_TOPIC_PREFIXES = + java.util.Arrays.asList("__beam_impulse_", "__beam_read_"); + + /** + * Prefixes of the topics whose partition count sets the pipeline's parallelism — the repartition + * topic a GroupByKey shuffles through. + */ + private static final List PARTITIONED_TOPIC_PREFIXES = + java.util.Arrays.asList("__beam_gbk_"); + + private KafkaStreamsTopicManager() {} + + /** + * Creates any runner-owned topic in {@code topology} that does not exist yet. + * + *

    Safe to run concurrently with another instance of the same job: a topic that appears between + * the existence check and the create request surfaces as {@link TopicExistsException}, which is + * treated as success. + */ + static void createMissingTopics(Topology topology, KafkaStreamsPipelineOptions options) { + Set runnerTopics = runnerOwnedTopics(topology); + if (runnerTopics.isEmpty()) { + return; + } + Properties adminConfig = new Properties(); + adminConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, options.getBootstrapServers()); + try (Admin admin = Admin.create(adminConfig)) { + Set existing = admin.listTopics().names().get(); + List toCreate = new ArrayList<>(); + for (String topic : runnerTopics) { + if (!existing.contains(topic)) { + toCreate.add( + new NewTopic( + topic, partitionsFor(topic, options), options.getTopicReplicationFactor())); + } + } + if (toCreate.isEmpty()) { + return; + } + LOG.info("Creating {} runner-owned topic(s): {}", toCreate.size(), toCreate); + createAll(admin, toCreate); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while creating the pipeline's Kafka topics", e); + } catch (ExecutionException e) { + throw new RuntimeException("Failed to create the pipeline's Kafka topics", e); + } + } + + private static void createAll(Admin admin, Collection topics) + throws InterruptedException, ExecutionException { + try { + admin.createTopics(topics).all().get(); + } catch (ExecutionException e) { + // Another instance of the same application may have created them first, which is fine. + if (!(e.getCause() instanceof TopicExistsException)) { + throw e; + } + LOG.debug("Some topics already existed; another instance created them first", e); + } + } + + /** The topics in the topology that the runner named, and so is responsible for creating. */ + private static Set runnerOwnedTopics(Topology topology) { + Set topics = new HashSet<>(); + for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { + for (TopologyDescription.Node node : subtopology.nodes()) { + if (node instanceof TopologyDescription.Source) { + Set sourceTopics = ((TopologyDescription.Source) node).topicSet(); + if (sourceTopics != null) { + topics.addAll(sourceTopics); + } + } else if (node instanceof TopologyDescription.Sink) { + String topic = ((TopologyDescription.Sink) node).topic(); + if (topic != null) { + topics.add(topic); + } + } + } + } + topics.removeIf(topic -> !isRunnerOwned(topic)); + return topics; + } + + /** + * The partition count a runner-owned topic is created with: one for a bootstrap topic, and the + * configured parallelism for a shuffle topic. + */ + private static int partitionsFor(String topic, KafkaStreamsPipelineOptions options) { + return hasAnyPrefix(topic, SINGLE_PARTITION_TOPIC_PREFIXES) + ? 1 + : options.getInternalParallelism(); + } + + private static boolean isRunnerOwned(String topic) { + return hasAnyPrefix(topic, SINGLE_PARTITION_TOPIC_PREFIXES) + || hasAnyPrefix(topic, PARTITIONED_TOPIC_PREFIXES); + } + + private static boolean hasAnyPrefix(String topic, List prefixes) { + for (String prefix : prefixes) { + if (topic.startsWith(prefix)) { + return true; + } + } + return false; + } +} 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 15265073dd21..63fefe3e15c4 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 @@ -286,10 +286,11 @@ private void closeBundleAndFlush(Record> record) { } private void forwardWatermark(Record> record, long watermarkMillis) { - // Stamped with this stage's own transform id; this stage is a single instance for now, so the - // report is for its only partition (0 of 1). Fanning the watermark out to every downstream - // partition — and producing it atomically with the offset commit so it is durable — lands with - // the topic-based shuffle work (#18479). + // Labelled as the only source a consumer will see. Forwarding here is in-process, to the + // stage's + // fused children, so exactly one instance of this stage reaches each of them. Where the output + // instead crosses a shuffle, ShuffleByKeyProcessor relabels the report with the real partition + // identity, because the broadcast then delivers every instance's report to every consumer. ProcessorContext> ctx = checkInitialized(context); ctx.forward( new Record>( 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 71e24f32c1f5..caefa6534fad 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 @@ -74,6 +74,8 @@ public void translate( // is unambiguous even before we add side-input support. String inputPCollectionId = stagePayload.getInput(); String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); + // A fused stage runs wherever its input runs: same task, so same partition identity. + int partitionCount = context.getPartitionCount(inputPCollectionId); // A multi-output stage (a DoFn with side outputs, or a Read whose SDF wrapper produces several // outputs) needs each output routed to the right downstream. Since downstream transforms are @@ -117,9 +119,11 @@ public void translate( topology.addProcessor( relayName, () -> new StageOutputProcessor(relayName), transformId); context.registerPCollectionProducer(outputPCollectionId, relayName); + context.registerPCollectionPartitionCount(outputPCollectionId, partitionCount); }); } else if (!outputPCollectionIds.isEmpty()) { context.registerPCollectionProducer(outputPCollectionIds.get(0), transformId); + context.registerPCollectionPartitionCount(outputPCollectionIds.get(0), partitionCount); } } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java index d09fed8185a1..e37da6774489 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java @@ -99,8 +99,7 @@ public void process(Record> record) { Instant advanced = watermarkAggregator.advance(); if (advanced.isAfter(lastForwardedWatermark)) { lastForwardedWatermark = advanced; - // Stamped with this Flatten's own transform id; Flatten is a single instance for now, so the - // report is for its only partition (0 of 1). + // Labelled as the only source a consumer will see; a shuffle downstream relabels it. ctx.forward( new Record>( record.key(), diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java index 3ac4c1304daa..a5c8ce05baeb 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java @@ -57,6 +57,9 @@ public void translate( Set seenInputs = new HashSet<>(); List parentProcessors = new ArrayList<>(); Set upstreamTransformIds = new HashSet<>(); + // Kafka Streams puts a processor and the parents it is wired to in one subtopology, so the + // inputs are co-partitioned and this Flatten runs at their partition count. + int partitionCount = 1; for (String inputPCollectionId : transform.getInputsMap().values()) { if (!seenInputs.add(inputPCollectionId)) { throw new UnsupportedOperationException( @@ -70,6 +73,7 @@ public void translate( String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); parentProcessors.add(parentProcessor); upstreamTransformIds.add(parentProcessor); + partitionCount = Math.max(partitionCount, context.getPartitionCount(inputPCollectionId)); } topology.addProcessor( @@ -78,5 +82,6 @@ public void translate( parentProcessors.toArray(new String[0])); context.registerPCollectionProducer(outputPCollectionId, transformId); + context.registerPCollectionPartitionCount(outputPCollectionId, partitionCount); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index c5327e28e069..c460436eedf8 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -95,6 +95,9 @@ public void translate( hydrateWindowingStrategy(pipeline, inputPCollectionId); String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId); + // The shuffle is what changes the parallelism: everything from the repartition topic onwards + // runs one task per partition of it. + int partitionCount = context.getPipelineOptions().getInternalParallelism(); String shuffleName = transformId + SHUFFLE_SUFFIX; String sinkName = transformId + SINK_SUFFIX; @@ -110,7 +113,13 @@ public void translate( Topology topology = context.getTopology(); // Re-key data records by the encoded Beam key; pass watermark reports through. - topology.addProcessor(shuffleName, () -> new ShuffleByKeyProcessor(keyCoder), parentProcessor); + // The shuffle runs in the upstream transform's task, so it relabels each report with that + // transform's instance identity before the sink broadcasts it to every partition. + int upstreamPartitionCount = context.getPartitionCount(inputPCollectionId); + topology.addProcessor( + shuffleName, + () -> new ShuffleByKeyProcessor(keyCoder, upstreamPartitionCount), + parentProcessor); // Shuffle through the repartition topic: data partitioned by key, watermark broadcast. topology.addSink( @@ -168,6 +177,7 @@ public void translate( transformId); context.registerPCollectionProducer(outputPCollectionId, transformId); + context.registerPCollectionPartitionCount(outputPCollectionId, partitionCount); } /** Hydrates the input PCollection's windowing strategy from the pipeline proto. */ 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 ec1b3f26aded..d03169615665 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 @@ -46,6 +46,12 @@ public class KafkaStreamsTranslationContext { private final KafkaStreamsPipelineOptions pipelineOptions; private final Topology topology; private final Map pCollectionIdToProcessorName; + + /** + * How many partitions the transform producing each PCollection runs across. A PCollection that + * has not been registered is produced by a single instance; only a shuffle raises the count. + */ + private final Map pCollectionIdToPartitionCount = new HashMap<>(); // 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 @@ -113,6 +119,30 @@ public void registerPCollectionProducer(String pCollectionId, String processorNa } } + /** + * Records how many partitions the transform producing {@code pCollectionId} runs across. + * + *

    This is the {@code totalSourcePartitions} its watermark reports carry, and what a downstream + * {@link WatermarkAggregator} waits to hear from before it lets the watermark advance. It changes + * only at a shuffle: everything fused downstream of one runs at the shuffle topic's partition + * count, and everything else runs as a single instance. + */ + public void registerPCollectionPartitionCount(String pCollectionId, int partitionCount) { + pCollectionIdToPartitionCount.put(pCollectionId, partitionCount); + } + + /** + * How many partitions the transform producing {@code pCollectionId} runs across; one unless a + * shuffle upstream raised it. + * + *

    Always at least one: an unregistered PCollection is produced by a single instance, and the + * only value ever registered is {@code --internalParallelism}, which the runner rejects below one + * before translating. + */ + public int getPartitionCount(String pCollectionId) { + return pCollectionIdToPartitionCount.getOrDefault(pCollectionId, 1); + } + /** Returns the processor node name producing the given PCollection. */ public String getProcessorNameForPCollection(String pCollectionId) { String name = pCollectionIdToProcessorName.get(pCollectionId); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java index 72c47db8d4b1..a44740887b3e 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java @@ -50,5 +50,8 @@ public void translate( // Passthrough: downstream lookups for the output PCollection resolve to the producer of the // input PCollection. No KS Processor / state store / source is added. context.registerPCollectionProducer(outputPCollectionId, parentProcessor); + // A pass-through: the output is produced by the same processor, so same partition identity. + context.registerPCollectionPartitionCount( + outputPCollectionId, context.getPartitionCount(inputPCollectionId)); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java index 774d36db0f2e..79595cba7c96 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java @@ -34,9 +34,16 @@ *

    This is not GroupByKey-specific: any transform that needs the values of a key co-located on * one partition uses it — GroupByKey today, and stateful ParDo later. For a data record it sets the * Kafka record key to the encoded Beam key (taken from the {@code KV}), so the downstream - * repartition sink co-locates every value of a key. Watermark reports are forwarded unchanged — the - * {@link GroupByKeyBroadcastPartitioner} fans them out to all partitions so every downstream task - * can fire. + * repartition sink co-locates every value of a key. + * + *

    Watermark reports are relabelled here with the reporting instance's real partition identity. + * This is the point at which a report stops being delivered in-process and starts crossing a topic: + * upstream of it a transform forwards to its fused children, which see exactly one instance of it, + * so the report names a single source. The {@link GroupByKeyBroadcastPartitioner} on the sink below + * fans each report out to every partition, so a downstream task instead sees a report from + * every instance of the upstream transform, and has to be able to tell them apart to know when it + * has heard from all of them. The transform id is left alone, so the report still names the + * transform that produced it. */ class ShuffleByKeyProcessor implements Processor, byte[], KStreamsPayload> { @@ -44,15 +51,25 @@ class ShuffleByKeyProcessor private static final Logger LOG = LoggerFactory.getLogger(ShuffleByKeyProcessor.class); private final Coder keyCoder; + + /** How many instances the transform being shuffled runs as, and which one this is. */ + private final int upstreamPartitionCount; + + private int upstreamPartition; + private @Nullable ProcessorContext> context; - ShuffleByKeyProcessor(Coder keyCoder) { + ShuffleByKeyProcessor(Coder keyCoder, int upstreamPartitionCount) { this.keyCoder = keyCoder; + this.upstreamPartitionCount = upstreamPartitionCount; } @Override public void init(ProcessorContext> context) { this.context = context; + // This processor runs in the upstream transform's task, so the task's partition is the + // identity of the instance whose reports it is forwarding. + this.upstreamPartition = context.taskId().partition(); } @Override @@ -82,8 +99,16 @@ public void process(Record> record) { } ctx.forward(record.withKey(encodedKey)); } else { - // Watermark report: forward as-is; the sink's partitioner broadcasts it to all partitions. - ctx.forward(record); + WatermarkPayload report = payload.asWatermark(); + ctx.forward( + new Record>( + record.key(), + KStreamsPayload.watermark( + report.getWatermarkMillis(), + report.getTransformId(), + upstreamPartition, + upstreamPartitionCount), + record.timestamp())); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java index 1b0ffae23fde..d00b01a3753d 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java @@ -284,7 +284,8 @@ private void forwardData( private void forwardWatermark(Record> trigger, long watermarkMillis) { ProcessorContext> ctx = checkInitialized(context); - // Stamped with this transform's own id; GroupByKey is a single instance for now (0 of 1). + // Labelled as the only source a consumer will see; see ExecutableStageProcessor for why an + // in-process edge reports a single source and a shuffle relabels. ctx.forward( new Record>( trigger.key(), diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java new file mode 100644 index 000000000000..bb82a403e557 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java @@ -0,0 +1,286 @@ +/* + * 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; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.nio.file.Files; +import java.util.UUID; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +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.Metrics; +import org.apache.beam.sdk.metrics.MetricsFilter; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.beam.sdk.util.construction.SplittableParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.joda.time.Duration; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.testcontainers.kafka.KafkaContainer; +import org.testcontainers.utility.DockerImageName; + +/** + * Runs a pipeline through the production {@link KafkaStreamsPipelineRunner} against a real Kafka + * broker, rather than through the {@code TopologyTestDriver} the rest of the suite uses. + * + *

    The test driver stands in for a broker well enough for translation and windowing logic, but it + * runs one instance in one thread and fakes the topics. Everything that only exists on a real + * cluster is untested by it: the runner creating its own bootstrap and repartition topics, records + * actually round-tripping through a repartition topic, exactly-once processing, the state stores' + * changelog, and the Kafka Streams application lifecycle. This test covers that path. + * + *

    It needs Docker and so is not part of the default build; the {@code brokerIntegrationTest} + * Gradle task runs it. + */ +@RunWith(JUnit4.class) +public class KafkaStreamsRunnerBrokerIT { + + private static final String NAMESPACE = "brokerIT"; + private static final String GROUPS_COUNTER = "groups"; + + /** How long to wait for the streaming application to work through the pipeline. */ + private static final Duration TIMEOUT = Duration.standardMinutes(2); + + private static KafkaContainer kafka; + + @BeforeClass + public static void startBroker() { + // The official Apache Kafka image. 4.0.0 rather than the 3.9.0 the runner's client is built + // against: Testcontainers' KafkaContainer cannot bring up the 3.9.0 image (it exits during + // startup), and a client talking to a newer broker is the compatibility direction Kafka + // supports anyway. + kafka = new KafkaContainer(DockerImageName.parse("apache/kafka:4.0.0")); + kafka.start(); + } + + @AfterClass + public static void stopBroker() { + if (kafka != null) { + kafka.stop(); + } + } + + /** Emits a fixed set of keyed elements, one per key group. */ + private static class EmitKvsFn extends DoFn> { + @ProcessElement + public void processElement(OutputReceiver> out) { + out.output(KV.of("a", 1)); + out.output(KV.of("a", 2)); + out.output(KV.of("b", 3)); + } + } + + /** Counts the groups that come out of the GroupByKey. */ + private static class CountGroupsFn extends DoFn>, Void> { + private final Counter groups = Metrics.counter(NAMESPACE, GROUPS_COUNTER); + + @ProcessElement + public void processElement() { + groups.inc(); + } + } + + private KafkaStreamsPipelineOptions options() { + return options(1); + } + + private KafkaStreamsPipelineOptions options(int topicPartitions) { + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + options.setRunner(CrashingRunner.class); + options.setBootstrapServers(kafka.getBootstrapServers()); + options.setApplicationId("ks-broker-it-" + UUID.randomUUID()); + options.setInternalParallelism(topicPartitions); + options + .as(PortablePipelineOptions.class) + .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); + try { + options.setStateDir(Files.createTempDirectory("ks-broker-it").toString()); + } catch (Exception e) { + throw new RuntimeException(e); + } + return options; + } + + @Test + public void groupByKeyRunsThroughARealBrokerAndReportsMetrics() throws Exception { + KafkaStreamsPipelineOptions options = options(); + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply(Impulse.create()) + .apply("emit", ParDo.of(new EmitKvsFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())) + .apply(GroupByKey.create()) + .apply("countGroups", ParDo.of(new CountGroupsFn())); + + // The same conversion the test runner does: translate Read-based sources as the primitive Read + // the runner supports rather than the splittable-DoFn expansion. + SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(pipeline); + RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); + JobInfo jobInfo = + JobInfo.create( + options.getApplicationId(), + options.getJobName(), + "", + PipelineOptionsTranslation.toProto(options)); + + PipelineResult result = new KafkaStreamsPipelineRunner(options).run(pipelineProto, jobInfo); + try { + // Two keys in, so two groups out once the elements have travelled through the repartition + // topic and the watermark has closed the global window. + assertThat(awaitCounter(result, 2L), is(2L)); + } finally { + result.cancel(); + } + } + + /** + * Collapses every group onto one key, so the next GroupByKey has to shuffle across partitions. + */ + private static class ToSingleKeyFn + extends DoFn>, KV> { + @ProcessElement + public void processElement( + @Element KV> group, OutputReceiver> out) { + int sum = 0; + for (int value : group.getValue()) { + sum += value; + } + out.output(KV.of("all", sum)); + } + } + + /** Builds the two-GroupByKey pipeline used by the chained tests. */ + private static void buildChainedPipeline(Pipeline pipeline) { + pipeline + .apply(Impulse.create()) + .apply("emit", ParDo.of(new EmitKvsFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())) + .apply("groupPerKey", GroupByKey.create()) + .apply("toSingleKey", ParDo.of(new ToSingleKeyFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())) + .apply("groupAll", GroupByKey.create()) + .apply("countGroups", ParDo.of(new CountGroupsFn())); + } + + private static PipelineResult runPipeline( + Pipeline pipeline, KafkaStreamsPipelineOptions options) { + SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(pipeline); + RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); + JobInfo jobInfo = + JobInfo.create( + options.getApplicationId(), + options.getJobName(), + "", + PipelineOptionsTranslation.toProto(options)); + return new KafkaStreamsPipelineRunner(options).run(pipelineProto, jobInfo); + } + + @Test + public void chainedGroupByKeysAreCorrectOnOnePartition() throws Exception { + // The control for the partitioned case below: the same shape, one partition throughout. + KafkaStreamsPipelineOptions options = options(1); + Pipeline pipeline = Pipeline.create(options); + buildChainedPipeline(pipeline); + + PipelineResult result = runPipeline(pipeline, options); + try { + assertThat(awaitCounter(result, 1L), is(1L)); + } finally { + result.cancel(); + } + } + + @Test + public void chainedGroupByKeysAreCorrectAcrossPartitions() throws Exception { + // Two GroupByKeys with a partitioned shuffle between them, which is what makes each task's + // watermark identity matter. The second GroupByKey aggregates the reports of every task of the + // first, so those tasks have to report under their own partition: if each claimed to be the + // only partition, the second would advance its watermark on the first report it saw and fire + // before the remaining partitions had contributed their groups. + KafkaStreamsPipelineOptions options = options(4); + Pipeline pipeline = Pipeline.create(options); + buildChainedPipeline(pipeline); + + PipelineResult result = runPipeline(pipeline, options); + try { + // Everything collapses onto one key, so the second GroupByKey emits exactly one group — and + // only once every partition of the first has contributed to it. + assertThat(awaitCounter(result, 1L), is(1L)); + // A premature firing would show up as a second group, so give one a chance to appear. + Thread.sleep(5_000L); + assertThat(counterValue(result), is(1L)); + } finally { + result.cancel(); + } + } + + /** + * Polls the pipeline's metrics until the counter reaches {@code expected} or the timeout hits. + */ + private static long awaitCounter(PipelineResult result, long expected) throws Exception { + long deadline = System.currentTimeMillis() + TIMEOUT.getMillis(); + long value = 0; + while (System.currentTimeMillis() < deadline) { + value = counterValue(result); + if (value >= expected) { + return value; + } + Thread.sleep(500L); + } + return value; + } + + private static long counterValue(PipelineResult result) { + MetricQueryResults query = + result + .metrics() + .queryMetrics( + MetricsFilter.builder() + .addNameFilter(MetricNameFilter.named(NAMESPACE, GROUPS_COUNTER)) + .build()); + if (Iterables.isEmpty(query.getCounters())) { + return 0L; + } + MetricResult counter = Iterables.getOnlyElement(query.getCounters()); + return counter.getAttempted(); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java new file mode 100644 index 000000000000..38669138075c --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java @@ -0,0 +1,123 @@ +/* + * 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 java.util.Properties; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.kafka.streams.processor.TaskId; +import org.apache.kafka.streams.processor.api.MockProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.junit.Test; + +/** + * Tests how {@link ShuffleByKeyProcessor} restamps a watermark report as it is about to cross a + * repartition topic. + * + *

    Upstream of the shuffle a transform forwards its watermark in process, to its fused children, + * which see exactly one instance of it — so the report names a single source. The sink below the + * shuffle broadcasts each report to every partition, so a downstream task sees a report from every + * instance of the upstream transform and has to tell them apart to know when it has heard from all + * of them. The shuffle is where that identity is attached. + */ +public class ShuffleByKeyProcessorTest { + + private static final String UPSTREAM_ID = "upstream"; + + @SuppressWarnings("unchecked") + private static ShuffleByKeyProcessor processorFor(int taskPartition, int upstreamPartitions) { + ShuffleByKeyProcessor processor = + new ShuffleByKeyProcessor( + (org.apache.beam.sdk.coders.Coder) + (org.apache.beam.sdk.coders.Coder) StringUtf8Coder.of(), + upstreamPartitions); + MockProcessorContext> ctx = + new MockProcessorContext<>(new Properties(), new TaskId(0, taskPartition), null); + processor.init(ctx); + lastContext = ctx; + return processor; + } + + private static MockProcessorContext> lastContext; + + private static Record> watermark(long millis) { + // As forwarded in process by the upstream transform: a single source, since a fused child sees + // exactly one instance of it. + return new Record<>(new byte[0], KStreamsPayload.watermark(millis, UPSTREAM_ID, 0, 1), 0L); + } + + @Test + public void restampsTheWatermarkWithTheUpstreamInstanceIdentity() { + // Instance 2 of a 4-instance upstream transform. + ShuffleByKeyProcessor processor = processorFor(2, 4); + + processor.process(watermark(500L)); + + assertThat(lastContext.forwarded().size(), is(1)); + WatermarkPayload out = lastContext.forwarded().get(0).record().value().asWatermark(); + assertThat(out.getWatermarkMillis(), is(500L)); + // The transform id still names the producer, so a downstream aggregator matches it to the + // upstream it expects; the partition identity is what it counts. + assertThat(out.getTransformId(), is(UPSTREAM_ID)); + assertThat(out.getSourcePartition(), is(2)); + assertThat(out.getTotalSourcePartitions(), is(4)); + } + + @Test + public void distinctUpstreamInstancesRestampDistinctly() { + processorFor(0, 4).process(watermark(100L)); + WatermarkPayload first = lastContext.forwarded().get(0).record().value().asWatermark(); + processorFor(3, 4).process(watermark(100L)); + WatermarkPayload second = lastContext.forwarded().get(0).record().value().asWatermark(); + + // Two instances of the same transform must be distinguishable downstream, or a consumer would + // treat one report as if every instance had already reported. + assertThat(first.getSourcePartition(), is(0)); + assertThat(second.getSourcePartition(), is(3)); + } + + @Test + public void anUnpartitionedUpstreamStillReportsASingleSource() { + ShuffleByKeyProcessor processor = processorFor(0, 1); + + processor.process(watermark(700L)); + + WatermarkPayload out = lastContext.forwarded().get(0).record().value().asWatermark(); + assertThat(out.getSourcePartition(), is(0)); + assertThat(out.getTotalSourcePartitions(), is(1)); + } + + @Test + public void dataIsRekeyedByTheBeamKeyAndNotRestamped() { + ShuffleByKeyProcessor processor = processorFor(1, 4); + + processor.process( + new Record<>( + new byte[0], + KStreamsPayload.data( + WindowedValues.valueInGlobalWindow( + org.apache.beam.sdk.values.KV.of("key", "value"))), + 0L)); + + assertThat(lastContext.forwarded().size(), is(1)); + assertThat(lastContext.forwarded().get(0).record().value().isData(), is(true)); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java new file mode 100644 index 000000000000..912d51d6b5d1 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java @@ -0,0 +1,152 @@ +/* + * 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.instanceOf; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.not; +import static org.hamcrest.MatcherAssert.assertThat; + +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.coders.VarIntCoder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; +import org.apache.beam.sdk.transforms.windowing.SlidingWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.transforms.windowing.WindowFn; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.beam.sdk.util.construction.RehydratedComponents; +import org.apache.beam.sdk.util.construction.WindowingStrategyTranslation; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.joda.time.Duration; +import org.junit.Test; + +/** + * Checks that the runner reconstructs the standard WindowFns from the language-neutral windowing + * strategy in the pipeline proto. + * + *

    The runner executes GroupAlsoByWindow itself (see {@link WindowedGroupByKeyProcessor}), so it + * has to rebuild the WindowFn from the proto rather than call into the SDK. Beam gives the standard + * WindowFns a URN and a parameter payload — {@code beam:window_fn:fixed_windows:v1} and friends — + * and those are what {@link org.apache.beam.runners.core.ReduceFnRunner} interprets directly. Every + * SDK emits the same URNs for them, so a pipeline built in another language that uses fixed, + * sliding, session or global windows produces a strategy this runner can rebuild; these tests pin + * that down. + * + *

    What this does not cover is a WindowFn the user wrote themselves. That cannot be + * interpreted runner-side at all: it is opaque to the runner and would have to be executed through + * the SDK harness that owns it. The runner does not support that today — {@code + * WindowingStrategyTranslation.windowFnFromProto} rejects an unrecognised URN — and it is the case + * that would really exercise cross-language windowing. + */ +public class StandardWindowFnTranslationTest { + + private static final Duration WINDOW_SIZE = Duration.millis(10); + + private static class EmitKvFn extends DoFn> { + @ProcessElement + public void processElement(OutputReceiver> out) { + out.output(KV.of("a", 1)); + } + } + + /** A windowed GroupByKey pipeline, as a proto — the form the runner is handed. */ + private static RunnerApi.Pipeline windowedPipelineProto(WindowFn windowFn) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + pipeline + .apply(Impulse.create()) + .apply(ParDo.of(new EmitKvFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())) + .apply(Window.into(windowFn)) + .apply(GroupByKey.create()); + return PipelineTranslation.toProto(pipeline); + } + + /** The windowing strategy of the GroupByKey's input, which is what the translator reads. */ + private static RunnerApi.WindowingStrategy nonGlobalStrategy(RunnerApi.Pipeline proto) { + for (RunnerApi.WindowingStrategy strategy : + proto.getComponents().getWindowingStrategiesMap().values()) { + if (!strategy + .getWindowFn() + .getUrn() + .equals(WindowingStrategyTranslation.GLOBAL_WINDOWS_URN)) { + return strategy; + } + } + throw new AssertionError("pipeline has no non-global windowing strategy"); + } + + @Test + public void fixedWindowsTravelAsTheStandardUrn() { + RunnerApi.WindowingStrategy strategy = + nonGlobalStrategy(windowedPipelineProto(FixedWindows.of(WINDOW_SIZE))); + + // The same URN and payload any SDK emits for fixed windows. + assertThat(strategy.getWindowFn().getUrn(), is(WindowingStrategyTranslation.FIXED_WINDOWS_URN)); + } + + @Test + public void slidingWindowsTravelAsTheStandardUrn() { + RunnerApi.WindowingStrategy strategy = + nonGlobalStrategy( + windowedPipelineProto(SlidingWindows.of(WINDOW_SIZE).every(Duration.millis(5)))); + + assertThat( + strategy.getWindowFn().getUrn(), is(WindowingStrategyTranslation.SLIDING_WINDOWS_URN)); + } + + @Test + public void aStandardWindowFnNeedsNoJavaSerialization() { + RunnerApi.Pipeline proto = windowedPipelineProto(FixedWindows.of(WINDOW_SIZE)); + + // Java serialization is the fallback for a WindowFn with no standard URN. A strategy that fell + // back to it here would only be reconstructable by a Java runner reading a Java pipeline. + for (RunnerApi.WindowingStrategy strategy : + proto.getComponents().getWindowingStrategiesMap().values()) { + assertThat( + strategy.getWindowFn().getUrn(), + is(not(WindowingStrategyTranslation.SERIALIZED_JAVA_WINDOWFN_URN))); + } + } + + @Test + public void hydratingTheStandardUrnRebuildsTheWindowFnTheRunnerRunsWith() throws Exception { + RunnerApi.Pipeline proto = windowedPipelineProto(FixedWindows.of(WINDOW_SIZE)); + RunnerApi.WindowingStrategy strategy = nonGlobalStrategy(proto); + + // The translator's own path: rebuild the strategy from the proto alone. + WindowingStrategy hydrated = + WindowingStrategyTranslation.fromProto( + strategy, RehydratedComponents.forComponents(proto.getComponents())); + + assertThat(hydrated.getWindowFn(), instanceOf(FixedWindows.class)); + assertThat(((FixedWindows) hydrated.getWindowFn()).getSize(), is(WINDOW_SIZE)); + // The window coder the runner encodes state and timers with comes from this WindowFn, so it is + // the standard interval-window coder rather than anything SDK-specific. + assertThat(hydrated.getWindowFn().windowCoder(), is(IntervalWindow.getCoder())); + } +} From cf1f10bdf825053ebe28c5ca8d4507419e5bb104 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Mon, 3 Aug 2026 12:45:55 +0500 Subject: [PATCH 24/37] [GSoC 2026] Kafka Streams runner: bound a bundle by element count (#39578) * [GSoC 2026] Kafka Streams runner: bound a bundle by element count A bundle stayed open until the next watermark, so on a stream that produces steadily it grew without limit and nothing it had already processed was emitted until a watermark happened to arrive. maxBundleSize was declared as a pipeline option but nothing read it. --- .../streams/KafkaStreamsPipelineOptions.java | 7 +- .../translation/ExecutableStageProcessor.java | 41 +++++- .../ExecutableStageTranslator.java | 3 +- .../translation/BundleBoundaryTest.java | 118 ++++++++++++++++++ ...ExecutableStageProcessorWatermarkTest.java | 4 +- .../translation/MetricsAcrossBundlesTest.java | 81 ++++++++++++ 6 files changed, 250 insertions(+), 4 deletions(-) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/BundleBoundaryTest.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/MetricsAcrossBundlesTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java index a5a8bb9328b1..e95268ac1308 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java @@ -49,7 +49,12 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { void setMaxBundleSize(int maxBundleSize); - @Description("Soft cap on bundle wall-clock duration in milliseconds.") + @Description( + "Intended cap on how long a bundle may stay open, in milliseconds. NOT APPLIED YET: closing a" + + " bundle from a wall-clock punctuator made a pipeline with two chained GroupByKeys" + + " across several partitions emit its groups repeatedly against a real broker, so only" + + " the element-count bound is enforced for now. See" + + " https://github.com/apache/beam/issues/18479.") @Default.Integer(1000) int getMaxBundleTimeMs(); 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 63fefe3e15c4..ef376606114b 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 @@ -65,6 +65,18 @@ * across the upstream transform's partitions actually advances. Until every partition has reported, * the watermark is held and nothing is forwarded — but data is still processed in the meantime. * + *

    A bundle is also bounded in size, by {@code --maxBundleSize}, and closed once that many + * elements have been fed to it. Without the bound a bundle stays open until the next watermark, + * which on a stream that produces steadily lets it grow without limit. The bound is checked as + * elements arrive. A time bound ({@code --maxBundleTimeMs}) is not applied yet — see the option's + * own documentation. + * + *

    Closing a bundle asks Kafka Streams to commit, so the elements a bundle consumed and the + * records it produced are committed together and a restart replays either all of the bundle or none + * of it. Note that this aligns commits to bundle boundaries but does not stop Kafka + * Streams from committing on its own interval part-way through a bundle; closing the bundle first + * from a pre-commit hook would be needed to rule that out entirely. + * *

    This is the Kafka Streams analogue of Flink's {@code ExecutableStageDoFnOperator} and Spark's * {@code SparkExecutableStageFunction}. State, timers, and side inputs are out of scope for this * first version: the stage is executed with {@link StateRequestHandler#unsupported()} and no timer @@ -107,6 +119,12 @@ class ExecutableStageProcessor private @Nullable StageBundleFactory stageBundleFactory; private @Nullable RemoteBundle currentBundle; + /** Bound on how many elements may be fed to one bundle. */ + private final int maxBundleSize; + + /** Elements fed to the open bundle, for the size bound above. */ + private int elementsInBundle; + /** * @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 @@ -120,13 +138,15 @@ class ExecutableStageProcessor String transformId, Set upstreamTransformIds, MetricsContainerImpl metricsContainer, - Map outputChildByPCollectionId) { + Map outputChildByPCollectionId, + int maxBundleSize) { this.stagePayload = stagePayload; this.jobInfo = jobInfo; this.transformId = transformId; this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); this.metricsContainer = metricsContainer; this.outputChildByPCollectionId = ImmutableMap.copyOf(outputChildByPCollectionId); + this.maxBundleSize = maxBundleSize; } /** A harness output element together with the id of the output PCollection it belongs to. */ @@ -185,9 +205,13 @@ public void process(Record> record) { try { ensureBundleOpen(); mainInputReceiver().accept(payload.getData()); + elementsInBundle++; } catch (Exception e) { throw new RuntimeException("Failed to process element through SDK harness", e); } + if (elementsInBundle >= maxBundleSize) { + closeBundleAndFlush(record); + } } private void ensureBundleOpen() throws Exception { @@ -241,6 +265,7 @@ public void onCompleted(ProcessBundleResponse response) { currentBundle = factory.getBundle( outputReceiverFactory, StateRequestHandler.unsupported(), progressHandler); + elementsInBundle = 0; } private FnDataReceiver> mainInputReceiver() { @@ -252,6 +277,18 @@ private FnDataReceiver> mainInputReceiver() { return receiver; } + /** + * Finishes the open bundle, forwards everything it produced, and asks Kafka Streams to commit. + * + *

    The commit request is what ties a bundle to a transaction: the elements the bundle consumed + * and the records it produced are then committed together, so a restart either replays the whole + * bundle or none of it. + * + *

    The outputs carry the key of the record that closed the bundle. An executable stage is + * unkeyed — it runs stateless, with no state or timers — so the Kafka record key means nothing to + * it and is only being carried along; where the key does matter, downstream sets it, as {@link + * ShuffleByKeyProcessor} does from the Beam key before a GroupByKey. + */ private void closeBundleAndFlush(Record> record) { RemoteBundle bundle = currentBundle; if (bundle == null) { @@ -265,6 +302,7 @@ private void closeBundleAndFlush(Record> record) { throw new RuntimeException("Failed to close SDK harness bundle", e); } finally { currentBundle = null; + elementsInBundle = 0; } ProcessorContext> ctx = checkInitialized(context); // The harness has finished the bundle (close() returned) so no further enqueues happen. @@ -283,6 +321,7 @@ private void closeBundleAndFlush(Record> record) { ctx.forward(outputRecord, childNode); } } + ctx.commit(); } private void forwardWatermark(Record> record, long watermarkMillis) { 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 caefa6534fad..c59e5b919aba 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 @@ -108,7 +108,8 @@ public void translate( transformId, ImmutableSet.of(parentProcessor), context.getMetricsContainerStepMap().getContainer(transformId), - outputChildByPCollectionId), + outputChildByPCollectionId, + context.getPipelineOptions().getMaxBundleSize()), parentProcessor); if (multiOutput) { diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/BundleBoundaryTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/BundleBoundaryTest.java new file mode 100644 index 000000000000..c057434f0b2a --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/BundleBoundaryTest.java @@ -0,0 +1,118 @@ +/* + * 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 static org.hamcrest.Matchers.greaterThanOrEqualTo; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.Before; +import org.junit.Test; + +/** + * Tests that a bundle is closed once it reaches {@code --maxBundleSize}, rather than staying open + * until the next watermark. + * + *

    The bound is observable through the DoFn's own lifecycle: {@code @FinishBundle} runs once per + * bundle the SDK harness processes, so feeding a known number of elements with a known bound tells + * us how many bundles the stage actually opened. + * + *

    The elements have to reach the stage as separate records for the bound to see them, since it + * counts what is fed to the stage rather than what the user's code emits inside it. A DoFn that + * fans one element out into many would be fused into the same stage and still be a single input, so + * these pipelines read the elements from a {@link Create} instead — the runner translates that to a + * primitive Read, which forwards one record per element. + */ +public class BundleBoundaryTest { + + private static final int ELEMENTS = 50; + private static final int MAX_BUNDLE_SIZE = 10; + + /** Records how many times {@code @FinishBundle} fired, i.e. how many bundles were processed. */ + private static final List FINISHED_BUNDLES = + Collections.synchronizedList(new ArrayList<>()); + + @Before + public void resetCounters() { + FINISHED_BUNDLES.clear(); + } + + /** Counts the bundles it is asked to process. */ + private static class CountBundlesFn extends DoFn { + @ProcessElement + public void processElement(@Element Integer element, OutputReceiver out) { + out.output(element); + } + + @FinishBundle + public void finishBundle() { + FINISHED_BUNDLES.add("bundle"); + } + } + + private static List elements() { + List elements = new ArrayList<>(); + for (int i = 0; i < ELEMENTS; i++) { + elements.add(i); + } + return elements; + } + + private static Pipeline buildPipeline(KafkaStreamsPipelineOptions options) { + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply("read", Create.of(elements())) + .apply("countBundles", ParDo.of(new CountBundlesFn())); + return pipeline; + } + + private static Pipeline pipelineWithBundleSize(int maxBundleSize) { + KafkaStreamsPipelineOptions options = + KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class); + options.setMaxBundleSize(maxBundleSize); + return buildPipeline(options); + } + + @Test + public void aBundleIsClosedOnceItReachesTheSizeBound() { + KafkaStreamsTestRunner.run(pipelineWithBundleSize(MAX_BUNDLE_SIZE)); + + // 50 elements bounded at 10 cannot have gone through in fewer than 5 bundles. Without the + // bound the whole run is one bundle, so this is what tells the two apart. The count is a lower + // bound rather than exact: a watermark arriving mid-bundle also closes one. + assertThat(FINISHED_BUNDLES.size(), is(greaterThanOrEqualTo(ELEMENTS / MAX_BUNDLE_SIZE))); + } + + @Test + public void aBoundLargerThanTheInputLeavesASingleBundle() { + // The control: with a bound nothing reaches, the stage keeps one bundle open until the + // terminal watermark closes it. + KafkaStreamsTestRunner.run(pipelineWithBundleSize(ELEMENTS * 10)); + + assertThat(FINISHED_BUNDLES.size(), is(1)); + } +} 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 010d98d52e88..290e109796a8 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 @@ -62,7 +62,9 @@ private static ExecutableStageProcessor newProcessor() { ImmutableSet.of(UPSTREAM_ID), new MetricsContainerImpl(STAGE_ID), // Single-output: no per-output routing (this test drives the watermark path directly). - ImmutableMap.of()); + ImmutableMap.of(), + // The bundle size bound is irrelevant to the watermark path this test drives. + 1000); } /** 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/MetricsAcrossBundlesTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/MetricsAcrossBundlesTest.java new file mode 100644 index 000000000000..a5b690c39ddb --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/MetricsAcrossBundlesTest.java @@ -0,0 +1,81 @@ +/* + * 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.KafkaStreamsPipelineOptions; +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.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; + +/** + * Checks that a user counter stays correct when the bundle size bound makes a stage run many small + * bundles instead of one large one. + * + *

    The runner folds the metrics the SDK harness reports into the job's step map as each bundle + * completes, and those updates add rather than replace. That is right only if each report covers + * its own bundle, so splitting the same input across more bundles must not change the total. + */ +public class MetricsAcrossBundlesTest { + private static class CountingFn extends DoFn { + private final Counter counter = Metrics.counter("probe", "elements"); + + @ProcessElement + public void processElement(@Element Integer in, OutputReceiver out) { + counter.inc(); + out.output(in); + } + } + + private static long run(int maxBundleSize) { + KafkaStreamsPipelineOptions options = + KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class); + options.setMaxBundleSize(maxBundleSize); + Pipeline p = Pipeline.create(options); + p.apply(Create.of(1, 2, 3, 4, 5, 6)).apply(ParDo.of(new CountingFn())); + MetricResults metrics = KafkaStreamsTestRunner.run(p); + MetricQueryResults q = + metrics.queryMetrics( + MetricsFilter.builder() + .addNameFilter(MetricNameFilter.named("probe", "elements")) + .build()); + return Iterables.getOnlyElement(q.getCounters()).getAttempted(); + } + + @Test + public void oneBundle() { + assertThat(run(1000), is(6L)); + } + + @Test + public void manyBundles() { + assertThat(run(1), is(6L)); + } +} From fc36301e6c2e6e98a9758eea815e7f4f8a03f966 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:54:53 +0500 Subject: [PATCH 25/37] [GSoC 2026] Kafka Streams runner: CombineTest coverage and two review follow-ups (#39610) * [GSoC 2026] Kafka Streams runner: CombineTest coverage and two review follow-ups Enables CombineTest in the ValidatesRunner suite, taking it from 49 to 59 tests. Combine was expected to work without a translator of its own, since the fuser expands Combine.perKey into a GroupByKey with the combining logic running as ordinary ParDos in the SDK harness, but nothing exercised that. BasicTests passes in full, including hot-key fanout and the accumulation-mode variant, and WindowingTests contributes the fixed-window and empty-window cases. The remainder falls out on category excludes the task already declares. testSessionsCombine is sickbayed alongside the existing merging windows entry, and it is the only Combine failure. Corrects the Flatten partition-count comment, which asserted that the inputs are co-partitioned and so implied the Math.max over them was redundant. Neither half held. The max is not a no-op in principle: Kafka Streams merges the subtopologies of every parent a processor is wired to and gives the result as many tasks as its largest source topic has partitions. But the mismatched shape does not reach this translator, because the fuser folds such a Flatten into the harness stage, and the runner Flattens that do arrive come from the fuser deduplicating partial outputs of one PCollection. FlattenParallelismTest records that, so a change letting the mismatched shape through starts failing there rather than producing a pipeline that stalls waiting for a watermark report that never comes. Guards the null record key in GroupByKeyBroadcastPartitioner.partitions(). partition() already guarded it, but partitions() is the method Kafka Streams calls and it hashed the key unguarded. --- runners/kafka-streams/build.gradle | 4 +- .../translation/FlattenTranslator.java | 10 +- .../GroupByKeyBroadcastPartitioner.java | 8 + .../translation/FlattenParallelismTest.java | 141 ++++++++++++++++++ 4 files changed, 160 insertions(+), 3 deletions(-) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenParallelismTest.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index d1326c079e5d..616257ddd788 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -113,8 +113,9 @@ def sickbayTests = [ // Merging (session) windows are not supported yet: ReduceFnRunner drives them through a merging // window set that moves per-window state as windows merge, which this first windowing pass does // not implement. Non-merging windows (fixed, sliding), the default trigger and timestamp - // combiners do work. Lands with the follow-up windowing PR. + // combiners do work, for both GroupByKey and Combine. Lands with the follow-up windowing PR. 'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests.testGroupByKeyMergingWindows', + 'org.apache.beam.sdk.transforms.CombineTest$WindowingTests.testSessionsCombine', // A DoFn whose @StartBundle throws never gets to report its error: SdkHarnessClient.newBundle // sends the ProcessBundleRequest and then blocks in GrpcDataService.createOutboundAggregator // waiting for the SDK harness to open its data stream, which a bundle that failed during setup @@ -177,6 +178,7 @@ tasks.register("validatesRunner", Test) { includeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest' includeTestsMatching 'org.apache.beam.sdk.transforms.GroupByKeyTest*' includeTestsMatching 'org.apache.beam.sdk.transforms.ParDoTest*' + includeTestsMatching 'org.apache.beam.sdk.transforms.CombineTest*' for (String test : sickbayTests) { excludeTestsMatching test } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java index a5c8ce05baeb..793c1e1dc530 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java @@ -57,8 +57,14 @@ public void translate( Set seenInputs = new HashSet<>(); List parentProcessors = new ArrayList<>(); Set upstreamTransformIds = new HashSet<>(); - // Kafka Streams puts a processor and the parents it is wired to in one subtopology, so the - // inputs are co-partitioned and this Flatten runs at their partition count. + // How many instances this Flatten runs as. Kafka Streams merges the subtopologies of every + // parent a processor is wired to and gives the merged subtopology as many tasks as its largest + // source topic has partitions, so the max is what that comes to. In practice the inputs agree: + // a Flatten whose branches could disagree — one through a GroupByKey, one straight from a + // source — is fused into the harness stage instead of becoming a node here, and the runner + // Flattens that do reach this translator come from the fuser deduplicating partial outputs of + // one PCollection. The max is kept as the cheap conservative choice rather than asserting that + // agreement, which is not enforced anywhere. int partitionCount = 1; for (String inputPCollectionId : transform.getInputsMap().values()) { if (!seenInputs.add(inputPCollectionId)) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java index b5ddcf2536a2..3c775c86f14c 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java @@ -56,6 +56,14 @@ public Optional> partitions( } return Optional.of(all); } + if (key == null) { + // A keyless record has no partition it must go to, so leave the choice to Kafka rather than + // hashing a null or pinning one partition: an empty Optional tells Kafka Streams no explicit + // partition was chosen, and the producer's default partitioner spreads keyless records over + // the topic instead of piling them onto one. This is the method Kafka Streams calls, so the + // null has to be handled here and not only in partition() above. + return Optional.empty(); + } int partition = Utils.toPositive(Utils.murmur2(key)) % numPartitions; return Optional.of(Collections.singleton(partition)); } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenParallelismTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenParallelismTest.java new file mode 100644 index 000000000000..601fa40e679c --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenParallelismTest.java @@ -0,0 +1,141 @@ +/* + * 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 java.util.ArrayList; +import java.util.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Flatten; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionList; +import org.apache.kafka.streams.TopologyDescription; +import org.junit.Test; + +/** + * Pins down what happens to a Flatten whose branches would run at different parallelisms — one + * through a GroupByKey and so at the shuffle's parallelism, one straight from a source and so a + * single instance. + * + *

    This matters because a Flatten runs as one set of tasks over all of its inputs. Kafka Streams + * merges the subtopologies of every parent a processor is wired to and gives the result as many + * tasks as its largest source topic has partitions, so a parent with fewer partitions would only + * produce on some of those tasks and the rest would wait forever for a watermark report from it. + * + *

    That does not arise, and these tests record why: the fuser folds such a Flatten into the SDK + * harness stages rather than leaving a node for the runner to translate, so the branches never + * share a subtopology and no Flatten node exists to run at a single parallelism. The Flattens that + * do reach {@link FlattenTranslator} come from the fuser deduplicating partial outputs of a single + * PCollection. If a change ever makes the mismatched shape reach the translator, these tests start + * failing and the partition-count handling there needs revisiting. + */ +public class FlattenParallelismTest { + + /** The name given to the Flatten below, which no topology node should be derived from. */ + private static final String FLATTEN_NAME = "merge"; + + private static class ToKvFn extends DoFn> { + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver> out) { + out.output(KV.of("k", input)); + } + } + + private static class UngroupFn extends DoFn>, Integer> { + @ProcessElement + public void processElement( + @Element KV> group, OutputReceiver out) { + for (int value : group.getValue()) { + out.output(value); + } + } + } + + private static Pipeline mixedParallelismFlatten(int internalParallelism) { + KafkaStreamsPipelineOptions options = + KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class); + options.setInternalParallelism(internalParallelism); + Pipeline pipeline = Pipeline.create(options); + + // Through a GroupByKey, so this branch runs at the shuffle's parallelism. + PCollection shuffled = + pipeline + .apply("createGrouped", Create.of(1, 2, 3)) + .apply("toKv", ParDo.of(new ToKvFn())) + .apply("group", GroupByKey.create()) + .apply("ungroup", ParDo.of(new UngroupFn())); + + // Straight from a source, so this branch is a single instance. + PCollection direct = pipeline.apply("createDirect", Create.of(4, 5, 6)); + + PCollectionList.of(shuffled).and(direct).apply("merge", Flatten.pCollections()); + return pipeline; + } + + /** Every processor node in the topology, across all subtopologies. */ + private static List processorNames(TopologyDescription description) { + List names = new ArrayList<>(); + for (TopologyDescription.Subtopology subtopology : description.subtopologies()) { + for (TopologyDescription.Node node : subtopology.nodes()) { + if (node instanceof TopologyDescription.Processor) { + names.add(node.name()); + } + } + } + return names; + } + + private static void assertFlattenWasFusedAway(TopologyDescription description) { + // No node stands for the Flatten. If one did, it would be wired to both branches and so would + // run over a merged subtopology whose smaller-parallelism parent could not reach all of its + // instances. + for (String name : processorNames(description)) { + assertThat( + "no processor node should stand for the Flatten, but found " + name, + name.contains(FLATTEN_NAME), + is(false)); + } + // The branches stay in separate subtopologies for the same reason: the source-fed branch, the + // one behind the shuffle, and the second source-fed branch. + assertThat(description.subtopologies().size(), is(3)); + } + + @Test + public void branchesAtDifferentParallelismsAreFusedRatherThanLeftToTheRunner() { + assertFlattenWasFusedAway( + KafkaStreamsTestRunner.translate(mixedParallelismFlatten(4)).getTopology().describe()); + } + + @Test + public void theSameHoldsAtASingleParallelism() { + // Whether the Flatten is fused is a property of the fused graph, not of the parallelism, so + // the shape is the same either way — which is why raising the parallelism cannot introduce a + // Flatten node over mismatched branches. + assertFlattenWasFusedAway( + KafkaStreamsTestRunner.translate(mixedParallelismFlatten(1)).getTopology().describe()); + } +} From 5861f31e8ac33749ea67061cd83ed213552c19ea Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Tue, 4 Aug 2026 21:09:06 +0500 Subject: [PATCH 26/37] [GSoC 2026] Kafka Streams runner: read unbounded sources (#39611) [GSoC 2026] Kafka Streams runner: read unbounded sources --- .../streams/KafkaStreamsPipelineOptions.java | 10 + .../streams/translation/ReadTranslator.java | 123 ++++++- .../translation/UnboundedReadProcessor.java | 304 ++++++++++++++++++ .../translation/UnboundedReadTest.java | 209 ++++++++++++ 4 files changed, 635 insertions(+), 11 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java index e95268ac1308..44abc8e5b34d 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java @@ -77,6 +77,16 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { void setTopicReplicationFactor(short topicReplicationFactor); + @Description( + "How many non-empty polls of an unbounded source to make before storing its checkpoint mark." + + " Taking a mark can be costly for some sources, so it is not worth doing on every poll;" + + " the cost of a larger value is that more elements are replayed after a restart, since" + + " the reader resumes from the last mark that was stored.") + @Default.Integer(10) + int getReadCheckpointNumBundles(); + + void setReadCheckpointNumBundles(int readCheckpointNumBundles); + @Description("Directory where Kafka Streams stores local state.") @Default.InstanceFactory(StateDirDefaultFactory.class) String getStateDir(); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java index 403499ca616c..f83442f97813 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java @@ -18,12 +18,14 @@ package org.apache.beam.runners.kafka.streams.translation; import java.io.IOException; +import java.util.List; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.WireCoderSetting; import org.apache.beam.runners.core.construction.SerializablePipelineOptions; import org.apache.beam.runners.fnexecution.wire.WireCoders; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.io.BoundedSource; +import org.apache.beam.sdk.io.UnboundedSource; import org.apache.beam.sdk.util.construction.ReadTranslation; import org.apache.beam.sdk.util.construction.RehydratedComponents; import org.apache.beam.sdk.util.construction.graph.PipelineNode; @@ -77,25 +79,92 @@ public void translate( // Read produces exactly one output PCollection; downstream consumers are separate PTransforms // whose inputs reference this PCollection id and are wired by their own translators. String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); - addReadNodes( - transformId, - boundedSource(transform), - pipeline.getComponents(), - outputPCollectionId, - context); - } - - private static BoundedSource boundedSource(RunnerApi.PTransform transform) { try { RunnerApi.ReadPayload payload = RunnerApi.ReadPayload.parseFrom(transform.getSpec().getPayload()); - return ReadTranslation.boundedSourceFromProto(payload); + // The same URN carries both kinds of source; the payload says which, and they need different + // processors. A bounded source is drained once and ends time; an unbounded one is polled + // repeatedly and moves the watermark as its reader reports progress. + if (payload.getIsBounded() == RunnerApi.IsBounded.Enum.UNBOUNDED) { + addUnboundedReadNodes( + transformId, + ReadTranslation.unboundedSourceFromProto(payload), + pipeline.getComponents(), + outputPCollectionId, + context); + } else { + addReadNodes( + transformId, + ReadTranslation.boundedSourceFromProto(payload), + pipeline.getComponents(), + outputPCollectionId, + context); + } } catch (IOException e) { throw new RuntimeException( - "Failed to read the BoundedSource from transform " + transform.getUniqueName(), e); + "Failed to read the source from transform " + transform.getUniqueName(), e); } } + /** + * Adds the source, {@link UnboundedReadProcessor}, and the store holding its checkpoint mark. + * + *

    The store keeps encoded bytes rather than the mark itself, since the mark's coder comes from + * the source and is only known here. + */ + private void addUnboundedReadNodes( + String transformId, + UnboundedSource source, + RunnerApi.Components components, + String outputPCollectionId, + KafkaStreamsTranslationContext context) { + PCollectionNode outputNode = + PipelineNode.pCollection( + outputPCollectionId, components.getPcollectionsOrThrow(outputPCollectionId)); + Coder> sdkWireCoder = sdkWireCoder(outputNode, components); + Coder> runnerWireCoder = runnerWireCoder(outputNode, components); + + Topology topology = context.getTopology(); + String sourceNodeName = transformId + SOURCE_SUFFIX; + String stateStoreName = transformId + STATE_STORE_SUFFIX; + String bootstrapTopic = context.getReadBootstrapTopic(transformId); + SerializablePipelineOptions options = + new SerializablePipelineOptions(context.getPipelineOptions()); + // Split here rather than in the processor: splitting belongs to translation, where it happens + // once for the pipeline instead of once per task instance, and the contract says nothing about + // splitting a source that has already been split. + UnboundedSource readableSource = singleSplitOf(source, context); + Coder checkpointCoder = readableSource.getCheckpointMarkCoder(); + int maxElementsPerPoll = context.getPipelineOptions().getMaxBundleSize(); + int checkpointEveryNPolls = context.getPipelineOptions().getReadCheckpointNumBundles(); + + topology.addSource( + sourceNodeName, + Serdes.ByteArray().deserializer(), + Serdes.ByteArray().deserializer(), + bootstrapTopic); + topology.addProcessor( + transformId, + () -> + new UnboundedReadProcessor<>( + readableSource, + options, + sdkWireCoder, + runnerWireCoder, + checkpointCoder, + stateStoreName, + transformId, + maxElementsPerPoll, + checkpointEveryNPolls), + sourceNodeName); + topology.addStateStore( + Stores.keyValueStoreBuilder( + Stores.persistentKeyValueStore(stateStoreName), Serdes.String(), Serdes.ByteArray()), + transformId); + + context.registerPCollectionProducer(outputPCollectionId, transformId); + } + /** * Adds the source, {@link ReadProcessor}, and state store for the read. The type variable {@code * T} captures the {@link BoundedSource}'s element type so the processor and its wire coders are @@ -138,6 +207,38 @@ private void addReadNodes( context.registerPCollectionProducer(outputPCollectionId, transformId); } + /** + * Splits an unbounded source into the single part this runner reads. + * + *

    A source is not obliged to be readable in its unsplit form — {@code split} is where several + * of them do their setup — so it is asked to split even though only one part is wanted. The count + * passed to {@code split} is only a hint, so what comes back has to be checked: taking the first + * of several splits would quietly drop whatever the others would have produced, which is data + * loss rather than a missing feature, so it fails instead. + */ + private static + UnboundedSource singleSplitOf( + UnboundedSource source, KafkaStreamsTranslationContext context) { + List> splits; + try { + splits = source.split(1, context.getPipelineOptions()); + } catch (Exception e) { + throw new RuntimeException("Failed to split unbounded source " + source, e); + } + if (splits.size() != 1) { + throw new UnsupportedOperationException( + "Unbounded source " + + source + + " split into " + + splits.size() + + " parts, but the Kafka Streams runner reads a source with a single reader and" + + " would therefore drop the data of every part but the first. Reading several" + + " splits in parallel is not supported yet; see" + + " https://github.com/apache/beam/issues/18479."); + } + return splits.get(0); + } + /** The coder the SDK harness would use on the wire, keeping unknown element coders intact. */ private static Coder> sdkWireCoder( PCollectionNode outputNode, RunnerApi.Components components) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java new file mode 100644 index 000000000000..b616e0fa8534 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java @@ -0,0 +1,304 @@ +/* + * 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 java.io.IOException; +import java.time.Duration; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark; +import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.kafka.streams.processor.Cancellable; +import org.apache.kafka.streams.processor.PunctuationType; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Reads an {@link UnboundedSource} and forwards its elements and watermark downstream. + * + *

    Where the bounded {@link ReadProcessor} drains its source once and jumps the watermark to the + * end of time, an unbounded source never finishes: it is polled repeatedly, and its watermark + * advances gradually as the reader reports progress. That difference is what makes this a streaming + * runner rather than a batch one — downstream windows close because the source says time has moved + * on, not because the input ran out. + * + *

    Polling happens on a wall-clock punctuator rather than in {@code process}, because the + * processor's bootstrap topic is empty and nothing else would drive it. Each turn reads at most + * {@link #maxElementsPerPoll} elements so a busy source cannot monopolise the Kafka Streams thread + * and starve the rest of the topology, then forwards the reader's watermark if it advanced. + * + *

    Restart is what the checkpoint mark is for. {@link UnboundedReader#getCheckpointMark()} + * describes the position the reader has consumed to; it is written to a persistent state store, and + * on {@link #init} the reader is created from the stored mark rather than from scratch, so a task + * that moves or restarts resumes where it left off instead of re-reading from the beginning. The + * store is changelogged and, under exactly-once, its writes commit atomically with the records the + * processor forwarded, so the mark can never be ahead of the data that was actually emitted. + * + *

    The source handed to this processor has already been split by {@link ReadTranslator}, which is + * where splitting belongs: it happens once for the pipeline rather than once per task instance, and + * the contract does not define splitting an already-split source. Reading several splits in + * parallel arrives with the topic-based shuffle work (#18479). As in the bounded processor, Kafka + * Streams disallows negative record timestamps, so each forwarded {@link Record} carries the Unix + * epoch and the Beam event time travels inside the {@link WindowedValue}. + */ +class UnboundedReadProcessor + implements Processor> { + + private static final Logger LOG = LoggerFactory.getLogger(UnboundedReadProcessor.class); + + /** Sole entry in the state store; the value is the encoded checkpoint mark. */ + static final String CHECKPOINT_KEY = "checkpoint"; + + /** How often the source is polled. */ + private static final Duration POLL_INTERVAL = Duration.ofMillis(50); + + private final UnboundedSource source; + private final SerializablePipelineOptions options; + // See ReadProcessor: a source produces decoded objects, but the downstream stage's harness input + // expects the runner-side wire form, so each element is transcoded through these two coders. + private final Coder> sdkWireCoder; + private final Coder> runnerWireCoder; + private final Coder checkpointCoder; + private final String stateStoreName; + private final String transformId; + private final int maxElementsPerPoll; + private final int checkpointEveryNPolls; + + private @Nullable ProcessorContext> context; + private @Nullable KeyValueStore checkpointStore; + private @Nullable UnboundedReader reader; + private boolean readerStarted; + private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + /** Set once the source's watermark reaches the end of time; it will produce nothing more. */ + private boolean exhausted; + + private @Nullable Cancellable scheduledPunctuator; + + UnboundedReadProcessor( + UnboundedSource source, + SerializablePipelineOptions options, + Coder> sdkWireCoder, + Coder> runnerWireCoder, + Coder checkpointCoder, + String stateStoreName, + String transformId, + int maxElementsPerPoll, + int checkpointEveryNPolls) { + this.source = source; + this.options = options; + this.sdkWireCoder = sdkWireCoder; + this.runnerWireCoder = runnerWireCoder; + this.checkpointCoder = checkpointCoder; + this.stateStoreName = stateStoreName; + this.transformId = transformId; + this.maxElementsPerPoll = maxElementsPerPoll; + this.checkpointEveryNPolls = checkpointEveryNPolls; + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + this.checkpointStore = context.getStateStore(stateStoreName); + this.scheduledPunctuator = + context.schedule(POLL_INTERVAL, PunctuationType.WALL_CLOCK_TIME, timestamp -> poll()); + } + + @Override + public void process(Record record) { + // The bootstrap topic carries no real data; a record arriving on it is just another chance to + // poll. The reader's own position decides what is actually emitted. + poll(); + } + + /** + * Drains what the source currently has, in batches, then publishes the watermark. + * + *

    A batch is capped at {@link #maxElementsPerPoll} so that the checkpoint mark and the + * watermark are updated as the reader progresses rather than only at the end. Batches run back to + * back while the source keeps filling them, since returning after every batch would cap + * throughput at one batch per punctuation interval. + * + *

    The run is bounded all the same. A source that always has data — which is the normal case + * for one that is keeping up — would otherwise never let this method return, and the Kafka + * Streams thread would never get back to committing or to the rest of the topology. So at most + * {@link #checkpointEveryNPolls} batches are taken before yielding, which is also where the + * checkpoint mark is stored, and the next punctuation carries on from there. + */ + private void poll() { + if (exhausted) { + return; + } + ProcessorContext> ctx = checkInitialized(context); + UnboundedReader currentReader = ensureReader(); + for (int batch = 0; batch < checkpointEveryNPolls; batch++) { + int emitted = readBatch(ctx, currentReader); + Instant watermark = currentReader.getWatermark(); + forwardWatermarkIfAdvanced(ctx, watermark); + if (!watermark.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) { + // The source has declared it will produce nothing further, so stop polling it. Store the + // final position first, since the loop will not come back to it. + storeCheckpoint(currentReader); + exhausted = true; + Cancellable punctuator = scheduledPunctuator; + if (punctuator != null) { + punctuator.cancel(); + scheduledPunctuator = null; + } + return; + } + if (emitted < maxElementsPerPoll) { + // Short batch: the source has nothing more for now, so store what was read and wait for + // the next punctuation rather than spinning on a reader that keeps returning false. + if (emitted > 0) { + storeCheckpoint(currentReader); + } + return; + } + } + // Yielded on the batch bound rather than on an empty source, so record the position reached. + storeCheckpoint(currentReader); + } + + /** Forwards up to {@link #maxElementsPerPoll} elements, returning how many were available. */ + private int readBatch( + ProcessorContext> ctx, UnboundedReader currentReader) { + int emitted = 0; + try { + while (emitted < maxElementsPerPoll) { + // start() positions the reader on its first element; advance() moves to the next. Either + // returning false means nothing is available right now — not that the source is finished, + // which is the difference from a bounded read. + boolean hasElement = readerStarted ? currentReader.advance() : currentReader.start(); + readerStarted = true; + if (!hasElement) { + break; + } + WindowedValue element = + WindowedValues.timestampedValueInGlobalWindow( + currentReader.getCurrent(), currentReader.getCurrentTimestamp()); + ctx.forward( + new Record>( + new byte[0], KStreamsPayload.data(toRunnerWire(element)), 0L)); + emitted++; + } + } catch (IOException e) { + throw new RuntimeException("Failed to read unbounded source for transform " + transformId, e); + } + return emitted; + } + + /** Publishes the reader's watermark, which is what lets downstream windows close. */ + private void forwardWatermarkIfAdvanced( + ProcessorContext> ctx, Instant watermark) { + if (!watermark.isAfter(lastForwardedWatermark)) { + return; + } + lastForwardedWatermark = watermark; + ctx.forward( + new Record>( + new byte[0], KStreamsPayload.watermark(watermark.getMillis(), transformId, 0, 1), 0L)); + } + + /** Creates the reader on first use, resuming from the stored checkpoint mark if there is one. */ + private UnboundedReader ensureReader() { + UnboundedReader existing = reader; + if (existing != null) { + return existing; + } + try { + UnboundedReader created = source.createReader(options.get(), restoreCheckpoint()); + reader = created; + return created; + } catch (Exception e) { + throw new RuntimeException( + "Failed to create a reader for unbounded source in transform " + transformId, e); + } + } + + private @Nullable CheckpointT restoreCheckpoint() { + KeyValueStore store = checkInitialized(checkpointStore); + byte[] encoded = store.get(CHECKPOINT_KEY); + if (encoded == null) { + return null; + } + try { + CheckpointT mark = CoderUtils.decodeFromByteArray(checkpointCoder, encoded); + LOG.info("Unbounded read {} resuming from a stored checkpoint mark", transformId); + return mark; + } catch (CoderException e) { + throw new RuntimeException( + "Failed to decode the checkpoint mark for transform " + transformId, e); + } + } + + private void storeCheckpoint(UnboundedReader currentReader) { + KeyValueStore store = checkInitialized(checkpointStore); + @SuppressWarnings("unchecked") + CheckpointT mark = (CheckpointT) currentReader.getCheckpointMark(); + try { + store.put(CHECKPOINT_KEY, CoderUtils.encodeToByteArray(checkpointCoder, mark)); + } catch (CoderException e) { + throw new RuntimeException( + "Failed to encode the checkpoint mark for transform " + transformId, e); + } + } + + /** Transcodes a raw element into the runner-side wire form the SDK harness input expects. */ + private WindowedValue toRunnerWire(WindowedValue element) { + try { + byte[] wireBytes = CoderUtils.encodeToByteArray(sdkWireCoder, element); + return CoderUtils.decodeFromByteArray(runnerWireCoder, wireBytes); + } catch (CoderException e) { + throw new RuntimeException( + "Failed to transcode an unbounded-read element to wire form for transform " + transformId, + e); + } + } + + @Override + public void close() { + UnboundedReader currentReader = reader; + if (currentReader != null) { + try { + currentReader.close(); + } catch (IOException e) { + LOG.warn("Error closing the reader for unbounded source {}", transformId, e); + } + reader = null; + } + } + + private static V checkInitialized(@Nullable V value) { + if (value == null) { + throw new IllegalStateException("UnboundedReadProcessor used before init()"); + } + return value; + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java new file mode 100644 index 000000000000..b668ee972d53 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java @@ -0,0 +1,209 @@ +/* + * 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.containsString; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.greaterThan; + +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.CountingSource; +import org.apache.beam.sdk.io.CountingSource.CounterMark; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.io.UnboundedSource; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.kafka.streams.TopologyTestDriver; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Before; +import org.junit.Test; + +/** + * Runs a pipeline whose source is unbounded, which is the shape the runner exists for: a Kafka + * Streams application is a long-running stream processor, and until now the runner could only read + * sources that finish. + * + *

    Two things separate this from the bounded read. The source is polled repeatedly rather than + * drained once, so elements arrive over several turns of the wall clock; and the watermark comes + * from the reader's own progress rather than jumping to the end of time when the input runs out, + * which is what lets downstream windows close on a stream that never ends. + */ +public class UnboundedReadTest { + + /** How many elements one poll of the source may take. */ + private static final int ELEMENTS_PER_POLL = 5; + + /** Elements the finite variant of the source produces before ending time. */ + private static final int ELEMENTS = 12; + + /** Elements the pipeline has seen, recorded in order. */ + private static final List RECEIVED = Collections.synchronizedList(new ArrayList<>()); + + @Before + public void reset() { + RECEIVED.clear(); + } + + private static class RecordFn extends DoFn { + @ProcessElement + public void processElement(@Element Long element, OutputReceiver out) { + RECEIVED.add(element); + out.output(element); + } + } + + /** + * A genuinely unbounded pipeline. Nothing caps the source — capping it with {@code + * withMaxNumRecords} would turn it back into a bounded read and test the wrong path — so the work + * is bounded instead by how many elements a single poll may take and how many turns the test + * drives. + */ + private static Pipeline unboundedPipeline() { + KafkaStreamsPipelineOptions options = + KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class); + options.setMaxBundleSize(ELEMENTS_PER_POLL); + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply("read", Read.from(CountingSource.unbounded())) + .apply("record", ParDo.of(new RecordFn())); + return pipeline; + } + + @Test + public void anUnboundedSourceIsPolledAndItsElementsReachTheHarness() { + Pipeline pipeline = unboundedPipeline(); + KafkaStreamsTranslationContext context = KafkaStreamsTestRunner.translate(pipeline); + + try (TopologyTestDriver driver = + new TopologyTestDriver( + context.getTopology(), KafkaStreamsTestRunner.streamsConfig(pipeline))) { + // Several turns, because an unbounded read yields what is available now rather than + // everything at once. + for (int turn = 0; turn < 4; turn++) { + driver.advanceWallClockTime(Duration.ofMillis(100)); + } + } + + // More than a single poll's worth, which is the point: a bounded read drains once, whereas + // this one has to be asked again on each turn of the clock and keep going from where it was. + assertThat( + "expected several polls' worth of elements, got " + RECEIVED.size(), + RECEIVED.size(), + is(greaterThan(ELEMENTS_PER_POLL))); + // The source counts from zero, so what arrived has to start there and be contiguous — no gap + // and no repeat, which is what the checkpoint mark between polls is for. + for (int i = 0; i < RECEIVED.size(); i++) { + assertThat(RECEIVED.get(i), is((long) i)); + } + } + + @Test + public void aSourceThatReachesTheEndOfTimeStopsBeingPolled() { + // CountingSource.unbounded() with a limit reports the terminal watermark once it has produced + // its elements, which is a source saying it will yield nothing further. Polling must stop + // there rather than spinning on a reader that can only return false. + KafkaStreamsPipelineOptions options = + KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class); + options.setMaxBundleSize(ELEMENTS_PER_POLL); + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply("read", Read.from(CountingSource.unbounded()).withMaxNumRecords(ELEMENTS)) + .apply("record", ParDo.of(new RecordFn())); + + KafkaStreamsTranslationContext context = KafkaStreamsTestRunner.translate(pipeline); + try (TopologyTestDriver driver = + new TopologyTestDriver( + context.getTopology(), KafkaStreamsTestRunner.streamsConfig(pipeline))) { + for (int turn = 0; turn < 10; turn++) { + driver.advanceWallClockTime(Duration.ofMillis(100)); + } + } + + // Every element exactly once: the source finished, and the turns after it finished added + // nothing. + assertThat(RECEIVED.size(), is(ELEMENTS)); + } + + /** A source that ignores the requested split count and always returns two parts. */ + private static class TwoSplitSource extends UnboundedSource { + private final UnboundedSource delegate = CountingSource.unbounded(); + + @Override + public List> split( + int desiredNumSplits, PipelineOptions options) throws Exception { + return delegate.split(2, options); + } + + @Override + public UnboundedReader createReader( + PipelineOptions options, @Nullable CounterMark checkpointMark) throws IOException { + return delegate.createReader(options, checkpointMark); + } + + @Override + public Coder getCheckpointMarkCoder() { + return delegate.getCheckpointMarkCoder(); + } + + @Override + public Coder getOutputCoder() { + return delegate.getOutputCoder(); + } + } + + @Test + public void aSourceThatSplitsIntoSeveralPartsIsRejectedRatherThanTruncated() { + // The count passed to split() is only a hint. Reading the first part of several and ignoring + // the rest would silently drop their data, so translation has to fail instead. + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + pipeline + .apply("read", Read.from(new TwoSplitSource())) + .apply("record", ParDo.of(new RecordFn())); + + try { + KafkaStreamsTestRunner.translate(pipeline); + throw new AssertionError("expected a multi-split source to be rejected"); + } catch (UnsupportedOperationException e) { + assertThat(e.getMessage(), containsString("split into 2 parts")); + assertThat(e.getMessage(), containsString("drop the data")); + } + } + + @Test + public void theSourceIsTranslatedAsAnUnboundedRead() { + // The bounded and unbounded reads share a URN and are told apart by the payload, so this pins + // down that the pipeline really did take the unbounded path. + KafkaStreamsTranslationContext context = KafkaStreamsTestRunner.translate(unboundedPipeline()); + + boolean hasReadProcessor = + context.getTopology().describe().subtopologies().stream() + .flatMap(subtopology -> subtopology.nodes().stream()) + .anyMatch(node -> node.name().contains("read")); + assertThat(hasReadProcessor, is(true)); + } +} From cb30afd092ef22a38e9d02d5544cfe208f0b183f Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Sun, 9 Aug 2026 19:05:27 +0500 Subject: [PATCH 27/37] [GSoC 2026] Kafka Streams runner: user documentation, marked experimental (#39627) * [GSoC 2026] Kafka Streams runner: user documentation, marked experimental Adds the runner's documentation page, linked from the runners menu: what the runner is and why someone would choose it, how to start the job server and submit a pipeline, every pipeline option with its default, the internal topics it creates, and what is and is not supported. The unsupported list is specific rather than a general disclaimer, since these are core parts of the Beam model rather than nice-to-haves: side inputs, stateful ParDo and user timers, merging windows, custom WindowFns, splittable DoFn, TestStream, reading a source in parallel, the bundle time bound, finalizeCheckpoint, and committed metrics. Each says what it means for a user. --- runners/kafka-streams/build.gradle | 10 + .../kafka/streams/KafkaStreamsRunner.java | 9 + .../en/documentation/runners/kafkastreams.md | 197 ++++++++++++++++++ website/www/site/data/capability_matrix.yaml | 155 ++++++++++++++ .../partials/section-menu/en/runners.html | 1 + 5 files changed, 372 insertions(+) create mode 100644 website/www/site/content/en/documentation/runners/kafkastreams.md diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 616257ddd788..203168d7ec40 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -87,6 +87,16 @@ dependencies { } +// Starts the job server a portable pipeline is submitted to. Pass driver arguments with +// -PjobServerArgs="--job-port=8099,--artifact-port=8098". +tasks.register("runJobServer", JavaExec) { + group = "Application" + description = "Runs the Kafka Streams job server." + mainClass = "org.apache.beam.runners.kafka.streams.KafkaStreamsJobServerDriver" + classpath = sourceSets.main.runtimeClasspath + args = project.hasProperty("jobServerArgs") ? project.property("jobServerArgs").split(",") : [] +} + // The broker integration test drives the production runner against a real Kafka in Docker, so it // is not part of the default build. Run it with :runners:kafka-streams:brokerIntegrationTest. test { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java index 6a8f105bb2ee..1b530fd22ab1 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java @@ -34,6 +34,15 @@ /** * A {@link PipelineRunner} that submits portable jobs to an in-process or external Beam job service * backed by the Kafka Streams translation path. + * + *

    This runner is experimental. It executes a subset of the Beam model correctly — the + * parts it supports are covered by Beam's {@code @ValidatesRunner} suite — but several capabilities + * that are core to the model are not implemented yet, among them side inputs, stateful {@code + * ParDo} and user timers, merging windows, custom {@code WindowFn}s and splittable {@code DoFn}. + * Its behaviour and its pipeline options may change. See the runner documentation for + * what is and is not supported, and #18479 for the work that remains. */ public class KafkaStreamsRunner extends PipelineRunner { diff --git a/website/www/site/content/en/documentation/runners/kafkastreams.md b/website/www/site/content/en/documentation/runners/kafkastreams.md new file mode 100644 index 000000000000..bef1f852094b --- /dev/null +++ b/website/www/site/content/en/documentation/runners/kafkastreams.md @@ -0,0 +1,197 @@ +--- +type: runners +title: "Kafka Streams Runner" +--- + + +# Kafka Streams Runner + +The Kafka Streams Runner executes Beam pipelines on [Kafka +Streams](https://kafka.apache.org/documentation/streams/), by translating a pipeline into a Kafka +Streams topology. + +What distinguishes it from the other runners is that Kafka Streams is a library rather than a +cluster. There is no job manager and no resource manager to operate: an application is an ordinary +JVM process that reads from and writes to Kafka, and scaling it means starting more copies of that +process. Fault tolerance, state, and exactly-once processing come from Kafka itself — from consumer +groups, changelog topics, and transactions. + +That makes it worth considering if you already run Kafka and want Beam's programming model without +introducing a second distributed system to operate. + +## The runner is experimental + +**The Kafka Streams Runner is experimental.** It executes a meaningful subset of the Beam model +correctly, and the parts it does support are covered by Beam's own `@ValidatesRunner` suite, but +several capabilities that are core to the model are not implemented yet. Read [what is not +supported](#what-is-not-supported-yet) before choosing it for anything real. + +It is also aimed squarely at streaming. A pipeline over bounded data will run, but there are more +efficient choices for batch work; this runner exists for pipelines that do not end. + +## Running a pipeline + +The runner is portable: it executes user code over the Fn API, in an SDK harness, so a pipeline goes +to a job server rather than being run directly. You do not have to start one yourself. + +### From Java + +Select `KafkaStreamsRunner` and point it at your Kafka cluster: + +``` +--runner=KafkaStreamsRunner \ +--bootstrapServers=localhost:9092 \ +--applicationId=my-beam-pipeline +``` + +With no `jobEndpoint` set, the runner starts a job server of its own on a dynamic port, submits to +it, and shuts it down when the pipeline finishes. Setting `--jobEndpoint` instead submits to a job +server you are already running. + +### From Python + +Select `KafkaStreamsRunner` there too. The Python runner builds the job server jar if it has to, +starts it, and stops it when the pipeline finishes: + +``` +python my_pipeline.py \ + --runner=KafkaStreamsRunner \ + --bootstrap_servers=localhost:9092 \ + --application_id=my-beam-pipeline +``` + +The SDK harness runs in `LOOPBACK` mode by default, so a local run needs no Docker. Building the jar +takes a while the first time; `--kafka_streams_job_server_jar` points at a prebuilt one, which +`./gradlew :runners:kafka-streams:job-server:shadowJar` produces. + +### Against a job server you are already running + +Start one, which listens on `localhost:8099` by default: + +``` +./gradlew :runners:kafka-streams:runJobServer +``` + +Then point a pipeline at it instead of letting the runner start its own. From Java: + +``` +--runner=KafkaStreamsRunner \ +--jobEndpoint=localhost:8099 \ +--bootstrapServers=localhost:9092 \ +--applicationId=my-beam-pipeline +``` + +and from Python, where the option names are the same in snake case: + +``` +--runner=PortableRunner \ +--job_endpoint=localhost:8099 \ +--bootstrap_servers=localhost:9092 \ +--application_id=my-beam-pipeline +``` + +The application id has no default and must be set. It becomes the Kafka Streams `application.id`, +which is the identity of the consumer group and of the runner's internal topics, so two different +pipelines sharing one would interfere with each other. + +## Pipeline options + +Named as Java spells them below; from Python the same options are in snake case, so +`internalParallelism` is `--internal_parallelism`. + +| Option | Default | Description | +| --- | --- | --- | +| `bootstrapServers` | `localhost:9092` | Kafka brokers the application connects to. | +| `applicationId` | *(required)* | Kafka Streams `application.id`. Must be unique per pipeline. | +| `internalParallelism` | `1` | Partitions for the internal topics the runner creates, which is the parallelism the shuffled parts of a pipeline can reach. | +| `topicReplicationFactor` | `1` | Replication factor for those topics. | +| `maxBundleSize` | `1000` | Elements per bundle, and elements taken per poll of an unbounded source. | +| `maxBundleTimeMs` | `1000` | Intended cap on how long a bundle may stay open. **Not applied yet** — see below. | +| `readCheckpointNumBundles` | `10` | Polls of an unbounded source between stores of its checkpoint mark. Larger values replay more after a restart. | +| `stateDir` | temp directory | Where Kafka Streams keeps local state. | + +### Topics the runner creates + +The runner shuffles through topics it names itself and creates before starting: a bootstrap topic +per `Impulse` and per source, and a repartition topic per `GroupByKey`. They carry a `__beam_` +prefix. Bootstrap topics always have one partition; repartition topics get `internalParallelism`, +which is what sets how many instances the parts of the pipeline behind a shuffle run across. + +Topics the pipeline itself reads or writes are never created implicitly. + +## What is supported + +* **Reading** — bounded and unbounded sources, through the primitive `Read`. +* **ParDo** — stateless, including multiple outputs. +* **GroupByKey**, and `Combine` through its GroupByKey expansion. +* **Windowing** — global, fixed and sliding windows, with the default trigger, allowed lateness, and + timestamp combiners. Windowing and triggering run through Beam's own `ReduceFnRunner`, backed by + Kafka Streams state and timers. +* **Flatten**, **Redistribute**. +* **Metrics** — user counters and distributions reported by the SDK harness surface as + `MetricResults`. +* **Exactly-once processing**, via Kafka transactions (`exactly_once_v2`). + +Because the runner is portable and reads the language-neutral pipeline proto, a pipeline built in +any Beam SDK should translate, provided it stays inside the subset above. Only the Java SDK has been +exercised so far. + +## What is not supported yet + +These are core parts of the Beam model that the runner does not implement. Each is a real gap rather +than a decision, and each is tracked: + +* **Side inputs** ([#39628](https://github.com/apache/beam/issues/39628)). +* **Stateful `ParDo` and user timers** ([#39629](https://github.com/apache/beam/issues/39629)) — + including timer families, looping timers + and processing-time timers. +* **Merging windows**, so session windows do not work, and **custom `WindowFn`s** + ([#39630](https://github.com/apache/beam/issues/39630)). The standard windows travel as URNs the + runner interprets directly; one the + user wrote themselves would have to run in the SDK harness, which is not wired up. +* **Splittable `DoFn`**, bounded or unbounded + ([#39631](https://github.com/apache/beam/issues/39631)). +* **`TestStream`** ([#39632](https://github.com/apache/beam/issues/39632)). +* **Reading a source in parallel** + ([#39626](https://github.com/apache/beam/issues/39626)). A source is split into exactly one part + and read by a single reader. A source that insists on splitting further is rejected at + translation rather than having its extra splits silently dropped. +* **A time bound on bundles** ([#39633](https://github.com/apache/beam/issues/39633)). + `maxBundleTimeMs` is accepted but has no effect: + closing a bundle from a wall-clock punctuator duplicated output against a real broker, and the + cause is not yet understood. Bundles are bounded by element count and closed on watermarks. +* **`finalizeCheckpoint`** ([#39634](https://github.com/apache/beam/issues/39634)) is not called on + an unbounded source's checkpoint mark, + so sources that rely on finalization to acknowledge data will not see it. +* **Committed metrics** ([#39635](https://github.com/apache/beam/issues/39635)) — only attempted + values are reported. + +## How it works + +A Beam pipeline arrives as a proto and is translated into a Kafka Streams `Topology`. Fused stages +of user code become processors that execute that code in an SDK harness over the Fn API; a +`GroupByKey` becomes a repartition topic plus a stateful processor; and the elements flowing between +them carry either data or a watermark report. + +Watermarks are the part with no direct Kafka Streams equivalent. Kafka Streams tracks stream-time, +which only advances when data arrives, whereas Beam needs a watermark that can advance on an idle +stream and that reflects every upstream instance. The runner therefore propagates its own watermark +reports alongside the data: a transform aggregates the reports of everything upstream of it, holds +until every partition of every upstream transform has reported, and only then lets its own watermark +advance. + +For the full design, see the [design +document](https://docs.google.com/document/d/1BBMURhSG4SxPcvvnKMTrmnKCr_jhXL6R4TBDBW7zsy8/edit) and +the tracking issue, [#18479](https://github.com/apache/beam/issues/18479). diff --git a/website/www/site/data/capability_matrix.yaml b/website/www/site/data/capability_matrix.yaml index a1afdc6f8abc..6f9b7256a2ff 100644 --- a/website/www/site/data/capability_matrix.yaml +++ b/website/www/site/data/capability_matrix.yaml @@ -26,6 +26,8 @@ capability-matrix: name: Apache Nemo - class: jet name: Hazelcast Jet + - class: kafka-streams + name: Kafka Streams - class: twister2 name: Twister2 - class: python direct @@ -80,6 +82,10 @@ capability-matrix: l1: "" l2: l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "Stateless ParDo runs in the SDK harness over the Fn API, including multiple outputs." - name: GroupByKey description: Grouping of key-value pairs per key, window, and pane. (See also other tabs.) values: @@ -119,6 +125,10 @@ capability-matrix: l1: "" l2: l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "Shuffles through a Kafka repartition topic keyed by the encoded Beam key." - name: Flatten description: Concatenates multiple homogenously typed collections together. values: @@ -158,6 +168,10 @@ capability-matrix: l1: "" l2: l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "" - name: Combine description: 'Application of an associative, commutative operation over all values ("globally") or over all values associated with each key ("per key"). Can be implemented using ParDo, but often more efficient implementations exist.' values: @@ -197,6 +211,10 @@ capability-matrix: l1: "" l2: l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "Executed through the GroupByKey expansion; there is no lifted pre-combine yet." - name: Composite Transforms description: Allows easy extensibility for library writers. In the near future, we expect there to be more information provided at this level -- customized metadata hooks for monitoring, additional runtime/environment hooks, etc. values: @@ -236,6 +254,10 @@ capability-matrix: l1: "" l2: l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "" - name: Side Inputs description: Side inputs are additional PCollections whose contents are computed during pipeline execution and then made accessible to DoFn code. The exact shape of the side input depends both on the PCollectionView used to describe the access pattern (interable, map, singleton) and the window of the element from the main input that is currently being processed. values: @@ -275,6 +297,10 @@ capability-matrix: l1: "" l2: l3: "" + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "Stages run without a side input handler." - name: Source API description: Allows users to provide additional input sources. Supports both bounded and unbounded data. Includes hooks necessary to provide efficient parallelization (size estimation, progress information, dynamic splitting, etc). values: @@ -314,6 +340,10 @@ capability-matrix: l1: "" l2: l3: "" + - class: kafka-streams + l1: "Partially" + l2: bounded and unbounded, read by a single reader + l3: "A source is split into exactly one part; reading several splits in parallel is not supported." - name: Metrics description: Allow transforms to gather simple metrics across bundles in a PTransform. Provide a mechanism to obtain both committed and attempted metrics. Semantically similar to using an additional output, but support partial results as the transform executes, and support both committed and attempted values. Will likely want to augment Metrics to be more useful for processing unbounded data by making them windowed. values: @@ -353,6 +383,10 @@ capability-matrix: l1: "" l2: l3: "" + - class: kafka-streams + l1: "Partially" + l2: attempted metrics only + l3: "User metrics reported by the SDK harness surface as attempted values; committed values are not available." - name: Stateful Processing description: Allows fine-grained access to per-key, per-window persistent state. Necessary for certain use cases (e.g. high-volume windows which store large amounts of data, but typically only access small portions of it; complex state machines; etc.) that are not easily or efficiently addressed via Combine or GroupByKey+ParDo. values: @@ -392,6 +426,10 @@ capability-matrix: l1: "" l2: l3: "" + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "User state is not wired to the harness." - description: Bounded Splittable DoFn Support Status anchor: what color-y: "fff" @@ -440,6 +478,10 @@ capability-matrix: l1: "Yes" l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Side Inputs description: "" values: @@ -479,6 +521,10 @@ capability-matrix: l1: l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Splittable DoFn Initiated Checkpointing description: "" values: @@ -518,6 +564,10 @@ capability-matrix: l1: "Yes" l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Dynamic Splitting description: "" values: @@ -557,6 +607,10 @@ capability-matrix: l1: "Yes" l2: Only with Python SDK l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Bundle Finalization description: "" values: @@ -596,6 +650,10 @@ capability-matrix: l1: "Yes" l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - description: Unbounded Splittable DoFn Support Status anchor: what color-y: "fff" @@ -644,6 +702,10 @@ capability-matrix: l1: "Yes" l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Side Inputs description: "" values: @@ -683,6 +745,10 @@ capability-matrix: l1: l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Splittable DoFn Initiated Checkpointing description: "" values: @@ -722,6 +788,10 @@ capability-matrix: l1: "Yes" l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Dynamic Splitting description: "" values: @@ -761,6 +831,10 @@ capability-matrix: l1: "No" l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Bundle Finalization description: "" values: @@ -800,6 +874,10 @@ capability-matrix: l1: "Yes" l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - description: Where in event time? anchor: where color-y: "fff" @@ -848,6 +926,10 @@ capability-matrix: l1: "Yes" l2: supported l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "" - name: Fixed windows description: Fixed-size, timestamp-based windows. (Hourly, Daily, etc) values: @@ -887,6 +969,10 @@ capability-matrix: l1: "Yes" l2: supported l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "Windowing runs through Beam's ReduceFnRunner over Kafka Streams state and timers." - name: Sliding windows description: Possibly overlapping fixed-size timestamp-based windows (Every minute, use the last ten minutes of data.) values: @@ -926,6 +1012,10 @@ capability-matrix: l1: "Yes" l2: supported l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "" - name: Session windows description: Based on bursts of activity separated by a gap size. Different per key. values: @@ -965,6 +1055,10 @@ capability-matrix: l1: "Yes" l2: supported l3: "" + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "Sessions are merging windows, which the windowing implementation does not handle yet." - name: Custom windows description: All windows must implement BoundedWindow, which specifies a max timestamp. Each WindowFn assigns elements to an associated window. values: @@ -1004,6 +1098,10 @@ capability-matrix: l1: "Yes" l2: supported l3: "" + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "Only the standard WindowFns, which travel as URNs the runner interprets directly." - name: Custom merging windows description: A custom WindowFn additionally specifies whether and how to merge windows. values: @@ -1043,6 +1141,10 @@ capability-matrix: l1: "Yes" l2: supported l3: "" + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Timestamp control description: For a grouping transform, such as GBK or Combine, an OutputTimeFn specifies (1) how to combine input timestamps within a window and (2) how to merge aggregated timestamps when windows merge. values: @@ -1082,6 +1184,10 @@ capability-matrix: l1: "Yes" l2: supported l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "Timestamp combiners are applied by ReduceFnRunner." - description: When in processing time? anchor: when @@ -1131,6 +1237,10 @@ capability-matrix: l1: "Yes" l2: fully supported l3: "" + - class: kafka-streams + l1: "No" + l2: default trigger only + l3: "" - name: Event-time triggers description: Triggers that fire in response to event-time completeness signals, such as watermarks progressing. values: @@ -1170,6 +1280,10 @@ capability-matrix: l1: "Yes" l2: fully supported l3: "" + - class: kafka-streams + l1: "Partially" + l2: the default trigger only + l3: "Panes fire when the watermark passes the end of the window; other event-time triggers are untested." - name: Processing-time triggers description: Triggers that fire in response to processing-time advancing. @@ -1210,6 +1324,10 @@ capability-matrix: l1: "Yes" l2: fully supported l3: "" + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "Processing-time timers are not wired up." - name: Count triggers description: Triggers that fire after seeing at least N elements. @@ -1250,6 +1368,10 @@ capability-matrix: l1: "Yes" l2: fully supported l3: "" + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Composite triggers description: Triggers which compose other triggers in more complex structures, such as logical AND, logical OR, early/on-time/late, etc. @@ -1290,6 +1412,10 @@ capability-matrix: l1: "Yes" l2: fully supported l3: "" + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Allowed lateness description: A way to bound the useful lifetime of a window (in event time), after which any unemitted results may be materialized, the window contents may be garbage collected, and any addtional late data that arrive for the window may be discarded. @@ -1330,6 +1456,10 @@ capability-matrix: l1: "Yes" l2: fully supported l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "Allowed lateness and late-data dropping are applied by ReduceFnRunner." - name: Timers description: A fine-grained mechanism for performing work at some point in the future, in either the event-time or processing-time domain. Useful for orchestrating delayed events, timeouts, etc in complex state per-key, per-window state machines. @@ -1370,6 +1500,10 @@ capability-matrix: l1: "Yes" l2: "Partially" l3: "" + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "User timers are not wired to the harness." - description: How do refinements relate? anchor: how @@ -1419,6 +1553,10 @@ capability-matrix: l1: "Yes" l2: fully supported l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "" - name: Accumulating description: Elements are accumulated in state across multiple pane firings for the same window. @@ -1459,6 +1597,10 @@ capability-matrix: l1: "Yes" l2: fully supported l3: "" + - class: kafka-streams + l1: "Yes" + l2: fully supported + l3: "" - description: Additional common features not yet part of the Beam model anchor: misc @@ -1508,6 +1650,10 @@ capability-matrix: l1: l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Checkpoint description: APIs and semantics for saving a pipeline checkpoint are under discussion. This would be a runner-specific materialization of the pipeline state required to resume or duplicate the pipeline. values: @@ -1547,6 +1693,10 @@ capability-matrix: l1: l2: l3: + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" - name: Key-ordered delivery description: The runner offers guarantees for the order in which elements are passed in between operations. See per-key ordering semantics. values: @@ -1586,3 +1736,8 @@ capability-matrix: l1: "Unverified" l2: l3: + + - class: kafka-streams + l1: "No" + l2: not implemented + l3: "" \ No newline at end of file diff --git a/website/www/site/layouts/partials/section-menu/en/runners.html b/website/www/site/layouts/partials/section-menu/en/runners.html index 337debf3ecec..6119debe8781 100644 --- a/website/www/site/layouts/partials/section-menu/en/runners.html +++ b/website/www/site/layouts/partials/section-menu/en/runners.html @@ -19,4 +19,5 @@

  • Apache Spark
  • Google Cloud Dataflow
  • Hazelcast Jet
    • +
  • Kafka Streams
  • Twister2
    • From e051e06bf83f34520f1c39cb26ddce9a68467bcd Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:18:37 +0500 Subject: [PATCH 28/37] [GSoC 2026] Kafka Streams runner: Python wrapper that starts its own job server (#39680) * [GSoC 2026] Kafka Streams runner: Python wrapper that starts its own job server The runner starting its own job server only helped Java, so a Python user still had to run one by hand. This adds the wrapper Flink and Spark provide, so a Python pipeline can select the runner and nothing else. --- runners/kafka-streams/build.gradle | 6 +- runners/kafka-streams/job-server/build.gradle | 88 +++++++++++++ .../apache_beam/options/pipeline_options.py | 20 +++ .../kafka_streams_java_job_server_test.py | 123 ++++++++++++++++++ .../portability/kafka_streams_runner.py | 106 +++++++++++++++ settings.gradle.kts | 1 + 6 files changed, 343 insertions(+), 1 deletion(-) create mode 100644 runners/kafka-streams/job-server/build.gradle create mode 100644 sdks/python/apache_beam/runners/portability/kafka_streams_java_job_server_test.py create mode 100644 sdks/python/apache_beam/runners/portability/kafka_streams_runner.py diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index 203168d7ec40..399550534092 100644 --- a/runners/kafka-streams/build.gradle +++ b/runners/kafka-streams/build.gradle @@ -21,7 +21,11 @@ import java.time.Duration plugins { id 'org.apache.beam.module' } -def kafka_version = '3.9.0' +// An extension property rather than a local, so the job server module can pin the same version +// instead of repeating it. Both have to pin: applyJavaNature forces every version in library.java, +// which includes an older kafka-clients, and a module that does not override it links against that +// one at runtime. +ext.kafka_version = '3.9.0' applyJavaNature( automaticModuleName: 'org.apache.beam.runners.kafka.streams', diff --git a/runners/kafka-streams/job-server/build.gradle b/runners/kafka-streams/job-server/build.gradle new file mode 100644 index 000000000000..aef2f271b7b1 --- /dev/null +++ b/runners/kafka-streams/job-server/build.gradle @@ -0,0 +1,88 @@ +/* + * 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. + */ + +/** + * Kafka Streams Runner JobServer build file. + * + * Packages the runner and everything it needs into one jar, so a pipeline from an SDK other than + * Java can start a job server without a Beam source tree. The Python KafkaStreamsRunner builds and + * launches this jar for the user. + */ + +apply plugin: 'org.apache.beam.module' +apply plugin: 'application' +// Must be set before the shadow plugin is applied. +mainClassName = "org.apache.beam.runners.kafka.streams.KafkaStreamsJobServerDriver" + +applyJavaNature( + automaticModuleName: 'org.apache.beam.runners.kafka.streams.jobserver', + validateShadowJar: false, + exportJavadoc: false, + shadowClosure: { + // Kafka's clients and Streams libraries ship reference.conf-style resources that have to be + // concatenated rather than overwritten when everything lands in one jar. + append "reference.conf" + }, +) + +def kafkaStreamsRunnerProject = ":runners:kafka-streams" + +description = "Apache Beam :: Runners :: Kafka Streams :: Job Server" + +evaluationDependsOn(kafkaStreamsRunnerProject) + +// The runner is compiled against this version, so the jar has to carry it. Without this the +// versions forced by applyJavaNature win, the shaded jar ships an older kafka-clients, and every +// pipeline fails at translation with a NoSuchMethodError rather than at build time. +def kafka_version = project(kafkaStreamsRunnerProject).kafka_version + +configurations.configureEach { + resolutionStrategy.eachDependency { details -> + if (details.requested.group == "org.apache.kafka") { + details.useVersion(kafka_version) + details.because("Kafka Streams runner is developed against Kafka ${kafka_version}.") + } + } +} + +dependencies { + implementation project(kafkaStreamsRunnerProject) + permitUnusedDeclared project(kafkaStreamsRunnerProject) + // A binding, or the job server starts but logs nothing at all, which is unhelpful for something + // a user runs in the foreground and reads to see what their pipeline is doing. + runtimeOnly library.java.slf4j_simple + runtimeOnly project(":sdks:java:extensions:google-cloud-platform-core") +} + +// The runner's classes only exist in the shadow jar, so the job server has to be started through +// runShadow rather than the plain run task. +runShadow { + args = [] + if (project.hasProperty('jobHost')) + args += ["--job-host=${project.property('jobHost')}"] + if (project.hasProperty('jobPort')) + args += ["--job-port=${project.property('jobPort')}"] + if (project.hasProperty('artifactPort')) + args += ["--artifact-port=${project.property('artifactPort')}"] + if (project.hasProperty('expansionPort')) + args += ["--expansion-port=${project.property('expansionPort')}"] + if (project.hasProperty('artifactsDir')) + args += ["--artifacts-dir=${project.property('artifactsDir')}"] + if (project.hasProperty('cleanArtifactsPerJob')) + args += ["--clean-artifacts-per-job=${project.property('cleanArtifactsPerJob')}"] +} diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 2533083f7e7e..b0d0bdfbaa64 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -732,6 +732,7 @@ class StandardOptions(PipelineOptions): 'apache_beam.runners.interactive.interactive_runner.InteractiveRunner', 'apache_beam.runners.portability.flink_runner.FlinkRunner', 'apache_beam.runners.portability.fn_api_runner.FnApiRunner', + 'apache_beam.runners.portability.kafka_streams_runner.KafkaStreamsRunner', 'apache_beam.runners.portability.portable_runner.PortableRunner', 'apache_beam.runners.portability.prism_runner.PrismRunner', 'apache_beam.runners.portability.spark_runner.SparkRunner', @@ -2064,6 +2065,25 @@ def _add_argparse_args(cls, parser): ' and the number of key groups used for partitioned state.') +class KafkaStreamsRunnerOptions(PipelineOptions): + @classmethod + def _add_argparse_args(cls, parser): + parser.add_argument( + '--bootstrap_servers', + default='localhost:9092', + help='Comma-separated list of host:port Kafka brokers the pipeline ' + 'connects to.') + parser.add_argument( + '--application_id', + help='Kafka Streams application.id for the pipeline. Must be unique ' + 'per pipeline, since it identifies the consumer group and the ' + 'runner\'s internal topics.') + parser.add_argument( + '--kafka_streams_job_server_jar', + help='Path or URL to a Beam Kafka Streams job server jar. If unset, ' + 'the jar is built from the Beam source tree.') + + class SparkRunnerOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): diff --git a/sdks/python/apache_beam/runners/portability/kafka_streams_java_job_server_test.py b/sdks/python/apache_beam/runners/portability/kafka_streams_java_job_server_test.py new file mode 100644 index 000000000000..fa785781ad49 --- /dev/null +++ b/sdks/python/apache_beam/runners/portability/kafka_streams_java_job_server_test.py @@ -0,0 +1,123 @@ +# +# 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. +# + +# pytype: skip-file + +import logging +import tempfile +import unittest + +import mock + +from apache_beam.options import pipeline_options +from apache_beam.runners.portability.kafka_streams_runner import KafkaStreamsJarJobServer +from apache_beam.runners.portability.kafka_streams_runner import KafkaStreamsRunner + + +class KafkaStreamsTestPipelineOptions(pipeline_options.PipelineOptions): + def view_as(self, cls): + # Ensure only KafkaStreamsRunnerOptions and JobServerOptions are used when + # calling default_job_server. If other options classes are needed, the + # cache key must include them to prevent incorrect hits. + assert ( + cls is pipeline_options.KafkaStreamsRunnerOptions or + cls is pipeline_options.JobServerOptions) + return super().view_as(cls) + + +class KafkaStreamsJavaJobServerTest(unittest.TestCase): + def test_job_server_cache(self): + # Multiple KafkaStreamsRunner instances may be created, so job servers have + # to be cached across runner instances: each one is an external Java + # process, and starting a second for the same configuration would fail to + # bind the same ports. + + # Options that do not affect job server configuration, such as + # sdk_worker_parallelism, should still hit the same cache entry. + job_server1 = KafkaStreamsRunner().default_job_server( + KafkaStreamsTestPipelineOptions(['--sdk_worker_parallelism=1'])) + job_server2 = KafkaStreamsRunner().default_job_server( + KafkaStreamsTestPipelineOptions(['--sdk_worker_parallelism=2'])) + self.assertIs(job_server2, job_server1) + + # JobServerOptions do affect it, so a different port is a different server. + job_server3 = KafkaStreamsRunner().default_job_server( + KafkaStreamsTestPipelineOptions(['--job_port=1234'])) + self.assertIsNot(job_server3, job_server1) + + # So do the runner's own options. + job_server4 = KafkaStreamsRunner().default_job_server( + KafkaStreamsTestPipelineOptions(['--bootstrap_servers=other:9092'])) + self.assertIsNot(job_server4, job_server1) + self.assertIsNot(job_server4, job_server3) + + job_server5 = KafkaStreamsRunner().default_job_server( + KafkaStreamsTestPipelineOptions(['--application_id=other-pipeline'])) + self.assertIsNot(job_server5, job_server1) + self.assertIsNot(job_server5, job_server4) + + def test_java_arguments(self): + # These are what the job server driver is launched with, so they have to be + # options it accepts. + job_server = KafkaStreamsJarJobServer( + pipeline_options.PipelineOptions(['--application_id=test-pipeline'])) + self.assertEqual([ + '--artifacts-dir', + '/tmp/artifacts', + '--job-port', + 8099, + '--artifact-port', + 8098, + '--expansion-port', + 8097 + ], + job_server.java_arguments( + 8099, 8098, 8097, '/tmp/artifacts')) + + def test_path_to_jar_defaults_to_the_job_server_module(self): + job_server = KafkaStreamsJarJobServer(pipeline_options.PipelineOptions([])) + # Without an explicit jar the runner resolves the one built by the job + # server module, which is what lets a user run a pipeline without having + # built or started anything first. Resolving it for real would either + # download or demand a built jar, so only the target is checked here. + with mock.patch.object(job_server, 'path_to_beam_jar') as path_to_beam_jar: + job_server.path_to_jar() + path_to_beam_jar.assert_called_once_with( + ':runners:kafka-streams:job-server:shadowJar') + + def test_path_to_jar_uses_an_explicit_jar(self): + with tempfile.NamedTemporaryFile(suffix='.jar') as jar: + job_server = KafkaStreamsJarJobServer( + pipeline_options.PipelineOptions( + ['--kafka_streams_job_server_jar=%s' % jar.name])) + self.assertEqual(jar.name, job_server.path_to_jar()) + + def test_path_to_jar_rejects_an_unusable_path(self): + job_server = KafkaStreamsJarJobServer( + pipeline_options.PipelineOptions( + ['--kafka_streams_job_server_jar=/no/such/jar.jar'])) + # A path that is neither an existing file nor a URL cannot be recovered + # from, so it fails with the command that would produce a jar rather than + # letting the job server fail to start later. + with self.assertRaises(ValueError) as context: + job_server.path_to_jar() + self.assertIn('job-server:shadowJar', str(context.exception)) + + +if __name__ == '__main__': + logging.getLogger().setLevel(logging.INFO) + unittest.main() diff --git a/sdks/python/apache_beam/runners/portability/kafka_streams_runner.py b/sdks/python/apache_beam/runners/portability/kafka_streams_runner.py new file mode 100644 index 000000000000..a2d043ca174e --- /dev/null +++ b/sdks/python/apache_beam/runners/portability/kafka_streams_runner.py @@ -0,0 +1,106 @@ +# +# 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. +# + +"""A runner for executing portable pipelines on Kafka Streams.""" + +# pytype: skip-file + +import os +import urllib + +from apache_beam.options import pipeline_options +from apache_beam.runners.portability import job_server +from apache_beam.runners.portability import portable_runner + +# A Java job server is a heavyweight external process, so reuse one across +# pipelines configured the same way. +JOB_SERVER_CACHE = {} + + +class KafkaStreamsRunner(portable_runner.PortableRunner): + """A runner for executing pipelines on Kafka Streams. + + Starts a job server automatically, so a pipeline can be submitted without + running one by hand: + + python my_pipeline.py \\ + --runner=KafkaStreamsRunner \\ + --bootstrap_servers=localhost:9092 \\ + --application_id=my-pipeline + + Pass --job_endpoint instead to submit to a job server that is already + running. + """ + + # Inherits run_portable_pipeline from PortableRunner. + + def default_environment(self, options): + portable_options = options.view_as(pipeline_options.PortableOptions) + if (not portable_options.environment_type and + not portable_options.output_executable_path): + # The job server runs on this machine, so the SDK harness can too, which + # saves the user from needing Docker for a local run. + portable_options.environment_type = 'LOOPBACK' + return super().default_environment(options) + + def default_job_server(self, options): + # Only these two option groups affect how the job server is configured, so + # they are what the cache is keyed on. + kafka_streams_options = options.view_as( + pipeline_options.KafkaStreamsRunnerOptions) + job_server_options = options.view_as(pipeline_options.JobServerOptions) + options_str = str(kafka_streams_options) + str(job_server_options) + if options_str not in JOB_SERVER_CACHE: + JOB_SERVER_CACHE[options_str] = job_server.StopOnExitJobServer( + KafkaStreamsJarJobServer(options)) + return JOB_SERVER_CACHE[options_str] + + +class KafkaStreamsJarJobServer(job_server.JavaJarJobServer): + def __init__(self, options): + super().__init__(options) + kafka_streams_options = options.view_as( + pipeline_options.KafkaStreamsRunnerOptions) + self._jar = kafka_streams_options.kafka_streams_job_server_jar + + def path_to_jar(self): + if self._jar: + if not os.path.exists(self._jar): + url = urllib.parse.urlparse(self._jar) + if not url.scheme: + raise ValueError( + 'Unable to parse jar URL "%s". If using a full URL, make sure ' + 'the scheme is specified. If using a local file path, make sure ' + 'the file exists; you may have to first build the job server ' + 'using `./gradlew runners:kafka-streams:job-server:shadowJar`.' % + self._jar) + return self._jar + return self.path_to_beam_jar( + ':runners:kafka-streams:job-server:shadowJar') + + def java_arguments( + self, job_port, artifact_port, expansion_port, artifacts_dir): + return [ + '--artifacts-dir', + artifacts_dir, + '--job-port', + job_port, + '--artifact-port', + artifact_port, + '--expansion-port', + expansion_port + ] diff --git a/settings.gradle.kts b/settings.gradle.kts index 9a9cb1dd1acb..cd1134d685c9 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -145,6 +145,7 @@ include(":runners:java-job-service") include(":runners:jet") include(":runners:kafka-streams") include(":runners:kafka-streams:proto") +include(":runners:kafka-streams:job-server") include(":runners:local-java") include(":runners:portability:java") include(":runners:prism") From 4ff618047c33b3bf9593adb66ad9ef918333ac9f Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Tue, 11 Aug 2026 13:14:08 +0500 Subject: [PATCH 29/37] [GSoC 2026] Kafka Streams runner: terminate a bounded pipeline when it is drained (#39700) * [GSoC 2026] Kafka Streams runner: terminate a bounded pipeline when it is drained Kafka Streams runs a topology until something closes the client, so a bounded pipeline produced its output and then ran for ever. Every processor already emits TIMESTAMP_MAX_VALUE once its input is exhausted, so each one now schedules its own termination when it emits that watermark, and the client is closed once they have all reported. Termination is scheduled rather than reported inline so that the work which follows the final watermark still runs. The callback waits for every processor rather than the first, because one instance can own both sides of a repartition topic, and it waits until the topology has finished starting, because processors register as their tasks are initialized. --- .../streams/KafkaStreamsPipelineRunner.java | 50 ++++- .../KafkaStreamsPortablePipelineResult.java | 7 +- .../translation/ExecutableStageProcessor.java | 13 +- .../ExecutableStageTranslator.java | 7 +- .../streams/translation/FlattenProcessor.java | 14 +- .../translation/FlattenTranslator.java | 4 +- .../translation/GroupByKeyTranslator.java | 27 ++- .../streams/translation/ImpulseProcessor.java | 13 +- .../translation/ImpulseTranslator.java | 4 +- .../KafkaStreamsTranslationContext.java | 14 ++ .../streams/translation/ReadProcessor.java | 13 +- .../streams/translation/ReadTranslator.java | 11 +- .../translation/ShuffleByKeyProcessor.java | 19 +- .../translation/StageOutputProcessor.java | 12 +- .../translation/TerminationReporter.java | 111 ++++++++++ .../translation/TerminationTracker.java | 191 ++++++++++++++++++ .../translation/UnboundedReadProcessor.java | 12 +- .../WindowedGroupByKeyProcessor.java | 15 +- .../streams/KafkaStreamsRunnerBrokerIT.java | 22 ++ ...ExecutableStageProcessorWatermarkTest.java | 3 +- .../ShuffleByKeyProcessorTest.java | 4 +- .../translation/StageOutputProcessorTest.java | 6 +- .../translation/TerminationTrackerTest.java | 172 ++++++++++++++++ 23 files changed, 717 insertions(+), 27 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/TerminationTrackerTest.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 96b4b2cd7f92..93e1058121ed 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 @@ -80,11 +80,59 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) // Build the result before starting: it registers a state listener, and Kafka Streams only // accepts one while the application is still in the CREATED state. KafkaStreamsPortablePipelineResult result = - new KafkaStreamsPortablePipelineResult(kafkaStreams, context.getMetricsContainerStepMap()); + new KafkaStreamsPortablePipelineResult( + kafkaStreams, + context.getMetricsContainerStepMap(), + // Only once every task is initialized are the processors that have registered the whole + // set, and only then can "all of them are finished" mean the pipeline is finished. + context.getTerminationTracker()::started); + // A bounded pipeline finishes; Kafka Streams has no notion of that, so the runner stops the + // client itself once every processor has reached the terminal watermark. Registered before + // start(), so a pipeline that drains quickly cannot finish before anything is listening. + context + .getTerminationTracker() + .onAllTerminated( + () -> closeInBackground(kafkaStreams, jobInfo.jobId(), "the pipeline is drained")); kafkaStreams.start(); + // The job service reads the result's state once, when this method returns, so returning while + // the pipeline is still running would leave the job reported as RUNNING for good. Blocking here + // is what FlinkPipelineRunner does too, by blocking in executor.execute(). + // + // A bounded pipeline unblocks this by draining: the processors report themselves terminated, + // the callback above stops the client, and the result's latch is released. A streaming pipeline + // never reaches the terminal watermark, so this blocks until the job is cancelled, which is the + // intended behaviour for a job that has no end. + result.waitUntilFinish(); + if (Thread.currentThread().isInterrupted()) { + // Cancelled: the job service interrupts this thread, and the invocation future it would + // otherwise have used to cancel the result has already been cancelled with it. Stop the + // client so it does not outlive the job — from another thread, since close() waits on the + // stream threads and the joins it does would throw straight back out of an interrupted one. + closeInBackground(kafkaStreams, jobInfo.jobId(), "the job was cancelled"); + } return result; } + /** + * Stops the Kafka Streams client from a thread of its own. + * + *

    Never called from a thread that {@code close()} itself waits for. When the pipeline drains, + * that is the task thread which reported the last termination; when the job is cancelled, it is + * the interrupted invocation thread. In both cases closing inline would either wait on the thread + * doing the closing or abandon the shutdown part-way. + */ + private static void closeInBackground(KafkaStreams kafkaStreams, String jobId, String reason) { + Thread closer = + new Thread( + () -> { + LOG.info("Stopping the Kafka Streams client for job {}: {}", jobId, reason); + kafkaStreams.close(); + }, + "kafka-streams-runner-shutdown-" + jobId); + closer.setDaemon(true); + closer.start(); + } + private static void checkRequiredOption(String name, @Nullable String value) { if (value == null || value.isEmpty()) { throw new IllegalArgumentException( 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 0d508b8189fd..81e80b66e303 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 @@ -53,11 +53,16 @@ class KafkaStreamsPortablePipelineResult implements PortablePipelineResult { * listener, and Kafka Streams rejects one once the application has left the CREATED state. */ KafkaStreamsPortablePipelineResult( - KafkaStreams kafkaStreams, MetricsContainerStepMap metricsContainerStepMap) { + KafkaStreams kafkaStreams, + MetricsContainerStepMap metricsContainerStepMap, + Runnable onRunning) { this.kafkaStreams = kafkaStreams; this.metricsContainerStepMap = metricsContainerStepMap; kafkaStreams.setStateListener( (newState, oldState) -> { + if (newState == KafkaStreams.State.RUNNING) { + onRunning.run(); + } if (newState == KafkaStreams.State.NOT_RUNNING || newState == KafkaStreams.State.ERROR) { terminated.countDown(); } 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 ef376606114b..fcb8b2ff27e6 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 @@ -111,6 +111,9 @@ class ExecutableStageProcessor // Computes this stage's input watermark from its upstream transform's reports, holding until // every partition of the upstream transform has reported (see WatermarkAggregator). private final WatermarkAggregator watermarkAggregator; + // Reports this stage instance as finished once it emits the terminal watermark, so a bounded + // pipeline can stop itself. + private final TerminationReporter terminationReporter; // The last watermark actually forwarded downstream, so we only forward when it advances. private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; @@ -139,7 +142,8 @@ class ExecutableStageProcessor Set upstreamTransformIds, MetricsContainerImpl metricsContainer, Map outputChildByPCollectionId, - int maxBundleSize) { + int maxBundleSize, + TerminationTracker terminationTracker) { this.stagePayload = stagePayload; this.jobInfo = jobInfo; this.transformId = transformId; @@ -147,6 +151,7 @@ class ExecutableStageProcessor this.metricsContainer = metricsContainer; this.outputChildByPCollectionId = ImmutableMap.copyOf(outputChildByPCollectionId); this.maxBundleSize = maxBundleSize; + this.terminationReporter = new TerminationReporter(terminationTracker, transformId); } /** A harness output element together with the id of the output PCollection it belongs to. */ @@ -163,6 +168,7 @@ private static final class PendingOutput { @Override public void init(ProcessorContext> context) { this.context = context; + terminationReporter.init(context); // The SDK harness (stage context + bundle factory) is created lazily on the first data // element, so a stage that only forwards watermarks never spins one up. This mirrors Spark's // SparkExecutableStageFunction, which likewise does not build a bundle factory when there are @@ -336,6 +342,7 @@ private void forwardWatermark(Record> record, long wa record.key(), KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), record.timestamp())); + terminationReporter.watermarkEmitted(ctx, watermarkMillis); } @Override @@ -364,6 +371,10 @@ public void close() { } catch (Exception e) { LOG.warn("Error closing executable stage context", e); } + // Last: this is what stops the pipeline waiting on this stage, and closing the bundle above can + // still forward records downstream. Releasing it first would let the pipeline be declared + // finished while this stage was flushing. + terminationReporter.close(); } private static T checkInitialized(@Nullable T value) { 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 c59e5b919aba..11f3192befdf 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 @@ -109,7 +109,8 @@ public void translate( ImmutableSet.of(parentProcessor), context.getMetricsContainerStepMap().getContainer(transformId), outputChildByPCollectionId, - context.getPipelineOptions().getMaxBundleSize()), + context.getPipelineOptions().getMaxBundleSize(), + context.getTerminationTracker()), parentProcessor); if (multiOutput) { @@ -118,7 +119,9 @@ public void translate( outputChildByPCollectionId.forEach( (outputPCollectionId, relayName) -> { topology.addProcessor( - relayName, () -> new StageOutputProcessor(relayName), transformId); + relayName, + () -> new StageOutputProcessor(relayName, context.getTerminationTracker()), + transformId); context.registerPCollectionProducer(outputPCollectionId, relayName); context.registerPCollectionPartitionCount(outputPCollectionId, partitionCount); }); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java index e37da6774489..5b36a53607e6 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java @@ -61,6 +61,9 @@ class FlattenProcessor // The last watermark actually forwarded downstream, so we only forward when it advances. private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; + // Reports this Flatten as finished once every branch it merges has gone terminal. + private final TerminationReporter terminationReporter; + private @Nullable ProcessorContext> context; /** @@ -68,14 +71,22 @@ class FlattenProcessor * @param upstreamTransformIds the producers of this Flatten's input PCollections (known from the * pipeline graph), whose reports the {@link WatermarkAggregator} waits for */ - FlattenProcessor(String transformId, Set upstreamTransformIds) { + FlattenProcessor( + String transformId, Set upstreamTransformIds, TerminationTracker terminationTracker) { this.transformId = transformId; this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); + this.terminationReporter = new TerminationReporter(terminationTracker, transformId); } @Override public void init(ProcessorContext> context) { this.context = context; + terminationReporter.init(context); + } + + @Override + public void close() { + terminationReporter.close(); } @Override @@ -105,6 +116,7 @@ public void process(Record> record) { record.key(), KStreamsPayload.watermark(advanced.getMillis(), transformId, 0, 1), record.timestamp())); + terminationReporter.watermarkEmitted(ctx, advanced.getMillis()); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java index 793c1e1dc530..998553d514d8 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java @@ -84,7 +84,9 @@ public void translate( topology.addProcessor( transformId, - () -> new FlattenProcessor(transformId, upstreamTransformIds), + () -> + new FlattenProcessor( + transformId, upstreamTransformIds, context.getTerminationTracker()), parentProcessors.toArray(new String[0])); context.registerPCollectionProducer(outputPCollectionId, transformId); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index c460436eedf8..c0e69302386a 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -106,7 +106,8 @@ public void translate( String holdsIndexStoreName = transformId + HOLDS_INDEX_STORE_SUFFIX; String timerStoreName = transformId + TIMER_STORE_SUFFIX; String timerIndexStoreName = transformId + TIMER_INDEX_STORE_SUFFIX; - String repartitionTopic = repartitionTopic(transformId); + String repartitionTopic = + repartitionTopic(transformId, context.getPipelineOptions().getApplicationId()); KStreamsPayloadSerde> payloadSerde = new KStreamsPayloadSerde<>(inputCoder); @@ -118,7 +119,9 @@ public void translate( int upstreamPartitionCount = context.getPartitionCount(inputPCollectionId); topology.addProcessor( shuffleName, - () -> new ShuffleByKeyProcessor(keyCoder, upstreamPartitionCount), + () -> + new ShuffleByKeyProcessor( + keyCoder, upstreamPartitionCount, shuffleName, context.getTerminationTracker()), parentProcessor); // Shuffle through the repartition topic: data partitioned by key, watermark broadcast. @@ -151,7 +154,8 @@ public void translate( keyCoder, valueCoder, windowingStrategy, - context.getPipelineOptions()), + context.getPipelineOptions(), + context.getTerminationTracker()), sourceName); topology.addStateStore( Stores.keyValueStoreBuilder( @@ -200,8 +204,19 @@ private static WindowingStrategy hydrateWindowingStrategy( } } - /** The internal repartition topic name for a GroupByKey transform. */ - static String repartitionTopic(String transformId) { - return REPARTITION_TOPIC_PREFIX + transformId.replaceAll("[^a-zA-Z0-9._-]", "_"); + /** + * The internal repartition topic name for a GroupByKey transform. + * + *

    Namespaced by application id, as the Impulse and Read bootstrap topics already are. + * Transform ids come from the pipeline's structure, so two jobs running the same pipeline would + * otherwise shuffle through the same topic and read each other's data — and, because the topic is + * created only if it does not already exist, the second job would silently inherit the first + * job's partition count rather than the one it asked for. + */ + static String repartitionTopic(String transformId, String applicationId) { + return REPARTITION_TOPIC_PREFIX + + applicationId.replaceAll("[^a-zA-Z0-9._-]", "_") + + "_" + + transformId.replaceAll("[^a-zA-Z0-9._-]", "_"); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java index bac91978a298..7c4590a0e5a1 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java @@ -69,24 +69,34 @@ class ImpulseProcessor implements Processor> context; private @Nullable KeyValueStore firedStore; private @Nullable Cancellable scheduledPunctuator; - ImpulseProcessor(String stateStoreName, String transformId) { + ImpulseProcessor( + String stateStoreName, String transformId, TerminationTracker terminationTracker) { this.stateStoreName = stateStoreName; this.transformId = transformId; + this.terminationReporter = new TerminationReporter(terminationTracker, transformId); } @Override public void init(ProcessorContext> context) { this.context = context; this.firedStore = context.getStateStore(stateStoreName); + terminationReporter.init(context); this.scheduledPunctuator = context.schedule(PUNCTUATION_DELAY, PunctuationType.WALL_CLOCK_TIME, ts -> maybeFire()); } + @Override + public void close() { + terminationReporter.close(); + } + @Override public void process(Record record) { // Records that happen to land on the bootstrap topic are not actual data; they just provide an @@ -132,6 +142,7 @@ private void forwardWatermarkMax(ProcessorContext>( new byte[0], KStreamsPayload.watermark(maxMillis, transformId, 0, 1), 0L)); + terminationReporter.watermarkEmitted(ctx, maxMillis); } /** Cancels the wall-clock punctuator after the impulse has fired to stop periodic wakeups. */ diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java index 79d8c3cf577f..3daf7782362e 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java @@ -79,7 +79,9 @@ public void translate( Serdes.ByteArray().deserializer(), bootstrapTopic); topology.addProcessor( - transformId, () -> new ImpulseProcessor(stateStoreName, transformId), sourceNodeName); + transformId, + () -> new ImpulseProcessor(stateStoreName, transformId, context.getTerminationTracker()), + sourceNodeName); topology.addStateStore( Stores.keyValueStoreBuilder( Stores.persistentKeyValueStore(stateStoreName), Serdes.String(), Serdes.Boolean()), 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 d03169615665..904bd8a71e3a 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 @@ -61,6 +61,11 @@ public class KafkaStreamsTranslationContext { // work. private final MetricsContainerStepMap metricsContainerStepMap = new MetricsContainerStepMap(); + // Decides when a bounded pipeline has finished. Owned by the context, so it is scoped to this one + // pipeline: the job server runs several jobs in a single process, and a tracker shared between + // them would let one pipeline finishing stop another. + private final TerminationTracker terminationTracker = new TerminationTracker(); + public static KafkaStreamsTranslationContext create( JobInfo jobInfo, KafkaStreamsPipelineOptions pipelineOptions) { return new KafkaStreamsTranslationContext(jobInfo, pipelineOptions, new Topology()); @@ -102,6 +107,15 @@ public MetricsContainerStepMap getMetricsContainerStepMap() { return metricsContainerStepMap; } + /** + * Returns the tracker that decides when this pipeline has finished. Processors report themselves + * to it as they reach the terminal watermark; the runner asks it to stop the Kafka Streams client + * once they all have. + */ + public TerminationTracker getTerminationTracker() { + return terminationTracker; + } + /** * 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/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java index eb20a8b83588..a6ef768ae2a1 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java @@ -95,6 +95,8 @@ class ReadProcessor implements Processor> runnerWireCoder; private final String stateStoreName; private final String transformId; + // Reports this source as finished once it emits the terminal watermark. + private final TerminationReporter terminationReporter; private @Nullable ProcessorContext> context; private @Nullable KeyValueStore firedStore; @@ -106,19 +108,27 @@ class ReadProcessor implements Processor> sdkWireCoder, Coder> runnerWireCoder, String stateStoreName, - String transformId) { + String transformId, + TerminationTracker terminationTracker) { this.source = source; this.options = options; this.sdkWireCoder = sdkWireCoder; this.runnerWireCoder = runnerWireCoder; this.stateStoreName = stateStoreName; this.transformId = transformId; + this.terminationReporter = new TerminationReporter(terminationTracker, transformId); + } + + @Override + public void close() { + terminationReporter.close(); } @Override public void init(ProcessorContext> context) { this.context = context; this.firedStore = context.getStateStore(stateStoreName); + terminationReporter.init(context); this.scheduledPunctuator = context.schedule(PUNCTUATION_DELAY, PunctuationType.WALL_CLOCK_TIME, ts -> maybeFire()); } @@ -193,6 +203,7 @@ private void forwardWatermarkMax(ProcessorContext> ct ctx.forward( new Record>( new byte[0], KStreamsPayload.watermark(maxMillis, transformId, 0, 1), 0L)); + terminationReporter.watermarkEmitted(ctx, maxMillis); } /** Cancels the wall-clock punctuator after the read has fired to stop periodic wakeups. */ diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java index f83442f97813..049f29651ebf 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java @@ -155,7 +155,8 @@ private void addUnbounde stateStoreName, transformId, maxElementsPerPoll, - checkpointEveryNPolls), + checkpointEveryNPolls, + context.getTerminationTracker()), sourceNodeName); topology.addStateStore( Stores.keyValueStoreBuilder( @@ -198,7 +199,13 @@ private void addReadNodes( transformId, () -> new ReadProcessor<>( - source, options, sdkWireCoder, runnerWireCoder, stateStoreName, transformId), + source, + options, + sdkWireCoder, + runnerWireCoder, + stateStoreName, + transformId, + context.getTerminationTracker()), sourceNodeName); KeyValueBytesStoreSupplier storeSupplier = Stores.persistentKeyValueStore(stateStoreName); topology.addStateStore( diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java index 79595cba7c96..638239406ed9 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java @@ -57,11 +57,21 @@ class ShuffleByKeyProcessor private int upstreamPartition; + // Reports this shuffle as finished once it has written the terminal watermark to the repartition + // topic. The downstream side reading that topic reports separately, which is why the pipeline + // waits for every processor rather than the first. + private final TerminationReporter terminationReporter; + private @Nullable ProcessorContext> context; - ShuffleByKeyProcessor(Coder keyCoder, int upstreamPartitionCount) { + ShuffleByKeyProcessor( + Coder keyCoder, + int upstreamPartitionCount, + String nodeName, + TerminationTracker terminationTracker) { this.keyCoder = keyCoder; this.upstreamPartitionCount = upstreamPartitionCount; + this.terminationReporter = new TerminationReporter(terminationTracker, nodeName); } @Override @@ -70,6 +80,12 @@ public void init(ProcessorContext> context) { // This processor runs in the upstream transform's task, so the task's partition is the // identity of the instance whose reports it is forwarding. this.upstreamPartition = context.taskId().partition(); + terminationReporter.init(context); + } + + @Override + public void close() { + terminationReporter.close(); } @Override @@ -109,6 +125,7 @@ public void process(Record> record) { upstreamPartition, upstreamPartitionCount), record.timestamp())); + terminationReporter.watermarkEmitted(ctx, report.getWatermarkMillis()); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java index eec3f2bae08a..6b8d62763453 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java @@ -48,15 +48,24 @@ class StageOutputProcessor private static final Logger LOG = LoggerFactory.getLogger(StageOutputProcessor.class); private final String transformId; + // Reports this output port as finished once the stage's terminal watermark reaches it. + private final TerminationReporter terminationReporter; private @Nullable ProcessorContext> context; - StageOutputProcessor(String transformId) { + StageOutputProcessor(String transformId, TerminationTracker terminationTracker) { this.transformId = transformId; + this.terminationReporter = new TerminationReporter(terminationTracker, transformId); } @Override public void init(ProcessorContext> context) { this.context = context; + terminationReporter.init(context); + } + + @Override + public void close() { + terminationReporter.close(); } @Override @@ -89,5 +98,6 @@ public void process(Record> record) { report.getSourcePartition(), report.getTotalSourcePartitions()), record.timestamp())); + terminationReporter.watermarkEmitted(ctx, report.getWatermarkMillis()); } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java new file mode 100644 index 000000000000..abc10500d211 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java @@ -0,0 +1,111 @@ +/* + * 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 java.time.Duration; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.kafka.streams.processor.Cancellable; +import org.apache.kafka.streams.processor.PunctuationType; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * The bit of every watermark-emitting processor that reports it has finished, so a bounded pipeline + * can stop itself. See {@link TerminationTracker} for why the runner has to work this out at all. + * + *

    A processor creates one of these, calls {@link #init} from {@code Processor#init}, passes + * every watermark it emits to {@link #watermarkEmitted}, and calls {@link #close} from {@code + * Processor#close}. + * + *

    Why termination is scheduled rather than reported inline

    + * + *

    Reporting from inside {@code process()} would announce the processor as finished while it is + * still in the middle of handling the record that carried the terminal watermark. Scheduling a + * punctuator instead defers the report until the current processing has completed, so anything that + * has to happen after the final watermark — flushing a bundle, forwarding downstream, committing — + * still runs first. + * + *

    The punctuator is {@link PunctuationType#WALL_CLOCK_TIME} rather than stream time: no further + * records arrive after the terminal watermark, so stream time would never advance and a stream-time + * punctuator would never fire. The interval is the smallest Kafka Streams accepts — it rejects + * anything below a millisecond with "The minimum supported scheduling interval is 1 millisecond." + */ +class TerminationReporter { + + /** Kafka Streams rejects any scheduling interval below this. */ + private static final Duration IMMEDIATELY = Duration.ofMillis(1); + + private final TerminationTracker tracker; + private final String transformId; + + private @Nullable String instanceId; + private @Nullable Cancellable scheduled; + private boolean reported; + + TerminationReporter(TerminationTracker tracker, String transformId) { + this.tracker = tracker; + this.transformId = transformId; + } + + /** Registers this processor instance as something the pipeline is waiting on. */ + void init(ProcessorContext context) { + // The task is what makes the id unique: one processor node runs as one instance per task. + this.instanceId = transformId + "#" + context.taskId(); + tracker.register(instanceId); + } + + /** + * Called with every watermark the processor emits. Once that watermark is terminal, schedules the + * report that this processor has no further work. + */ + void watermarkEmitted(ProcessorContext context, long watermarkMillis) { + if (reported || watermarkMillis < BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()) { + return; + } + reported = true; + scheduled = + context.schedule( + IMMEDIATELY, + PunctuationType.WALL_CLOCK_TIME, + timestamp -> { + cancelSchedule(); + String id = instanceId; + if (id != null) { + tracker.terminate(id); + } + }); + } + + /** Stops the pipeline waiting on this processor, e.g. when its task migrates on a rebalance. */ + void close() { + cancelSchedule(); + String id = instanceId; + if (id != null) { + tracker.unregister(id); + instanceId = null; + } + } + + private void cancelSchedule() { + Cancellable handle = scheduled; + if (handle != null) { + handle.cancel(); + scheduled = null; + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java new file mode 100644 index 000000000000..6c570db1aff1 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java @@ -0,0 +1,191 @@ +/* + * 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 java.util.HashSet; +import java.util.Set; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Decides when a bounded pipeline has finished, so the Kafka Streams client can be stopped. + * + *

    Kafka Streams has no notion of a processor being finished: a topology runs until something + * closes the client. A bounded Beam pipeline does finish, though, and the runner already knows + * when: every processor emits a watermark of {@link + * org.apache.beam.sdk.transforms.windowing.BoundedWindow#TIMESTAMP_MAX_VALUE} once its input is + * exhausted. This class collects those reports and fires a callback when there is nothing left to + * do. + * + *

    Why no coordination between instances is needed

    + * + *

    A watermark that crosses a repartition topic is broadcast to every partition (see + * {@link GroupByKeyBroadcastPartitioner}), so every task of every downstream transform observes the + * terminal watermark on its own, whichever instance it happens to run on. Each instance can + * therefore decide to stop from what it sees locally, and they all reach the same conclusion + * without talking to each other. + * + *

    Why every local processor has to be counted, not just the first

    + * + *

    One instance can own tasks from both sides of a repartition topic. The upstream side goes + * terminal as soon as it has written its data to the topic, while the downstream side still has to + * consume it. Stopping the client when the first processor finishes would cut that downstream work + * off and report the pipeline as done having silently dropped it. So the callback only fires once + * every processor instance registered here has terminated. + * + *

    An instance that happens to own only upstream tasks still terminates on its own, which is + * correct: what it wrote is durable in the topic for whichever instance reads it. + * + *

    Scope

    + * + *

    One tracker belongs to one pipeline, not to the JVM. The job server runs many jobs in a single + * process, so a shared static tracker would let one pipeline finishing tear down another. + * + *

    A pipeline with an unbounded source never produces a terminal watermark, so the callback never + * fires and the client keeps running — which is the intended behaviour for a streaming job. + */ +public class TerminationTracker { + + private static final Logger LOG = LoggerFactory.getLogger(TerminationTracker.class); + + /** Processor instances currently running here, by {@code transformId#taskId}. */ + private final Set live = new HashSet<>(); + + /** Those of {@link #live} that have emitted the terminal watermark. */ + private final Set terminated = new HashSet<>(); + + /** + * What to do when the pipeline is finished, cleared as it is taken. + * + *

    Clearing it is what stops it running twice: a pipeline only finishes once, but processors go + * on reporting afterwards — the callback stops the client, and closing it makes every remaining + * task close its processors, each of which unregisters and asks again. + */ + private @Nullable Runnable onAllTerminated; + + /** + * Whether the topology is fully up, and so whether the registered processors are the whole set. + * + *

    Processors register as their task is initialized, which happens gradually while the client + * starts. Deciding before that is finished reads "every processor is done" off a set that is + * merely incomplete: on a short pipeline the source can drain before the task downstream of the + * repartition topic exists, and stopping there discards the rest of the pipeline and reports a + * successful run that produced nothing. + */ + private boolean started; + + /** + * Sets what to do when the pipeline is finished. Must be called before the topology starts, so + * that no processor can terminate before there is anything to call. + * + *

    The callback runs on whichever thread completes the picture: usually the Kafka Streams task + * thread reporting the last termination, but the thread reporting startup when the pipeline + * drained before it finished starting. Both are threads {@code KafkaStreams.close()} waits for, + * so stopping the client is the callback's job to hand off to a thread of its own. + */ + public synchronized void onAllTerminated(Runnable callback) { + this.onAllTerminated = callback; + } + + /** + * Marks the topology as fully started, after which the registered processors are taken to be the + * whole set. Called when Kafka Streams reports {@code RUNNING}, which it does once every assigned + * task has been initialized. + * + *

    A pipeline short enough to drain during startup will already have reported terminations by + * then, so this re-checks rather than only gating what comes later. + */ + public void started() { + Runnable callback; + synchronized (this) { + started = true; + callback = takeCallbackIfDone(); + } + run(callback); + } + + /** Registers a processor instance, called from {@code Processor#init}. */ + synchronized void register(String instanceId) { + live.add(instanceId); + } + + /** + * Removes a processor instance, called from {@code Processor#close}, so that a task migrating + * away during a rebalance is not waited on forever. + */ + void unregister(String instanceId) { + Runnable callback; + synchronized (this) { + live.remove(instanceId); + terminated.remove(instanceId); + callback = takeCallbackIfDone(); + } + run(callback); + } + + /** + * Records that a processor instance has emitted the terminal watermark and has no further work. + */ + void terminate(String instanceId) { + Runnable callback; + synchronized (this) { + if (!live.contains(instanceId)) { + // Terminated after being unregistered, or never registered: nothing is waiting on it. + return; + } + if (terminated.add(instanceId)) { + LOG.debug( + "Processor {} reached the terminal watermark ({}/{})", + instanceId, + terminated.size(), + live.size()); + } + callback = takeCallbackIfDone(); + } + run(callback); + } + + private static void run(@Nullable Runnable callback) { + // Deliberately outside the lock: the callback shuts the pipeline down, and holding the monitor + // while calling into shutdown makes this class part of that path for anyone who changes what + // the callback does later. + if (callback != null) { + callback.run(); + } + } + + /** + * Returns the callback to run if the pipeline is finished, having claimed the right to run it. + */ + private @Nullable Runnable takeCallbackIfDone() { + if (!started || live.isEmpty() || !terminated.containsAll(live)) { + return null; + } + Runnable callback = onAllTerminated; + if (callback == null) { + // Never set, or already taken — either way there is nothing left to do. + return null; + } + onAllTerminated = null; + LOG.info( + "All {} processor instances reached the terminal watermark; stopping the pipeline", + live.size()); + return callback; + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java index b616e0fa8534..96aa1bf0c69e 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java @@ -99,6 +99,10 @@ class UnboundedReadProcessor /** Set once the source's watermark reaches the end of time; it will produce nothing more. */ private boolean exhausted; + // An unbounded source normally never reaches the terminal watermark, so this normally never + // reports anything. It does matter for a source that is drained or is bounded in practice. + private final TerminationReporter terminationReporter; + private @Nullable Cancellable scheduledPunctuator; UnboundedReadProcessor( @@ -110,7 +114,9 @@ class UnboundedReadProcessor String stateStoreName, String transformId, int maxElementsPerPoll, - int checkpointEveryNPolls) { + int checkpointEveryNPolls, + TerminationTracker terminationTracker) { + this.terminationReporter = new TerminationReporter(terminationTracker, transformId); this.source = source; this.options = options; this.sdkWireCoder = sdkWireCoder; @@ -126,6 +132,7 @@ class UnboundedReadProcessor public void init(ProcessorContext> context) { this.context = context; this.checkpointStore = context.getStateStore(stateStoreName); + terminationReporter.init(context); this.scheduledPunctuator = context.schedule(POLL_INTERVAL, PunctuationType.WALL_CLOCK_TIME, timestamp -> poll()); } @@ -224,6 +231,7 @@ private void forwardWatermarkIfAdvanced( ctx.forward( new Record>( new byte[0], KStreamsPayload.watermark(watermark.getMillis(), transformId, 0, 1), 0L)); + terminationReporter.watermarkEmitted(ctx, watermark.getMillis()); } /** Creates the reader on first use, resuming from the stored checkpoint mark if there is one. */ @@ -293,6 +301,8 @@ public void close() { } reader = null; } + // Last, so the pipeline is not declared finished while this source is still closing down. + terminationReporter.close(); } private static V checkInitialized(@Nullable V value) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java index d00b01a3753d..4c328542a322 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java @@ -92,6 +92,10 @@ class WindowedGroupByKeyProcessor private Instant inputWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; private @Nullable ProcessorContext> context; + // Reports this GroupByKey instance as finished once it emits the terminal watermark, which it + // only does after firing every pane it was holding. + private final TerminationReporter terminationReporter; + private @Nullable KeyValueStore stateStore; private @Nullable KeyValueStore holdsIndexStore; private @Nullable KeyValueStore timerStore; @@ -107,7 +111,9 @@ class WindowedGroupByKeyProcessor Coder keyCoder, Coder valueCoder, WindowingStrategy windowingStrategy, - PipelineOptions options) { + PipelineOptions options, + TerminationTracker terminationTracker) { + this.terminationReporter = new TerminationReporter(terminationTracker, transformId); this.stateStoreName = stateStoreName; this.holdsIndexStoreName = holdsIndexStoreName; this.timerStoreName = timerStoreName; @@ -129,6 +135,12 @@ public void init(ProcessorContext> context) { this.holdsIndexStore = context.getStateStore(holdsIndexStoreName); this.timerStore = context.getStateStore(timerStoreName); this.timerIndexStore = context.getStateStore(timerIndexStoreName); + terminationReporter.init(context); + } + + @Override + public void close() { + terminationReporter.close(); } @Override @@ -291,6 +303,7 @@ private void forwardWatermark(Record> trigger, long w trigger.key(), KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1), trigger.timestamp())); + terminationReporter.watermarkEmitted(ctx, watermarkMillis); } private @NonNull K decodeKey(byte[] bytes) { diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java index bb82a403e557..b060561eaab7 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java @@ -253,6 +253,28 @@ public void chainedGroupByKeysAreCorrectAcrossPartitions() throws Exception { } } + @Test + public void aBoundedPipelineTerminatesOnItsOwn() throws Exception { + // Kafka Streams runs a topology until something stops the client, so a bounded pipeline used to + // run for ever against a real broker: it produced the right answer and then sat there. The + // other tests here cannot see that, because they cancel rather than wait, and the + // ValidatesRunner + // suite cannot either, because TopologyTestDriver is synchronous and always reports DONE. + // + // Nothing cancels this one. Returning from run() at all is the assertion. + KafkaStreamsPipelineOptions options = options(4); + Pipeline pipeline = Pipeline.create(options); + buildChainedPipeline(pipeline); + + PipelineResult result = runPipeline(pipeline, options); + + assertThat(result.getState(), is(PipelineResult.State.DONE)); + // And it stopped for the right reason — having produced its output exactly once. Termination is + // driven from a wall-clock punctuator, which is the same mechanism that duplicates output when + // it is used to close bundles on time (#39633), so the count matters as much as the state. + assertThat(counterValue(result), is(1L)); + } + /** * Polls the pipeline's metrics until the counter reaches {@code expected} or the timeout hits. */ 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 290e109796a8..22728500bf0a 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 @@ -64,7 +64,8 @@ private static ExecutableStageProcessor newProcessor() { // Single-output: no per-output routing (this test drives the watermark path directly). ImmutableMap.of(), // The bundle size bound is irrelevant to the watermark path this test drives. - 1000); + 1000, + new TerminationTracker()); } /** 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/ShuffleByKeyProcessorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java index 38669138075c..b6716764da7f 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java @@ -48,7 +48,9 @@ private static ShuffleByKeyProcessor processorFor(int taskPartition, int upstrea new ShuffleByKeyProcessor( (org.apache.beam.sdk.coders.Coder) (org.apache.beam.sdk.coders.Coder) StringUtf8Coder.of(), - upstreamPartitions); + upstreamPartitions, + "shuffle-node", + new TerminationTracker()); MockProcessorContext> ctx = new MockProcessorContext<>(new Properties(), new TaskId(0, taskPartition), null); processor.init(ctx); diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java index 83653f13c891..9d17ae73fcb6 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessorTest.java @@ -56,7 +56,7 @@ private static Record> watermark( @Test public void watermarkKeepsPartitionIdentityAndRelabelsTransformId() { MockProcessorContext> ctx = new MockProcessorContext<>(); - StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID); + StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID, new TerminationTracker()); processor.init(ctx); // A report from partition 1 of a 3-instance stage. @@ -76,7 +76,7 @@ public void watermarkKeepsPartitionIdentityAndRelabelsTransformId() { @Test public void distinctStagePartitionsStayDistinctDownstream() { MockProcessorContext> ctx = new MockProcessorContext<>(); - StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID); + StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID, new TerminationTracker()); processor.init(ctx); processor.process(watermark(100L, 0, 3)); @@ -93,7 +93,7 @@ public void distinctStagePartitionsStayDistinctDownstream() { @Test public void dataIsForwardedUnchanged() { MockProcessorContext> ctx = new MockProcessorContext<>(); - StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID); + StageOutputProcessor processor = new StageOutputProcessor(RELAY_ID, new TerminationTracker()); processor.init(ctx); WindowedValue element = WindowedValues.valueInGlobalWindow(new byte[] {7}); diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/TerminationTrackerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/TerminationTrackerTest.java new file mode 100644 index 000000000000..72ef28d8bef1 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/TerminationTrackerTest.java @@ -0,0 +1,172 @@ +/* + * 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.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.is; + +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link TerminationTracker}. */ +@RunWith(JUnit4.class) +public class TerminationTrackerTest { + + private final AtomicInteger calls = new AtomicInteger(); + + /** A tracker for a topology that has finished starting up. */ + private TerminationTracker tracker() { + TerminationTracker tracker = new TerminationTracker(); + tracker.onAllTerminated(calls::incrementAndGet); + tracker.started(); + return tracker; + } + + @Test + public void doesNotFireBeforeTheTopologyHasFinishedStarting() { + // Processors register as their task is initialized, so mid-startup the registered set is only + // part of the pipeline. A source that drains that quickly would otherwise look like a finished + // pipeline, and stopping there discards every stage that had not started yet — which shows up + // as a run that reports success and produces no output. + TerminationTracker tracker = new TerminationTracker(); + tracker.onAllTerminated(calls::incrementAndGet); + + tracker.register("source#0_0"); + tracker.terminate("source#0_0"); + assertThat("the rest of the topology may not exist yet", calls.get(), is(0)); + + // The stage downstream of the repartition topic comes up late and has real work to do. + tracker.register("downstream#1_0"); + tracker.started(); + assertThat(calls.get(), is(0)); + + tracker.terminate("downstream#1_0"); + assertThat(calls.get(), is(1)); + } + + @Test + public void firesOnStartupIfEverythingAlreadyTerminated() { + // A pipeline short enough to drain entirely during startup still has to be noticed. + TerminationTracker tracker = new TerminationTracker(); + tracker.onAllTerminated(calls::incrementAndGet); + tracker.register("source#0_0"); + tracker.terminate("source#0_0"); + + tracker.started(); + + assertThat(calls.get(), is(1)); + } + + @Test + public void firesOnceEveryRegisteredProcessorHasTerminated() { + TerminationTracker tracker = tracker(); + tracker.register("stage#0_0"); + tracker.register("stage#0_1"); + + tracker.terminate("stage#0_0"); + assertThat("one of two done is not the whole pipeline", calls.get(), is(0)); + + tracker.terminate("stage#0_1"); + assertThat(calls.get(), is(1)); + } + + @Test + public void doesNotFireWhileAProcessorIsStillRunning() { + TerminationTracker tracker = tracker(); + // The shape that makes counting every processor necessary: one instance owning both sides of a + // repartition topic. The upstream goes terminal as soon as it has written to the topic, while + // the downstream still has to consume it. + tracker.register("upstream#0_0"); + tracker.register("downstream#1_0"); + + tracker.terminate("upstream#0_0"); + + assertThat("stopping here would cut the downstream off", calls.get(), is(0)); + } + + @Test + public void doesNotFireWithNothingRegistered() { + TerminationTracker tracker = tracker(); + tracker.terminate("never-registered#0_0"); + assertThat(calls.get(), is(0)); + } + + @Test + public void firesOnlyOnce() { + TerminationTracker tracker = tracker(); + tracker.register("stage#0_0"); + + tracker.terminate("stage#0_0"); + tracker.terminate("stage#0_0"); + + assertThat(calls.get(), is(1)); + } + + @Test + public void shuttingDownDoesNotFireAgain() { + // What the callback does is stop the client, which closes every task's processors, and each of + // those unregisters on the way out. So the tracker is asked again several times after the + // pipeline has already been declared finished. + TerminationTracker tracker = tracker(); + tracker.register("source#0_0"); + tracker.register("stage#1_0"); + tracker.terminate("source#0_0"); + tracker.terminate("stage#1_0"); + assertThat(calls.get(), is(1)); + + tracker.unregister("source#0_0"); + tracker.unregister("stage#1_0"); + + assertThat("stopping the pipeline must not stop it a second time", calls.get(), is(1)); + } + + @Test + public void aProcessorThatMigratesAwayIsNoLongerWaitedOn() { + TerminationTracker tracker = tracker(); + tracker.register("stage#0_0"); + tracker.register("stage#0_1"); + tracker.terminate("stage#0_0"); + + // Task 0_1 is reassigned to another instance during a rebalance. What is left here is done, so + // this instance has nothing to keep it alive. + tracker.unregister("stage#0_1"); + + assertThat(calls.get(), is(1)); + } + + @Test + public void unregisteringTheLastProcessorDoesNotFire() { + TerminationTracker tracker = tracker(); + tracker.register("stage#0_0"); + + tracker.unregister("stage#0_0"); + + assertThat("nothing registered means nothing finished", calls.get(), is(0)); + } + + @Test + public void withoutACallbackNothingHappens() { + TerminationTracker tracker = new TerminationTracker(); + tracker.register("stage#0_0"); + tracker.terminate("stage#0_0"); + // No callback set: the point is that this does not throw. + assertThat(calls.get(), is(0)); + } +} From 5f1658f8b5e0842aaedd12612a0cb44ae5fd3219 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:34:54 +0500 Subject: [PATCH 30/37] [GSoC 2026] Kafka Streams runner: portable ValidatesRunner suite for Python (#39736) * [GSoC 2026] Kafka Streams runner: portable ValidatesRunner suite for Python Runs Beam's portable ValidatesRunner suite against the runner, which is what shows it works for a pipeline that was not written in Java. 29 tests pass and 46 are skipped, each skip naming the issue for the feature it needs. * [GSoC 2026] Kafka Streams runner: integration test for two instances over three partitions --- .../streams/KafkaStreamsPipelineRunner.java | 32 +- .../KafkaStreamsPortablePipelineResult.java | 17 +- .../translation/GroupByKeyTranslator.java | 12 +- .../translation/ImpulseTranslator.java | 3 +- .../KafkaStreamsTranslationContext.java | 17 ++ .../streams/translation/ReadTranslator.java | 6 +- ...afkaStreamsPortablePipelineResultTest.java | 96 ++++++ .../streams/KafkaStreamsRunnerBrokerIT.java | 66 +++- .../portability/kafka_streams_runner_test.py | 285 ++++++++++++++++++ .../python/test-suites/portable/common.gradle | 35 +++ sdks/python/tox.ini | 5 + 11 files changed, 558 insertions(+), 16 deletions(-) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResultTest.java create mode 100644 sdks/python/apache_beam/runners/portability/kafka_streams_runner_test.py 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 93e1058121ed..de08480a61d4 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 @@ -18,6 +18,8 @@ package org.apache.beam.runners.kafka.streams; import java.util.Properties; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicReference; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.jobsubmission.PortablePipelineResult; @@ -27,6 +29,7 @@ import org.apache.kafka.streams.KafkaStreams; import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.errors.StreamsUncaughtExceptionHandler; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -77,6 +80,19 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) topology.describe()); KafkaStreams kafkaStreams = new KafkaStreams(topology, streamsConfig(jobInfo)); + // Kafka Streams reports a failed task by moving the client to ERROR and keeping the exception + // to itself, which left a failed job with nothing to say beyond "unknown error". Hold on to the + // first failure so this method can rethrow it: the job service turns what run() throws into the + // job's error message. + AtomicReference<@Nullable Throwable> failure = new AtomicReference<>(); + kafkaStreams.setUncaughtExceptionHandler( + throwable -> { + failure.compareAndSet(null, throwable); + LOG.error("Pipeline {} failed", jobInfo.jobId(), throwable); + // The pipeline is a job with an owner waiting on it, not a service to keep alive, so a + // failure stops the client rather than replacing the thread and carrying on. + return StreamsUncaughtExceptionHandler.StreamThreadExceptionResponse.SHUTDOWN_CLIENT; + }); // Build the result before starting: it registers a state listener, and Kafka Streams only // accepts one while the application is still in the CREATED state. KafkaStreamsPortablePipelineResult result = @@ -110,6 +126,12 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) // stream threads and the joins it does would throw straight back out of an interrupted one. closeInBackground(kafkaStreams, jobInfo.jobId(), "the job was cancelled"); } + Throwable thrown = failure.get(); + if (thrown != null) { + // Thrown rather than returned as a failed result: the job service reads the state of what is + // returned, but only what is thrown carries a reason the user can act on. + throw new RuntimeException("Pipeline " + jobInfo.jobId() + " failed", thrown); + } return result; } @@ -146,7 +168,15 @@ private Properties streamsConfig(JobInfo jobInfo) { props.put(StreamsConfig.APPLICATION_ID_CONFIG, pipelineOptions.getApplicationId()); props.put(StreamsConfig.STATE_DIR_CONFIG, pipelineOptions.getStateDir()); props.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE_V2); - props.put(StreamsConfig.CLIENT_ID_CONFIG, jobInfo.jobId()); + // The job id identifies the pipeline, which every instance of it shares, so on its own it does + // not identify an instance. Kafka Streams names threads, consumers and metrics after the client + // id, so two workers running the same job would produce logs and JMX metrics that cannot be + // told + // apart — in a deployment whose whole point is that you add workers. Keeping the job id as the + // prefix leaves the pipeline recognizable; the suffix is what makes each worker distinct, and + // is + // what Kafka Streams does by default when no client id is set. + props.put(StreamsConfig.CLIENT_ID_CONFIG, jobInfo.jobId() + "-" + UUID.randomUUID()); return props; } } 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 81e80b66e303..29c07917986f 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 @@ -26,8 +26,6 @@ import org.apache.beam.sdk.metrics.MetricResults; import org.apache.kafka.streams.KafkaStreams; import org.joda.time.Duration; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; /** * Result of executing a portable pipeline as a {@link KafkaStreams} application. @@ -38,9 +36,6 @@ */ class KafkaStreamsPortablePipelineResult implements PortablePipelineResult { - private static final Logger LOG = - 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. @@ -126,8 +121,16 @@ public MetricResults metrics() { @Override public JobApi.MetricResults portableMetrics() throws UnsupportedOperationException { - LOG.debug("portableMetrics() not yet implemented in the Kafka Streams runner"); - return JobApi.MetricResults.newBuilder().build(); + // How a pipeline from another SDK reads its metrics. The job service asks for these once the + // job is terminal and returns them over the job API. Without it a Python pipeline saw no + // metrics at all, even though the same values were already available to a Java one. + // + // Reported as attempted only, and deliberately not also as committed: the values are what the + // SDK harness reported per bundle, which is not tied to the commit of the records that produced + // them. Committed metrics are https://github.com/apache/beam/issues/39635. + return JobApi.MetricResults.newBuilder() + .addAllAttempted(metricsContainerStepMap.getMonitoringInfos()) + .build(); } private static State mapState(KafkaStreams.State state) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index c0e69302386a..a62d937170e9 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -102,10 +102,14 @@ public void translate( String shuffleName = transformId + SHUFFLE_SUFFIX; String sinkName = transformId + SINK_SUFFIX; String sourceName = transformId + SOURCE_SUFFIX; - String stateStoreName = transformId + STATE_STORE_SUFFIX; - String holdsIndexStoreName = transformId + HOLDS_INDEX_STORE_SUFFIX; - String timerStoreName = transformId + TIMER_STORE_SUFFIX; - String timerIndexStoreName = transformId + TIMER_INDEX_STORE_SUFFIX; + String stateStoreName = + KafkaStreamsTranslationContext.getStoreName(transformId, STATE_STORE_SUFFIX); + String holdsIndexStoreName = + KafkaStreamsTranslationContext.getStoreName(transformId, HOLDS_INDEX_STORE_SUFFIX); + String timerStoreName = + KafkaStreamsTranslationContext.getStoreName(transformId, TIMER_STORE_SUFFIX); + String timerIndexStoreName = + KafkaStreamsTranslationContext.getStoreName(transformId, TIMER_INDEX_STORE_SUFFIX); String repartitionTopic = repartitionTopic(transformId, context.getPipelineOptions().getApplicationId()); diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java index 3daf7782362e..29bd6e6bd9d3 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java @@ -70,7 +70,8 @@ public void translate( Topology topology = context.getTopology(); String sourceNodeName = transformId + SOURCE_SUFFIX; - String stateStoreName = transformId + STATE_STORE_SUFFIX; + String stateStoreName = + KafkaStreamsTranslationContext.getStoreName(transformId, STATE_STORE_SUFFIX); String bootstrapTopic = context.getImpulseBootstrapTopic(transformId); topology.addSource( 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 904bd8a71e3a..4a5463677163 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 @@ -192,4 +192,21 @@ public String getReadBootstrapTopic(String transformId) { + "_" + sanitizedTransformId; } + + /** + * Returns the name of a state store belonging to a transform. + * + *

    The transform id is sanitized to Kafka's legal topic-name characters even though a store + * name is not itself a topic: Kafka Streams names a persistent store's changelog topic after the + * store, so a transform whose name contains a character a topic may not — which is ordinary, + * {@code CombinePerKey(MeanCombineFn)/Group} is a Beam transform name — would fail at runtime + * when the changelog is created. + * + *

    Two transform ids differing only in characters that are replaced would sanitize to one name. + * Kafka Streams rejects a store name that is already taken when the topology is built, so that + * surfaces as a failure to start rather than as two transforms quietly sharing state. + */ + public static String getStoreName(String transformId, String suffix) { + return ILLEGAL_TOPIC_CHARS.matcher(transformId).replaceAll("_") + suffix; + } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java index 049f29651ebf..69006661ee65 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java @@ -126,7 +126,8 @@ private void addUnbounde Topology topology = context.getTopology(); String sourceNodeName = transformId + SOURCE_SUFFIX; - String stateStoreName = transformId + STATE_STORE_SUFFIX; + String stateStoreName = + KafkaStreamsTranslationContext.getStoreName(transformId, STATE_STORE_SUFFIX); String bootstrapTopic = context.getReadBootstrapTopic(transformId); SerializablePipelineOptions options = new SerializablePipelineOptions(context.getPipelineOptions()); @@ -185,7 +186,8 @@ private void addReadNodes( Topology topology = context.getTopology(); String sourceNodeName = transformId + SOURCE_SUFFIX; - String stateStoreName = transformId + STATE_STORE_SUFFIX; + String stateStoreName = + KafkaStreamsTranslationContext.getStoreName(transformId, STATE_STORE_SUFFIX); String bootstrapTopic = context.getReadBootstrapTopic(transformId); SerializablePipelineOptions options = new SerializablePipelineOptions(context.getPipelineOptions()); diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResultTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResultTest.java new file mode 100644 index 000000000000..2c44765151b4 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResultTest.java @@ -0,0 +1,96 @@ +/* + * 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; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.hasItem; +import static org.hamcrest.Matchers.hasProperty; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.not; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import org.apache.beam.model.jobmanagement.v1.JobApi; +import org.apache.beam.runners.core.metrics.MetricsContainerImpl; +import org.apache.beam.runners.core.metrics.MetricsContainerStepMap; +import org.apache.beam.sdk.metrics.MetricName; +import org.apache.kafka.streams.KafkaStreams; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests for {@link KafkaStreamsPortablePipelineResult}, in particular the metrics an SDK other than + * Java reads its results through. + */ +@RunWith(JUnit4.class) +public class KafkaStreamsPortablePipelineResultTest { + + private static final String STEP = "a-stage"; + private static final String NAMESPACE = "ns"; + private static final String COUNTER = "elements"; + + private static KafkaStreams idleClient() { + KafkaStreams kafkaStreams = mock(KafkaStreams.class); + // The result registers a state listener and checks the current state, so it has to have one. + when(kafkaStreams.state()).thenReturn(KafkaStreams.State.CREATED); + return kafkaStreams; + } + + @Test + public void portableMetricsReportWhatTheHarnessMeasured() { + MetricsContainerStepMap stepMap = new MetricsContainerStepMap(); + MetricsContainerImpl container = stepMap.getContainer(STEP); + container.getCounter(MetricName.named(NAMESPACE, COUNTER)).inc(7); + + KafkaStreamsPortablePipelineResult result = + new KafkaStreamsPortablePipelineResult(idleClient(), stepMap, () -> {}); + + JobApi.MetricResults metrics = result.portableMetrics(); + + // A pipeline from another SDK reads these over the job API; before they were reported the list + // was empty and a Python pipeline saw no metrics at all. + assertThat( + metrics.getAttemptedList(), + hasItem(hasProperty("urn", is("beam:metric:user:sum_int64:v1")))); + assertThat(metrics.getAttemptedCount(), is(not(0))); + } + + @Test + public void portableMetricsAreNotReportedAsCommitted() { + // The values are what the SDK harness reported per bundle, which is not tied to the commit of + // the records that produced them, so claiming them as committed would be wrong. + // See https://github.com/apache/beam/issues/39635. + MetricsContainerStepMap stepMap = new MetricsContainerStepMap(); + stepMap.getContainer(STEP).getCounter(MetricName.named(NAMESPACE, COUNTER)).inc(1); + + KafkaStreamsPortablePipelineResult result = + new KafkaStreamsPortablePipelineResult(idleClient(), stepMap, () -> {}); + + assertThat(result.portableMetrics().getCommittedCount(), is(0)); + } + + @Test + public void aPipelineThatMeasuredNothingReportsNothing() { + KafkaStreamsPortablePipelineResult result = + new KafkaStreamsPortablePipelineResult( + idleClient(), new MetricsContainerStepMap(), () -> {}); + + assertThat(result.portableMetrics().getAttemptedCount(), is(0)); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java index b060561eaab7..acc4130b8216 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java @@ -21,7 +21,14 @@ import static org.hamcrest.MatcherAssert.assertThat; import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.UUID; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.sdk.Pipeline; @@ -123,11 +130,22 @@ private KafkaStreamsPipelineOptions options() { } private KafkaStreamsPipelineOptions options(int topicPartitions) { + return options(topicPartitions, "ks-broker-it-" + UUID.randomUUID()); + } + + /** + * Options for one runner instance. + * + *

    Two instances of the same job share an application id — that is what puts them in one + * consumer group and so splits the work between them — but each needs its own state directory, + * since the local stores are per instance. + */ + private KafkaStreamsPipelineOptions options(int topicPartitions, String applicationId) { KafkaStreamsPipelineOptions options = PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); options.setRunner(CrashingRunner.class); options.setBootstrapServers(kafka.getBootstrapServers()); - options.setApplicationId("ks-broker-it-" + UUID.randomUUID()); + options.setApplicationId(applicationId); options.setInternalParallelism(topicPartitions); options .as(PortablePipelineOptions.class) @@ -275,6 +293,52 @@ public void aBoundedPipelineTerminatesOnItsOwn() throws Exception { assertThat(counterValue(result), is(1L)); } + @Test + public void twoInstancesShareThreePartitions() throws Exception { + // Everything else here runs one instance, which leaves the thing the runner exists for + // untested: the work being split between instances by Kafka's own group membership. + // + // Three partitions across two instances is deliberate. It does not divide, so the instances + // take an unequal share, and a watermark aggregator on either of them has to hear from all + // three upstream partitions — some of which are being produced by the other instance — before + // it may let its watermark advance. If the reports were tied to the instance that produced + // them rather than to the partition, this is the shape that would break. + String applicationId = "ks-broker-it-" + UUID.randomUUID(); + List results = Collections.synchronizedList(new ArrayList<>()); + ExecutorService instances = Executors.newFixedThreadPool(2); + try { + List> running = new ArrayList<>(); + for (int instance = 0; instance < 2; instance++) { + KafkaStreamsPipelineOptions options = options(3, applicationId); + Pipeline pipeline = Pipeline.create(options); + buildChainedPipeline(pipeline); + running.add( + instances.submit( + () -> { + // run() blocks until its instance has finished, so each needs its own thread. + results.add(runPipeline(pipeline, options)); + })); + } + for (Future future : running) { + // Fails rather than hangs if an instance never finishes — which is the interesting way for + // this to go wrong, since an instance only stops once every processor it owns is done. + future.get(TIMEOUT.getMillis(), TimeUnit.MILLISECONDS); + } + } finally { + instances.shutdownNow(); + } + + // The pipeline collapses everything onto one key, so exactly one group comes out of the second + // GroupByKey however the partitions were shared. Each instance counts what it processed, so the + // total across both is what has to be one: a group counted twice would mean the instances had + // both claimed the same partition's data. + long groups = 0; + for (PipelineResult result : results) { + groups += counterValue(result); + } + assertThat(groups, is(1L)); + } + /** * Polls the pipeline's metrics until the counter reaches {@code expected} or the timeout hits. */ diff --git a/sdks/python/apache_beam/runners/portability/kafka_streams_runner_test.py b/sdks/python/apache_beam/runners/portability/kafka_streams_runner_test.py new file mode 100644 index 000000000000..9285e0fd3d31 --- /dev/null +++ b/sdks/python/apache_beam/runners/portability/kafka_streams_runner_test.py @@ -0,0 +1,285 @@ +# +# 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. +# + +# pytype: skip-file + +import argparse +import logging +import shlex +import unittest +import uuid +from shutil import rmtree +from tempfile import mkdtemp + +import pytest + +from apache_beam.options.pipeline_options import KafkaStreamsRunnerOptions +from apache_beam.options.pipeline_options import PortableOptions +from apache_beam.runners.portability import job_server +from apache_beam.runners.portability import portable_runner +from apache_beam.runners.portability import portable_runner_test +from apache_beam.utils import subprocess_server + +# Runs Beam's portable ValidatesRunner suite against the Kafka Streams runner, which is what shows +# the runner works for a pipeline that was not written in Java. +# +# Needs a Kafka broker, unlike the Flink and Spark suites, because the runner executes on Kafka +# rather than on a cluster of its own. Point it at one with --bootstrap_servers. +# +# Run as +# +# pytest kafka_streams_runner_test.py[::TestClass::test_case] \ +# --test-pipeline-options="--bootstrap_servers=localhost:9092" + +_LOGGER = logging.getLogger(__name__) + + +class KafkaStreamsRunnerTest(portable_runner_test.PortableRunnerTest): + _use_grpc = True + _use_subprocesses = True + + expansion_port = None + kafka_streams_job_server_jar = None + bootstrap_servers = 'localhost:9092' + environment_type = 'LOOPBACK' + environment_options = None + + @pytest.fixture(autouse=True) + def parse_options(self, request): + if not request.config.option.test_pipeline_options: + raise unittest.SkipTest( + 'Skipping because --test-pipeline-options is not specified.') + test_pipeline_options = request.config.option.test_pipeline_options + parser = argparse.ArgumentParser(add_help=True) + parser.add_argument( + '--kafka_streams_job_server_jar', + help='Job server jar to submit jobs.', + action='store') + parser.add_argument( + '--bootstrap_servers', + default='localhost:9092', + help='Kafka the runner executes on, and creates its own topics in.') + parser.add_argument( + '--environment_type', + default='LOOPBACK', + choices=['DOCKER', 'PROCESS', 'LOOPBACK'], + help='Set the environment type for running user code. DOCKER runs ' + 'user code in a container. PROCESS runs user code in ' + 'automatically started processes. LOOPBACK runs user code on ' + 'the same process that originally submitted the job.') + parser.add_argument( + '--environment_option', + '--environment_options', + dest='environment_options', + action='append', + default=None, + help=( + 'Environment configuration for running the user code. ' + 'Recognized options depend on --environment_type.')) + known_args, unknown_args = parser.parse_known_args( + shlex.split(test_pipeline_options)) + if unknown_args: + _LOGGER.warning('Discarding unrecognized arguments %s' % unknown_args) + self.set_kafka_streams_job_server_jar( + known_args.kafka_streams_job_server_jar or + job_server.JavaJarJobServer.path_to_beam_jar( + ':runners:kafka-streams:job-server:shadowJar')) + type(self).bootstrap_servers = known_args.bootstrap_servers + self.environment_type = known_args.environment_type + self.environment_options = known_args.environment_options + + @classmethod + def _subprocess_command(cls, job_port, expansion_port): + # Created and used by the job server; removed here so the job server makes it itself. + tmp_dir = mkdtemp(prefix='kafkastreamstest') + + cls.expansion_port = expansion_port + + try: + return [ + subprocess_server.JavaHelper.get_java(), + '-jar', + cls.kafka_streams_job_server_jar, + '--artifacts-dir', + tmp_dir, + '--job-port', + str(job_port), + '--artifact-port', + '0', + '--expansion-port', + str(expansion_port), + ] + finally: + rmtree(tmp_dir) + + @classmethod + def get_runner(cls): + return portable_runner.PortableRunner() + + @classmethod + def get_expansion_service(cls): + return 'localhost:%s' % cls.expansion_port + + @classmethod + def set_kafka_streams_job_server_jar(cls, kafka_streams_job_server_jar): + cls.kafka_streams_job_server_jar = kafka_streams_job_server_jar + + def create_options(self): + options = super().create_options() + options.view_as(PortableOptions).environment_type = self.environment_type + options.view_as( + PortableOptions).environment_options = self.environment_options + + kafka_streams_options = options.view_as(KafkaStreamsRunnerOptions) + kafka_streams_options.bootstrap_servers = self.bootstrap_servers + # A fresh application id per pipeline. The id names the consumer group and the runner's own + # topics, so reusing one would have a test resume another test's offsets and read its data. + kafka_streams_options.application_id = 'beam-vr-%s' % uuid.uuid4() + return options + + # --------------------------------------------------------------------------- + # Features the runner does not support yet. Each skip points at the issue that + # would implement it, so this list doubles as the runner's capability gaps. + # --------------------------------------------------------------------------- + + def test_pardo_side_inputs(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_pardo_windowed_side_inputs(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_flattened_side_input(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_multimap_side_input(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_multimap_multiside_input(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_multimap_side_input_type_coercion(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_pardo_unfusable_side_inputs(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_pardo_state_only(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_pardo_timers(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_pardo_timers_clear(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_pardo_state_timers(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_pardo_state_timers_non_standard_coder(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_windowed_pardo_state_timers(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_pardo_dynamic_timer(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_custom_merging_window(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39630") + + def test_custom_window_type(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39630") + + def test_sdf_with_watermark_tracking(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39631") + + def test_sdf_with_sdf_initiated_checkpointing(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39631") + + def test_sdf_synthetic_source(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39631") + + def test_sdf_with_dofn_as_watermark_estimator(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39631") + + def test_callbacks_with_exception(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/18479") + + def test_register_finalizations(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/18479") + + def test_batch_pardo_fusion_break(self): + # CombineGlobally expands to a stage with side inputs. + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_batch_to_element_pardo(self): + # CombineGlobally expands to a stage with side inputs. + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_gbk_side_input(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_pack_combiners(self): + # The packed combiners are CombineGlobally, which needs side inputs. + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_pardo_side_input_dependencies(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_pardo_unfusable_side_inputs_with_separation(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39628") + + def test_pardo_state_with_custom_key_coder(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_pardo_et_timer_with_no_firing(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_pardo_et_timer_with_no_reset(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_pardo_et_timer_with_no_reset_and_no_clear(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39629") + + def test_windowing(self): + # Sessions, which are merging windows. + raise unittest.SkipTest("https://github.com/apache/beam/issues/39630") + + def test_windowed_combine_per_key(self): + # The fixed and sliding parts pass; the sessions part does not, sessions being merging windows. + raise unittest.SkipTest("https://github.com/apache/beam/issues/39630") + + def test_reshuffle_after_custom_window(self): + raise unittest.SkipTest("https://github.com/apache/beam/issues/39630") + + def test_metrics(self): + # The runner reports attempted values, which do reach a Python pipeline; this asserts on + # committed ones. + raise unittest.SkipTest("https://github.com/apache/beam/issues/39635") + + def test_read(self): + # The runner reads sources through the deprecated primitive Read, which carries a serialized + # Java source; a Python pipeline's source is a Python object, so reading one needs splittable + # DoFn support rather than anything specific to Read. + raise unittest.SkipTest("https://github.com/apache/beam/issues/39631") + + # Inherits all other tests from PortableRunnerTest. + + +if __name__ == '__main__': + logging.getLogger().setLevel(logging.INFO) + unittest.main() diff --git a/sdks/python/test-suites/portable/common.gradle b/sdks/python/test-suites/portable/common.gradle index 17bf9989f28b..c8c789ca96df 100644 --- a/sdks/python/test-suites/portable/common.gradle +++ b/sdks/python/test-suites/portable/common.gradle @@ -174,6 +174,41 @@ tasks.register("sparkValidatesRunner") { dependsOn 'sparkCompatibilityMatrixLOOPBACK' } +// Unlike the Flink and Spark suites, this one needs a Kafka broker: the runner executes on Kafka +// rather than on a cluster of its own. It is therefore not wired into any aggregate build. Run it +// against a broker, e.g. +// +// docker run -d -p 9092:9092 --name beam-kafka apache/kafka:4.0.0 +// ./gradlew :sdks:python:test-suites:portable:py312:kafkaStreamsValidatesRunner +// +// The broker defaults to localhost:9092; point it elsewhere with +// -PkafkaStreamsBootstrapServers=host:port. +// +// LOOPBACK only, unlike the other runners: it is the environment the suite has been run in, and a +// task for an environment nobody has tried would claim more than is known. +def createKafkaStreamsRunnerTestTask() { + def taskName = "kafkaStreamsCompatibilityMatrixLOOPBACK" + // Not resolvable until runtime, as for the Spark job server above, so the path is spelled out. + def jobServerJar = + "${rootDir}/runners/kafka-streams/job-server/build/libs/beam-runners-kafka-streams-job-server-${version}.jar" + def bootstrapServers = + project.findProperty('kafkaStreamsBootstrapServers') ?: 'localhost:9092' + def options = + "--kafka_streams_job_server_jar=${jobServerJar} --environment_type=LOOPBACK" + + " --bootstrap_servers=${bootstrapServers}" + def task = toxTask(taskName, 'kafka-streams-runner-test', options) + task.configure { + dependsOn ':runners:kafka-streams:job-server:shadowJar' + } + return task +} + +createKafkaStreamsRunnerTestTask() + +tasks.register("kafkaStreamsValidatesRunner") { + dependsOn 'kafkaStreamsCompatibilityMatrixLOOPBACK' +} + def createPrismRunnerTestTask(String workerType) { def taskName = "prismCompatibilityMatrix${workerType}" diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index 6dab85083a02..beab0399d50c 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -330,6 +330,11 @@ extras = test commands = bash {toxinidir}/scripts/pytest_validates_runner.sh {envname} {toxinidir}/apache_beam/runners/portability/spark_runner_test.py {posargs} +[testenv:kafka-streams-runner-test] +extras = test +commands = + bash {toxinidir}/scripts/pytest_validates_runner.sh {envname} {toxinidir}/apache_beam/runners/portability/kafka_streams_runner_test.py {posargs} + [testenv:prism-runner-test] extras = test commands = From 511a40e4f20739645943ab8eff69a1279591b813 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Fri, 14 Aug 2026 14:19:59 +0500 Subject: [PATCH 31/37] [GSoC 2026] Kafka Streams runner: separate the source's poll size from the bundle size, and expose the session timeout (#39748) --- .../streams/KafkaStreamsPipelineOptions.java | 21 +++++ .../streams/KafkaStreamsPipelineRunner.java | 14 ++- .../streams/translation/ReadTranslator.java | 2 +- .../KafkaStreamsPipelineRunnerConfigTest.java | 86 +++++++++++++++++++ 4 files changed, 121 insertions(+), 2 deletions(-) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunnerConfigTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java index 44abc8e5b34d..180e1ef028f6 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java @@ -49,6 +49,27 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { void setMaxBundleSize(int maxBundleSize); + @Description( + "How many elements an unbounded source may yield per poll. Separate from --maxBundleSize:" + + " a small bundle is how you get output promptly, while how much a source reads at a" + + " time is about throughput, and tying them together means a pipeline cannot have both.") + @Default.Integer(1000) + int getReadMaxElementsPerPoll(); + + void setReadMaxElementsPerPoll(int readMaxElementsPerPoll); + + @Description( + "How long the consumer group waits before deciding an instance has gone, in milliseconds." + + " This is the floor on how quickly work moves to another instance after one is lost," + + " since a departed instance is not noticed any sooner. Kafka's default of 45s is kept," + + " but a pipeline that values recovery over tolerance of a slow or briefly paused" + + " instance can lower it — a broker will not accept a value below its" + + " group.min.session.timeout.ms, which itself defaults to 6s.") + @Default.Integer(45_000) + int getSessionTimeoutMs(); + + void setSessionTimeoutMs(int sessionTimeoutMs); + @Description( "Intended cap on how long a bundle may stay open, in milliseconds. NOT APPLIED YET: closing a" + " bundle from a wall-clock punctuator made a pipeline with two chained GroupByKeys" 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 de08480a61d4..baa62f31aa4a 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 @@ -26,6 +26,7 @@ import org.apache.beam.runners.jobsubmission.PortablePipelineRunner; import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsPipelineTranslator; import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsTranslationContext; +import org.apache.kafka.clients.consumer.ConsumerConfig; import org.apache.kafka.streams.KafkaStreams; import org.apache.kafka.streams.StreamsConfig; import org.apache.kafka.streams.Topology; @@ -162,7 +163,8 @@ private static void checkRequiredOption(String name, @Nullable String value) { } } - private Properties streamsConfig(JobInfo jobInfo) { + // Visible for testing: the session timeout and the heartbeat derived from it. + Properties streamsConfig(JobInfo jobInfo) { Properties props = new Properties(); props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, pipelineOptions.getBootstrapServers()); props.put(StreamsConfig.APPLICATION_ID_CONFIG, pipelineOptions.getApplicationId()); @@ -177,6 +179,16 @@ private Properties streamsConfig(JobInfo jobInfo) { // is // what Kafka Streams does by default when no client id is set. props.put(StreamsConfig.CLIENT_ID_CONFIG, jobInfo.jobId() + "-" + UUID.randomUUID()); + // How quickly a lost instance is noticed, which is the floor on how quickly its work moves + // elsewhere. The heartbeat must be shorter than the timeout, or a healthy instance would be + // declared dead between beats; a third is the ratio Kafka's own defaults use. Deriving it + // rather than exposing it keeps the pair consistent whatever the timeout is set to. + int sessionTimeoutMs = pipelineOptions.getSessionTimeoutMs(); + props.put( + StreamsConfig.consumerPrefix(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG), sessionTimeoutMs); + props.put( + StreamsConfig.consumerPrefix(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG), + Math.max(1, sessionTimeoutMs / 3)); return props; } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java index 69006661ee65..b6620203b2d4 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java @@ -136,7 +136,7 @@ private void addUnbounde // splitting a source that has already been split. UnboundedSource readableSource = singleSplitOf(source, context); Coder checkpointCoder = readableSource.getCheckpointMarkCoder(); - int maxElementsPerPoll = context.getPipelineOptions().getMaxBundleSize(); + int maxElementsPerPoll = context.getPipelineOptions().getReadMaxElementsPerPoll(); int checkpointEveryNPolls = context.getPipelineOptions().getReadCheckpointNumBundles(); topology.addSource( diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunnerConfigTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunnerConfigTest.java new file mode 100644 index 000000000000..c98d50c75a07 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunnerConfigTest.java @@ -0,0 +1,86 @@ +/* + * 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; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.lessThan; + +import java.util.Properties; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.streams.StreamsConfig; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for the Kafka Streams configuration the runner builds from its pipeline options. */ +@RunWith(JUnit4.class) +public class KafkaStreamsPipelineRunnerConfigTest { + + private static final String SESSION_TIMEOUT = + StreamsConfig.consumerPrefix(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG); + private static final String HEARTBEAT = + StreamsConfig.consumerPrefix(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG); + + private static Properties configFor(Integer sessionTimeoutMs) { + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + options.setApplicationId("an-application"); + options.setBootstrapServers("localhost:9092"); + if (sessionTimeoutMs != null) { + options.setSessionTimeoutMs(sessionTimeoutMs); + } + JobInfo jobInfo = + JobInfo.create("a-job", "a-job", "", PipelineOptionsTranslation.toProto(options)); + return new KafkaStreamsPipelineRunner(options).streamsConfig(jobInfo); + } + + @Test + public void theSessionTimeoutDefaultsToKafkasOwn() { + // Changing this would change how long a lost instance goes unnoticed, so it is deliberate that + // the runner keeps Kafka's default rather than choosing its own. + assertThat(configFor(null).get(SESSION_TIMEOUT), is(45_000)); + } + + @Test + public void theSessionTimeoutIsWhateverThePipelineAskedFor() { + assertThat(configFor(6_000).get(SESSION_TIMEOUT), is(6_000)); + } + + @Test + public void theHeartbeatIsAThirdOfTheSessionTimeout() { + assertThat(configFor(6_000).get(HEARTBEAT), is(2_000)); + } + + @Test + public void theHeartbeatStaysShorterThanEvenATinySessionTimeout() { + // Kafka rejects a heartbeat that is not shorter than the session timeout, so deriving it has + // to hold for small values too — a fixed floor would not. One millisecond is excluded because + // no positive heartbeat is shorter than it, and a broker would refuse such a timeout anyway. + for (int sessionTimeoutMs : new int[] {2, 10, 100, 200, 1_000, 6_000, 45_000}) { + Properties config = configFor(sessionTimeoutMs); + assertThat( + "heartbeat must stay under the session timeout for " + sessionTimeoutMs + "ms", + (Integer) config.get(HEARTBEAT), + lessThan((Integer) config.get(SESSION_TIMEOUT))); + } + } +} From 65a2e400c736ce207a072918f5bd3899237e5ca3 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Sun, 16 Aug 2026 22:09:28 +0500 Subject: [PATCH 32/37] [GSoC 2026] Kafka Streams runner: bound a source poll in time, not only in elements (#39761) * [GSoC 2026] Kafka Streams runner: bound a source poll in time, not only in elements --readMaxPollTimeMs bounds the turn in time as well; whichever bound comes first ends it. --- .../streams/KafkaStreamsPipelineOptions.java | 17 ++++++ .../streams/translation/ReadTranslator.java | 2 + .../translation/UnboundedReadProcessor.java | 49 +++++++++++++++-- .../translation/UnboundedReadTest.java | 54 +++++++++++++++++++ .../en/documentation/runners/kafkastreams.md | 1 + 5 files changed, 118 insertions(+), 5 deletions(-) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java index 180e1ef028f6..99454b3a6588 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java @@ -58,6 +58,23 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { void setReadMaxElementsPerPoll(int readMaxElementsPerPoll); + @Description( + "How long one turn of reading an unbounded source may take, in milliseconds, before the" + + " source yields the Kafka Streams thread. A source is polled from a punctuator" + + " scheduled every 50ms, and the same thread runs the rest of the topology, so a turn" + + " that overruns that interval is already due again when it returns and fires straight" + + " away: the source then holds the thread and the stages below it are never scheduled," + + " which shows up as a pipeline that reads steadily and emits nothing at all rather" + + " than one that falls behind. Roughly, the source takes this fraction of a 50ms" + + " interval, so the default of 10ms leaves the thread four fifths of its time." + + " --readMaxElementsPerPoll bounds the same turn by count; whichever bound is reached" + + " first ends it, and a count alone cannot bound the time because how long an element" + + " takes depends on the pipeline below.") + @Default.Integer(10) + int getReadMaxPollTimeMs(); + + void setReadMaxPollTimeMs(int readMaxPollTimeMs); + @Description( "How long the consumer group waits before deciding an instance has gone, in milliseconds." + " This is the floor on how quickly work moves to another instance after one is lost," diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java index b6620203b2d4..9a727e82a70d 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java @@ -138,6 +138,7 @@ private void addUnbounde Coder checkpointCoder = readableSource.getCheckpointMarkCoder(); int maxElementsPerPoll = context.getPipelineOptions().getReadMaxElementsPerPoll(); int checkpointEveryNPolls = context.getPipelineOptions().getReadCheckpointNumBundles(); + int maxPollTimeMs = context.getPipelineOptions().getReadMaxPollTimeMs(); topology.addSource( sourceNodeName, @@ -157,6 +158,7 @@ private void addUnbounde transformId, maxElementsPerPoll, checkpointEveryNPolls, + maxPollTimeMs, context.getTerminationTracker()), sourceNodeName); topology.addStateStore( diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java index 96aa1bf0c69e..4f8c5f105566 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java @@ -79,6 +79,9 @@ class UnboundedReadProcessor /** How often the source is polled. */ private static final Duration POLL_INTERVAL = Duration.ofMillis(50); + /** Elements read between two checks of whether the turn is out of time. */ + private static final int ELEMENTS_BETWEEN_DEADLINE_CHECKS = 64; + private final UnboundedSource source; private final SerializablePipelineOptions options; // See ReadProcessor: a source produces decoded objects, but the downstream stage's harness input @@ -90,6 +93,7 @@ class UnboundedReadProcessor private final String transformId; private final int maxElementsPerPoll; private final int checkpointEveryNPolls; + private final int maxPollTimeMs; private @Nullable ProcessorContext> context; private @Nullable KeyValueStore checkpointStore; @@ -115,6 +119,7 @@ class UnboundedReadProcessor String transformId, int maxElementsPerPoll, int checkpointEveryNPolls, + int maxPollTimeMs, TerminationTracker terminationTracker) { this.terminationReporter = new TerminationReporter(terminationTracker, transformId); this.source = source; @@ -126,6 +131,7 @@ class UnboundedReadProcessor this.transformId = transformId; this.maxElementsPerPoll = maxElementsPerPoll; this.checkpointEveryNPolls = checkpointEveryNPolls; + this.maxPollTimeMs = maxPollTimeMs; } @Override @@ -157,6 +163,16 @@ public void process(Record record) { * Streams thread would never get back to committing or to the rest of the topology. So at most * {@link #checkpointEveryNPolls} batches are taken before yielding, which is also where the * checkpoint mark is stored, and the next punctuation carries on from there. + * + *

    That batch bound is a count, and a count cannot bound the time: how long an element takes is + * decided by the pipeline underneath it, which the source knows nothing about. A punctuator is + * expected to be quick, and this one runs on the thread that also serves the rest of the + * topology, so a turn that overruns its own {@link #POLL_INTERVAL} is due again as soon as it + * returns and runs once more instead of the tasks below it. Measured on a grouping pipeline, a + * turn of 200 elements took 3ms and held the thread 6% of the time, while a turn of 5000 took + * 57ms and held it 89%, and the pipeline read tens of millions of elements while emitting none. + * {@link #maxPollTimeMs} bounds the turn in time as well, and whichever bound is reached first + * ends it. */ private void poll() { if (exhausted) { @@ -164,8 +180,9 @@ private void poll() { } ProcessorContext> ctx = checkInitialized(context); UnboundedReader currentReader = ensureReader(); + long deadline = System.currentTimeMillis() + maxPollTimeMs; for (int batch = 0; batch < checkpointEveryNPolls; batch++) { - int emitted = readBatch(ctx, currentReader); + int emitted = readBatch(ctx, currentReader, deadline); Instant watermark = currentReader.getWatermark(); forwardWatermarkIfAdvanced(ctx, watermark); if (!watermark.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) { @@ -181,24 +198,46 @@ private void poll() { return; } if (emitted < maxElementsPerPoll) { - // Short batch: the source has nothing more for now, so store what was read and wait for - // the next punctuation rather than spinning on a reader that keeps returning false. + // Short batch: either the source has nothing more for now, or the turn ran out of time. + // Either way, store what was read and wait for the next punctuation rather than spinning + // on a reader that keeps returning false. if (emitted > 0) { storeCheckpoint(currentReader); } return; } + if (System.currentTimeMillis() >= deadline) { + // A full batch and the turn is out of time: yield with the position recorded, so the next + // punctuation carries on rather than this one running the thread out from under the rest + // of the topology. + storeCheckpoint(currentReader); + return; + } } // Yielded on the batch bound rather than on an empty source, so record the position reached. storeCheckpoint(currentReader); } - /** Forwards up to {@link #maxElementsPerPoll} elements, returning how many were available. */ + /** + * Forwards up to {@link #maxElementsPerPoll} elements, returning how many were available. + * + *

    Stops early if the turn's deadline passes, since one batch can be long enough on its own to + * overrun it. The clock is read every {@link #ELEMENTS_BETWEEN_DEADLINE_CHECKS} elements rather + * than every element, which bounds the overshoot to that many elements without putting a clock + * read in front of each one. + */ private int readBatch( - ProcessorContext> ctx, UnboundedReader currentReader) { + ProcessorContext> ctx, + UnboundedReader currentReader, + long deadline) { int emitted = 0; try { while (emitted < maxElementsPerPoll) { + if (emitted % ELEMENTS_BETWEEN_DEADLINE_CHECKS == 0 + && emitted > 0 + && System.currentTimeMillis() >= deadline) { + break; + } // start() positions the reader on its first element; advance() moves to the next. Either // returning false means nothing is available right now — not that the source is finished, // which is the difference from a bounded read. diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java index b668ee972d53..01e1a3146e3a 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadTest.java @@ -21,6 +21,7 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.lessThan; import java.io.IOException; import java.time.Duration; @@ -122,6 +123,59 @@ public void anUnboundedSourceIsPolledAndItsElementsReachTheHarness() { } } + /** A pipeline whose source may read {@code elementsPerPoll} elements and run for {@code ms}. */ + private static Pipeline pipelineWithPollBounds(int elementsPerPoll, int maxPollTimeMs) { + KafkaStreamsPipelineOptions options = + KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class); + options.setReadMaxElementsPerPoll(elementsPerPoll); + options.setReadMaxPollTimeMs(maxPollTimeMs); + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply("read", Read.from(CountingSource.unbounded())) + .apply("record", ParDo.of(new RecordFn())); + return pipeline; + } + + /** Drives one turn of the wall clock and reports how many elements the source produced. */ + private static int elementsInOneTurn(Pipeline pipeline) { + KafkaStreamsTranslationContext context = KafkaStreamsTestRunner.translate(pipeline); + try (TopologyTestDriver driver = + new TopologyTestDriver( + context.getTopology(), KafkaStreamsTestRunner.streamsConfig(pipeline))) { + driver.advanceWallClockTime(Duration.ofMillis(100)); + } + return RECEIVED.size(); + } + + /** + * The element bound cannot bound the time a turn takes, because how long an element takes is + * decided by the pipeline below the source. Left on the count alone, a source with data always + * available runs the full count every turn, overruns the punctuation interval, and is due again + * the moment it returns — so it keeps the thread and the rest of the topology never runs. + */ + @Test + public void aPollOutOfTimeYieldsBeforeReachingItsElementBound() { + // A turn that is out of time before it starts, so what stops it can only be the time bound. + int elements = elementsInOneTurn(pipelineWithPollBounds(1_000, 0)); + + assertThat( + "the source should have yielded, not run to its element bound", + elements, + is(lessThan(1_000))); + assertThat("the source should still have made progress", elements, is(greaterThan(0))); + } + + /** The time bound only cuts a turn short; with time to spare the element bound still applies. */ + @Test + public void aPollWithTimeToSpareReachesItsElementBound() { + int elements = elementsInOneTurn(pipelineWithPollBounds(100, 60_000)); + + assertThat( + "a turn with time to spare should read at least a full batch", + elements, + is(greaterThan(99))); + } + @Test public void aSourceThatReachesTheEndOfTimeStopsBeingPolled() { // CountingSource.unbounded() with a limit reports the terminal watermark once it has produced diff --git a/website/www/site/content/en/documentation/runners/kafkastreams.md b/website/www/site/content/en/documentation/runners/kafkastreams.md index bef1f852094b..058e85b6d461 100644 --- a/website/www/site/content/en/documentation/runners/kafkastreams.md +++ b/website/www/site/content/en/documentation/runners/kafkastreams.md @@ -119,6 +119,7 @@ Named as Java spells them below; from Python the same options are in snake case, | `topicReplicationFactor` | `1` | Replication factor for those topics. | | `maxBundleSize` | `1000` | Elements per bundle, and elements taken per poll of an unbounded source. | | `maxBundleTimeMs` | `1000` | Intended cap on how long a bundle may stay open. **Not applied yet** — see below. | +| `readMaxPollTimeMs` | `10` | How long one turn of reading an unbounded source may take before it yields the Kafka Streams thread. A source is polled every 50ms and shares its thread with the rest of the topology, so a turn that overruns that interval leaves the stages below it unscheduled; a bound on elements alone cannot bound the time. | | `readCheckpointNumBundles` | `10` | Polls of an unbounded source between stores of its checkpoint mark. Larger values replay more after a restart. | | `stateDir` | temp directory | Where Kafka Streams keeps local state. | From 104dc272e8d4e00ac51211dd3bbe5577bb67087e Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Sun, 16 Aug 2026 22:14:25 +0500 Subject: [PATCH 33/37] [GSoC 2026] Kafka Streams runner: ask for primitive reads in the Java wrapper (#39766) A Read expands into a splittable DoFn by default and the runner cannot translate one, so a pipeline that merely reads failed to translate unless it knew to convert the reads itself. The wrapper now sets use_deprecated_read and converts the pipeline before handing it on, so a pipeline does not have to know, and a pipeline that asks for splittable reads still gets primitive ones rather than something that cannot run. --- .../kafka/streams/KafkaStreamsRunner.java | 30 +++++++++- .../kafka/streams/KafkaStreamsRunnerTest.java | 59 +++++++++++++++++++ 2 files changed, 88 insertions(+), 1 deletion(-) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java index 1b530fd22ab1..924c3ac01fe8 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunner.java @@ -26,6 +26,8 @@ import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.SplittableParDo; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; @@ -60,7 +62,7 @@ protected KafkaStreamsRunner(KafkaStreamsPipelineOptions pipelineOptions) { @Override public PipelineResult run(Pipeline pipeline) { - assignPortableDefaults(pipelineOptions); + prepareForTranslation(pipeline, pipelineOptions); @Nullable KafkaStreamsJobServerDriver jobServerDriver = null; try { if (Strings.isNullOrEmpty(pipelineOptions.getJobEndpoint())) { @@ -89,6 +91,22 @@ public PipelineResult run(Pipeline pipeline) { } } + /** + * Settles the options the runner needs and rewrites the pipeline into what it can translate. + * + *

    The runner does not translate splittable DoFns, and a {@link org.apache.beam.sdk.io.Read} + * expands into one by default, so a pipeline that merely reads would otherwise fail to translate. + * Beam keeps the primitive read for exactly this case, behind an experiment that {@link + * #assignPortableDefaults} sets, so a pipeline does not have to ask for it and the proto that + * reaches the job server already holds primitive reads. + */ + @VisibleForTesting + static void prepareForTranslation( + Pipeline pipeline, KafkaStreamsPipelineOptions pipelineOptions) { + assignPortableDefaults(pipelineOptions); + SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReadsIfNecessary(pipeline); + } + private static void assignPortableDefaults(KafkaStreamsPipelineOptions pipelineOptions) { if (Strings.isNullOrEmpty(pipelineOptions.getDefaultEnvironmentType())) { pipelineOptions.setDefaultEnvironmentType(Environments.ENVIRONMENT_LOOPBACK); @@ -97,8 +115,18 @@ private static void assignPortableDefaults(KafkaStreamsPipelineOptions pipelineO @Nullable List existingExperiments = experimentalOptions.getExperiments(); List experiments = existingExperiments == null ? new ArrayList<>() : new ArrayList<>(existingExperiments); + boolean changed = false; if (!experiments.contains("beam_fn_api")) { experiments.add("beam_fn_api"); + changed = true; + } + // Splittable DoFns are not translated, so the Read that expands into one has to stay the + // primitive it used to be. This is the experiment Beam looks for when deciding that. + if (!experiments.contains("use_deprecated_read")) { + experiments.add("use_deprecated_read"); + changed = true; + } + if (changed) { experimentalOptions.setExperiments(experiments); } } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java index 02bc1a887c73..321b16f303c5 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java @@ -17,18 +17,28 @@ */ package org.apache.beam.runners.kafka.streams; +import static org.hamcrest.CoreMatchers.hasItem; import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.MatcherAssert.assertThat; import java.time.Duration; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; import org.apache.beam.model.pipeline.v1.RunnerApi; import org.apache.beam.runners.kafka.streams.translation.KStreamsPayload; import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.io.CountingSource; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.options.ExperimentalOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.Impulse; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.PTransformTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; import org.apache.kafka.streams.Topology; import org.apache.kafka.streams.TopologyTestDriver; import org.apache.kafka.streams.processor.api.Processor; @@ -50,6 +60,55 @@ */ public class KafkaStreamsRunnerTest { + /** The transform urns the pipeline would hand to the job server. */ + private static Set translatedUrns(Pipeline pipeline) { + return PipelineTranslation.toProto(pipeline).getComponents().getTransformsMap().values() + .stream() + .map(transform -> transform.getSpec().getUrn()) + .collect(Collectors.toSet()); + } + + /** + * A {@code Read} expands into a splittable DoFn by default, which this runner cannot translate. + * The runner asks for the primitive read instead, so that a pipeline does not have to know to. + */ + @Test + public void aReadReachesTheJobServerAsAPrimitiveReadRatherThanASplittableDoFn() { + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + options.setApplicationId("read-conversion-test"); + // Pipeline.create insists on a runner; it does not run here, the pipeline is only translated. + options.setRunner(KafkaStreamsRunner.class); + Pipeline pipeline = Pipeline.create(options); + pipeline.apply("read", Read.from(CountingSource.unbounded())); + + // Left alone, the read is a splittable DoFn expansion and no primitive read is present. + assertThat(translatedUrns(pipeline), not(hasItem(PTransformTranslation.READ_TRANSFORM_URN))); + + KafkaStreamsRunner.prepareForTranslation(pipeline, options); + + assertThat(translatedUrns(pipeline), hasItem(PTransformTranslation.READ_TRANSFORM_URN)); + } + + /** + * A pipeline may ask for splittable reads outright. The runner cannot translate them, so it asks + * for the primitive read anyway rather than letting a pipeline choose something that cannot run. + */ + @Test + public void aPipelineAskingForSplittableReadsStillGetsPrimitiveOnes() { + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + options.setApplicationId("sdf-read-override-test"); + options.setRunner(KafkaStreamsRunner.class); + options.as(ExperimentalOptions.class).setExperiments(new ArrayList<>(List.of("use_sdf_read"))); + Pipeline pipeline = Pipeline.create(options); + pipeline.apply("read", Read.from(CountingSource.unbounded())); + + KafkaStreamsRunner.prepareForTranslation(pipeline, options); + + assertThat(translatedUrns(pipeline), hasItem(PTransformTranslation.READ_TRANSFORM_URN)); + } + @Test public void impulseOnlyPipelineEmitsDataAndTerminalWatermark() { Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); From 49d459a243eb613e97e6dae5e935638d8e4f4677 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Sun, 16 Aug 2026 22:43:37 +0500 Subject: [PATCH 34/37] [GSoC 2026] Kafka Streams runner: put the runner behind an opt-in build flag (#39762) --- .../beam_KafkaStreamsRunner_FeatureBranch.yml | 4 +- ...am_PreCommit_Java_Kafka_Streams_Runner.yml | 1 + build.gradle.kts | 5 ++- .../python/test-suites/portable/common.gradle | 10 +++-- settings.gradle.kts | 14 +++++-- .../en/documentation/runners/kafkastreams.md | 40 +++++++++++++++++-- 6 files changed, 61 insertions(+), 13 deletions(-) diff --git a/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml b/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml index e290c9beee1a..8b558865c218 100644 --- a/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml +++ b/.github/workflows/beam_KafkaStreamsRunner_FeatureBranch.yml @@ -69,6 +69,6 @@ jobs: restore-keys: | ${{ runner.os }}-gradle-kafka-streams- - name: Build and test Kafka Streams runner - run: ./gradlew :runners:kafka-streams:build --no-daemon --stacktrace + run: ./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:build --no-daemon --stacktrace - name: Run ValidatesRunner suite - run: ./gradlew :runners:kafka-streams:validatesRunner --no-daemon --stacktrace + run: ./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:validatesRunner --no-daemon --stacktrace diff --git a/.github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml b/.github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml index 564b2bbc4bc8..005764b4d33c 100644 --- a/.github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml +++ b/.github/workflows/beam_PreCommit_Java_Kafka_Streams_Runner.yml @@ -89,6 +89,7 @@ jobs: uses: ./.github/actions/gradle-command-self-hosted-action with: gradle-command: :runners:kafka-streams:build + arguments: -Pwith-kafka-streams-runner max-workers: 4 - name: Archive JUnit Test Results uses: actions/upload-artifact@v7 diff --git a/build.gradle.kts b/build.gradle.kts index b3b9fdd7fdf0..3f5d39d99347 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -278,7 +278,10 @@ tasks.register("javaPreCommit") { dependsOn(":runners:java-fn-execution:build") dependsOn(":runners:java-job-service:build") dependsOn(":runners:jet:build") - dependsOn(":runners:kafka-streams:build") + // Only when the opt-in flag put it in the build; see settings.gradle.kts. + if (findProject(":runners:kafka-streams") != null) { + dependsOn(":runners:kafka-streams:build") + } dependsOn(":runners:local-java:build") dependsOn(":runners:portability:java:build") dependsOn(":runners:prism:java:build") diff --git a/sdks/python/test-suites/portable/common.gradle b/sdks/python/test-suites/portable/common.gradle index c8c789ca96df..6cc30fcadab6 100644 --- a/sdks/python/test-suites/portable/common.gradle +++ b/sdks/python/test-suites/portable/common.gradle @@ -203,10 +203,14 @@ def createKafkaStreamsRunnerTestTask() { return task } -createKafkaStreamsRunnerTestTask() +// The Kafka Streams runner is opt-in (-Pwith-kafka-streams-runner), so its job server is only a +// project when it was asked for; without it there is nothing for these tasks to run against. +if (project.findProject(':runners:kafka-streams:job-server') != null) { + createKafkaStreamsRunnerTestTask() -tasks.register("kafkaStreamsValidatesRunner") { - dependsOn 'kafkaStreamsCompatibilityMatrixLOOPBACK' + tasks.register("kafkaStreamsValidatesRunner") { + dependsOn 'kafkaStreamsCompatibilityMatrixLOOPBACK' + } } def createPrismRunnerTestTask(String workerType) { diff --git a/settings.gradle.kts b/settings.gradle.kts index cd1134d685c9..050d97dc600e 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -143,9 +143,17 @@ include(":runners:google-cloud-dataflow-java:examples-streaming") include(":runners:java-fn-execution") include(":runners:java-job-service") include(":runners:jet") -include(":runners:kafka-streams") -include(":runners:kafka-streams:proto") -include(":runners:kafka-streams:job-server") +// The Kafka Streams runner is opt-in, and is left out of the build unless it is asked for with +// -Pwith-kafka-streams-runner. It is being developed in the open so that others can build it and +// work on it, but it is not ready to be released: bundles are not yet closed after a bounded time +// (https://github.com/apache/beam/issues/39633), among other things. Keeping it out of the default +// build means it reaches nobody who did not ask for it, and the flag can be dropped when the runner +// is stable enough - or the runner can be dropped, without either affecting users. +if (startParameter.projectProperties.containsKey("with-kafka-streams-runner")) { + include(":runners:kafka-streams") + include(":runners:kafka-streams:proto") + include(":runners:kafka-streams:job-server") +} include(":runners:local-java") include(":runners:portability:java") include(":runners:prism") diff --git a/website/www/site/content/en/documentation/runners/kafkastreams.md b/website/www/site/content/en/documentation/runners/kafkastreams.md index 058e85b6d461..1fe5d4e86cad 100644 --- a/website/www/site/content/en/documentation/runners/kafkastreams.md +++ b/website/www/site/content/en/documentation/runners/kafkastreams.md @@ -31,7 +31,7 @@ groups, changelog topics, and transactions. That makes it worth considering if you already run Kafka and want Beam's programming model without introducing a second distributed system to operate. -## The runner is experimental +## The runner is experimental, and is not built by default **The Kafka Streams Runner is experimental.** It executes a meaningful subset of the Beam model correctly, and the parts it does support are covered by Beam's own `@ValidatesRunner` suite, but @@ -41,6 +41,36 @@ supported](#what-is-not-supported-yet) before choosing it for anything real. It is also aimed squarely at streaming. A pipeline over bounded data will run, but there are more efficient choices for batch work; this runner exists for pipelines that do not end. +It is **not part of a Beam release, and not part of the default build**. It is developed in the open +so that people can build it, use it and work on it, but it is not ready to be released: there are +known bugs, not only missing features — bundles are not yet closed after a bounded time +([#39633](https://github.com/apache/beam/issues/39633)), for one. Building it takes an opt-in flag: + +``` +./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:build +``` + +Without `-Pwith-kafka-streams-runner` the runner's projects are left out of the build entirely, so +it reaches nobody who has not asked for it. The intent is to give the runner somewhere to be +developed and maintained by whoever is interested in it. If it becomes stable enough the flag will +be dropped and the runner built like any other; if it does not, it can be removed again without +affecting anyone, since no release ever contained it. + +## Building it + +Every command in this page needs the opt-in flag. To build the runner and run its tests: + +``` +./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:build +./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:validatesRunner +``` + +To build the job server jar that the Python SDK submits to: + +``` +./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:job-server:shadowJar +``` + ## Running a pipeline The runner is portable: it executes user code over the Fn API, in an SDK harness, so a pipeline goes @@ -74,14 +104,14 @@ python my_pipeline.py \ The SDK harness runs in `LOOPBACK` mode by default, so a local run needs no Docker. Building the jar takes a while the first time; `--kafka_streams_job_server_jar` points at a prebuilt one, which -`./gradlew :runners:kafka-streams:job-server:shadowJar` produces. +`./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:job-server:shadowJar` produces. ### Against a job server you are already running Start one, which listens on `localhost:8099` by default: ``` -./gradlew :runners:kafka-streams:runJobServer +./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:runJobServer ``` Then point a pipeline at it instead of letting the runner start its own. From Java: @@ -117,10 +147,12 @@ Named as Java spells them below; from Python the same options are in snake case, | `applicationId` | *(required)* | Kafka Streams `application.id`. Must be unique per pipeline. | | `internalParallelism` | `1` | Partitions for the internal topics the runner creates, which is the parallelism the shuffled parts of a pipeline can reach. | | `topicReplicationFactor` | `1` | Replication factor for those topics. | -| `maxBundleSize` | `1000` | Elements per bundle, and elements taken per poll of an unbounded source. | +| `maxBundleSize` | `1000` | Elements per bundle. | | `maxBundleTimeMs` | `1000` | Intended cap on how long a bundle may stay open. **Not applied yet** — see below. | +| `readMaxElementsPerPoll` | `1000` | Elements an unbounded source may take per poll. Separate from `maxBundleSize`, so a pipeline can have small bundles without throttling its source. | | `readMaxPollTimeMs` | `10` | How long one turn of reading an unbounded source may take before it yields the Kafka Streams thread. A source is polled every 50ms and shares its thread with the rest of the topology, so a turn that overruns that interval leaves the stages below it unscheduled; a bound on elements alone cannot bound the time. | | `readCheckpointNumBundles` | `10` | Polls of an unbounded source between stores of its checkpoint mark. Larger values replay more after a restart. | +| `sessionTimeoutMs` | `45000` | How long the consumer group waits before deciding an instance has gone, which is the floor on how quickly its work moves elsewhere. A broker refuses a value below its own `group.min.session.timeout.ms`. | | `stateDir` | temp directory | Where Kafka Streams keeps local state. | ### Topics the runner creates From 10ff557d616bcc108bbf7eb659528ae2c20d12b2 Mon Sep 17 00:00:00 2001 From: M Junaid Shaukat <154750865+junaiddshaukat@users.noreply.github.com> Date: Sun, 16 Aug 2026 23:17:06 +0500 Subject: [PATCH 35/37] [GSoC 2026] Kafka Streams runner: an application for measuring instances coming and going (#39752) * [GSoC 2026] Kafka Streams runner: an application for measuring instances coming and going * [GSoC 2026] Kafka Streams runner: count groups in the pipeline rather than beside it SpotBugs is turned off for this module. It runs the pipeline in process, so the SDK harness and its dependencies are on the classpath and SpotBugs reports on those instead of on the four classes here. The it/ modules do the same for the same reason. --- .../kafka-streams/measurement/build.gradle | 78 ++++++ .../measurement/docker-compose.yml | 27 ++ .../measurement/RescalingMeasurement.java | 261 ++++++++++++++++++ .../streams/measurement/package-info.java | 26 ++ settings.gradle.kts | 1 + 5 files changed, 393 insertions(+) create mode 100644 runners/kafka-streams/measurement/build.gradle create mode 100644 runners/kafka-streams/measurement/docker-compose.yml create mode 100644 runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/RescalingMeasurement.java create mode 100644 runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/package-info.java diff --git a/runners/kafka-streams/measurement/build.gradle b/runners/kafka-streams/measurement/build.gradle new file mode 100644 index 000000000000..5e7992a0916e --- /dev/null +++ b/runners/kafka-streams/measurement/build.gradle @@ -0,0 +1,78 @@ +/* + * 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. + */ + +/** + * An application for measuring the Kafka Streams runner's behaviour when instances come and go. + * + * Not part of the build's verification: it is something a person runs against a Kafka, several + * copies at once, and watches. + */ + +apply plugin: 'org.apache.beam.module' +apply plugin: 'application' +mainClassName = "org.apache.beam.runners.kafka.streams.measurement.RescalingMeasurement" + +applyJavaNature( + automaticModuleName: 'org.apache.beam.runners.kafka.streams.measurement', + publish: false, + exportJavadoc: false, + // This module runs the pipeline in its own process, so the SDK harness and its dependencies are + // on the classpath, and SpotBugs reports on those rather than on the four classes here — some + // eleven thousand warnings, none of them in this source tree. The same is done in the it/ + // modules, which are on the classpath of what they exercise for the same reason. Checkstyle, + // ErrorProne, spotless and the nullness checker all still run. + enableSpotbugs: false, +) + +description = "Apache Beam :: Runners :: Kafka Streams :: Measurement" + +def kafkaStreamsRunnerProject = ":runners:kafka-streams" + +evaluationDependsOn(kafkaStreamsRunnerProject) + +// Same pin as the runner and the job server: applyJavaNature forces the versions in library.java, +// which includes an older kafka-clients than the runner is compiled against. +def kafka_version = project(kafkaStreamsRunnerProject).kafka_version + +configurations.configureEach { + resolutionStrategy.eachDependency { details -> + if (details.requested.group == "org.apache.kafka") { + details.useVersion(kafka_version) + details.because("Kafka Streams runner is developed against Kafka ${kafka_version}.") + } + } +} + +dependencies { + implementation project(kafkaStreamsRunnerProject) + implementation project(path: ":sdks:java:core", configuration: "shadow") + implementation project(path: ":model:pipeline", configuration: "shadow") + implementation project(":runners:java-fn-execution") + // On the compile classpath to resolve PortablePipelineRunner, which KafkaStreamsPipelineRunner + // implements; no class of it is named here, so the dependency analysis does not see it used. + implementation project(":runners:java-job-service") + permitUnusedDeclared project(":runners:java-job-service") + implementation project(":runners:core-java") + permitUnusedDeclared project(":runners:core-java") + // The pipeline's own code runs in this process, so the Java SDK harness has to be present. + runtimeOnly project(":sdks:java:harness") + implementation library.java.joda_time + // Without a binding the application starts and says nothing, which is unhelpful for something + // whose whole purpose is to be watched while it runs. + runtimeOnly library.java.slf4j_simple +} diff --git a/runners/kafka-streams/measurement/docker-compose.yml b/runners/kafka-streams/measurement/docker-compose.yml new file mode 100644 index 000000000000..836f121ca331 --- /dev/null +++ b/runners/kafka-streams/measurement/docker-compose.yml @@ -0,0 +1,27 @@ +# One Kafka for the measurement application. One broker is enough: what gets run several times is +# the runner instance, not the broker. +# +# docker compose -f runners/kafka-streams/measurement/docker-compose.yml up -d +# +# group.min.session.timeout.ms is lowered because a broker refuses a session timeout below it, and +# how quickly the group notices a departed instance is the floor on how quickly its work moves. The +# default of 6s would put a floor under every measurement of recovery. +services: + kafka: + image: apache/kafka:4.0.0 + container_name: ks-measurement-kafka + ports: + - "9092:9092" + environment: + KAFKA_NODE_ID: 1 + KAFKA_PROCESS_ROLES: broker,controller + KAFKA_LISTENERS: PLAINTEXT://:9092,CONTROLLER://:9093 + KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://localhost:9092 + KAFKA_CONTROLLER_QUORUM_VOTERS: 1@localhost:9093 + KAFKA_CONTROLLER_LISTENER_NAMES: CONTROLLER + KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: CONTROLLER:PLAINTEXT,PLAINTEXT:PLAINTEXT + KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1 + KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR: 1 + KAFKA_TRANSACTION_STATE_LOG_MIN_ISR: 1 + KAFKA_GROUP_MIN_SESSION_TIMEOUT_MS: 1000 + KAFKA_GROUP_INITIAL_REBALANCE_DELAY_MS: 0 diff --git a/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/RescalingMeasurement.java b/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/RescalingMeasurement.java new file mode 100644 index 000000000000..3ba9b22ca0ee --- /dev/null +++ b/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/RescalingMeasurement.java @@ -0,0 +1,261 @@ +/* + * 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.measurement; + +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineRunner; +import org.apache.beam.runners.kafka.streams.KafkaStreamsRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.io.GenerateSequence; +import org.apache.beam.sdk.options.Default; +import org.apache.beam.sdk.options.Description; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.FixedWindows; +import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.beam.sdk.util.construction.SplittableParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.TypeDescriptors; +import org.joda.time.Duration; + +/** + * One instance of a streaming pipeline, run as an ordinary application, for measuring what happens + * when instances are added and removed. + * + *

    Run several of these against one Kafka. They share an application id, so Kafka's consumer + * group divides the work between them, and stopping one hands its share to the others. + * + *

    This is an application rather than a test on purpose. The numbers only mean something if the + * pipeline is doing a realistic amount of work — a grouping over thousands of keys, fed fast enough + * that every partition has something to do. A pipeline that trickles produces idle partitions, and + * an idle partition holds a watermark back for reasons that have nothing to do with rescaling. + * + *

    The source produces a fixed number of elements per second over a fixed set of keys, so what a + * complete window looks like is known before the run starts: every window should report the same + * number of groups. That is what makes a shortfall legible as a shortfall, rather than as one of + * the many rates a pipeline could happen to be running at. + * + * 

    + *   docker compose -f runners/kafka-streams/measurement/docker-compose.yml up -d
+ *   ./gradlew :runners:kafka-streams:measurement:installDist
+ *
    + * + *

    Then start two instances, sharing an application id and differing in everything local to the + * instance. Each needs its own {@code --stateDir}: two instances sharing one directory fail with a + * {@code LockException}, because Kafka Streams locks the state it keeps on disk. + * + * 

    + *   BIN=runners/kafka-streams/measurement/build/install/measurement/bin/measurement
+ *   $BIN --applicationId=demo --instanceName=one --stateDir=/tmp/ks-one &
+ *   $BIN --applicationId=demo --instanceName=two --stateDir=/tmp/ks-two &
+ *
    + * + *

    The pipeline logs one line per key per window. Nothing is counted beside the pipeline: the + * groups in a window are its own output, so the tally does not depend on how many instances are + * running or on which of them happens to be doing the work. + * + * 

    + *   <millis> <instance> window_end=<millis> key=<key> count=<n> skew_ms=<n>
+ *
    + * + *

    Because the rate and the key space are both fixed, a complete window has one line per key and + * the same count on each, so counting the lines for a window says whether the window was complete. + * + *

    {@code skew_ms} is the gap between the window's event time and the wall clock when the group + * came out. It is what falling behind should look like: a pipeline that cannot keep up ought to + * report its groups later and later while still reporting all of them, so a climbing skew with + * complete windows is congestion, and missing groups are something else. + * + *

    To watch a handover, kill one instance and watch the other's lines. The delay before the + * survivor reports the killed instance's share again is dominated by {@code --sessionTimeoutMs}, + * which is how long the consumer group waits before deciding the instance is gone. + */ +public final class RescalingMeasurement { + + private RescalingMeasurement() {} + + /** Whether the command line mentioned an option, so that a default is not applied over it. */ + private static boolean given(String[] args, String name) { + for (String arg : args) { + if (arg.equals("--" + name) || arg.startsWith("--" + name + "=")) { + return true; + } + } + return false; + } + + /** + * Applies the defaults this measurement needs, where they differ from the runner's own. + * + *

    The runner's defaults are meant for a pipeline, not for this. Left alone, the source is read + * in large enough turns that the read starves the rest of the topology and no groups come out at + * all. The parallelism is raised for a related reason: a measurement of work moving between + * instances needs more than the single partition the runner defaults to, since with one partition + * there is nothing to divide. + */ + private static void applyMeasurementDefaults(String[] args, MeasurementOptions options) { + if (!given(args, "readMaxElementsPerPoll")) { + options.setReadMaxElementsPerPoll(200); + } + if (!given(args, "internalParallelism")) { + options.setInternalParallelism(3); + } + } + + /** Options of the measurement itself, on top of the runner's own. */ + public interface MeasurementOptions extends KafkaStreamsPipelineOptions { + + @Description("Name for this instance in the output, so several can be told apart.") + @Default.String("instance") + String getInstanceName(); + + void setInstanceName(String instanceName); + + @Description( + "How many distinct keys the grouping runs over. Thousands, so that every partition of the" + + " shuffle has work and no partition sits idle holding a watermark back. With a" + + " window long enough to contain them all, this is also how many groups a complete" + + " window has.") + @Default.Integer(2_000) + int getNumKeys(); + + void setNumKeys(int numKeys); + + @Description( + "How many elements the source produces per second. Fixed rather than as-fast-as-possible so" + + " that a window's contents are known in advance and a shortfall is visible.") + @Default.Integer(20_000) + int getElementsPerSecond(); + + void setElementsPerSecond(int elementsPerSecond); + + @Description("Window size in milliseconds; how often the groups are counted and reported.") + @Default.Integer(1_000) + int getWindowMs(); + + void setWindowMs(int windowMs); + } + + /** + * Logs each group the pipeline produces, with how far behind the wall clock its window was. + * + *

    One line per key per window. With a fixed rate over a fixed key space every window holds the + * same groups, so counting the lines for a window says whether the window was complete, and no + * counter has to be kept anywhere for that to be true — the count is the pipeline's own output + * rather than a tally maintained beside it, which is what makes it independent of how many + * instances are running. + * + *

    The skew is the point of the timestamp. A pipeline that cannot keep up should report its + * groups later and later rather than stop reporting them, so a skew that climbs while the groups + * stay complete is the pipeline falling behind, and groups going missing is something else. + */ + private static class ReportGroupFn extends DoFn, Void> { + private final String instanceName; + + ReportGroupFn(String instanceName) { + this.instanceName = instanceName; + } + + @ProcessElement + public void processElement(@Element KV group, BoundedWindow window) { + long windowEnd = window.maxTimestamp().getMillis(); + long now = System.currentTimeMillis(); + System.out.printf( + "%d %s window_end=%d key=%s count=%d skew_ms=%d%n", + now, instanceName, windowEnd, group.getKey(), group.getValue(), now - windowEnd); + } + } + + public static void main(String[] args) throws Exception { + PipelineOptionsFactory.register(MeasurementOptions.class); + // Deliberately not withValidation(): that enforces the options a pipeline needs when it is + // submitted to a job server, and --jobEndpoint above all, which means nothing here because this + // application runs the pipeline itself. + MeasurementOptions options = PipelineOptionsFactory.fromArgs(args).as(MeasurementOptions.class); + if (options.getApplicationId() == null || options.getApplicationId().isEmpty()) { + throw new IllegalArgumentException( + "--applicationId is required, and every instance of one measurement must share it: it is" + + " what puts them in the same consumer group and so divides the work between them."); + } + applyMeasurementDefaults(args, options); + // Pipeline.create needs a runner class even though this application never calls pipeline.run() + // — it builds the pipeline proto and hands it to the runner below itself. + options.setRunner(KafkaStreamsRunner.class); + // The user code runs in this same process, so no container or separate worker is needed. + options + .as(PortablePipelineOptions.class) + .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); + + // A window holds every key as long as it is long enough for the rate to reach them all; below + // that the source has not got round to each key once and the window is short by construction. + long elementsPerWindow = (long) options.getElementsPerSecond() * options.getWindowMs() / 1_000L; + long expectedGroups = Math.min(options.getNumKeys(), elementsPerWindow); + + int numKeys = options.getNumKeys(); + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "read", + GenerateSequence.from(0) + .withRate(options.getElementsPerSecond(), Duration.standardSeconds(1))) + .apply( + "key", + // numKeys is read here rather than inside the lambda: reaching for it through options + // would capture the PipelineOptions in the transform, which cannot be serialized. + MapElements.into(TypeDescriptors.strings()).via((Long n) -> "key-" + (n % numKeys))) + .apply("window", Window.into(FixedWindows.of(Duration.millis(options.getWindowMs())))) + .apply("countPerKey", Count.perElement()) + .apply("report", ParDo.of(new ReportGroupFn(options.getInstanceName()))); + + SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(pipeline); + RunnerApi.Pipeline proto = PipelineTranslation.toProto(pipeline); + JobInfo jobInfo = + JobInfo.create( + options.getApplicationId(), + options.getJobName(), + "", + PipelineOptionsTranslation.toProto(options)); + + System.out.printf( + "starting %s: application=%s keys=%d rate=%d/s parallelism=%d window=%dms" + + " session_timeout=%dms read_per_poll=%d bundle=%d expected_groups_per_window=%d%n", + options.getInstanceName(), + options.getApplicationId(), + options.getNumKeys(), + options.getElementsPerSecond(), + options.getInternalParallelism(), + options.getWindowMs(), + options.getSessionTimeoutMs(), + options.getReadMaxElementsPerPoll(), + options.getMaxBundleSize(), + expectedGroups); + + // Blocks until the instance is stopped; a streaming pipeline has no end of its own. + new KafkaStreamsPipelineRunner(options).run(proto, jobInfo); + } +} diff --git a/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/package-info.java b/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/package-info.java new file mode 100644 index 000000000000..41579676fb8d --- /dev/null +++ b/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/package-info.java @@ -0,0 +1,26 @@ +/* + * 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. + */ + +/** + * An application for measuring the Kafka Streams runner's behaviour when instances come and go. + * + *

    Not part of the build's verification: it is something a person runs against a Kafka, several + * copies at once, and watches. See {@link + * org.apache.beam.runners.kafka.streams.measurement.RescalingMeasurement} for how to run it. + */ +package org.apache.beam.runners.kafka.streams.measurement; diff --git a/settings.gradle.kts b/settings.gradle.kts index 050d97dc600e..de1e5ea6533d 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -153,6 +153,7 @@ if (startParameter.projectProperties.containsKey("with-kafka-streams-runner")) { include(":runners:kafka-streams") include(":runners:kafka-streams:proto") include(":runners:kafka-streams:job-server") + include(":runners:kafka-streams:measurement") } include(":runners:local-java") include(":runners:portability:java") From 75392abc047f5d63d6de53567a58ac6d99b65b22 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Mon, 17 Aug 2026 04:59:49 +0500 Subject: [PATCH 36/37] [GSoC 2026] Kafka Streams runner: license header and Python formatting for master CI The measurement docker-compose file had no Apache license header, which RAT rejects, and one line in the Python wrapper was not as yapf formats it. Neither ran on the feature branch, whose CI only built the runner. --- .../measurement/docker-compose.yml | 17 +++++++++++++++++ .../runners/portability/kafka_streams_runner.py | 3 +-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/runners/kafka-streams/measurement/docker-compose.yml b/runners/kafka-streams/measurement/docker-compose.yml index 836f121ca331..aa9102b97424 100644 --- a/runners/kafka-streams/measurement/docker-compose.yml +++ b/runners/kafka-streams/measurement/docker-compose.yml @@ -1,3 +1,20 @@ +# +# 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. +# + # One Kafka for the measurement application. One broker is enough: what gets run several times is # the runner instance, not the broker. # diff --git a/sdks/python/apache_beam/runners/portability/kafka_streams_runner.py b/sdks/python/apache_beam/runners/portability/kafka_streams_runner.py index a2d043ca174e..890ec41164a6 100644 --- a/sdks/python/apache_beam/runners/portability/kafka_streams_runner.py +++ b/sdks/python/apache_beam/runners/portability/kafka_streams_runner.py @@ -89,8 +89,7 @@ def path_to_jar(self): 'using `./gradlew runners:kafka-streams:job-server:shadowJar`.' % self._jar) return self._jar - return self.path_to_beam_jar( - ':runners:kafka-streams:job-server:shadowJar') + return self.path_to_beam_jar(':runners:kafka-streams:job-server:shadowJar') def java_arguments( self, job_port, artifact_port, expansion_port, artifacts_dir): From c5d800609cce353bc7423d165b87bcc6894ddeeb Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Mon, 17 Aug 2026 13:11:44 +0500 Subject: [PATCH 37/37] [GSoC 2026] Kafka Streams runner: shorten the explanation comments Several class comments explained more than they needed to, which makes them less likely to be read rather than more. Shortened the longest, keeping the reasoning and dropping the retelling. Also corrects four that had gone stale: two translators claiming topics are not created automatically, the payload claiming its serde does not exist yet, and the read translator pointing at the test runner for a conversion the runner now does itself. --- .../measurement/RescalingMeasurement.java | 51 ++++------ .../streams/KafkaStreamsPipelineRunner.java | 35 +++---- .../translation/ExecutableStageProcessor.java | 93 +++++++------------ .../streams/translation/FlattenProcessor.java | 28 +++--- .../translation/GroupByKeyTranslator.java | 31 ++----- .../streams/translation/ImpulseProcessor.java | 29 ++---- .../translation/ImpulseTranslator.java | 32 ++----- .../streams/translation/KStreamsPayload.java | 25 ++--- .../KafkaStreamsTimerInternals.java | 31 +++---- .../KafkaStreamsTranslationContext.java | 61 +++--------- .../streams/translation/ReadProcessor.java | 45 ++++----- .../streams/translation/ReadTranslator.java | 37 +++----- .../translation/TerminationReporter.java | 28 +++--- .../translation/TerminationTracker.java | 43 +++------ .../translation/UnboundedReadProcessor.java | 63 +++++-------- .../translation/WatermarkAggregator.java | 31 +++---- .../streams/translation/WatermarkManager.java | 47 +++------- 17 files changed, 239 insertions(+), 471 deletions(-) diff --git a/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/RescalingMeasurement.java b/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/RescalingMeasurement.java index 3ba9b22ca0ee..34ef1654345d 100644 --- a/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/RescalingMeasurement.java +++ b/runners/kafka-streams/measurement/src/main/java/org/apache/beam/runners/kafka/streams/measurement/RescalingMeasurement.java @@ -45,56 +45,41 @@ /** * One instance of a streaming pipeline, run as an ordinary application, for measuring what happens - * when instances are added and removed. + * when instances come and go. * - *

    Run several of these against one Kafka. They share an application id, so Kafka's consumer - * group divides the work between them, and stopping one hands its share to the others. + *

    Run several against one Kafka. They share an application id, so the consumer group divides the + * work between them and stopping one hands its share to the others. It is an application rather + * than a test because the numbers only mean something under a realistic load: a grouping over + * thousands of keys, fed fast enough that no partition sits idle holding a watermark back. * - *

    This is an application rather than a test on purpose. The numbers only mean something if the - * pipeline is doing a realistic amount of work — a grouping over thousands of keys, fed fast enough - * that every partition has something to do. A pipeline that trickles produces idle partitions, and - * an idle partition holds a watermark back for reasons that have nothing to do with rescaling. - * - *

    The source produces a fixed number of elements per second over a fixed set of keys, so what a - * complete window looks like is known before the run starts: every window should report the same - * number of groups. That is what makes a shortfall legible as a shortfall, rather than as one of - * the many rates a pipeline could happen to be running at. + *

    The source runs at a fixed rate over a fixed key space, so a complete window is known before + * the run starts — one line per key, the same count on each — which is what makes a shortfall + * legible as one. * * 

      *   docker compose -f runners/kafka-streams/measurement/docker-compose.yml up -d
- *   ./gradlew :runners:kafka-streams:measurement:installDist
- *
    + * ./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:measurement:installDist * - *

    Then start two instances, sharing an application id and differing in everything local to the - * instance. Each needs its own {@code --stateDir}: two instances sharing one directory fail with a - * {@code LockException}, because Kafka Streams locks the state it keeps on disk. - * - * 

      *   BIN=runners/kafka-streams/measurement/build/install/measurement/bin/measurement
  *   $BIN --applicationId=demo --instanceName=one --stateDir=/tmp/ks-one &
  *   $BIN --applicationId=demo --instanceName=two --stateDir=/tmp/ks-two &
  *
    * - *

    The pipeline logs one line per key per window. Nothing is counted beside the pipeline: the - * groups in a window are its own output, so the tally does not depend on how many instances are - * running or on which of them happens to be doing the work. + *

    Each instance needs its own {@code --stateDir}; sharing one fails with a {@code + * LockException}. Output is one line per key per window, counted by the pipeline itself rather than + * beside it, so the tally does not depend on how many instances are running: * * 

      *   <millis> <instance> window_end=<millis> key=<key> count=<n> skew_ms=<n>
  *
    * - *

    Because the rate and the key space are both fixed, a complete window has one line per key and - * the same count on each, so counting the lines for a window says whether the window was complete. - * - *

    {@code skew_ms} is the gap between the window's event time and the wall clock when the group - * came out. It is what falling behind should look like: a pipeline that cannot keep up ought to - * report its groups later and later while still reporting all of them, so a climbing skew with - * complete windows is congestion, and missing groups are something else. - * - *

    To watch a handover, kill one instance and watch the other's lines. The delay before the - * survivor reports the killed instance's share again is dominated by {@code --sessionTimeoutMs}, - * which is how long the consumer group waits before deciding the instance is gone. + *

    {@code skew_ms} is the gap between the window's event time and the wall clock when it came + * out. A pipeline that cannot keep up should report its groups later and later while still + * reporting all of them, so climbing skew with complete windows is congestion and missing groups + * are something else. To watch a handover, kill one instance and watch the other; the delay before + * it reports the dead instance's share is dominated by {@code --sessionTimeoutMs}. */ + public final class RescalingMeasurement { private RescalingMeasurement() {} 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 baa62f31aa4a..486cc05d8b54 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 @@ -48,17 +48,13 @@ public KafkaStreamsPipelineRunner(KafkaStreamsPipelineOptions pipelineOptions) { @Override public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) { - // Surface a clear error if an option this runner needs is missing, instead of letting - // Properties.put fail with a raw NullPointerException further down. Only the options that are - // meaningful here are checked, rather than validating the whole interface: this runs on the job - // server, executing a pipeline that has already been submitted, so the client-side options - // PortablePipelineOptions marks required — jobEndpoint above all — do not apply. Flink's - // equivalent PortablePipelineRunner does not validate here either. + // Only the options meaningful here are checked, not the whole interface: this runs on the job + // server, so the client-side options PortablePipelineOptions marks required — jobEndpoint above + // all — do not apply. Flink's PortablePipelineRunner does not validate here either. checkRequiredOption("applicationId", pipelineOptions.getApplicationId()); checkRequiredOption("bootstrapServers", pipelineOptions.getBootstrapServers()); - // A topic cannot have fewer than one partition, and the value is also the number of watermark - // reports a shuffle's consumer waits for, so a non-positive value would leave it waiting - // forever rather than failing. + // Also the number of watermark reports a shuffle's consumer waits for, so a non-positive value + // would leave it waiting forever rather than failing. if (pipelineOptions.getInternalParallelism() < 1) { throw new IllegalArgumentException( "--internalParallelism must be at least 1, but was " @@ -81,31 +77,26 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) topology.describe()); KafkaStreams kafkaStreams = new KafkaStreams(topology, streamsConfig(jobInfo)); - // Kafka Streams reports a failed task by moving the client to ERROR and keeping the exception - // to itself, which left a failed job with nothing to say beyond "unknown error". Hold on to the - // first failure so this method can rethrow it: the job service turns what run() throws into the - // job's error message. + // Kafka Streams moves the client to ERROR and keeps the exception to itself, which left failed + // jobs saying only "unknown error". Keep the first failure so run() can rethrow it. AtomicReference<@Nullable Throwable> failure = new AtomicReference<>(); kafkaStreams.setUncaughtExceptionHandler( throwable -> { failure.compareAndSet(null, throwable); LOG.error("Pipeline {} failed", jobInfo.jobId(), throwable); - // The pipeline is a job with an owner waiting on it, not a service to keep alive, so a - // failure stops the client rather than replacing the thread and carrying on. + // A job with an owner waiting on it, not a service: a failure stops the client. return StreamsUncaughtExceptionHandler.StreamThreadExceptionResponse.SHUTDOWN_CLIENT; }); - // Build the result before starting: it registers a state listener, and Kafka Streams only - // accepts one while the application is still in the CREATED state. + // Before start(): Kafka Streams only accepts a state listener while still in CREATED. KafkaStreamsPortablePipelineResult result = new KafkaStreamsPortablePipelineResult( kafkaStreams, context.getMetricsContainerStepMap(), - // Only once every task is initialized are the processors that have registered the whole - // set, and only then can "all of them are finished" mean the pipeline is finished. + // Only once every task is initialized is the registered set complete, so that "all + // finished" can mean the pipeline is finished. context.getTerminationTracker()::started); - // A bounded pipeline finishes; Kafka Streams has no notion of that, so the runner stops the - // client itself once every processor has reached the terminal watermark. Registered before - // start(), so a pipeline that drains quickly cannot finish before anything is listening. + // Kafka Streams has no notion of a finished pipeline, so the runner stops the client once every + // processor reaches the terminal watermark. Registered before start() so a fast drain is seen. context .getTerminationTracker() .onAllTerminated( 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 fcb8b2ff27e6..5cb0672ff9fc 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 @@ -48,39 +48,32 @@ import org.slf4j.LoggerFactory; /** - * Kafka Streams {@link Processor} that executes a fused {@link ExecutableStage} (stateless user - * code such as ParDo) in the Beam SDK harness over the Fn API. + * Kafka Streams {@link Processor} that executes a fused {@link ExecutableStage} — stateless user + * code such as ParDo — in the Beam SDK harness over the Fn API. * - *

    For each {@link KStreamsPayload#isData() data} payload it unwraps the {@link WindowedValue} - * and feeds it to the harness through the stage's main input {@link FnDataReceiver}. Harness - * outputs are collected on the harness threads into {@link #pendingOutputs} and then flushed - * downstream on the Kafka Streams processing thread when the bundle closes — Kafka Streams' {@link - * ProcessorContext#forward} must only be called from the processing thread, so outputs are never - * forwarded directly from a harness callback. + *

    Each {@link KStreamsPayload#isData() data} payload is unwrapped and fed to the harness through + * the stage's main input {@link FnDataReceiver}. Harness outputs are collected on the harness + * threads into {@link #pendingOutputs} and flushed downstream when the bundle closes, because + * {@link ProcessorContext#forward} may only be called from the processing thread. * - *

    A {@link KStreamsPayload#isWatermark() watermark} payload is a report from one partition of - * one upstream transform and marks a bundle boundary: the open bundle (if any) is closed (flushing - * outputs), the report is fed to the {@link WatermarkAggregator}, and the stage's output watermark - * is forwarded downstream — stamped with this stage's own transform id — only when the aggregate - * across the upstream transform's partitions actually advances. Until every partition has reported, - * the watermark is held and nothing is forwarded — but data is still processed in the meantime. + *

    A {@link KStreamsPayload#isWatermark() watermark} payload marks a bundle boundary: the open + * bundle is closed and flushed, the report goes to the {@link WatermarkAggregator}, and the stage's + * output watermark is forwarded — stamped with this stage's transform id — only once the aggregate + * across the upstream partitions advances. Until every partition has reported the watermark is + * held, though data is still processed meanwhile. * - *

    A bundle is also bounded in size, by {@code --maxBundleSize}, and closed once that many - * elements have been fed to it. Without the bound a bundle stays open until the next watermark, - * which on a stream that produces steadily lets it grow without limit. The bound is checked as - * elements arrive. A time bound ({@code --maxBundleTimeMs}) is not applied yet — see the option's - * own documentation. + *

    A bundle is also bounded by {@code --maxBundleSize}, checked as elements arrive; without it a + * bundle would stay open until the next watermark and grow without limit on a steady stream. The + * time bound {@code --maxBundleTimeMs} is not applied yet, see that option's documentation. * - *

    Closing a bundle asks Kafka Streams to commit, so the elements a bundle consumed and the - * records it produced are committed together and a restart replays either all of the bundle or none - * of it. Note that this aligns commits to bundle boundaries but does not stop Kafka - * Streams from committing on its own interval part-way through a bundle; closing the bundle first - * from a pre-commit hook would be needed to rule that out entirely. + *

    Closing a bundle asks Kafka Streams to commit, so the elements consumed and the records + * produced commit together and a restart replays all of a bundle or none. This aligns commits to + * bundle boundaries but does not stop Kafka Streams committing on its own interval mid-bundle; + * ruling that out needs a pre-commit hook. * - *

    This is the Kafka Streams analogue of Flink's {@code ExecutableStageDoFnOperator} and Spark's - * {@code SparkExecutableStageFunction}. State, timers, and side inputs are out of scope for this - * first version: the stage is executed with {@link StateRequestHandler#unsupported()} and no timer - * receivers. + *

    The analogue of Flink's {@code ExecutableStageDoFnOperator} and Spark's {@code + * SparkExecutableStageFunction}. State, timers and side inputs are out of scope here: the stage + * runs with {@link StateRequestHandler#unsupported()} and no timer receivers. */ class ExecutableStageProcessor implements Processor, byte[], KStreamsPayload> { @@ -89,30 +82,21 @@ class ExecutableStageProcessor private final RunnerApi.ExecutableStagePayload stagePayload; private final JobInfo jobInfo; - // 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. + // Stamped on every watermark forwarded, so downstream aggregators know which transform reported. 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. + // Updated from the MonitoringInfos the harness reports as each bundle completes. 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. - // Each entry carries the output PCollection id so it can be routed to that output's downstream on - // flush. The element type is intentionally wildcarded: the runner does not need to know the - // runtime value type — the bundle factory handles all coder application at the Fn-API boundary - // using the PCollection coders from the ExecutableStagePayload. + // Enqueued by harness threads and drained by the processing thread on bundle close, so it must + // be thread-safe. Each entry carries its output PCollection id for routing on flush. The element + // type is wildcarded: coders are applied by the bundle factory at the Fn-API boundary. private final Queue pendingOutputs = new ConcurrentLinkedQueue<>(); - // Output PCollection id -> the child node (a StageOutputProcessor relay) to forward that output - // to. Empty for a single-output stage, which forwards to its one downstream directly. + // Output PCollection id -> relay child node. Empty for a single-output stage. private final Map outputChildByPCollectionId; - // Computes this stage's input watermark from its upstream transform's reports, holding until - // every partition of the upstream transform has reported (see WatermarkAggregator). + // Holds until every partition of the upstream transform has reported; see WatermarkAggregator. private final WatermarkAggregator watermarkAggregator; - // Reports this stage instance as finished once it emits the terminal watermark, so a bounded - // pipeline can stop itself. + // Reports this stage finished at the terminal watermark, so a bounded pipeline can stop. private final TerminationReporter terminationReporter; // The last watermark actually forwarded downstream, so we only forward when it advances. private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE; @@ -169,10 +153,8 @@ private static final class PendingOutput { public void init(ProcessorContext> context) { this.context = context; terminationReporter.init(context); - // The SDK harness (stage context + bundle factory) is created lazily on the first data - // element, so a stage that only forwards watermarks never spins one up. This mirrors Spark's - // SparkExecutableStageFunction, which likewise does not build a bundle factory when there are - // no inputs to process. + // Created lazily on the first data element, so a stage that only forwards watermarks never + // spins up a harness. Spark's SparkExecutableStageFunction does the same. } private void ensureStageBundleFactory() { @@ -188,18 +170,16 @@ private void ensureStageBundleFactory() { public void process(Record> record) { KStreamsPayload payload = record.value(); if (payload == null) { - // A topic feeding the runner can always be written to from outside (or carry a tombstone), - // so recover from the obvious error instead of crashing the task: warn and drop. + // A topic can always be written to from outside, so warn and drop rather than crash. LOG.warn( "Stage {} dropping record with null payload (external write or tombstone)", transformId); return; } if (payload.isWatermark()) { - // Emit any buffered outputs before the watermark. Data is processed regardless of watermark - // readiness; only the watermark itself is held until every source partition has reported. + // Flush buffered outputs before the watermark. Data is processed regardless of readiness; + // only the watermark waits for every source partition. closeBundleAndFlush(record); - // Feed the report into the aggregator and forward the stage's output watermark only when the - // aggregate across the upstream transform's partitions actually advances. + // Forward the output watermark only when the aggregate across upstream partitions advances. watermarkAggregator.observe(payload.asWatermark()); Instant advanced = watermarkAggregator.advance(); if (advanced.isAfter(lastForwardedWatermark)) { @@ -230,8 +210,7 @@ private void ensureBundleOpen() throws Exception { new OutputReceiverFactory() { @Override public FnDataReceiver create(String pCollectionId) { - // Outputs are queued here on harness threads, tagged with their output PCollection id, - // and drained on the processing thread after the bundle closes. + // Queued on harness threads, drained on the processing thread after the bundle closes. return receivedElement -> { if (receivedElement != null) { pendingOutputs.add( diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java index 5b36a53607e6..56ae3a5d4550 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java @@ -28,25 +28,19 @@ import org.slf4j.LoggerFactory; /** - * Kafka Streams {@link Processor} implementing Beam's {@code Flatten} primitive ({@code - * beam:transform:flatten:v1}): the union of N input PCollections into one output PCollection. + * Kafka Streams {@link Processor} implementing Beam's {@code Flatten} primitive: the union of N + * input PCollections into one. * - *

    Data records are forwarded straight through unchanged — the merge of the N parents' - * data streams is the flatten. + *

    Data records pass straight through — merging the parents' streams is the flatten. The work is + * in the watermark, which Flatten owns as GroupByKey does: a {@link WatermarkAggregator} over its + * inputs, forwarding its own watermark only when the minimum across them advances and stamping it + * as a single source. That holds the output back until every branch has reported, so a downstream + * GroupByKey cannot fire before all branches are drained. * - *

    Watermark reports are where Flatten does real work, and it owns its output watermark - * the same way GroupByKey does: it runs a {@link WatermarkAggregator} over its inputs, forwards its - * own watermark only when the {@code min()} across them advances, and stamps that as a single - * source ({@code 0 of 1}) to its downstream. This holds the output watermark back until - * every input branch has reported, so a downstream GroupByKey does not fire before all - * flattened branches are drained. - * - *

    The {@link WatermarkAggregator} tells the input branches apart by the transform id each - * branch's producer stamps on its watermark (Kafka Streams does not tell a processor which parent - * forwarded a record). Each producer stamps its own identity regardless of who consumes it, so a - * PCollection feeding several Flattens reports one identity and every Flatten still waits only for - * the upstream transforms it expects — the set handed to it at construction from the pipeline - * graph. + *

    Branches are told apart by the transform id each producer stamps, since Kafka Streams does not + * say which parent forwarded a record. A producer stamps its own identity regardless of who + * consumes it, so a PCollection feeding several Flattens reports one identity and each Flatten + * still waits only for the upstream transforms handed to it at construction. */ class FlattenProcessor implements Processor, byte[], KStreamsPayload> { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java index a62d937170e9..8438e6af2a97 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java @@ -39,28 +39,17 @@ * Translates the {@code beam:transform:group_by_key:v1} URN — the runner's first stateful, * shuffle-bearing transform. * - *

    Windowing and triggering are executed by Beam's {@link - * org.apache.beam.runners.core.ReduceFnRunner} inside {@link WindowedGroupByKeyProcessor}, the same - * way the Flink and Spark portable runners do it — so fixed/sliding windows, the default trigger, - * allowed lateness and timestamp combiners all work. The input PCollection's windowing strategy is - * hydrated from the pipeline proto and handed to the processor. + *

    Windowing and triggering run through Beam's {@link + * org.apache.beam.runners.core.ReduceFnRunner} inside {@link WindowedGroupByKeyProcessor}, as the + * Flink and Spark portable runners do, so fixed and sliding windows, the default trigger, allowed + * lateness and timestamp combiners all work. The input's windowing strategy is hydrated from the + * pipeline proto and handed to the processor. * - *

    Topology added (the Beam key becomes the Kafka record key so Kafka Streams shuffles by it): - * - *

      - *
    • a {@link ShuffleByKeyProcessor} wired to the input's producer, which sets the Kafka record - * key to the encoded Beam key for data records and passes watermark reports through; - *
    • a {@link Topology#addSink sink} to an internal repartition topic, with the payload encoded - * via {@link KStreamsPayloadSerde} and a {@link GroupByKeyBroadcastPartitioner} that hashes - * data by key and fans watermark reports out to every partition; - *
    • a {@link Topology#addSource source} reading the repartition topic back; - *
    • the {@link WindowedGroupByKeyProcessor} plus persistent state and timer stores, wired to - * the source. - *
    - * - *

    The repartition topic is expected to exist on the broker before the job starts (same - * pre-create assumption as the Impulse bootstrap topic); auto-creation lands with the AdminClient - * wiring in a follow-up. + *

    The Beam key becomes the Kafka record key so Kafka Streams shuffles by it. The topology is a + * {@link ShuffleByKeyProcessor} that sets that key and passes watermark reports through, a sink to + * an internal repartition topic using a {@link GroupByKeyBroadcastPartitioner} that hashes data by + * key and fans watermarks out to every partition, a source reading that topic back, and the {@link + * WindowedGroupByKeyProcessor} with its state and timer stores. */ class GroupByKeyTranslator implements PTransformTranslator { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java index 7c4590a0e5a1..9a84035bf8bb 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java @@ -34,28 +34,17 @@ /** * Kafka Streams {@link Processor} implementing Beam's {@code Impulse} transform. * - *

    For each task instance, emits exactly two {@link KStreamsPayload}s downstream: + *

    Each task emits exactly two payloads: a {@link KStreamsPayload#data data} payload wrapping an + * empty {@code byte[]} in the {@link org.apache.beam.sdk.transforms.windowing.GlobalWindow} at + * {@link BoundedWindow#TIMESTAMP_MIN_VALUE}, then a {@link KStreamsPayload#watermark watermark} at + * {@link BoundedWindow#TIMESTAMP_MAX_VALUE} to say the source is done. * - *

      - *
    1. A {@link KStreamsPayload#data data} payload wrapping a {@link WindowedValue} of an empty - * {@code byte[]} in the {@link org.apache.beam.sdk.transforms.windowing.GlobalWindow}, with - * event-time {@link BoundedWindow#TIMESTAMP_MIN_VALUE}. - *
    2. A {@link KStreamsPayload#watermark watermark} payload at {@link - * BoundedWindow#TIMESTAMP_MAX_VALUE} that tells downstream transforms the source is done. - *
    + *

    A persistent state store records whether the data element was already emitted, so a restart + * does not duplicate it. The terminal watermark is re-emitted on every restart instead, so + * downstream watermark holds still release after recovery. * - *

    A persistent state store records whether the data element has already been emitted so that - * task restarts do not duplicate the data. The terminal watermark, on the other hand, is re-emitted - * on every restart so downstream watermark holds release correctly after recovery (per Jan's review - * on PR #38689). - * - *

    The trigger comes from a wall-clock punctuator scheduled on {@link #init} — this lets the - * processor fire even when the dedicated bootstrap source topic is empty, which is the expected - * production state. - * - *

    Kafka Streams disallows negative record timestamps, so the forwarded {@link Record} carries - * the Unix epoch ({@code 0L}). The Beam event-time lives inside the {@link KStreamsPayload} - * variant: inside the {@link WindowedValue} for data, or as the explicit watermark millis. + *

    A wall-clock punctuator scheduled in {@link #init} drives it, so the processor fires even + * though its bootstrap topic is empty, which is the normal production state. */ class ImpulseProcessor implements Processor> { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java index 29bd6e6bd9d3..21f2647e0986 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslator.java @@ -21,36 +21,20 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.kafka.common.serialization.Serdes; import org.apache.kafka.streams.Topology; -import org.apache.kafka.streams.state.KeyValueBytesStoreSupplier; import org.apache.kafka.streams.state.Stores; /** * Translates the {@code beam:transform:impulse:v1} URN. * - *

    Adds three nodes to the Kafka Streams {@link Topology}: + *

    Adds three nodes: a {@code byte[]} source bound to a per-transform bootstrap topic, which + * exists only because Kafka Streams refuses to start a topology with no source topic and whose + * records {@link ImpulseProcessor} ignores; the processor itself, which fires a one-shot wall-clock + * punctuator and emits one empty data payload followed by a terminal watermark; and a persistent + * state store recording whether it already fired, so a restart does not duplicate the impulse. * - *

      - *
    • A {@code byte[]} source bound to a dedicated per-transform bootstrap topic (see {@link - * KafkaStreamsTranslationContext#getImpulseBootstrapTopic}). Kafka Streams refuses to start a - * topology that has no real source topic, so the bootstrap topic exists purely to satisfy - * that requirement — records published to it are ignored by {@link ImpulseProcessor}. - *
    • The {@link ImpulseProcessor} itself, which schedules a one-shot wall-clock punctuator on - * {@code init} and emits a single empty data {@link KStreamsPayload} followed by a terminal - * watermark payload at {@link - * org.apache.beam.sdk.transforms.windowing.BoundedWindow#TIMESTAMP_MAX_VALUE}. - *
    • A per-processor {@link KeyValueBytesStoreSupplier persistent state store} that records - * whether the impulse has already fired so task restarts do not duplicate it. - *
    - * - *

    The processor's output PCollection is registered with the translation context so subsequent - * translators can wire themselves to this node by id. - * - *

    Bootstrap topic lifecycle: this translator does not auto-create the bootstrap - * topic. The topic is expected to exist on the broker before the job starts; otherwise Kafka - * Streams raises {@code MissingSourceTopicException} on startup. The auto-create-vs-pre-create - * decision (design doc §12.1) is deferred to a follow-up sub-issue along with the {@code - * AdminClient} wiring; pre-creation is sufficient for the {@code TopologyTestDriver}-based unit - * tests in this PR. + *

    The output PCollection is registered with the translation context so later translators can + * wire to this node by id. The bootstrap topic itself is created before startup by {@link + * org.apache.beam.runners.kafka.streams.KafkaStreamsTopicManager}. */ class ImpulseTranslator implements PTransformTranslator { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java index c165f0e875d4..647c64953b06 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java @@ -24,26 +24,13 @@ import org.checkerframework.checker.nullness.qual.Nullable; /** - * Sum-type envelope flowing between Kafka Streams processors in the Beam Kafka Streams runner. + * Envelope for every record value passed between the runner's processors. It is either a {@link + * #isData() data} element wrapping a {@link WindowedValue}, or a {@link #isWatermark() watermark} + * report carrying an event time plus the partition fields the downstream {@link WatermarkManager} + * needs. * - *

    Every record value emitted by a runner-introduced processor is one of: - * - *

      - *
    • A {@link #isData() data} element wrapping a {@link WindowedValue}, or - *
    • A {@link #isWatermark() watermark} report carrying an event-time milliseconds value plus - * the in-band coordination fields (source partition and total source partition count) the - * downstream {@link WatermarkManager} needs. - *
    - * - *

    The envelope lets a single Kafka Streams output channel carry both Beam data and the watermark - * / synchronization primitives that Kafka Streams does not natively support. Future control - * messages (e.g. the {@code (epoch, assigned_partitions)} propagation from design doc §5) can be - * added here as additional variants. - * - *

    This class is intentionally in-JVM only for now; serialization across topic boundaries - * (repartition or sink topics) will be introduced when the first translator that emits to a topic - * lands, at which point a corresponding Kafka {@link org.apache.kafka.common.serialization.Serde} - * will be added. + *

    One channel therefore carries both Beam data and the watermark coordination Kafka Streams has + * no notion of. Across topic boundaries it is encoded by {@link KStreamsPayloadSerde}. * * @param element type carried by data variants */ diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java index dddc28eb4381..aaf6d52b1d33 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTimerInternals.java @@ -33,26 +33,19 @@ * A {@link TimerInternals} for one key, backed by two Kafka Streams stores shared by a GroupByKey. * *

    Kafka Streams has no per-key timer service, so timers are persisted like any other state, in - * two stores that serve the two ways a timer is looked up: + * two stores serving the two ways a timer is looked up. The identity store, keyed by {@code key | + * domain | timerFamily | timerId | namespace}, is how {@link #setTimer} overwrites and {@link + * #deleteTimer} removes exactly one timer as the contract requires; its value is the index key, so + * an overwritten timer's index entry can be removed without knowing what time it was set for. The + * index store, keyed by {@code domain | fireTimestamp | identity}, is how due timers are found: the + * timestamp is in the sortable form described on {@link StoreKeys}, so every event-time timer due + * at a watermark is one range scan rather than a scan of every timer of every key, and its value is + * the {@link TimerData} so firing needs no second lookup. * - *

      - *
    • the identity store, keyed by {@code key | domain | timerFamily | timerId | - * namespace}, is how {@link #setTimer} overwrites and {@link #deleteTimer} removes exactly - * one timer, as {@link TimerInternals}' contract requires. Its value is the index key below, - * so a timer that is overwritten or deleted can have its index entry removed without knowing - * what time it had been set for. - *
    • the index store, keyed by {@code domain | fireTimestamp | identity}, is how due - * timers are found. Because the timestamp is written in the sortable form described on {@link - * StoreKeys}, all event-time timers due at a watermark are one range scan — {@link - * #dueEventTimeRangeStart} to {@link #dueEventTimeRangeEnd} — rather than a scan of every - * timer of every key. Its value is the {@link TimerData}, so firing needs no second lookup. - *
    - * - *

    Firing is driven by {@link WindowedGroupByKeyProcessor}: on a watermark advance it range-scans - * the index for event-time timers that are due and replays them through {@link - * org.apache.beam.runners.core.ReduceFnRunner#onTimers}. - * - *

    This instance reports the times it was constructed with; it never fires timers itself. + *

    Firing is driven by {@link WindowedGroupByKeyProcessor}, which range-scans the index on a + * watermark advance and replays due timers through {@link + * org.apache.beam.runners.core.ReduceFnRunner#onTimers}. This instance only reports the times it + * was constructed with; it never fires timers itself. */ class KafkaStreamsTimerInternals implements TimerInternals { 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 4a5463677163..cc8f57e0370c 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 @@ -52,18 +52,13 @@ public class KafkaStreamsTranslationContext { * has not been registered is produced by a single instance; only a shuffle raises the count. */ private final Map pCollectionIdToPartitionCount = new HashMap<>(); - // 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. + // Beam metrics from the SDK harness, one container per stage. Sharing a container across a + // stage's parallel tasks is safe: the cells are thread-safe and updates add rather than + // overwrite. Aggregating across runner JVMs is out of scope for now. private final MetricsContainerStepMap metricsContainerStepMap = new MetricsContainerStepMap(); - // Decides when a bounded pipeline has finished. Owned by the context, so it is scoped to this one - // pipeline: the job server runs several jobs in a single process, and a tracker shared between - // them would let one pipeline finishing stop another. + // Scoped to this pipeline rather than the JVM: the job server runs several jobs in one process, + // and a shared tracker would let one job's completion stop another. private final TerminationTracker terminationTracker = new TerminationTracker(); public static KafkaStreamsTranslationContext create( @@ -92,34 +87,23 @@ public KafkaStreamsPipelineOptions getPipelineOptions() { return pipelineOptions; } - /** Returns the {@link Topology} being built by the translation. */ 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}. - */ + /** One container per stage, updated by the processors and read by the pipeline result. */ public MetricsContainerStepMap getMetricsContainerStepMap() { return metricsContainerStepMap; } /** - * Returns the tracker that decides when this pipeline has finished. Processors report themselves - * to it as they reach the terminal watermark; the runner asks it to stop the Kafka Streams client - * once they all have. + * Processors report to it at the terminal watermark; the runner stops the client once all have. */ public TerminationTracker getTerminationTracker() { return terminationTracker; } - /** - * Registers the processor node that produces the given Beam PCollection. Downstream translators - * resolve their parent processor names by looking up the input PCollection id. - */ + /** Downstream translators resolve their parent node by looking up the input PCollection id. */ public void registerPCollectionProducer(String pCollectionId, String processorName) { String existing = pCollectionIdToProcessorName.putIfAbsent(pCollectionId, processorName); if (existing != null && !existing.equals(processorName)) { @@ -134,25 +118,14 @@ public void registerPCollectionProducer(String pCollectionId, String processorNa } /** - * Records how many partitions the transform producing {@code pCollectionId} runs across. - * - *

    This is the {@code totalSourcePartitions} its watermark reports carry, and what a downstream - * {@link WatermarkAggregator} waits to hear from before it lets the watermark advance. It changes - * only at a shuffle: everything fused downstream of one runs at the shuffle topic's partition - * count, and everything else runs as a single instance. + * The {@code totalSourcePartitions} this PCollection's watermark reports carry, which a + * downstream {@link WatermarkAggregator} waits on. It changes only at a shuffle. */ public void registerPCollectionPartitionCount(String pCollectionId, int partitionCount) { pCollectionIdToPartitionCount.put(pCollectionId, partitionCount); } - /** - * How many partitions the transform producing {@code pCollectionId} runs across; one unless a - * shuffle upstream raised it. - * - *

    Always at least one: an unregistered PCollection is produced by a single instance, and the - * only value ever registered is {@code --internalParallelism}, which the runner rejects below one - * before translating. - */ + /** One unless a shuffle upstream raised it; never less, as --internalParallelism is validated. */ public int getPartitionCount(String pCollectionId) { return pCollectionIdToPartitionCount.getOrDefault(pCollectionId, 1); } @@ -167,10 +140,8 @@ public String getProcessorNameForPCollection(String pCollectionId) { } /** - * Returns the dedicated bootstrap topic name for one Impulse transform. Keyed by transform id - * (sanitized to Kafka's legal topic-name character set) because a pipeline can contain several - * Impulses (e.g. an empty {@code Create} plus the dummy branch {@code PAssert} adds), and Kafka - * Streams rejects registering the same topic on two source nodes. + * Keyed by transform id, sanitized to Kafka's legal topic characters: a pipeline can hold several + * Impulses, and Kafka Streams rejects the same topic on two source nodes. */ public String getImpulseBootstrapTopic(String transformId) { String sanitizedTransformId = ILLEGAL_TOPIC_CHARS.matcher(transformId).replaceAll("_"); @@ -180,11 +151,7 @@ public String getImpulseBootstrapTopic(String transformId) { + sanitizedTransformId; } - /** - * Returns the dedicated bootstrap topic name a primitive Read reads from. Keyed by transform id - * (sanitized to Kafka's legal topic-name character set) so multiple Reads — and Impulse — never - * register the same topic on two source nodes, which Kafka Streams rejects. - */ + /** Keyed by transform id, for the same reason as {@link #getImpulseBootstrapTopic}. */ public String getReadBootstrapTopic(String transformId) { String sanitizedTransformId = ILLEGAL_TOPIC_CHARS.matcher(transformId).replaceAll("_"); return READ_BOOTSTRAP_TOPIC_PREFIX diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java index a6ef768ae2a1..a80376cdefc9 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java @@ -42,37 +42,26 @@ * Kafka Streams {@link Processor} implementing Beam's deprecated primitive {@code Read} * (beam:transform:read:v1) over a {@link BoundedSource}. * - *

    For each task instance, reads the whole {@link BoundedSource} once and emits, in order: + *

    Each task reads the whole source once, emitting one {@link KStreamsPayload#data data} payload + * per element in the {@link org.apache.beam.sdk.transforms.windowing.GlobalWindow} at its own event + * time, then a {@link KStreamsPayload#watermark watermark} at {@link + * BoundedWindow#TIMESTAMP_MAX_VALUE} to say the source is done. * - *

      - *
    1. One {@link KStreamsPayload#data data} payload per source element, each wrapping a {@link - * WindowedValue} in the {@link org.apache.beam.sdk.transforms.windowing.GlobalWindow} at the - * element's own event time (from {@link BoundedReader#getCurrentTimestamp()}). - *
    2. A {@link KStreamsPayload#watermark watermark} payload at {@link - * BoundedWindow#TIMESTAMP_MAX_VALUE} telling downstream transforms the source is done. - *
    + *

    Wire form. A Read produces decoded Java objects, but the harness's main-input receiver + * expects the runner-side wire form: a raw object for a model coder, a length-prefixed {@code + * byte[]} for a coder the runner does not know. Stage-to-stage edges already carry that form, so + * each element is transcoded here, encoded with the SDK-side wire coder and decoded with the + * runner-side one. The two are byte-compatible by construction, so this yields exactly what the + * receiver expects, nesting and all. * - *

    Wire form. Unlike Impulse (whose element is already an opaque {@code byte[]}), a Read - * produces decoded Java objects. Downstream {@link ExecutableStageProcessor} feeds - * whatever it receives straight into the SDK harness, whose main-input receiver expects each - * element in the runner-side wire form — a raw object for a model coder, but a length-prefixed - * {@code byte[]} for a coder the runner does not know (e.g. {@code VarIntCoder}). Stage-to-stage - * edges already carry that wire form because harness outputs are decoded with the runner-side wire - * coder; this processor reproduces it for the source edge by transcoding each element through the - * SDK-side wire coder (encode) and back through the runner-side wire coder (decode). The two are - * byte-compatible by construction, so the transcode yields exactly the object the receiver expects, - * nesting and all. + *

    As in {@link ImpulseProcessor}, a state store records whether the elements were already + * emitted so a restart does not duplicate them, while the terminal watermark is re-emitted on every + * restart so downstream holds still release. A wall-clock punctuator scheduled in {@link #init} + * drives it, since the bootstrap topic is empty. * - *

    This mirrors {@link ImpulseProcessor}: a persistent state store records whether the elements - * have already been emitted so task restarts do not duplicate them, while the terminal watermark is - * re-emitted on every restart so downstream watermark holds still release after recovery. The - * trigger is a wall-clock punctuator scheduled on {@link #init} so the processor fires even though - * its bootstrap source topic is empty. - * - *

    The source is read in a single instance with no splitting — parallelism across the source's - * splits arrives with the topic-based shuffle work (#18479). Kafka Streams disallows negative - * record timestamps, so each forwarded {@link Record} carries the Unix epoch ({@code 0L}); the Beam - * event time lives inside the {@link WindowedValue}. + *

    The source is read single-instance without splitting; parallel reads arrive with #18479. Kafka + * Streams rejects negative record timestamps, so each {@link Record} carries the Unix epoch and the + * Beam event time travels inside the {@link WindowedValue}. */ class ReadProcessor implements Processor> { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java index 9a727e82a70d..76d4649c0442 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java @@ -38,34 +38,21 @@ import org.apache.kafka.streams.state.Stores; /** - * Translates the deprecated primitive {@code Read} URN ({@code beam:transform:read:v1}) over a - * {@link BoundedSource}. + * Translates the deprecated primitive {@code Read} URN ({@code beam:transform:read:v1}). * - *

    The runner forces every {@code Read.Bounded} (including the one {@code Create} of two or more - * elements expands to) into this primitive read before translation — see {@code - * KafkaStreamsTestRunner.translate}, which applies {@code - * SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads}. This deliberately avoids the - * default {@code BoundedSourceAsSDFWrapperFn} splittable-DoFn expansion, which the runner cannot - * execute yet (no SDF restriction protocol), as agreed with the mentor. + *

    The runner converts every {@code Read} into this primitive before translation, rather than + * letting it expand into the default splittable-DoFn wrapper, which it cannot execute; see {@code + * KafkaStreamsRunner.prepareForTranslation}. * - *

    Adds the same three-node shape as {@link ImpulseTranslator}: + *

    The topology is the same three-node shape as {@link ImpulseTranslator}: a {@code byte[]} + * source on a per-transform bootstrap topic, since Kafka Streams will not start a topology with no + * source topic and the records on it are ignored; the {@link ReadProcessor}; and a persistent state + * store recording whether the read already fired, so a restart does not duplicate elements. * - *

      - *
    • A {@code byte[]} source bound to a dedicated per-transform bootstrap topic (see {@link - * KafkaStreamsTranslationContext#getReadBootstrapTopic(String)}). Kafka Streams refuses to - * start a topology with no real source topic; records published to it are ignored by {@link - * ReadProcessor}. - *
    • The {@link ReadProcessor}, which reads the {@link BoundedSource} on a one-shot wall-clock - * punctuator and emits one data payload per element followed by a terminal watermark. - *
    • A per-processor persistent state store recording whether the read has already fired so task - * restarts do not duplicate elements. - *
    - * - *

    The processor emits elements in the runner-side wire form the downstream stage's SDK harness - * expects, so it is handed the SDK-side and runner-side wire coders for the read's output - * PCollection (see {@link ReadProcessor} for why). Only {@link - * org.apache.beam.model.pipeline.v1.RunnerApi.IsBounded.Enum#BOUNDED bounded} sources are - * supported; {@link ReadTranslation#boundedSourceFromProto} rejects an unbounded payload. + *

    Elements are emitted in the runner-side wire form the downstream harness expects, so the + * processor is handed both wire coders for the output PCollection — see {@link ReadProcessor}. + * Bounded and unbounded sources are both supported, by {@link ReadProcessor} and {@link + * UnboundedReadProcessor} respectively. */ class ReadTranslator implements PTransformTranslator { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java index abc10500d211..6597b277770c 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationReporter.java @@ -25,25 +25,19 @@ import org.checkerframework.checker.nullness.qual.Nullable; /** - * The bit of every watermark-emitting processor that reports it has finished, so a bounded pipeline - * can stop itself. See {@link TerminationTracker} for why the runner has to work this out at all. + * The part of every watermark-emitting processor that reports it has finished, so a bounded + * pipeline can stop itself; see {@link TerminationTracker} for why that has to be worked out at + * all. A processor calls {@link #init} from {@code Processor#init}, passes every watermark it emits + * to {@link #watermarkEmitted}, and calls {@link #close} from {@code Processor#close}. * - *

    A processor creates one of these, calls {@link #init} from {@code Processor#init}, passes - * every watermark it emits to {@link #watermarkEmitted}, and calls {@link #close} from {@code - * Processor#close}. + *

    The report is scheduled rather than made inline, because reporting from inside {@code + * process()} would announce the processor finished while it is still handling the record that + * carried the terminal watermark; deferring it lets flushing, forwarding and committing happen + * first. * - *

    Why termination is scheduled rather than reported inline

    - * - *

    Reporting from inside {@code process()} would announce the processor as finished while it is - * still in the middle of handling the record that carried the terminal watermark. Scheduling a - * punctuator instead defers the report until the current processing has completed, so anything that - * has to happen after the final watermark — flushing a bundle, forwarding downstream, committing — - * still runs first. - * - *

    The punctuator is {@link PunctuationType#WALL_CLOCK_TIME} rather than stream time: no further - * records arrive after the terminal watermark, so stream time would never advance and a stream-time - * punctuator would never fire. The interval is the smallest Kafka Streams accepts — it rejects - * anything below a millisecond with "The minimum supported scheduling interval is 1 millisecond." + *

    It uses {@link PunctuationType#WALL_CLOCK_TIME}: no records arrive after the terminal + * watermark, so stream time would never advance and a stream-time punctuator would never fire. The + * interval is 1ms, the smallest Kafka Streams accepts. */ class TerminationReporter { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java index 6c570db1aff1..d0a04e38126a 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/TerminationTracker.java @@ -26,39 +26,24 @@ /** * Decides when a bounded pipeline has finished, so the Kafka Streams client can be stopped. * - *

    Kafka Streams has no notion of a processor being finished: a topology runs until something - * closes the client. A bounded Beam pipeline does finish, though, and the runner already knows - * when: every processor emits a watermark of {@link + *

    Kafka Streams has no notion of a finished processor; a topology runs until something closes + * the client. A bounded pipeline does finish, and the runner already knows when, because every + * processor emits {@link * org.apache.beam.sdk.transforms.windowing.BoundedWindow#TIMESTAMP_MAX_VALUE} once its input is - * exhausted. This class collects those reports and fires a callback when there is nothing left to - * do. + * exhausted. This collects those reports and fires a callback when nothing is left to do. * - *

    Why no coordination between instances is needed

    + *

    No coordination between instances is needed: a watermark crossing a repartition topic is + * broadcast to every partition (see {@link GroupByKeyBroadcastPartitioner}), so every task observes + * the terminal watermark itself and all instances reach the same conclusion independently. * - *

    A watermark that crosses a repartition topic is broadcast to every partition (see - * {@link GroupByKeyBroadcastPartitioner}), so every task of every downstream transform observes the - * terminal watermark on its own, whichever instance it happens to run on. Each instance can - * therefore decide to stop from what it sees locally, and they all reach the same conclusion - * without talking to each other. + *

    Every local processor is counted, not just the first. One instance can own tasks from both + * sides of a repartition topic, and the upstream side goes terminal as soon as it has written to + * the topic while the downstream side still has to consume it. Stopping at the first would drop + * that work and still report success. * - *

    Why every local processor has to be counted, not just the first

    - * - *

    One instance can own tasks from both sides of a repartition topic. The upstream side goes - * terminal as soon as it has written its data to the topic, while the downstream side still has to - * consume it. Stopping the client when the first processor finishes would cut that downstream work - * off and report the pipeline as done having silently dropped it. So the callback only fires once - * every processor instance registered here has terminated. - * - *

    An instance that happens to own only upstream tasks still terminates on its own, which is - * correct: what it wrote is durable in the topic for whichever instance reads it. - * - *

    Scope

    - * - *

    One tracker belongs to one pipeline, not to the JVM. The job server runs many jobs in a single - * process, so a shared static tracker would let one pipeline finishing tear down another. - * - *

    A pipeline with an unbounded source never produces a terminal watermark, so the callback never - * fires and the client keeps running — which is the intended behaviour for a streaming job. + *

    A tracker belongs to one pipeline rather than to the JVM: the job server runs many jobs in one + * process, and a static tracker would let one job stop another. An unbounded pipeline never + * produces a terminal watermark, so the callback never fires and the client keeps running. */ public class TerminationTracker { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java index 4f8c5f105566..21fc392990ea 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/UnboundedReadProcessor.java @@ -44,29 +44,24 @@ * Reads an {@link UnboundedSource} and forwards its elements and watermark downstream. * *

    Where the bounded {@link ReadProcessor} drains its source once and jumps the watermark to the - * end of time, an unbounded source never finishes: it is polled repeatedly, and its watermark - * advances gradually as the reader reports progress. That difference is what makes this a streaming - * runner rather than a batch one — downstream windows close because the source says time has moved - * on, not because the input ran out. + * end of time, an unbounded source never finishes: it is polled repeatedly and its watermark + * advances as the reader reports progress. That is what makes downstream windows close because time + * moved on rather than because the input ran out. * - *

    Polling happens on a wall-clock punctuator rather than in {@code process}, because the - * processor's bootstrap topic is empty and nothing else would drive it. Each turn reads at most - * {@link #maxElementsPerPoll} elements so a busy source cannot monopolise the Kafka Streams thread - * and starve the rest of the topology, then forwards the reader's watermark if it advanced. + *

    Polling runs on a wall-clock punctuator, since the bootstrap topic is empty and nothing else + * would drive it. A turn is bounded by {@link #maxElementsPerPoll} and by {@link #maxPollTimeMs} so + * a busy source cannot hold the Kafka Streams thread and starve the rest of the topology. * - *

    Restart is what the checkpoint mark is for. {@link UnboundedReader#getCheckpointMark()} - * describes the position the reader has consumed to; it is written to a persistent state store, and - * on {@link #init} the reader is created from the stored mark rather than from scratch, so a task - * that moves or restarts resumes where it left off instead of re-reading from the beginning. The - * store is changelogged and, under exactly-once, its writes commit atomically with the records the - * processor forwarded, so the mark can never be ahead of the data that was actually emitted. + *

    The checkpoint mark is what makes restart work. {@link UnboundedReader#getCheckpointMark()} is + * written to a persistent state store and the reader is recreated from it in {@link #init}, so a + * task that moves or restarts resumes where it left off. The store is changelogged and, under + * exactly-once, commits atomically with the records forwarded, so the mark cannot run ahead of the + * data actually emitted. * - *

    The source handed to this processor has already been split by {@link ReadTranslator}, which is - * where splitting belongs: it happens once for the pipeline rather than once per task instance, and - * the contract does not define splitting an already-split source. Reading several splits in - * parallel arrives with the topic-based shuffle work (#18479). As in the bounded processor, Kafka - * Streams disallows negative record timestamps, so each forwarded {@link Record} carries the Unix - * epoch and the Beam event time travels inside the {@link WindowedValue}. + *

    The source is split once by {@link ReadTranslator} rather than per task; reading several + * splits in parallel arrives with #18479. Kafka Streams rejects negative record timestamps, so each + * {@link Record} carries the Unix epoch and the event time travels inside the {@link + * WindowedValue}. */ class UnboundedReadProcessor implements Processor> { @@ -153,26 +148,16 @@ public void process(Record record) { /** * Drains what the source currently has, in batches, then publishes the watermark. * - *

    A batch is capped at {@link #maxElementsPerPoll} so that the checkpoint mark and the - * watermark are updated as the reader progresses rather than only at the end. Batches run back to - * back while the source keeps filling them, since returning after every batch would cap - * throughput at one batch per punctuation interval. + *

    A batch is capped at {@link #maxElementsPerPoll} so the checkpoint mark and the watermark + * move as the reader progresses, and batches run back to back while the source keeps filling + * them, since returning after each would cap throughput at one batch per punctuation. * - *

    The run is bounded all the same. A source that always has data — which is the normal case - * for one that is keeping up — would otherwise never let this method return, and the Kafka - * Streams thread would never get back to committing or to the rest of the topology. So at most - * {@link #checkpointEveryNPolls} batches are taken before yielding, which is also where the - * checkpoint mark is stored, and the next punctuation carries on from there. - * - *

    That batch bound is a count, and a count cannot bound the time: how long an element takes is - * decided by the pipeline underneath it, which the source knows nothing about. A punctuator is - * expected to be quick, and this one runs on the thread that also serves the rest of the - * topology, so a turn that overruns its own {@link #POLL_INTERVAL} is due again as soon as it - * returns and runs once more instead of the tasks below it. Measured on a grouping pipeline, a - * turn of 200 elements took 3ms and held the thread 6% of the time, while a turn of 5000 took - * 57ms and held it 89%, and the pipeline read tens of millions of elements while emitting none. - * {@link #maxPollTimeMs} bounds the turn in time as well, and whichever bound is reached first - * ends it. + *

    Both bounds exist because this runs on the thread that also serves the rest of the topology. + * At most {@link #checkpointEveryNPolls} batches are taken before yielding, and {@link + * #maxPollTimeMs} bounds the turn in time — a count cannot, since how long an element takes is + * decided by the pipeline below the source. A turn that overruns {@link #POLL_INTERVAL} is due + * again the moment it returns and runs instead of the tasks beneath it, which shows up as a + * pipeline that reads steadily and emits nothing. */ private void poll() { if (exhausted) { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java index c9af03df26a9..e081d2e9db6b 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java @@ -26,29 +26,20 @@ import org.joda.time.Instant; /** - * Computes a transform's input watermark from the watermark reports of its upstream transforms. + * Computes a transform's input watermark from the reports of its upstream transforms. * - *

    A watermark report carries three orthogonal pieces of information (see {@link - * WatermarkPayload}): which transform produced it, which partition (physical - * instance) of that transform it is for, and how many partitions that transform has. A - * producer stamps its own identity without regard to who consumes the report. This aggregator is - * the consuming side, used by every transform that aggregates a watermark — ExecutableStage, - * GroupByKey, Flatten (and CombinePerKey later): + *

    A report says which transform produced it, which partition of that transform it is for, and + * how many partitions that transform has (see {@link WatermarkPayload}); a producer stamps its own + * identity without regard to who consumes it. This is the consuming side, used by every transform + * that aggregates a watermark — ExecutableStage, GroupByKey, Flatten. * - *

      - *
    • It is constructed with the set of upstream transform ids the consumer expects, known from - * the pipeline graph at translation time (a single-input transform passes its one parent; a - * Flatten passes the producers of all of its input PCollections). - *
    • Per upstream transform it tracks partitions with a dedicated {@link WatermarkManager}, - * which holds until every partition of that transform has reported and keeps each partition - * monotonic. - *
    • The aggregate input watermark is the {@code min()} across the upstream transforms' - * watermarks, defined only once every expected upstream transform is ready; until - * then {@link #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} and the caller - * emits nothing. - *
    + *

    It is constructed with the upstream transform ids the consumer expects, known from the + * pipeline graph at translation time, and tracks each with its own {@link WatermarkManager}. The + * input watermark is the minimum across them, defined only once every expected upstream is ready; + * until then {@link #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} and the caller + * emits nothing. * - *

    Not thread-safe; the caller (a single Kafka Streams processor thread) serializes access. + *

    Not thread-safe; the calling Kafka Streams processor thread serializes access. */ final class WatermarkAggregator { diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java index cd0fca3654ae..39b9f11bb57d 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java @@ -24,44 +24,23 @@ import org.joda.time.Instant; /** - * In-memory tracker of a single fused stage's input watermark, computed from the committed - * watermarks reported by the upstream source partitions that feed it (the output / - * repartition-topic partitions of the parent stage). + * Tracks one fused stage's input watermark from the committed watermarks reported by the upstream + * source partitions feeding it. Kept free of Kafka wiring so it can be unit-tested on its own. * - *

    This is the core of the Kafka Streams runner's watermark propagation, decoupled from the Kafka - * wiring so it can be unit-tested in isolation. The wiring that produces the reports (flushing - * {@code (sourcePartition, committedWatermark, totalSourcePartitions)} atomically with each offset - * commit and fanning it out to every downstream partition) and consumes them lands in a follow-up. + *

    It counts source partitions rather than producer instances. A partition count is fixed, known, + * and travels in-band with every report, whereas instances come and go on every rebalance and can + * die without notice; a dead instance's partitions are reassigned and the new owner keeps + * reporting. * - *

    Why source partitions, not producer instances

    + *

    Until every source partition has reported, the stage's input watermark is undefined and {@link + * #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE}. A change in the partition count + * clears the reports and re-opens that hold, which subsumes an explicit epoch rule. * - *

    The question a stage has to answer is "have I received the watermark from every upstream - * producer, so that {@code min()} across them is meaningful?". Counting producer instances - * is hard: an instance can be killed without notice, leaving stale state, and the number changes on - * every rebalance. Counting source partitions is robust instead, because the partition count - * is fixed and known: it travels in-band with every report ({@code totalSourcePartitions}), a - * partition is always owned by exactly one live instance, and when an instance dies its partitions - * are reassigned and the new owner keeps reporting. So the manager only ever reasons about - * partitions, never about instances. (Design agreed with the mentor; see the watermark - * coordination-channel PoC findings.) + *

    Watermarks must not go backwards, so each partition's watermark is held monotonic and the + * emitted one is clamped against the last emitted — a newly appeared partition may report an older + * watermark than the stage has already reached. * - *

    Holding until ready

    - * - *

    Until a committed watermark has been seen for every source partition, the stage's input - * watermark is undefined and {@link #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} — - * i.e. the stage emits no meaningful watermark downstream. A change in {@code - * totalSourcePartitions} (e.g. a repartition) clears the accumulated reports and re-opens this hold - * until the new full set has reported, which subsumes the "new epoch / revert" rule without an - * explicit epoch. - * - *

    Monotonicity

    - * - *

    Beam watermarks must be non-decreasing. Each source partition's watermark is held monotonic (a - * lower report is ignored), and the emitted stage watermark is additionally clamped so it never - * regresses below the previously emitted value — relevant if a newly appeared partition reports an - * older watermark after the stage had already advanced. - * - *

    Not thread-safe; the caller (a single Kafka Streams processor thread) serializes access. + *

    Not thread-safe; the calling Kafka Streams processor thread serializes access. */ public final class WatermarkManager {