Kafka Streams runner: target flush markers across repartition topics

[junaiddshaukat]

Summary

Second PR toward bundles bounded by time (#39633), after #40068. It makes the shuffle and the repartition partitioner handle a flush marker. Nothing emits markers yet, so there is no behaviour change.

What changed

• GroupByKeyBroadcastPartitioner is renamed to KStreamsPayloadPartitioner, since it no longer only broadcasts. Data is still hashed by key and watermarks still go to every partition. A flush goes to exactly the partitions it names.
• ShuffleByKeyProcessor picks those partitions. Upstream partition i of U addresses downstream partitions [floor(i*D/U), floor((i+1)*D/U)) of D. Over all i this covers every downstream partition exactly once, for fan-out, fan-in and equal counts, so each downstream partition gets one flush per interval and not one per upstream partition. It only depends on partition numbers, so a rebalance does not change it.
• When the shuffle fans in, some upstream partitions have nothing to address (partition 0 of 10 into 8, for example). Those forward no marker.
• The partitioner fails on a flush if the topic does not have internalParallelism partitions. The targets are computed for that count, so on a topic left by an earlier run with a different parallelism they would point at partitions that don't exist, or skip some. The topic manager reuses an existing topic without checking its partition count, so nothing catches this earlier today. Handling stale topics properly belongs to Kafka Streams runner: lifecycle management for runner-created topics #39566.

Order of the remaining PRs

I swapped the last two compared to #40068. Today ExecutableStageProcessor and WindowedGroupByKeyProcessor would throw on a flush, because each treats anything that is not one kind as the other kind. So consumers have to handle the marker before any source emits it:

3. Consumers handle a flush: ExecutableStageProcessor closes its bundle and forwards the marker, WindowedGroupByKeyProcessor forwards it. StageOutputProcessor and FlattenProcessor already pass it through.
4. Sources emit the marker on their existing punctuator, at maxBundleTimeMs.

One open question for step 4: a marker created by a source has no edge yet, but KStreamsPayload.flush does not accept empty targets. The shuffle recomputes targets anyway, so what a source puts there is ignored. I'd rather settle that in step 4 than change the payload now.

Testing

./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:build
./gradlew -Pwith-kafka-streams-runner :runners:kafka-streams:validatesRunner

Both pass, 119 unit tests and 59 ValidatesRunner tests. New tests: KStreamsPayloadPartitionerTest (watermark broadcast, flush targets, data by key, keyless data, partition count mismatch) and four in ShuffleByKeyProcessorTest, including one that checks every downstream partition is hit exactly once for several upstream/downstream shapes. I checked they catch real mistakes: an off-by-one in the range rule plus removing the partition count check fails five of them, and forwarding a flush with no targets fails the fan-in test.

[junaiddshaukat]

R: @je-ik

[github-actions]

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

[je-ik]

I'm a little confused about how we do actually ensure that the elements are really actually emitted to the target partitions? I.e. if there are multiple partitions, how do we make sure the element is correctly duplicated across them?

[je-ik]

Should this actually expand the target partitions and emit multiple elements that target different partitions? Or is this done elsewhere?

[junaiddshaukat]

It is done in the sink, by Kafka Streams. The shuffle forwards the flush once, KStreamsPayloadPartitioner.partitions() returns the target set, and RecordCollectorImpl sends one copy to each partition in it (RecordCollectorImpl.java#L163-L177, 3.9.0):

for (final int multicastPartition: multicastPartitions) {
    send(topic, key, value, headers, multicastPartition, timestamp, ...);
}

This is the multicast from KIP-837, and it is the same path the watermark broadcast already uses, only with a subset instead of every partition. An empty set is dropped there with a warning, which is why the shuffle does not forward one.

Expanding it here would also work, but it would do the same thing with more records, since a processor cannot pick the partition itself; only the sink's partitioner can.

It was not obvious from the code, so I added a short comment here, and a broker test (KStreamsPayloadPartitionerBrokerIT) that sends a flush for partitions 1 and 2 of a 4 partition topic and checks it arrives once on each and nowhere else. The unit tests cannot show this, because TopologyTestDriver reports every topic as having one partition.

[junaiddshaukat]

It is done in the sink, by Kafka Streams. The shuffle forwards the flush once, KStreamsPayloadPartitioner.partitions() returns the target set, and RecordCollectorImpl sends one copy to each partition in it (RecordCollectorImpl.java#L163-L177, 3.9.0):

for (final int multicastPartition: multicastPartitions) {
    send(topic, key, value, headers, multicastPartition, timestamp, ...);
}

This is the multicast from KIP-837, and it is the same path the watermark broadcast already uses, only with a subset instead of every partition. An empty set is dropped there with a warning, which is why the shuffle does not forward one.

Expanding it here would also work, but it would do the same thing with more records, since a processor cannot pick the partition itself; only the sink's partitioner can.

It was not obvious from the code, so I added a short comment here, and a broker test (KStreamsPayloadPartitionerBrokerIT) that sends a flush for partitions 1 and 2 of a 4 partition topic and checks it arrives once on each and nowhere else. The unit tests cannot show this, because TopologyTestDriver reports every topic as having one partition.