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