diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java index b8eb5413b27e..00f85032d041 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java @@ -28,6 +28,7 @@ import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.fnexecution.state.StateRequestHandler; import org.apache.beam.sdk.fn.data.FnDataReceiver; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.util.construction.graph.ExecutableStage; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; @@ -35,6 +36,7 @@ import org.apache.kafka.streams.processor.api.ProcessorContext; import org.apache.kafka.streams.processor.api.Record; import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -49,9 +51,12 @@ * ProcessorContext#forward} must only be called from the processing thread, so outputs are never * forwarded directly from a harness callback. * - *

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

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

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

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

* *

The envelope lets a single Kafka Streams output channel carry both Beam data and the watermark @@ -54,21 +57,45 @@ private enum Kind { private final Kind kind; private final @Nullable WindowedValue data; private final long watermarkMillis; + private final int sourcePartition; + private final int totalSourcePartitions; - private KStreamsPayload(Kind kind, @Nullable WindowedValue data, long watermarkMillis) { + private KStreamsPayload( + Kind kind, + @Nullable WindowedValue data, + long watermarkMillis, + int sourcePartition, + int totalSourcePartitions) { this.kind = kind; this.data = data; this.watermarkMillis = watermarkMillis; + this.sourcePartition = sourcePartition; + this.totalSourcePartitions = totalSourcePartitions; } /** Returns a data payload wrapping the given {@link WindowedValue}. */ public static KStreamsPayload data(WindowedValue value) { - return new KStreamsPayload<>(Kind.DATA, value, 0L); + return new KStreamsPayload<>(Kind.DATA, value, 0L, 0, 0); } - /** Returns a watermark payload carrying the given event-time milliseconds. */ - public static KStreamsPayload watermark(long watermarkMillis) { - return new KStreamsPayload<>(Kind.WATERMARK, null, watermarkMillis); + /** + * Returns a watermark report payload: the event-time milliseconds together with the in-band + * coordination fields the downstream stage's {@link WatermarkManager} needs — which source + * partition this report is for and how many source partitions feed the stage in total. + */ + public static KStreamsPayload watermark( + long watermarkMillis, int sourcePartition, int totalSourcePartitions) { + Preconditions.checkArgument( + totalSourcePartitions > 0, + "totalSourcePartitions must be positive: %s", + totalSourcePartitions); + Preconditions.checkArgument( + sourcePartition >= 0 && sourcePartition < totalSourcePartitions, + "sourcePartition %s out of range for totalSourcePartitions %s", + sourcePartition, + totalSourcePartitions); + return new KStreamsPayload<>( + Kind.WATERMARK, null, watermarkMillis, sourcePartition, totalSourcePartitions); } public boolean isData() { @@ -91,14 +118,31 @@ public WindowedValue getData() { } /** - * Returns the watermark event-time milliseconds. Caller must check {@link #isWatermark()} first; - * calling this on a data payload throws. + * Narrows this payload to its {@link WatermarkPayload} view, through which the watermark report + * fields are read. Caller must check {@link #isWatermark()} first; calling this on a data payload + * throws. */ - public long getWatermarkMillis() { - if (kind != Kind.WATERMARK) { - throw new IllegalStateException("Payload is not a watermark: kind=" + kind); + public WatermarkPayload asWatermark() { + Preconditions.checkState(isWatermark(), "Payload is not a watermark: kind=%s", kind); + return new WatermarkView(); + } + + /** {@link WatermarkPayload} view backed by this payload's fields. */ + private final class WatermarkView implements WatermarkPayload { + @Override + public long getWatermarkMillis() { + return watermarkMillis; + } + + @Override + public int getSourcePartition() { + return sourcePartition; + } + + @Override + public int getTotalSourcePartitions() { + return totalSourcePartitions; } - return watermarkMillis; } @Override @@ -112,12 +156,14 @@ public boolean equals(@Nullable Object o) { KStreamsPayload that = (KStreamsPayload) o; return kind == that.kind && watermarkMillis == that.watermarkMillis + && sourcePartition == that.sourcePartition + && totalSourcePartitions == that.totalSourcePartitions && Objects.equals(data, that.data); } @Override public int hashCode() { - return Objects.hash(kind, data, watermarkMillis); + return Objects.hash(kind, data, watermarkMillis, sourcePartition, totalSourcePartitions); } @Override @@ -126,7 +172,10 @@ public String toString() { if (kind == Kind.DATA) { helper.add("data", data); } else { - helper.add("watermarkMillis", watermarkMillis); + helper + .add("watermarkMillis", watermarkMillis) + .add("sourcePartition", sourcePartition) + .add("totalSourcePartitions", totalSourcePartitions); } return helper.toString(); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java new file mode 100644 index 000000000000..194be701ecec --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java @@ -0,0 +1,41 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.kafka.streams.translation; + +/** + * The watermark-only view of a {@link KStreamsPayload}, obtained via {@link + * KStreamsPayload#asWatermark()}. Keeping the watermark accessors on this interface — rather than + * on {@link KStreamsPayload} itself — means they are only reachable after the caller has checked + * {@link KStreamsPayload#isWatermark()} and narrowed the payload, so there is no kind check to do + * on each accessor. + * + *

A watermark report is the in-band coordination message a downstream stage's {@link + * WatermarkManager} consumes: the watermark value plus which source partition reported it and how + * many source partitions feed the stage in total. + */ +public interface WatermarkPayload { + + /** The reported watermark, in event-time milliseconds. */ + long getWatermarkMillis(); + + /** The source partition this report is for, in {@code [0, getTotalSourcePartitions())}. */ + int getSourcePartition(); + + /** The total number of source partitions feeding the downstream stage. */ + int getTotalSourcePartitions(); +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java index 30000d77d864..0c576d27122f 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerTest.java @@ -97,7 +97,7 @@ public void impulseOnlyPipelineEmitsDataAndTerminalWatermark() { assertThat(capture.received.get(0).getData().getValue().length, is(0)); assertThat(capture.received.get(1).isWatermark(), is(true)); assertThat( - capture.received.get(1).getWatermarkMillis(), + capture.received.get(1).asWatermark().getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java new file mode 100644 index 000000000000..fc5797a12c26 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java @@ -0,0 +1,124 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.kafka.streams.translation; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.kafka.streams.processor.api.MockProcessorContext; +import org.apache.kafka.streams.processor.api.Record; +import org.junit.Test; + +/** + * Tests the watermark wiring of {@link ExecutableStageProcessor}: how it feeds incoming watermark + * reports to the {@link WatermarkManager} and forwards the stage's output watermark. + * + *

Only the watermark path is exercised, so the SDK harness is never started (it is created + * lazily on the first data element). A {@link MockProcessorContext} captures what the processor + * forwards downstream. + */ +public class ExecutableStageProcessorWatermarkTest { + + private static ExecutableStageProcessor newProcessor() { + JobInfo jobInfo = + JobInfo.create( + "job-id", + "job-name", + "", + PipelineOptionsTranslation.toProto(PipelineOptionsFactory.create())); + return new ExecutableStageProcessor( + RunnerApi.ExecutableStagePayload.getDefaultInstance(), jobInfo); + } + + private static Record> watermark( + long millis, int sourcePartition, int totalSourcePartitions) { + KStreamsPayload payload = + KStreamsPayload.watermark(millis, sourcePartition, totalSourcePartitions); + return new Record<>(new byte[0], payload, 0L); + } + + private static KStreamsPayload onlyForwarded( + MockProcessorContext> ctx) { + assertThat(ctx.forwarded().size(), is(1)); + return ctx.forwarded().get(0).record().value(); + } + + @Test + public void singleSourcePartitionForwardsImmediatelyStampedAsItsOwnSource() { + MockProcessorContext> ctx = new MockProcessorContext<>(); + ExecutableStageProcessor processor = newProcessor(); + processor.init(ctx); + + processor.process(watermark(100L, 0, 1)); + + KStreamsPayload out = onlyForwarded(ctx); + assertThat(out.isWatermark(), is(true)); + WatermarkPayload report = out.asWatermark(); + assertThat(report.getWatermarkMillis(), is(100L)); + // The stage forwards as its own single source (0 of 1), not the upstream's identity. + assertThat(report.getSourcePartition(), is(0)); + assertThat(report.getTotalSourcePartitions(), is(1)); + } + + @Test + public void holdsUntilAllSourcePartitionsReportThenForwardsMin() { + MockProcessorContext> ctx = new MockProcessorContext<>(); + ExecutableStageProcessor processor = newProcessor(); + processor.init(ctx); + + processor.process(watermark(300L, 0, 3)); + processor.process(watermark(100L, 1, 3)); + // Two of three source partitions reported — still holding, nothing forwarded. + assertThat(ctx.forwarded().isEmpty(), is(true)); + + processor.process(watermark(500L, 2, 3)); + // All three reported — forward min(300, 100, 500) = 100. + assertThat(onlyForwarded(ctx).asWatermark().getWatermarkMillis(), is(100L)); + } + + @Test + public void doesNotReforwardWhenWatermarkDoesNotAdvance() { + MockProcessorContext> ctx = new MockProcessorContext<>(); + ExecutableStageProcessor processor = newProcessor(); + processor.init(ctx); + + processor.process(watermark(100L, 0, 1)); + assertThat(ctx.forwarded().size(), is(1)); + + // A repeated, non-advancing watermark must not be forwarded again. + processor.process(watermark(100L, 0, 1)); + assertThat(ctx.forwarded().size(), is(1)); + } + + @Test + public void forwardsTerminalMaxWatermark() { + long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); + MockProcessorContext> ctx = new MockProcessorContext<>(); + ExecutableStageProcessor processor = newProcessor(); + processor.init(ctx); + + processor.process(watermark(maxMillis, 0, 1)); + + assertThat(onlyForwarded(ctx).asWatermark().getWatermarkMillis(), is(maxMillis)); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java index a092b10e02c4..f15e818f43ed 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ImpulseTranslatorTest.java @@ -77,7 +77,8 @@ public void impulseEmitsDataElementFollowedByTerminalWatermark() { KStreamsPayload watermarkPayload = capture.received.get(1); assertThat(watermarkPayload.isWatermark(), is(true)); assertThat( - watermarkPayload.getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); + watermarkPayload.asWatermark().getWatermarkMillis(), + is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); } @Test diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java new file mode 100644 index 000000000000..e4bdbe4d9733 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPropagationTest.java @@ -0,0 +1,155 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.runners.kafka.streams.translation; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Properties; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.runners.fnexecution.provisioning.JobInfo; +import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.PortablePipelineOptions; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.Impulse; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; +import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.TopologyDescription; +import org.apache.kafka.streams.TopologyTestDriver; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.Record; +import org.junit.Test; + +/** + * End-to-end test that a watermark propagates through the topology: {@code Impulse -> + * ExecutableStage -> a recording sink}. The Impulse source emits a terminal {@code + * TIMESTAMP_MAX_VALUE} watermark report; the ExecutableStage routes it through its {@link + * WatermarkManager} and forwards it on, stamped as its own single source partition. A sink + * processor attached to the leaf captures the forwarded watermark and the test asserts on it. + */ +public class WatermarkPropagationTest { + + private static final String APPLICATION_ID = "ks-watermark-propagation-test"; + + /** Identity DoFn so the pipeline contains a fused ExecutableStage. */ + private static class IdentityFn extends DoFn { + @ProcessElement + public void processElement(@Element byte[] input, OutputReceiver out) { + out.output(input); + } + } + + /** Sink processor that records the watermark payloads it is forwarded. */ + private static final class WatermarkCapture + implements Processor, Void, Void> { + private final List> watermarks; + + WatermarkCapture(List> watermarks) { + this.watermarks = watermarks; + } + + @Override + public void process(Record> record) { + if (record.value().isWatermark()) { + watermarks.add(record.value()); + } + } + } + + @Test + public void terminalWatermarkPropagatesToDownstreamStampedAsSingleSource() throws Exception { + Pipeline pipeline = Pipeline.create(pipelineOptions()); + pipeline.apply("impulse", Impulse.create()).apply("identity", ParDo.of(new IdentityFn())); + + RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); + KafkaStreamsPipelineOptions options = + pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); + JobInfo jobInfo = + JobInfo.create( + APPLICATION_ID, options.getJobName(), "", PipelineOptionsTranslation.toProto(options)); + KafkaStreamsTranslationContext context = translator.createTranslationContext(jobInfo, options); + translator.translate(context, translator.prepareForTranslation(pipelineProto)); + + // Attach a sink to the leaf ExecutableStage processor to capture the watermark it forwards. + Topology topology = context.getTopology(); + List> captured = new ArrayList<>(); + topology.addProcessor( + "watermark-capture", () -> new WatermarkCapture(captured), findLeafProcessor(topology)); + + try (TopologyTestDriver driver = new TopologyTestDriver(topology, streamsConfig())) { + driver.advanceWallClockTime(Duration.ofSeconds(1)); + driver.advanceWallClockTime(Duration.ofSeconds(1)); + } + + assertThat("a watermark reached the downstream sink", captured.isEmpty(), is(false)); + WatermarkPayload terminal = captured.get(captured.size() - 1).asWatermark(); + assertThat(terminal.getWatermarkMillis(), is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); + assertThat(terminal.getSourcePartition(), is(0)); + assertThat(terminal.getTotalSourcePartitions(), is(1)); + } + + /** + * Returns the name of the single processor node with no successors (the leaf of the topology). + */ + private static String findLeafProcessor(Topology topology) { + for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { + for (TopologyDescription.Node node : subtopology.nodes()) { + if (node instanceof TopologyDescription.Processor && node.successors().isEmpty()) { + return node.name(); + } + } + } + throw new IllegalStateException("no leaf processor found in topology"); + } + + private static PipelineOptions pipelineOptions() { + PipelineOptions options = + PipelineOptionsFactory.fromArgs("--applicationId=" + APPLICATION_ID).create(); + options.setRunner(CrashingRunner.class); + options.as(KafkaStreamsPipelineOptions.class).setApplicationId(APPLICATION_ID); + options + .as(PortablePipelineOptions.class) + .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); + return options; + } + + private static Properties streamsConfig() { + Properties props = new Properties(); + props.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID); + props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092"); + props.put( + StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + props.put( + StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.ByteArray().getClass().getName()); + return props; + } +}