[GSoC 2026] Kafka Streams runner — WatermarkManager part 2: wire into ExecutableStageProcessor

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

@@ -28,13 +28,15 @@
 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;
 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 @@
  * 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 @@ 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<byte[], KStreamsPayload<?>> record) {
   }
 
   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

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.
+   *
+   * <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. */

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 @@
  *
  * <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 @@ private enum Kind {
   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 WindowedValue<T> 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();
   }

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

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