[GSoC 2026] Kafka Streams runner — WatermarkManager part 2: wire into ExecutableStageProcessor#38987
Conversation
Part 2 of the WatermarkManager work: drive watermark propagation through the part-1 WatermarkManager instead of forwarding on every received watermark. - KStreamsPayload: the watermark variant now carries the in-band report fields (sourcePartition, totalSourcePartitions) the downstream WatermarkManager needs, alongside the watermark millis. - ExecutableStageProcessor: on a watermark payload, flush the open bundle, feed the report to a WatermarkManager, and forward the stage's output watermark only when min() across the source partitions advances (stamped as the stage's own single source, 0 of 1). Data is still processed while the watermark is held. The SDK harness is now created lazily on the first data element, so a stage that only forwards watermarks never starts one. - ImpulseProcessor: stamp its terminal TIMESTAMP_MAX_VALUE watermark as (sourcePartition=0, totalSourcePartitions=1). The durable/distributed produce side — flushing the report atomically with the EOS offset commit, fanning it out to all downstream partitions, and a serde for it to cross topic boundaries — is deferred until the topology gains topic-based shuffle and there are real source partitions to track. Tests: ExecutableStageProcessorWatermarkTest drives the watermark path with a MockProcessorContext (hold-until-all-report, min, monotonic, single-source stamping); WatermarkPropagationTest is an end-to-end TopologyTestDriver test that a terminal watermark propagates Impulse -> ExecutableStage -> a recording sink. ./gradlew :runners:kafka-streams:check green. Closes apache#38977 Refs apache#18479
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request implements the second part of the WatermarkManager for the Kafka Streams runner. It transitions the runner from a provisional 'flush-on-every-watermark' behavior to a more robust model where stages track and compute the minimum watermark across source partitions. This change ensures correct watermark propagation while maintaining data processing throughput, and lays the groundwork for future distributed watermark persistence and topic-based shuffle integration. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request integrates the WatermarkManager into ExecutableStageProcessor to handle watermarks per source partition, holding the output watermark until all source partitions have reported. It also introduces lazy initialization of the SDK harness and updates KStreamsPayload to carry partition coordination fields. Feedback on these changes suggests adding a null check for the record payload in ExecutableStageProcessor to avoid potential NullPointerExceptions, and validating partition coordinates in KStreamsPayload to prevent invalid values.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
Assigning reviewers: R: @shunping added as fallback since no labels match configuration Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
je-ik
left a comment
je-ik
left a comment
There was a problem hiding this comment.
Great job! I like the tests. 👍
I left a few comments.
| StageBundleFactory factory = stageBundleFactory; | ||
| if (factory == null) { | ||
| ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload); | ||
| ExecutableStageContext sc = |
There was a problem hiding this comment.
Are the variables factory and sc necessary? Could we use directly the target fields?
There was a problem hiding this comment.
Done ensureStageBundleFactory is void now and assigns the fields directly; ensureBundleOpen reads the field via checkInitialized().
| * #isWatermark()} first; calling this on a data payload throws. | ||
| */ | ||
| public int getSourcePartition() { | ||
| if (kind != Kind.WATERMARK) { |
There was a problem hiding this comment.
The idiomatic approach would be to extract interface WatermarkPayload from the KStreamsPayload and keep the kind == Kind.WATERMARK methods only on that interface.
It would then be cast like
KStreamsPayload payload = ...;
if (payload.isWatermark()) {
WatermarkPayload watermark = payload.asWatermark()l;
watermark.getSourcePartition(); // specific methods.
}
There was a problem hiding this comment.
Added Preconditions.checkArgument in the watermark() factory to validate sourcePartition / totalSourcePartitions.
Done — pulled the watermark accessors onto a WatermarkPayload interface; you reach them via payload.asWatermark() (which checkStates isWatermark()), and I updated the callers.
- Extract a WatermarkPayload interface holding the watermark-only accessors (getWatermarkMillis / getSourcePartition / getTotalSourcePartitions). They no longer live on KStreamsPayload; callers narrow via payload.asWatermark(), which Preconditions.checkState(isWatermark()) once instead of every accessor doing its own kind check. - Validate the watermark() factory arguments with Preconditions.checkArgument. - ExecutableStageProcessor.ensureStageBundleFactory: assign the fields directly instead of going through local variables.
je-ik
merged commit 1058e94
into
apache:feat/18479-kafka-streams-runner-skeleton
2 participants
Summary
Part 2 of the WatermarkManager (part 1 = #38957). Wires the in-memory WatermarkManager into the data path so each fused stage computes and forwards its input watermark through it, replacing the provisional "flush on every received watermark" behavior.
This is the in-JVM wiring. The stage still forwards its output watermark downstream via
ctx.forward(now gated by the WatermarkManager so it only forwards whenmin()advances), so watermarks keep propagating and are observed in tests. Deferred is only the durable/distributed produce side — flushing the report atomically with the EOS offset commit, fanning it out to all downstream partitions, and a serde for it to cross topic boundaries — which needs the topic-based shuffle (GroupByKey) infra that isn't built yet.Changes
KStreamsPayload: watermark variant carries(sourcePartition, totalSourcePartitions)in-band with the millis.ExecutableStageProcessor: feeds reports to aWatermarkManager, forwardsthe output watermark only when it advances, stamped as the stage's own single
source (0 of 1); data still processed while holding. SDK harness created
lazily on first data element.
ImpulseProcessor: stamps its terminalTIMESTAMP_MAX_VALUEas(0, 1).Out of scope (later, depend on topic-based shuffle)
persistence and downstream timer firing.
Testing
ExecutableStageProcessorWatermarkTest(MockProcessorContext): hold-until-all-report, min, monotonic non-re-forward, single-source stamping.
WatermarkPropagationTest(TopologyTestDriver): terminal watermarkpropagates Impulse -> ExecutableStage -> recording sink.
./gradlew :runners:kafka-streams:checkgreen; 30 runner tests.Closes #38977
Refs #18479
cc @je-ik