Kafka Streams runner: add a flush marker payload variant

runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto

@@ -50,9 +50,19 @@ message KafkaStreamsPayload {
     bytes value = 1;
   }
 
-  // Exactly one variant is set; the oneof case discriminates data vs watermark.
+  // A request to close the open bundle and flush its output. See
+  // https://github.com/apache/beam/issues/39633.
+  message FlushPayload {
+    // Which partitions of the repartition topic this marker is for. The producer picks them so
+    // that each downstream partition gets exactly one flush per interval; broadcasting would give
+    // it one per upstream partition instead.
+    repeated uint32 target_partitions = 1;
+  }
+
+  // Exactly one variant is set; the oneof case discriminates data vs watermark vs flush.
   oneof payload {
     WatermarkPayload watermark = 1;
     DataPayload data = 2;
+    FlushPayload flush = 3;
   }
 }

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

@@ -0,0 +1,37 @@
+/*
+ * 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 java.util.Set;
+
+/**
+ * The flush-only view of a {@link KStreamsPayload}, obtained via {@link KStreamsPayload#asFlush()}.
+ *
+ * <p>A flush marker asks the stage that receives it to close its bundle, which is how a bundle is
+ * bounded in time. It arrives as a record so the bundle is closed from {@code process()} rather
+ * than from a punctuator; see https://github.com/apache/beam/issues/39633.
+ */
+public interface FlushPayload {
+
+  /**
+   * The repartition-topic partitions this marker is addressed to. Never empty: a producer with
+   * nothing to address emits no marker. The receiving stage ignores this; only the partitioner
+   * reads it.
+   */
+  Set<Integer> getTargetPartitions();
+}

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

@@ -18,16 +18,18 @@
 package org.apache.beam.runners.kafka.streams.translation;
 
 import java.util.Objects;
+import java.util.Set;
 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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
- * Envelope for every record value passed between the runner's processors. It is either a {@link
- * #isData() data} element wrapping a {@link WindowedValue}, or a {@link #isWatermark() watermark}
- * report carrying an event time plus the partition fields the downstream {@link WatermarkManager}
- * needs.
+ * Envelope for every record value passed between the runner's processors. It is a {@link #isData()
+ * data} element wrapping a {@link WindowedValue}, a {@link #isWatermark() watermark} report
+ * carrying an event time plus the partition fields the downstream {@link WatermarkManager} needs,
+ * or a {@link #isFlush() flush} marker asking the receiving stage to close its bundle.
  *
  * <p>One channel therefore carries both Beam data and the watermark coordination Kafka Streams has
  * no notion of. Across topic boundaries it is encoded by {@link KStreamsPayloadSerde}.
@@ -38,7 +40,8 @@ public final class KStreamsPayload<T> {
 
   private enum Kind {
     DATA,
-    WATERMARK
+    WATERMARK,
+    FLUSH
   }
 
   private final Kind kind;
@@ -47,25 +50,28 @@ private enum Kind {
   private final String transformId;
   private final int sourcePartition;
   private final int totalSourcePartitions;
+  private final Set<Integer> targetPartitions;
 
   private KStreamsPayload(
       Kind kind,
       @Nullable WindowedValue<T> data,
       long watermarkMillis,
       String transformId,
       int sourcePartition,
-      int totalSourcePartitions) {
+      int totalSourcePartitions,
+      Set<Integer> targetPartitions) {
     this.kind = kind;
     this.data = data;
     this.watermarkMillis = watermarkMillis;
     this.transformId = transformId;
     this.sourcePartition = sourcePartition;
     this.totalSourcePartitions = totalSourcePartitions;
+    this.targetPartitions = targetPartitions;
   }
 
   /** 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, "", 0, 0);
+    return new KStreamsPayload<>(Kind.DATA, value, 0L, "", 0, 0, ImmutableSet.of());
   }
 
   /**
@@ -89,7 +95,24 @@ public static <T> KStreamsPayload<T> watermark(
         sourcePartition,
         totalSourcePartitions);
     return new KStreamsPayload<>(
-        Kind.WATERMARK, null, watermarkMillis, transformId, sourcePartition, totalSourcePartitions);
+        Kind.WATERMARK,
+        null,
+        watermarkMillis,
+        transformId,
+        sourcePartition,
+        totalSourcePartitions,
+        ImmutableSet.of());
+  }
+
+  /**
+   * Returns a flush marker addressed to the given repartition-topic partitions. A producer with no
+   * partitions to address emits no marker, so the set must not be empty.
+   */
+  public static <T> KStreamsPayload<T> flush(Set<Integer> targetPartitions) {
+    Preconditions.checkArgument(
+        !targetPartitions.isEmpty(), "flush marker must target at least one partition");
+    return new KStreamsPayload<>(
+        Kind.FLUSH, null, 0L, "", 0, 0, ImmutableSet.copyOf(targetPartitions));
   }
 
   public boolean isData() {
@@ -100,6 +123,10 @@ public boolean isWatermark() {
     return kind == Kind.WATERMARK;
   }
 
+  public boolean isFlush() {
+    return kind == Kind.FLUSH;
+  }
+
   /**
    * Returns the wrapped data element. Caller must check {@link #isData()} first; calling this on a
    * watermark payload throws.
@@ -121,6 +148,23 @@ public WatermarkPayload asWatermark() {
     return new WatermarkView();
   }
 
+  /**
+   * Narrows this payload to its {@link FlushPayload} view. Caller must check {@link #isFlush()}
+   * first; calling this on any other payload throws.
+   */
+  public FlushPayload asFlush() {
+    Preconditions.checkState(isFlush(), "Payload is not a flush marker: kind=%s", kind);
+    return new FlushView();
+  }
+
+  /** {@link FlushPayload} view backed by this payload's fields. */
+  private final class FlushView implements FlushPayload {
+    @Override
+    public Set<Integer> getTargetPartitions() {
+      return targetPartitions;
+    }
+  }
+
   /** {@link WatermarkPayload} view backed by this payload's fields. */
   private final class WatermarkView implements WatermarkPayload {
     @Override
@@ -158,20 +202,31 @@ public boolean equals(@Nullable Object o) {
         && transformId.equals(that.transformId)
         && sourcePartition == that.sourcePartition
         && totalSourcePartitions == that.totalSourcePartitions
+        && targetPartitions.equals(that.targetPartitions)
         && Objects.equals(data, that.data);
   }
 
   @Override
   public int hashCode() {
     return Objects.hash(
-        kind, data, watermarkMillis, transformId, sourcePartition, totalSourcePartitions);
+        kind,
+        data,
+        watermarkMillis,
+        transformId,
+        sourcePartition,
+        totalSourcePartitions,
+        targetPartitions);
   }
 
   @Override
   public String toString() {
     MoreObjects.ToStringHelper helper = MoreObjects.toStringHelper(this).add("kind", kind);
     if (kind == Kind.DATA) {
       helper.add("data", data);
+    } else if (kind == Kind.FLUSH) {
+      helper
+          .add("sourcePartition", sourcePartition)
+          .add("totalSourcePartitions", totalSourcePartitions);
     } else {
       helper
           .add("watermarkMillis", watermarkMillis)

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

@@ -24,6 +24,7 @@
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
 import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.Serde;
@@ -36,9 +37,9 @@
  *
  * <p>The wire form is the {@link KafkaStreamsPayload} protobuf message — protobuf gives compatible
  * schema evolution and compact varint encoding. The data variant carries the {@link WindowedValue}
- * encoded with the {@link Coder} supplied for the topic's PCollection; the watermark variant
- * carries the coder-independent watermark report. A {@link KStreamsPayloadSerde} is therefore
- * parameterized by the data {@link Coder} (different topics carry different element types).
+ * encoded with the {@link Coder} supplied for the topic's PCollection; the watermark and flush
+ * variants are coder-independent. A {@link KStreamsPayloadSerde} is therefore parameterized by the
+ * data {@link Coder} (different topics carry different element types).
  *
  * <p>The serde assumes non-null payloads: the topics it is used on (repartition and watermark
  * fan-out) are not log-compacted, so no tombstone (null-valued) records occur.
@@ -77,6 +78,10 @@ public byte[] serialize(String topic, KStreamsPayload<T> payload) {
         proto.setData(
             KafkaStreamsPayload.DataPayload.newBuilder()
                 .setValue(ByteString.copyFrom(encoded.toByteArray())));
+      } else if (payload.isFlush()) {
+        proto.setFlush(
+            KafkaStreamsPayload.FlushPayload.newBuilder()
+                .addAllTargetPartitions(payload.asFlush().getTargetPartitions()));
       } else {
         WatermarkPayload watermark = payload.asWatermark();
         proto.setWatermark(
@@ -113,6 +118,9 @@ public KStreamsPayload<T> deserialize(String topic, byte[] bytes) {
               watermark.getTransformId(),
               watermark.getSourcePartition(),
               watermark.getTotalPartitions());
+        case FLUSH:
+          return KStreamsPayload.flush(
+              ImmutableSet.copyOf(proto.getFlush().getTargetPartitionsList()));
         case PAYLOAD_NOT_SET:
         default:
           throw new SerializationException(

runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java

@@ -27,6 +27,7 @@
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.Serializer;
@@ -47,6 +48,33 @@ private KStreamsPayload<Integer> roundTrip(KStreamsPayload<Integer> payload) {
     return deserializer.deserialize(TOPIC, serializer.serialize(TOPIC, payload));
   }
 
+  @Test
+  public void roundTripsFlushPayload() {
+    KStreamsPayload<Integer> payload = KStreamsPayload.flush(ImmutableSet.of(0, 3, 7));
+    KStreamsPayload<Integer> out = roundTrip(payload);
+    assertThat(out.isFlush(), is(true));
+    assertThat(out.isData(), is(false));
+    assertThat(out.isWatermark(), is(false));
+    assertThat(out.asFlush().getTargetPartitions(), is(ImmutableSet.of(0, 3, 7)));
+    assertThat(out, is(payload));
+  }
+
+  @Test
+  public void aFlushPayloadMustTargetSomething() {
+    // A producer with nothing to address emits no marker at all, so an empty set is a bug rather
+    // than a case to encode.
+    assertThrows(IllegalArgumentException.class, () -> KStreamsPayload.flush(ImmutableSet.of()));
+  }
+
+  @Test
+  public void aFlushPayloadIsNotAWatermarkOrData() {
+    KStreamsPayload<Integer> flush = KStreamsPayload.flush(ImmutableSet.of(0));
+    assertThrows(IllegalStateException.class, flush::asWatermark);
+    assertThrows(IllegalStateException.class, flush::getData);
+    KStreamsPayload<Integer> data = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(1));
+    assertThrows(IllegalStateException.class, data::asFlush);
+  }
+
   @Test
   public void roundTripsDataPayload() {
     KStreamsPayload<Integer> payload = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(42));