diff --git a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto index 4dadf2d14b70..43428ffb7f52 100644 --- a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto +++ b/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; } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java new file mode 100644 index 000000000000..b00902cbdfdc --- /dev/null +++ b/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()}. + * + *

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 getTargetPartitions(); +} 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 647c64953b06..e31e3719aaf3 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 @@ -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. * *

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 { private enum Kind { DATA, - WATERMARK + WATERMARK, + FLUSH } private final Kind kind; @@ -47,6 +50,7 @@ private enum Kind { private final String transformId; private final int sourcePartition; private final int totalSourcePartitions; + private final Set targetPartitions; private KStreamsPayload( Kind kind, @@ -54,18 +58,20 @@ private KStreamsPayload( long watermarkMillis, String transformId, int sourcePartition, - int totalSourcePartitions) { + int totalSourcePartitions, + Set 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 KStreamsPayload data(WindowedValue 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 KStreamsPayload 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 KStreamsPayload flush(Set 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 getTargetPartitions() { + return targetPartitions; + } + } + /** {@link WatermarkPayload} view backed by this payload's fields. */ private final class WatermarkView implements WatermarkPayload { @Override @@ -158,13 +202,20 @@ 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 @@ -172,6 +223,10 @@ 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) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java index 1363740d58bd..1f39a9c424ce 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java +++ b/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 @@ * *

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). * *

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 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 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( diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java index 95ce70b88578..f0cd19a781c9 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java +++ b/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 roundTrip(KStreamsPayload payload) { return deserializer.deserialize(TOPIC, serializer.serialize(TOPIC, payload)); } + @Test + public void roundTripsFlushPayload() { + KStreamsPayload payload = KStreamsPayload.flush(ImmutableSet.of(0, 3, 7)); + KStreamsPayload 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 flush = KStreamsPayload.flush(ImmutableSet.of(0)); + assertThrows(IllegalStateException.class, flush::asWatermark); + assertThrows(IllegalStateException.class, flush::getData); + KStreamsPayload data = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(1)); + assertThrows(IllegalStateException.class, data::asFlush); + } + @Test public void roundTripsDataPayload() { KStreamsPayload payload = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(42));