Skip to content

Commit 2719876

Browse files
junaiddshaukatjunaiddshaukat
authored
[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.
1 parent 47614a9 commit 2719876

17 files changed

Lines changed: 900 additions & 20 deletions

‎runners/kafka-streams/build.gradle‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
*/
1818

1919
import groovy.json.JsonOutput
20+
import java.time.Duration
2021

2122
plugins { id 'org.apache.beam.module' }
2223

@@ -74,6 +75,7 @@ dependencies {
7475
testImplementation library.java.junit
7576
testImplementation library.java.mockito_core
7677
testImplementation "org.apache.kafka:kafka-streams-test-utils:$kafka_version"
78+
testImplementation library.java.testcontainers_kafka
7779

7880
// Beam's @ValidatesRunner suite: the test classes come from the SDK core test jar; the runner
7981
// (TestKafkaStreamsRunner) and its TopologyTestDriver harness come from this module's test
@@ -85,6 +87,27 @@ dependencies {
8587
}
8688

8789

90+
// The broker integration test drives the production runner against a real Kafka in Docker, so it
91+
// is not part of the default build. Run it with :runners:kafka-streams:brokerIntegrationTest.
92+
test {
93+
filter {
94+
excludeTestsMatching 'org.apache.beam.runners.kafka.streams.*IT'
95+
}
96+
}
97+
98+
tasks.register("brokerIntegrationTest", Test) {
99+
group = "Verification"
100+
description = "Runs the Kafka Streams runner against a real broker (requires Docker)."
101+
outputs.upToDateWhen { false }
102+
testClassesDirs = sourceSets.test.output.classesDirs
103+
classpath = sourceSets.test.runtimeClasspath
104+
filter {
105+
includeTestsMatching 'org.apache.beam.runners.kafka.streams.*IT'
106+
}
107+
// A container start plus a streaming run is well past the default per-test expectations.
108+
timeout = Duration.ofMinutes(15)
109+
}
110+
88111
// Known-failing @ValidatesRunner tests, excluded until the feature they need lands.
89112
def sickbayTests = [
90113
// Merging (session) windows are not supported yet: ReduceFnRunner drives them through a merging

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,23 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions {
5555

5656
void setMaxBundleTimeMs(int maxBundleTimeMs);
5757

58+
@Description(
59+
"How many partitions the runner gives the internal topics it creates to shuffle a pipeline"
60+
+ " through, which is the parallelism the shuffled parts of that pipeline can reach. A"
61+
+ " GroupByKey runs one task per partition of its repartition topic, so this is the"
62+
+ " number of instances its state and its downstream stages are spread over. Must be at"
63+
+ " least 1.")
64+
@Default.Integer(1)
65+
int getInternalParallelism();
66+
67+
void setInternalParallelism(int internalParallelism);
68+
69+
@Description("Replication factor for the internal topics the runner creates for a pipeline.")
70+
@Default.Short(1)
71+
short getTopicReplicationFactor();
72+
73+
void setTopicReplicationFactor(short topicReplicationFactor);
74+
5875
@Description("Directory where Kafka Streams stores local state.")
5976
@Default.InstanceFactory(StateDirDefaultFactory.class)
6077
String getStateDir();

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

Lines changed: 32 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,10 @@
2424
import org.apache.beam.runners.jobsubmission.PortablePipelineRunner;
2525
import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsPipelineTranslator;
2626
import org.apache.beam.runners.kafka.streams.translation.KafkaStreamsTranslationContext;
27-
import org.apache.beam.sdk.options.PipelineOptionsValidator;
2827
import org.apache.kafka.streams.KafkaStreams;
2928
import org.apache.kafka.streams.StreamsConfig;
3029
import org.apache.kafka.streams.Topology;
30+
import org.checkerframework.checker.nullness.qual.Nullable;
3131
import org.slf4j.Logger;
3232
import org.slf4j.LoggerFactory;
3333

@@ -44,9 +44,22 @@ public KafkaStreamsPipelineRunner(KafkaStreamsPipelineOptions pipelineOptions) {
4444

4545
@Override
4646
public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) {
47-
// Surface a clear error if a required option (e.g. applicationId) is missing instead of
48-
// letting Properties.put fail with a raw NullPointerException further down.
49-
PipelineOptionsValidator.validate(KafkaStreamsPipelineOptions.class, pipelineOptions);
47+
// Surface a clear error if an option this runner needs is missing, instead of letting
48+
// Properties.put fail with a raw NullPointerException further down. Only the options that are
49+
// meaningful here are checked, rather than validating the whole interface: this runs on the job
50+
// server, executing a pipeline that has already been submitted, so the client-side options
51+
// PortablePipelineOptions marks required — jobEndpoint above all — do not apply. Flink's
52+
// equivalent PortablePipelineRunner does not validate here either.
53+
checkRequiredOption("applicationId", pipelineOptions.getApplicationId());
54+
checkRequiredOption("bootstrapServers", pipelineOptions.getBootstrapServers());
55+
// A topic cannot have fewer than one partition, and the value is also the number of watermark
56+
// reports a shuffle's consumer waits for, so a non-positive value would leave it waiting
57+
// forever rather than failing.
58+
if (pipelineOptions.getInternalParallelism() < 1) {
59+
throw new IllegalArgumentException(
60+
"--internalParallelism must be at least 1, but was "
61+
+ pipelineOptions.getInternalParallelism());
62+
}
5063

5164
KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator();
5265
KafkaStreamsTranslationContext context =
@@ -55,15 +68,28 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo)
5568
translator.translate(context, prepared);
5669

5770
Topology topology = context.getTopology();
71+
// The runner names its own bootstrap and repartition topics, which Kafka Streams treats as
72+
// user topics and will not create; it refuses to start if a source topic is missing.
73+
KafkaStreamsTopicManager.createMissingTopics(topology, pipelineOptions);
5874
LOG.info(
5975
"Translated pipeline {} into Kafka Streams topology:\n{}",
6076
jobInfo.jobId(),
6177
topology.describe());
6278

6379
KafkaStreams kafkaStreams = new KafkaStreams(topology, streamsConfig(jobInfo));
80+
// Build the result before starting: it registers a state listener, and Kafka Streams only
81+
// accepts one while the application is still in the CREATED state.
82+
KafkaStreamsPortablePipelineResult result =
83+
new KafkaStreamsPortablePipelineResult(kafkaStreams, context.getMetricsContainerStepMap());
6484
kafkaStreams.start();
65-
return new KafkaStreamsPortablePipelineResult(
66-
kafkaStreams, context.getMetricsContainerStepMap());
85+
return result;
86+
}
87+
88+
private static void checkRequiredOption(String name, @Nullable String value) {
89+
if (value == null || value.isEmpty()) {
90+
throw new IllegalArgumentException(
91+
"Missing required pipeline option --" + name + " for the Kafka Streams runner");
92+
}
6793
}
6894

6995
private Properties streamsConfig(JobInfo jobInfo) {

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,10 @@ class KafkaStreamsPortablePipelineResult implements PortablePipelineResult {
4848
private final CountDownLatch terminated = new CountDownLatch(1);
4949
private volatile boolean cancelled = false;
5050

51+
/**
52+
* Must be constructed before {@link KafkaStreams#start()} is called: it registers a state
53+
* listener, and Kafka Streams rejects one once the application has left the CREATED state.
54+
*/
5155
KafkaStreamsPortablePipelineResult(
5256
KafkaStreams kafkaStreams, MetricsContainerStepMap metricsContainerStepMap) {
5357
this.kafkaStreams = kafkaStreams;

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

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.runners.kafka.streams;
19+
20+
import java.util.ArrayList;
21+
import java.util.Collection;
22+
import java.util.HashSet;
23+
import java.util.List;
24+
import java.util.Properties;
25+
import java.util.Set;
26+
import java.util.concurrent.ExecutionException;
27+
import org.apache.kafka.clients.admin.Admin;
28+
import org.apache.kafka.clients.admin.AdminClientConfig;
29+
import org.apache.kafka.clients.admin.NewTopic;
30+
import org.apache.kafka.common.errors.TopicExistsException;
31+
import org.apache.kafka.streams.Topology;
32+
import org.apache.kafka.streams.TopologyDescription;
33+
import org.slf4j.Logger;
34+
import org.slf4j.LoggerFactory;
35+
36+
/**
37+
* Creates the topics a translated pipeline needs before the Kafka Streams application starts.
38+
*
39+
* <p>The runner shuffles data through topics it names itself: a bootstrap topic per Impulse and per
40+
* primitive Read, and a repartition topic per GroupByKey. Kafka Streams does create the internal
41+
* topics it manages on its own, but these are declared with explicit names through {@code
42+
* addSource} and {@code addSink}, so to Kafka Streams they are ordinary user topics — it will not
43+
* create them, and refuses to start with {@code MissingSourceTopicException} if a source topic is
44+
* absent. Relying on the broker's {@code auto.create.topics.enable} is not an option either: it is
45+
* off on many clusters, and a topic auto-created on first fetch gets the broker's default partition
46+
* count rather than the pipeline's.
47+
*
48+
* <p>Only topics carrying one of the runner's own prefixes are created. Any other topic in the
49+
* topology belongs to the user (a source or sink they named), and creating those implicitly would
50+
* hide a misconfiguration behind an empty topic.
51+
*/
52+
class KafkaStreamsTopicManager {
53+
54+
private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsTopicManager.class);
55+
56+
/**
57+
* Prefixes of the bootstrap topics, which must have exactly one partition.
58+
*
59+
* <p>An Impulse or a primitive Read emits its elements once per task, gated by a state store that
60+
* is itself per task. Kafka Streams creates one task per partition of the source topic, so a
61+
* bootstrap topic with several partitions would make the same Impulse fire once per partition and
62+
* the same source be read once per partition.
63+
*/
64+
private static final List<String> SINGLE_PARTITION_TOPIC_PREFIXES =
65+
java.util.Arrays.asList("__beam_impulse_", "__beam_read_");
66+
67+
/**
68+
* Prefixes of the topics whose partition count sets the pipeline's parallelism — the repartition
69+
* topic a GroupByKey shuffles through.
70+
*/
71+
private static final List<String> PARTITIONED_TOPIC_PREFIXES =
72+
java.util.Arrays.asList("__beam_gbk_");
73+
74+
private KafkaStreamsTopicManager() {}
75+
76+
/**
77+
* Creates any runner-owned topic in {@code topology} that does not exist yet.
78+
*
79+
* <p>Safe to run concurrently with another instance of the same job: a topic that appears between
80+
* the existence check and the create request surfaces as {@link TopicExistsException}, which is
81+
* treated as success.
82+
*/
83+
static void createMissingTopics(Topology topology, KafkaStreamsPipelineOptions options) {
84+
Set<String> runnerTopics = runnerOwnedTopics(topology);
85+
if (runnerTopics.isEmpty()) {
86+
return;
87+
}
88+
Properties adminConfig = new Properties();
89+
adminConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, options.getBootstrapServers());
90+
try (Admin admin = Admin.create(adminConfig)) {
91+
Set<String> existing = admin.listTopics().names().get();
92+
List<NewTopic> toCreate = new ArrayList<>();
93+
for (String topic : runnerTopics) {
94+
if (!existing.contains(topic)) {
95+
toCreate.add(
96+
new NewTopic(
97+
topic, partitionsFor(topic, options), options.getTopicReplicationFactor()));
98+
}
99+
}
100+
if (toCreate.isEmpty()) {
101+
return;
102+
}
103+
LOG.info("Creating {} runner-owned topic(s): {}", toCreate.size(), toCreate);
104+
createAll(admin, toCreate);
105+
} catch (InterruptedException e) {
106+
Thread.currentThread().interrupt();
107+
throw new RuntimeException("Interrupted while creating the pipeline's Kafka topics", e);
108+
} catch (ExecutionException e) {
109+
throw new RuntimeException("Failed to create the pipeline's Kafka topics", e);
110+
}
111+
}
112+
113+
private static void createAll(Admin admin, Collection<NewTopic> topics)
114+
throws InterruptedException, ExecutionException {
115+
try {
116+
admin.createTopics(topics).all().get();
117+
} catch (ExecutionException e) {
118+
// Another instance of the same application may have created them first, which is fine.
119+
if (!(e.getCause() instanceof TopicExistsException)) {
120+
throw e;
121+
}
122+
LOG.debug("Some topics already existed; another instance created them first", e);
123+
}
124+
}
125+
126+
/** The topics in the topology that the runner named, and so is responsible for creating. */
127+
private static Set<String> runnerOwnedTopics(Topology topology) {
128+
Set<String> topics = new HashSet<>();
129+
for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) {
130+
for (TopologyDescription.Node node : subtopology.nodes()) {
131+
if (node instanceof TopologyDescription.Source) {
132+
Set<String> sourceTopics = ((TopologyDescription.Source) node).topicSet();
133+
if (sourceTopics != null) {
134+
topics.addAll(sourceTopics);
135+
}
136+
} else if (node instanceof TopologyDescription.Sink) {
137+
String topic = ((TopologyDescription.Sink) node).topic();
138+
if (topic != null) {
139+
topics.add(topic);
140+
}
141+
}
142+
}
143+
}
144+
topics.removeIf(topic -> !isRunnerOwned(topic));
145+
return topics;
146+
}
147+
148+
/**
149+
* The partition count a runner-owned topic is created with: one for a bootstrap topic, and the
150+
* configured parallelism for a shuffle topic.
151+
*/
152+
private static int partitionsFor(String topic, KafkaStreamsPipelineOptions options) {
153+
return hasAnyPrefix(topic, SINGLE_PARTITION_TOPIC_PREFIXES)
154+
? 1
155+
: options.getInternalParallelism();
156+
}
157+
158+
private static boolean isRunnerOwned(String topic) {
159+
return hasAnyPrefix(topic, SINGLE_PARTITION_TOPIC_PREFIXES)
160+
|| hasAnyPrefix(topic, PARTITIONED_TOPIC_PREFIXES);
161+
}
162+
163+
private static boolean hasAnyPrefix(String topic, List<String> prefixes) {
164+
for (String prefix : prefixes) {
165+
if (topic.startsWith(prefix)) {
166+
return true;
167+
}
168+
}
169+
return false;
170+
}
171+
}

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -286,10 +286,11 @@ private void closeBundleAndFlush(Record<byte[], KStreamsPayload<?>> record) {
286286
}
287287

288288
private void forwardWatermark(Record<byte[], KStreamsPayload<?>> record, long watermarkMillis) {
289-
// Stamped with this stage's own transform id; this stage is a single instance for now, so the
290-
// report is for its only partition (0 of 1). Fanning the watermark out to every downstream
291-
// partition — and producing it atomically with the offset commit so it is durable — lands with
292-
// the topic-based shuffle work (#18479).
289+
// Labelled as the only source a consumer will see. Forwarding here is in-process, to the
290+
// stage's
291+
// fused children, so exactly one instance of this stage reaches each of them. Where the output
292+
// instead crosses a shuffle, ShuffleByKeyProcessor relabels the report with the real partition
293+
// identity, because the broadcast then delivers every instance's report to every consumer.
293294
ProcessorContext<byte[], KStreamsPayload<?>> ctx = checkInitialized(context);
294295
ctx.forward(
295296
new Record<byte[], KStreamsPayload<?>>(

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,8 @@ public void translate(
7474
// is unambiguous even before we add side-input support.
7575
String inputPCollectionId = stagePayload.getInput();
7676
String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId);
77+
// A fused stage runs wherever its input runs: same task, so same partition identity.
78+
int partitionCount = context.getPartitionCount(inputPCollectionId);
7779

7880
// A multi-output stage (a DoFn with side outputs, or a Read whose SDF wrapper produces several
7981
// outputs) needs each output routed to the right downstream. Since downstream transforms are
@@ -117,9 +119,11 @@ public void translate(
117119
topology.addProcessor(
118120
relayName, () -> new StageOutputProcessor(relayName), transformId);
119121
context.registerPCollectionProducer(outputPCollectionId, relayName);
122+
context.registerPCollectionPartitionCount(outputPCollectionId, partitionCount);
120123
});
121124
} else if (!outputPCollectionIds.isEmpty()) {
122125
context.registerPCollectionProducer(outputPCollectionIds.get(0), transformId);
126+
context.registerPCollectionPartitionCount(outputPCollectionIds.get(0), partitionCount);
123127
}
124128
}
125129
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,8 +99,7 @@ public void process(Record<byte[], KStreamsPayload<?>> record) {
9999
Instant advanced = watermarkAggregator.advance();
100100
if (advanced.isAfter(lastForwardedWatermark)) {
101101
lastForwardedWatermark = advanced;
102-
// Stamped with this Flatten's own transform id; Flatten is a single instance for now, so the
103-
// report is for its only partition (0 of 1).
102+
// Labelled as the only source a consumer will see; a shuffle downstream relabels it.
104103
ctx.forward(
105104
new Record<byte[], KStreamsPayload<?>>(
106105
record.key(),

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,9 @@ public void translate(
5757
Set<String> seenInputs = new HashSet<>();
5858
List<String> parentProcessors = new ArrayList<>();
5959
Set<String> upstreamTransformIds = new HashSet<>();
60+
// Kafka Streams puts a processor and the parents it is wired to in one subtopology, so the
61+
// inputs are co-partitioned and this Flatten runs at their partition count.
62+
int partitionCount = 1;
6063
for (String inputPCollectionId : transform.getInputsMap().values()) {
6164
if (!seenInputs.add(inputPCollectionId)) {
6265
throw new UnsupportedOperationException(
@@ -70,6 +73,7 @@ public void translate(
7073
String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId);
7174
parentProcessors.add(parentProcessor);
7275
upstreamTransformIds.add(parentProcessor);
76+
partitionCount = Math.max(partitionCount, context.getPartitionCount(inputPCollectionId));
7377
}
7478

7579
topology.addProcessor(
@@ -78,5 +82,6 @@ public void translate(
7882
parentProcessors.toArray(new String[0]));
7983

8084
context.registerPCollectionProducer(outputPCollectionId, transformId);
85+
context.registerPCollectionPartitionCount(outputPCollectionId, partitionCount);
8186
}
8287
}

0 commit comments

Comments
 (0)