This is an automated email from the ASF dual-hosted git repository.
je-ik pushed a commit to branch feat/18479-kafka-streams-runner-skeleton
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to
refs/heads/feat/18479-kafka-streams-runner-skeleton by this push:
new 1058e94027e [GSoC 2026] Kafka Streams runner #38987: Wire
WatermarkManager into ExecutableStageProcessor
1058e94027e is described below
commit 1058e94027e5925a6324c0cbd9f3013a5a870218
Author: M Junaid Shaukat <[email protected]>
AuthorDate: Thu Jun 18 16:30:35 2026 +0500
[GSoC 2026] Kafka Streams runner #38987: Wire WatermarkManager into
ExecutableStageProcessor
---
.../translation/ExecutableStageProcessor.java | 62 ++++++---
.../streams/translation/ImpulseProcessor.java | 10 +-
.../kafka/streams/translation/KStreamsPayload.java | 77 ++++++++--
.../streams/translation/WatermarkPayload.java | 41 ++++++
.../kafka/streams/KafkaStreamsRunnerTest.java | 2 +-
.../ExecutableStageProcessorWatermarkTest.java | 124 +++++++++++++++++
.../streams/translation/ImpulseTranslatorTest.java | 3 +-
.../translation/WatermarkPropagationTest.java | 155 +++++++++++++++++++++
8 files changed, 440 insertions(+), 34 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 b8eb5413b27..00f85032d04 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
@@ -28,6 +28,7 @@ import
org.apache.beam.runners.fnexecution.control.StageBundleFactory;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
import org.apache.beam.sdk.values.WindowedValue;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
@@ -35,6 +36,7 @@ import org.apache.kafka.streams.processor.api.Processor;
import org.apache.kafka.streams.processor.api.ProcessorContext;
import org.apache.kafka.streams.processor.api.Record;
import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -49,9 +51,12 @@ import org.slf4j.LoggerFactory;
* ProcessorContext#forward} must only be called from the processing thread,
so outputs are never
* forwarded directly from a harness callback.
*
- * <p>A {@link KStreamsPayload#isWatermark() watermark} payload marks a bundle
boundary: the open
- * bundle (if any) is closed (flushing outputs), and the watermark is then
forwarded downstream so
- * that subsequent stages observe it after all data of the bundle.
+ * <p>A {@link KStreamsPayload#isWatermark() watermark} payload is a
per-source-partition report and
+ * marks a bundle boundary: the open bundle (if any) is closed (flushing
outputs), the report is fed
+ * to the {@link WatermarkManager}, and the stage's output watermark is
forwarded downstream only
+ * when the {@code min()} across its source partitions actually advances.
Until every source
+ * partition has reported, the watermark is held and nothing is forwarded —
but data is still
+ * processed in the meantime.
*
* <p>This is the Kafka Streams analogue of Flink's {@code
ExecutableStageDoFnOperator} and Spark's
* {@code SparkExecutableStageFunction}. State, timers, and side inputs are
out of scope for this
@@ -74,6 +79,12 @@ class ExecutableStageProcessor
// only safe because the Impulse output coder happens to be ByteArrayCoder.
private final Queue<WindowedValue<?>> pendingOutputs = new
ConcurrentLinkedQueue<>();
+ // Computes this stage's output watermark as min() over its source
partitions' reported
+ // watermarks, holding until every source partition has reported (see
WatermarkManager).
+ private final WatermarkManager watermarkManager = new WatermarkManager();
+ // The last watermark actually forwarded downstream, so we only forward when
it advances.
+ private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
+
private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
private @Nullable ExecutableStageContext stageContext;
private @Nullable StageBundleFactory stageBundleFactory;
@@ -87,22 +98,40 @@ class ExecutableStageProcessor
@Override
public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
this.context = context;
+ // The SDK harness (stage context + bundle factory) is created lazily on
the first data
+ // element, so a stage that only forwards watermarks never spins one up.
This mirrors Spark's
+ // SparkExecutableStageFunction, which likewise does not build a bundle
factory when there are
+ // no inputs to process.
+ }
+
+ private void ensureStageBundleFactory() {
+ if (stageBundleFactory != null) {
+ return;
+ }
ExecutableStage executableStage =
ExecutableStage.fromPayload(stagePayload);
- this.stageContext =
KafkaStreamsExecutableStageContextFactory.getInstance().get(jobInfo);
- this.stageBundleFactory =
stageContext.getStageBundleFactory(executableStage);
+ stageContext =
KafkaStreamsExecutableStageContextFactory.getInstance().get(jobInfo);
+ stageBundleFactory = stageContext.getStageBundleFactory(executableStage);
}
@Override
public void process(Record<byte[], KStreamsPayload<?>> record) {
KStreamsPayload<?> payload = record.value();
if (payload.isWatermark()) {
- // NOTE: flushing the bundle on every received watermark is provisional.
Once the
- // WatermarkManager lands, a stage will receive watermarks from multiple
parent instances and
- // the output watermark becomes min() across them — the bundle should
flush / the output
- // watermark advance only when that minimum actually moves forward, not
on every received
- // watermark. Tracked in #38743.
+ // Emit any buffered outputs before the watermark. Data is processed
regardless of watermark
+ // readiness; only the watermark itself is held until every source
partition has reported.
closeBundleAndFlush(record);
- forwardWatermark(record, payload.getWatermarkMillis());
+ // Feed the report into the WatermarkManager and forward the stage's
output watermark only
+ // when min() across the source partitions actually advances, not on
every received watermark.
+ WatermarkPayload report = payload.asWatermark();
+ watermarkManager.observe(
+ report.getSourcePartition(),
+ new Instant(report.getWatermarkMillis()),
+ report.getTotalSourcePartitions());
+ Instant advanced = watermarkManager.advance();
+ if (advanced.isAfter(lastForwardedWatermark)) {
+ lastForwardedWatermark = advanced;
+ forwardWatermark(record, advanced.getMillis());
+ }
return;
}
try {
@@ -117,6 +146,7 @@ class ExecutableStageProcessor
if (currentBundle != null) {
return;
}
+ ensureStageBundleFactory();
StageBundleFactory factory = checkInitialized(stageBundleFactory);
OutputReceiverFactory outputReceiverFactory =
new OutputReceiverFactory() {
@@ -173,14 +203,14 @@ class ExecutableStageProcessor
}
private void forwardWatermark(Record<byte[], KStreamsPayload<?>> record,
long watermarkMillis) {
- // TODO(#38743 / WatermarkManager): a watermark must reach every parallel
instance of every
- // downstream processor, but ctx.forward routes to one downstream
partition per Kafka Streams'
- // partitioning. The simplest correct approach is to fan the watermark out
to all downstream
- // partitions; that wiring lands with the WatermarkManager sub-issue (per
Jan on PR #38764).
+ // This stage is a single instance for now, so it forwards its watermark
as the only source
+ // partition (0 of 1). Fanning the watermark out to every downstream
partition — and producing
+ // it atomically with the offset commit so it is durable — lands with the
topic-based shuffle
+ // work, when there are real source partitions to track (#18479).
ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
ctx.forward(
new Record<byte[], KStreamsPayload<?>>(
- record.key(), KStreamsPayload.watermark(watermarkMillis),
record.timestamp()));
+ record.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1),
record.timestamp()));
}
@Override
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
index e9fc8ddb36a..675bee4d859 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
@@ -121,12 +121,18 @@ class ImpulseProcessor implements Processor<byte[],
byte[], byte[], KStreamsPayl
LOG.debug("Impulse {} emitted single element and terminal watermark",
transformId);
}
- /** Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to
downstream processors. */
+ /**
+ * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to
downstream processors.
+ *
+ * <p>Impulse is a single-instance source, so the report is stamped as the
only source partition:
+ * {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real
per-partition identities
+ * arrive once the topology gains topic-based shuffle.
+ */
private static void forwardWatermarkMax(ProcessorContext<byte[],
KStreamsPayload<byte[]>> ctx) {
long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis();
ctx.forward(
new Record<byte[], KStreamsPayload<byte[]>>(
- new byte[0], KStreamsPayload.watermark(maxMillis), 0L));
+ new byte[0], KStreamsPayload.watermark(maxMillis, 0, 1), 0L));
}
/** Cancels the wall-clock punctuator after the impulse has fired to stop
periodic wakeups. */
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java
index 53e47b1216b..93b40346761 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java
@@ -20,6 +20,7 @@ package org.apache.beam.runners.kafka.streams.translation;
import java.util.Objects;
import org.apache.beam.sdk.values.WindowedValue;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.checkerframework.checker.nullness.qual.Nullable;
/**
@@ -29,7 +30,9 @@ import org.checkerframework.checker.nullness.qual.Nullable;
*
* <ul>
* <li>A {@link #isData() data} element wrapping a {@link WindowedValue}, or
- * <li>A {@link #isWatermark() watermark} signal carrying an event-time
milliseconds value.
+ * <li>A {@link #isWatermark() watermark} report carrying an event-time
milliseconds value plus
+ * the in-band coordination fields (source partition and total source
partition count) the
+ * downstream {@link WatermarkManager} needs.
* </ul>
*
* <p>The envelope lets a single Kafka Streams output channel carry both Beam
data and the watermark
@@ -54,21 +57,45 @@ public final class KStreamsPayload<T> {
private final Kind kind;
private final @Nullable WindowedValue<T> data;
private final long watermarkMillis;
+ private final int sourcePartition;
+ private final int totalSourcePartitions;
- private KStreamsPayload(Kind kind, @Nullable WindowedValue<T> data, long
watermarkMillis) {
+ private KStreamsPayload(
+ Kind kind,
+ @Nullable WindowedValue<T> data,
+ long watermarkMillis,
+ int sourcePartition,
+ int totalSourcePartitions) {
this.kind = kind;
this.data = data;
this.watermarkMillis = watermarkMillis;
+ this.sourcePartition = sourcePartition;
+ this.totalSourcePartitions = totalSourcePartitions;
}
/** Returns a data payload wrapping the given {@link WindowedValue}. */
public static <T> KStreamsPayload<T> data(WindowedValue<T> value) {
- return new KStreamsPayload<>(Kind.DATA, value, 0L);
+ return new KStreamsPayload<>(Kind.DATA, value, 0L, 0, 0);
}
- /** Returns a watermark payload carrying the given event-time milliseconds.
*/
- public static <T> KStreamsPayload<T> watermark(long watermarkMillis) {
- return new KStreamsPayload<>(Kind.WATERMARK, null, watermarkMillis);
+ /**
+ * Returns a watermark report payload: the event-time milliseconds together
with the in-band
+ * coordination fields the downstream stage's {@link WatermarkManager} needs
— which source
+ * partition this report is for and how many source partitions feed the
stage in total.
+ */
+ public static <T> KStreamsPayload<T> watermark(
+ long watermarkMillis, int sourcePartition, int totalSourcePartitions) {
+ Preconditions.checkArgument(
+ totalSourcePartitions > 0,
+ "totalSourcePartitions must be positive: %s",
+ totalSourcePartitions);
+ Preconditions.checkArgument(
+ sourcePartition >= 0 && sourcePartition < totalSourcePartitions,
+ "sourcePartition %s out of range for totalSourcePartitions %s",
+ sourcePartition,
+ totalSourcePartitions);
+ return new KStreamsPayload<>(
+ Kind.WATERMARK, null, watermarkMillis, sourcePartition,
totalSourcePartitions);
}
public boolean isData() {
@@ -91,14 +118,31 @@ public final class KStreamsPayload<T> {
}
/**
- * Returns the watermark event-time milliseconds. Caller must check {@link
#isWatermark()} first;
- * calling this on a data payload throws.
+ * Narrows this payload to its {@link WatermarkPayload} view, through which
the watermark report
+ * fields are read. Caller must check {@link #isWatermark()} first; calling
this on a data payload
+ * throws.
*/
- public long getWatermarkMillis() {
- if (kind != Kind.WATERMARK) {
- throw new IllegalStateException("Payload is not a watermark: kind=" +
kind);
+ public WatermarkPayload asWatermark() {
+ Preconditions.checkState(isWatermark(), "Payload is not a watermark:
kind=%s", kind);
+ return new WatermarkView();
+ }
+
+ /** {@link WatermarkPayload} view backed by this payload's fields. */
+ private final class WatermarkView implements WatermarkPayload {
+ @Override
+ public long getWatermarkMillis() {
+ return watermarkMillis;
+ }
+
+ @Override
+ public int getSourcePartition() {
+ return sourcePartition;
+ }
+
+ @Override
+ public int getTotalSourcePartitions() {
+ return totalSourcePartitions;
}
- return watermarkMillis;
}
@Override
@@ -112,12 +156,14 @@ public final class KStreamsPayload<T> {
KStreamsPayload<?> that = (KStreamsPayload<?>) o;
return kind == that.kind
&& watermarkMillis == that.watermarkMillis
+ && sourcePartition == that.sourcePartition
+ && totalSourcePartitions == that.totalSourcePartitions
&& Objects.equals(data, that.data);
}
@Override
public int hashCode() {
- return Objects.hash(kind, data, watermarkMillis);
+ return Objects.hash(kind, data, watermarkMillis, sourcePartition,
totalSourcePartitions);
}
@Override
@@ -126,7 +172,10 @@ public final class KStreamsPayload<T> {
if (kind == Kind.DATA) {
helper.add("data", data);
} else {
- helper.add("watermarkMillis", watermarkMillis);
+ helper
+ .add("watermarkMillis", watermarkMillis)
+ .add("sourcePartition", sourcePartition)
+ .add("totalSourcePartitions", totalSourcePartitions);
}
return helper.toString();
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java
new file mode 100644
index 00000000000..194be701ece
--- /dev/null
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+/**
+ * The watermark-only view of a {@link KStreamsPayload}, obtained via {@link
+ * KStreamsPayload#asWatermark()}. Keeping the watermark accessors on this
interface — rather than
+ * on {@link KStreamsPayload} itself — means they are only reachable after the
caller has checked
+ * {@link KStreamsPayload#isWatermark()} and narrowed the payload, so there is
no kind check to do
+ * on each accessor.
+ *
+ * <p>A watermark report is the in-band coordination message a downstream
stage's {@link
+ * WatermarkManager} consumes: the watermark value plus which source partition
reported it and how
+ * many source partitions feed the stage in total.
+ */
+public interface WatermarkPayload {
+
+ /** The reported watermark, in event-time milliseconds. */
+ long getWatermarkMillis();
+
+ /** The source partition this report is for, in {@code [0,
getTotalSourcePartitions())}. */
+ int getSourcePartition();
+
+ /** The total number of source partitions feeding the downstream stage. */
+ int getTotalSourcePartitions();
+}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java
index 30000d77d86..0c576d27122 100644
---
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java
@@ -97,7 +97,7 @@ public class KafkaStreamsRunnerTest {
assertThat(capture.received.get(0).getData().getValue().length, is(0));
assertThat(capture.received.get(1).isWatermark(), is(true));
assertThat(
- capture.received.get(1).getWatermarkMillis(),
+ capture.received.get(1).asWatermark().getWatermarkMillis(),
is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
new file mode 100644
index 00000000000..fc5797a12c2
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
+import org.apache.kafka.streams.processor.api.MockProcessorContext;
+import org.apache.kafka.streams.processor.api.Record;
+import org.junit.Test;
+
+/**
+ * Tests the watermark wiring of {@link ExecutableStageProcessor}: how it
feeds incoming watermark
+ * reports to the {@link WatermarkManager} and forwards the stage's output
watermark.
+ *
+ * <p>Only the watermark path is exercised, so the SDK harness is never
started (it is created
+ * lazily on the first data element). A {@link MockProcessorContext} captures
what the processor
+ * forwards downstream.
+ */
+public class ExecutableStageProcessorWatermarkTest {
+
+ private static ExecutableStageProcessor newProcessor() {
+ JobInfo jobInfo =
+ JobInfo.create(
+ "job-id",
+ "job-name",
+ "",
+
PipelineOptionsTranslation.toProto(PipelineOptionsFactory.create()));
+ return new ExecutableStageProcessor(
+ RunnerApi.ExecutableStagePayload.getDefaultInstance(), jobInfo);
+ }
+
+ private static Record<byte[], KStreamsPayload<?>> watermark(
+ long millis, int sourcePartition, int totalSourcePartitions) {
+ KStreamsPayload<?> payload =
+ KStreamsPayload.watermark(millis, sourcePartition,
totalSourcePartitions);
+ return new Record<>(new byte[0], payload, 0L);
+ }
+
+ private static KStreamsPayload<?> onlyForwarded(
+ MockProcessorContext<byte[], KStreamsPayload<?>> ctx) {
+ assertThat(ctx.forwarded().size(), is(1));
+ return ctx.forwarded().get(0).record().value();
+ }
+
+ @Test
+ public void singleSourcePartitionForwardsImmediatelyStampedAsItsOwnSource() {
+ MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
+ ExecutableStageProcessor processor = newProcessor();
+ processor.init(ctx);
+
+ processor.process(watermark(100L, 0, 1));
+
+ KStreamsPayload<?> out = onlyForwarded(ctx);
+ assertThat(out.isWatermark(), is(true));
+ WatermarkPayload report = out.asWatermark();
+ assertThat(report.getWatermarkMillis(), is(100L));
+ // The stage forwards as its own single source (0 of 1), not the
upstream's identity.
+ assertThat(report.getSourcePartition(), is(0));
+ assertThat(report.getTotalSourcePartitions(), is(1));
+ }
+
+ @Test
+ public void holdsUntilAllSourcePartitionsReportThenForwardsMin() {
+ MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
+ ExecutableStageProcessor processor = newProcessor();
+ processor.init(ctx);
+
+ processor.process(watermark(300L, 0, 3));
+ processor.process(watermark(100L, 1, 3));
+ // Two of three source partitions reported — still holding, nothing
forwarded.
+ assertThat(ctx.forwarded().isEmpty(), is(true));
+
+ processor.process(watermark(500L, 2, 3));
+ // All three reported — forward min(300, 100, 500) = 100.
+ assertThat(onlyForwarded(ctx).asWatermark().getWatermarkMillis(),
is(100L));
+ }
+
+ @Test
+ public void doesNotReforwardWhenWatermarkDoesNotAdvance() {
+ MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
+ ExecutableStageProcessor processor = newProcessor();
+ processor.init(ctx);
+
+ processor.process(watermark(100L, 0, 1));
+ assertThat(ctx.forwarded().size(), is(1));
+
+ // A repeated, non-advancing watermark must not be forwarded again.
+ processor.process(watermark(100L, 0, 1));
+ assertThat(ctx.forwarded().size(), is(1));
+ }
+
+ @Test
+ public void forwardsTerminalMaxWatermark() {
+ long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis();
+ MockProcessorContext<byte[], KStreamsPayload<?>> ctx = new
MockProcessorContext<>();
+ ExecutableStageProcessor processor = newProcessor();
+ processor.init(ctx);
+
+ processor.process(watermark(maxMillis, 0, 1));
+
+ assertThat(onlyForwarded(ctx).asWatermark().getWatermarkMillis(),
is(maxMillis));
+ }
+}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java
index a092b10e02c..f15e818f43e 100644
---
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java
@@ -77,7 +77,8 @@ public class ImpulseTranslatorTest {
KStreamsPayload<byte[]> watermarkPayload = capture.received.get(1);
assertThat(watermarkPayload.isWatermark(), is(true));
assertThat(
- watermarkPayload.getWatermarkMillis(),
is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
+ watermarkPayload.asWatermark().getWatermarkMillis(),
+ is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
}
@Test
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java
new file mode 100644
index 00000000000..e4bdbe4d973
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Properties;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.PortablePipelineOptions;
+import org.apache.beam.sdk.testing.CrashingRunner;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.Impulse;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.util.construction.Environments;
+import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
+import org.apache.beam.sdk.util.construction.PipelineTranslation;
+import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.Topology;
+import org.apache.kafka.streams.TopologyDescription;
+import org.apache.kafka.streams.TopologyTestDriver;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.Record;
+import org.junit.Test;
+
+/**
+ * End-to-end test that a watermark propagates through the topology: {@code
Impulse ->
+ * ExecutableStage -> a recording sink}. The Impulse source emits a terminal
{@code
+ * TIMESTAMP_MAX_VALUE} watermark report; the ExecutableStage routes it
through its {@link
+ * WatermarkManager} and forwards it on, stamped as its own single source
partition. A sink
+ * processor attached to the leaf captures the forwarded watermark and the
test asserts on it.
+ */
+public class WatermarkPropagationTest {
+
+ private static final String APPLICATION_ID = "ks-watermark-propagation-test";
+
+ /** Identity DoFn so the pipeline contains a fused ExecutableStage. */
+ private static class IdentityFn extends DoFn<byte[], byte[]> {
+ @ProcessElement
+ public void processElement(@Element byte[] input, OutputReceiver<byte[]>
out) {
+ out.output(input);
+ }
+ }
+
+ /** Sink processor that records the watermark payloads it is forwarded. */
+ private static final class WatermarkCapture
+ implements Processor<byte[], KStreamsPayload<?>, Void, Void> {
+ private final List<KStreamsPayload<?>> watermarks;
+
+ WatermarkCapture(List<KStreamsPayload<?>> watermarks) {
+ this.watermarks = watermarks;
+ }
+
+ @Override
+ public void process(Record<byte[], KStreamsPayload<?>> record) {
+ if (record.value().isWatermark()) {
+ watermarks.add(record.value());
+ }
+ }
+ }
+
+ @Test
+ public void terminalWatermarkPropagatesToDownstreamStampedAsSingleSource()
throws Exception {
+ Pipeline pipeline = Pipeline.create(pipelineOptions());
+ pipeline.apply("impulse", Impulse.create()).apply("identity", ParDo.of(new
IdentityFn()));
+
+ RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline);
+ KafkaStreamsPipelineOptions options =
+ pipeline.getOptions().as(KafkaStreamsPipelineOptions.class);
+ KafkaStreamsPipelineTranslator translator = new
KafkaStreamsPipelineTranslator();
+ JobInfo jobInfo =
+ JobInfo.create(
+ APPLICATION_ID, options.getJobName(), "",
PipelineOptionsTranslation.toProto(options));
+ KafkaStreamsTranslationContext context =
translator.createTranslationContext(jobInfo, options);
+ translator.translate(context,
translator.prepareForTranslation(pipelineProto));
+
+ // Attach a sink to the leaf ExecutableStage processor to capture the
watermark it forwards.
+ Topology topology = context.getTopology();
+ List<KStreamsPayload<?>> captured = new ArrayList<>();
+ topology.addProcessor(
+ "watermark-capture", () -> new WatermarkCapture(captured),
findLeafProcessor(topology));
+
+ try (TopologyTestDriver driver = new TopologyTestDriver(topology,
streamsConfig())) {
+ driver.advanceWallClockTime(Duration.ofSeconds(1));
+ driver.advanceWallClockTime(Duration.ofSeconds(1));
+ }
+
+ assertThat("a watermark reached the downstream sink", captured.isEmpty(),
is(false));
+ WatermarkPayload terminal = captured.get(captured.size() -
1).asWatermark();
+ assertThat(terminal.getWatermarkMillis(),
is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
+ assertThat(terminal.getSourcePartition(), is(0));
+ assertThat(terminal.getTotalSourcePartitions(), is(1));
+ }
+
+ /**
+ * Returns the name of the single processor node with no successors (the
leaf of the topology).
+ */
+ private static String findLeafProcessor(Topology topology) {
+ for (TopologyDescription.Subtopology subtopology :
topology.describe().subtopologies()) {
+ for (TopologyDescription.Node node : subtopology.nodes()) {
+ if (node instanceof TopologyDescription.Processor &&
node.successors().isEmpty()) {
+ return node.name();
+ }
+ }
+ }
+ throw new IllegalStateException("no leaf processor found in topology");
+ }
+
+ private static PipelineOptions pipelineOptions() {
+ PipelineOptions options =
+ PipelineOptionsFactory.fromArgs("--applicationId=" +
APPLICATION_ID).create();
+ options.setRunner(CrashingRunner.class);
+
options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID);
+ options
+ .as(PortablePipelineOptions.class)
+ .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED);
+ return options;
+ }
+
+ private static Properties streamsConfig() {
+ Properties props = new Properties();
+ props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID);
+ props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
+ props.put(
+ StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG,
Serdes.ByteArray().getClass().getName());
+ props.put(
+ StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG,
Serdes.ByteArray().getClass().getName());
+ return props;
+ }
+}
