From 90bc4f4b6ea7d7fedfb857e18f9d127b50043344 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Sat, 1 Aug 2026 20:30:16 +0500 Subject: [PATCH 1/2] [GSoC 2026] Kafka Streams runner: bound a bundle by element count MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. The stage now counts the elements fed to the open bundle and closes it once that many have gone in, and 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. That aligns commits to bundle boundaries; it does not stop Kafka Streams from committing on its own interval part-way through a bundle, which would need a pre-commit hook, and the class documents that. maxBundleTimeMs is still not applied. Closing a bundle from a wall-clock punctuator made the pipeline with two chained GroupByKeys across four partitions emit its group repeatedly against a real broker, with the count still climbing after the input stopped. The same bundles closed from the record path are fine, and requesting the commit is not the cause — the duplication happens with the commit request removed. Rather than ship behaviour whose failure mode is not understood, the option documents that it has no effect yet. BundleBoundaryTest covers the size bound and the case of a bound the input never reaches. MetricsAcrossBundlesTest pins down that splitting the same input across many bundles does not change a user counter, since the runner folds each bundle's metrics by adding them. --- .../streams/KafkaStreamsPipelineOptions.java | 7 +- .../translation/ExecutableStageProcessor.java | 50 +++++++- .../ExecutableStageTranslator.java | 3 +- .../translation/BundleBoundaryTest.java | 118 ++++++++++++++++++ ...ExecutableStageProcessorWatermarkTest.java | 4 +- .../translation/MetricsAcrossBundlesTest.java | 81 ++++++++++++ 6 files changed, 258 insertions(+), 5 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..694ef24af3c1 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,15 @@ 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; + + /** The key of the last record processed, used when a close has no record to take one from. */ + private byte[] lastKey = new byte[0]; + /** * @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 +141,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. */ @@ -182,12 +205,20 @@ public void process(Record> record) { } return; } + byte[] key = record.key(); + if (key != null) { + lastKey = key; + } 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 +272,7 @@ public void onCompleted(ProcessBundleResponse response) { currentBundle = factory.getBundle( outputReceiverFactory, StateRequestHandler.unsupported(), progressHandler); + elementsInBundle = 0; } private FnDataReceiver> mainInputReceiver() { @@ -253,6 +285,18 @@ private FnDataReceiver> mainInputReceiver() { } private void closeBundleAndFlush(Record> record) { + byte[] key = record.key(); + closeBundleAndFlush(key == null ? lastKey : key, record.timestamp()); + } + + /** + * 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. + */ + private void closeBundleAndFlush(byte[] key, long timestamp) { RemoteBundle bundle = currentBundle; if (bundle == null) { return; @@ -265,6 +309,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. @@ -275,7 +320,7 @@ private void closeBundleAndFlush(Record> record) { while ((output = pendingOutputs.poll()) != null) { Record> outputRecord = new Record>( - record.key(), KStreamsPayload.data(output.value), record.timestamp()); + key, KStreamsPayload.data(output.value), timestamp); String childNode = outputChildByPCollectionId.get(output.pCollectionId); if (childNode == null) { ctx.forward(outputRecord); @@ -283,6 +328,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 53054b717bebaadc95cea8852019dc2841031cae Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Mon, 3 Aug 2026 00:26:44 +0500 Subject: [PATCH 2/2] Address review: do not substitute a key when closing a bundle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bundle close took a key and a timestamp separately so that a wall-clock punctuator, which has no record of its own, could pass the last key it had seen. That punctuator is not in this PR, so the substitution only ever applied to a record whose own key was null — and standing in a different record's key there would be wrong rather than merely unnecessary. The close now takes the record again and carries that record's own key and timestamp onto the outputs, as it did before this PR, and the field holding the last key is gone. Documents why the key is incidental here: an executable stage is unkeyed, running stateless with no state or timers, so the Kafka record key is only carried along, and where it matters downstream sets it -- ShuffleByKeyProcessor derives it from the Beam key before a GroupByKey. --- .../translation/ExecutableStageProcessor.java | 21 +++++++------------ 1 file changed, 7 insertions(+), 14 deletions(-) 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 694ef24af3c1..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 @@ -125,9 +125,6 @@ class ExecutableStageProcessor /** Elements fed to the open bundle, for the size bound above. */ private int elementsInBundle; - /** The key of the last record processed, used when a close has no record to take one from. */ - private byte[] lastKey = new byte[0]; - /** * @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 @@ -205,10 +202,6 @@ public void process(Record> record) { } return; } - byte[] key = record.key(); - if (key != null) { - lastKey = key; - } try { ensureBundleOpen(); mainInputReceiver().accept(payload.getData()); @@ -284,19 +277,19 @@ private FnDataReceiver> mainInputReceiver() { return receiver; } - private void closeBundleAndFlush(Record> record) { - byte[] key = record.key(); - closeBundleAndFlush(key == null ? lastKey : key, record.timestamp()); - } - /** * 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(byte[] key, long timestamp) { + private void closeBundleAndFlush(Record> record) { RemoteBundle bundle = currentBundle; if (bundle == null) { return; @@ -320,7 +313,7 @@ private void closeBundleAndFlush(byte[] key, long timestamp) { while ((output = pendingOutputs.poll()) != null) { Record> outputRecord = new Record>( - key, KStreamsPayload.data(output.value), timestamp); + record.key(), KStreamsPayload.data(output.value), record.timestamp()); String childNode = outputChildByPCollectionId.get(output.pCollectionId); if (childNode == null) { ctx.forward(outputRecord);