Kafka Streams runner: target flush markers across repartition topics

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

@@ -47,8 +47,8 @@
  *
  * <p>The Beam key becomes the Kafka record key so Kafka Streams shuffles by it. The topology is a
  * {@link ShuffleByKeyProcessor} that sets that key and passes watermark reports through, a sink to
- * an internal repartition topic using a {@link GroupByKeyBroadcastPartitioner} that hashes data by
- * key and fans watermarks out to every partition, a source reading that topic back, and the {@link
+ * an internal repartition topic using a {@link KStreamsPayloadPartitioner} that hashes data by key
+ * and fans watermarks out to every partition, a source reading that topic back, and the {@link
  * WindowedGroupByKeyProcessor} with its state and timer stores.
  */
 class GroupByKeyTranslator implements PTransformTranslator {
@@ -114,16 +114,20 @@ public void translate(
         shuffleName,
         () ->
             new ShuffleByKeyProcessor(
-                keyCoder, upstreamPartitionCount, shuffleName, context.getTerminationTracker()),
+                keyCoder,
+                upstreamPartitionCount,
+                partitionCount,
+                shuffleName,
+                context.getTerminationTracker()),
         parentProcessor);
 
-    // Shuffle through the repartition topic: data partitioned by key, watermark broadcast.
+    // Shuffle through the repartition topic: data by key, watermarks broadcast, flushes targeted.
     topology.addSink(
         sinkName,
         repartitionTopic,
         Serdes.ByteArray().serializer(),
         payloadSerde.serializer(),
-        new GroupByKeyBroadcastPartitioner<>(),
+        new KStreamsPayloadPartitioner<>(partitionCount),
         shuffleName);
     topology.addSource(
         sourceName,

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyBroadcastPartitioner.java → runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadPartitioner.java

@@ -25,19 +25,27 @@
 import org.apache.kafka.streams.processor.StreamPartitioner;
 
 /**
- * Partitions records on the GroupByKey repartition topic.
+ * Partitions records on a repartition topic.
  *
  * <ul>
  *   <li><b>data</b> records go to the single partition selected by hashing the (already encoded
  *       Beam key) Kafka record key — the same scheme Kafka's default partitioner uses — so every
  *       value of a key lands together;
  *   <li><b>watermark</b> reports are broadcast to <i>every</i> partition, so each downstream
- *       GroupByKey task observes the terminal watermark and fires its keys.
+ *       GroupByKey task observes the terminal watermark and fires its keys;
+ *   <li><b>flush</b> markers go to the partitions they name, one copy to each.
  * </ul>
  *
  * @param <T> the data element type carried by data payloads
  */
-class GroupByKeyBroadcastPartitioner<T> implements StreamPartitioner<byte[], KStreamsPayload<T>> {
+class KStreamsPayloadPartitioner<T> implements StreamPartitioner<byte[], KStreamsPayload<T>> {
+
+  /** The partition count the flush targets were computed for. */
+  private final int expectedPartitions;
+
+  KStreamsPayloadPartitioner(int expectedPartitions) {
+    this.expectedPartitions = expectedPartitions;
+  }
 
   @Override
   public Integer partition(String topic, byte[] key, KStreamsPayload<T> value, int numPartitions) {
@@ -56,6 +64,17 @@ public Optional<Set<Integer>> partitions(
       }
       return Optional.of(all);
     }
+    if (value.isFlush()) {
+      // Targets assume internalParallelism partitions; a topic from an earlier run may differ.
+      if (numPartitions != expectedPartitions) {
+        throw new IllegalStateException(
+            String.format(
+                "Repartition topic %s has %d partitions but the pipeline expects %d. It was"
+                    + " probably created by an earlier run with a different --internalParallelism.",
+                topic, numPartitions, expectedPartitions));
+      }
+      return Optional.of(value.asFlush().getTargetPartitions());
+    }
     if (key == null) {
       // A keyless record has no partition it must go to, so leave the choice to Kafka rather than
       // hashing a null or pinning one partition: an empty Optional tells Kafka Streams no explicit

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

@@ -17,10 +17,12 @@
  */
 package org.apache.beam.runners.kafka.streams.translation;
 
+import java.util.Set;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.CoderException;
 import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.sdk.values.KV;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 import org.apache.kafka.streams.processor.api.Processor;
 import org.apache.kafka.streams.processor.api.ProcessorContext;
 import org.apache.kafka.streams.processor.api.Record;
@@ -39,11 +41,13 @@
  * <p>Watermark reports are relabelled here with the reporting instance's real partition identity.
  * This is the point at which a report stops being delivered in-process and starts crossing a topic:
  * upstream of it a transform forwards to its fused children, which see exactly one instance of it,
- * so the report names a single source. The {@link GroupByKeyBroadcastPartitioner} on the sink below
+ * so the report names a single source. The {@link KStreamsPayloadPartitioner} on the sink below
  * fans each report out to <em>every</em> partition, so a downstream task instead sees a report from
  * every instance of the upstream transform, and has to be able to tell them apart to know when it
  * has heard from all of them. The transform id is left alone, so the report still names the
  * transform that produced it.
+ *
+ * <p>Flush markers are readdressed here too; see {@link #flushTargets}.
  */
 class ShuffleByKeyProcessor
     implements Processor<byte[], KStreamsPayload<?>, byte[], KStreamsPayload<?>> {
@@ -57,6 +61,9 @@ class ShuffleByKeyProcessor
 
   private int upstreamPartition;
 
+  /** Partitions of the repartition topic. */
+  private final int downstreamPartitionCount;
+
   // Reports this shuffle as finished once it has written the terminal watermark to the repartition
   // topic. The downstream side reading that topic reports separately, which is why the pipeline
   // waits for every processor rather than the first.
@@ -67,10 +74,12 @@ class ShuffleByKeyProcessor
   ShuffleByKeyProcessor(
       Coder<Object> keyCoder,
       int upstreamPartitionCount,
+      int downstreamPartitionCount,
       String nodeName,
       TerminationTracker terminationTracker) {
     this.keyCoder = keyCoder;
     this.upstreamPartitionCount = upstreamPartitionCount;
+    this.downstreamPartitionCount = downstreamPartitionCount;
     this.terminationReporter = new TerminationReporter(terminationTracker, nodeName);
   }
 
@@ -114,6 +123,16 @@ public void process(Record<byte[], KStreamsPayload<?>> record) {
         throw new RuntimeException("Failed to encode shuffle key", e);
       }
       ctx.forward(record.withKey(encodedKey));
+    } else if (payload.isFlush()) {
+      // Forwarded once: the sink's partitioner makes Kafka Streams write a copy to each target.
+      // Fanning in, an instance may have nothing to address.
+      Set<Integer> targets =
+          flushTargets(upstreamPartition, upstreamPartitionCount, downstreamPartitionCount);
+      if (!targets.isEmpty()) {
+        ctx.forward(
+            new Record<byte[], KStreamsPayload<?>>(
+                record.key(), KStreamsPayload.flush(targets), record.timestamp()));
+      }
     } else {
       WatermarkPayload report = payload.asWatermark();
       ctx.forward(
@@ -129,6 +148,20 @@ public void process(Record<byte[], KStreamsPayload<?>> record) {
     }
   }
 
+  /**
+   * Downstream partitions {@code [i*D/U, (i+1)*D/U)} for upstream partition {@code i}. Over all
+   * {@code i} these cover each downstream partition once, and do not change on a rebalance.
+   */
+  static Set<Integer> flushTargets(int upstreamPartition, int upstreamCount, int downstreamCount) {
+    long from = (long) upstreamPartition * downstreamCount / upstreamCount;
+    long to = ((long) upstreamPartition + 1) * downstreamCount / upstreamCount;
+    ImmutableSet.Builder<Integer> targets = ImmutableSet.builder();
+    for (long partition = from; partition < to; partition++) {
+      targets.add((int) partition);
+    }
+    return targets.build();
+  }
+
   private static <T> T checkInitialized(@Nullable T value) {
     if (value == null) {
       throw new IllegalStateException("ShuffleByKeyProcessor used before init()");

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

@@ -33,8 +33,8 @@
  * exhausted. This collects those reports and fires a callback when nothing is left to do.
  *
  * <p>No coordination between instances is needed: a watermark crossing a repartition topic is
- * broadcast to every partition (see {@link GroupByKeyBroadcastPartitioner}), so every task observes
- * the terminal watermark itself and all instances reach the same conclusion independently.
+ * broadcast to every partition (see {@link KStreamsPayloadPartitioner}), so every task observes the
+ * terminal watermark itself and all instances reach the same conclusion independently.
  *
  * <p>Every local processor is counted, not just the first. One instance can own tasks from both
  * sides of a repartition topic, and the upstream side goes terminal as soon as it has written to

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

@@ -0,0 +1,176 @@
+/*
+ * 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.nio.file.Files;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.TreeMap;
+import java.util.UUID;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AdminClientConfig;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.Topology;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.testcontainers.kafka.KafkaContainer;
+import org.testcontainers.utility.DockerImageName;
+
+/**
+ * Checks on a real broker that a flush is written once to each partition it names. The {@code
+ * TopologyTestDriver} cannot show this, since it reports every topic as having one partition.
+ */
+@RunWith(JUnit4.class)
+public class KStreamsPayloadPartitionerBrokerIT {
+
+  private static final int PARTITIONS = 4;
+
+  private static KafkaContainer kafka;
+
+  @BeforeClass
+  public static void startBroker() {
+    kafka = new KafkaContainer(DockerImageName.parse("apache/kafka:4.0.0"));
+    kafka.start();
+  }
+
+  @AfterClass
+  public static void stopBroker() {
+    if (kafka != null) {
+      kafka.stop();
+    }
+  }
+
+  @Test
+  public void aFlushIsWrittenOnceToEachPartitionItNames() throws Exception {
+    String id = UUID.randomUUID().toString();
+    String input = "flush-in-" + id;
+    String output = "flush-out-" + id;
+    try (Admin admin =
+        Admin.create(
+            ImmutableMap.<String, Object>of(
+                AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, kafka.getBootstrapServers()))) {
+      admin
+          .createTopics(
+              Arrays.asList(
+                  new NewTopic(input, 1, (short) 1), new NewTopic(output, PARTITIONS, (short) 1)))
+          .all()
+          .get();
+    }
+
+    KStreamsPayloadSerde<Integer> serde =
+        new KStreamsPayloadSerde<>(
+            WindowedValues.getFullCoder(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE));
+    Topology topology = new Topology();
+    topology.addSource("in", new ByteArrayDeserializer(), serde.deserializer(), input);
+    topology.addSink(
+        "out",
+        output,
+        new ByteArraySerializer(),
+        serde.serializer(),
+        new KStreamsPayloadPartitioner<>(PARTITIONS),
+        "in");
+
+    Properties config = new Properties();
+    config.put(StreamsConfig.APPLICATION_ID_CONFIG, "flush-it-" + id);
+    config.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, kafka.getBootstrapServers());
+    config.put(StreamsConfig.STATE_DIR_CONFIG, Files.createTempDirectory("flush-it").toString());
+    KafkaStreams streams = new KafkaStreams(topology, config);
+    streams.start();
+    try {
+      try (KafkaProducer<byte[], KStreamsPayload<Integer>> producer =
+          new KafkaProducer<>(
+              ImmutableMap.<String, Object>of(
+                  ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka.getBootstrapServers()),
+              new ByteArraySerializer(),
+              serde.serializer())) {
+        producer
+            .send(
+                new ProducerRecord<>(
+                    input, new byte[0], KStreamsPayload.<Integer>flush(ImmutableSet.of(1, 2))))
+            .get();
+      }
+
+      Map<Integer, Integer> expected = ImmutableMap.of(1, 1, 2, 1);
+      assertThat(copiesPerPartition(output, serde, 2), is(expected));
+    } finally {
+      streams.close(Duration.ofSeconds(30));
+      streams.cleanUp();
+    }
+  }
+
+  /** Reads the topic until {@code expected} records arrive, then a little longer for extras. */
+  private static Map<Integer, Integer> copiesPerPartition(
+      String topic, KStreamsPayloadSerde<Integer> serde, int expected) {
+    List<TopicPartition> partitions = new ArrayList<>();
+    for (int partition = 0; partition < PARTITIONS; partition++) {
+      partitions.add(new TopicPartition(topic, partition));
+    }
+    Map<Integer, Integer> copies = new TreeMap<>();
+    try (KafkaConsumer<byte[], KStreamsPayload<Integer>> consumer =
+        new KafkaConsumer<>(
+            ImmutableMap.<String, Object>of(
+                ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka.getBootstrapServers()),
+            new ByteArrayDeserializer(),
+            serde.deserializer())) {
+      consumer.assign(partitions);
+      consumer.seekToBeginning(partitions);
+      int seen = 0;
+      boolean draining = false;
+      long deadline = System.currentTimeMillis() + 60_000;
+      while (System.currentTimeMillis() < deadline) {
+        for (ConsumerRecord<byte[], KStreamsPayload<Integer>> record :
+            consumer.poll(Duration.ofMillis(500))) {
+          assertThat(record.value().isFlush(), is(true));
+          copies.merge(record.partition(), 1, Integer::sum);
+          seen++;
+        }
+        if (!draining && seen >= expected) {
+          draining = true;
+          deadline = System.currentTimeMillis() + 3_000;
+        }
+      }
+    }
+    return copies;
+  }
+}