Skip to content

Commit 4636b1f

Browse files
junaiddshaukatjunaiddshaukat
authored
#38957: Add in-memory WatermarkManager core (per-source-partition tracking
1 parent 4c02a77 commit 4636b1f

2 files changed

Lines changed: 312 additions & 0 deletions

File tree

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

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.runners.kafka.streams.translation;
19+
20+
import java.util.HashMap;
21+
import java.util.Map;
22+
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
23+
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
24+
import org.joda.time.Instant;
25+
26+
/**
27+
* In-memory tracker of a single fused stage's input watermark, computed from the committed
28+
* watermarks reported by the <i>upstream source partitions</i> that feed it (the output /
29+
* repartition-topic partitions of the parent stage).
30+
*
31+
* <p>This is the core of the Kafka Streams runner's watermark propagation, decoupled from the Kafka
32+
* wiring so it can be unit-tested in isolation. The wiring that produces the reports (flushing
33+
* {@code (sourcePartition, committedWatermark, totalSourcePartitions)} atomically with each offset
34+
* commit and fanning it out to every downstream partition) and consumes them lands in a follow-up.
35+
*
36+
* <h3>Why source partitions, not producer instances</h3>
37+
*
38+
* <p>The question a stage has to answer is "have I received the watermark from every upstream
39+
* producer, so that {@code min()} across them is meaningful?". Counting <i>producer instances</i>
40+
* is hard: an instance can be killed without notice, leaving stale state, and the number changes on
41+
* every rebalance. Counting <i>source partitions</i> is robust instead, because the partition count
42+
* is fixed and known: it travels in-band with every report ({@code totalSourcePartitions}), a
43+
* partition is always owned by exactly one live instance, and when an instance dies its partitions
44+
* are reassigned and the new owner keeps reporting. So the manager only ever reasons about
45+
* partitions, never about instances. (Design agreed with the mentor; see the watermark
46+
* coordination-channel PoC findings.)
47+
*
48+
* <h3>Holding until ready</h3>
49+
*
50+
* <p>Until a committed watermark has been seen for <i>every</i> source partition, the stage's input
51+
* watermark is undefined and {@link #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} —
52+
* i.e. the stage emits no meaningful watermark downstream. A change in {@code
53+
* totalSourcePartitions} (e.g. a repartition) clears the accumulated reports and re-opens this hold
54+
* until the new full set has reported, which subsumes the "new epoch / revert" rule without an
55+
* explicit epoch.
56+
*
57+
* <h3>Monotonicity</h3>
58+
*
59+
* <p>Beam watermarks must be non-decreasing. Each source partition's watermark is held monotonic (a
60+
* lower report is ignored), and the emitted stage watermark is additionally clamped so it never
61+
* regresses below the previously emitted value — relevant if a newly appeared partition reports an
62+
* older watermark after the stage had already advanced.
63+
*
64+
* <p>Not thread-safe; the caller (a single Kafka Streams processor thread) serializes access.
65+
*/
66+
public final class WatermarkManager {
67+
68+
/** Total source partitions feeding this stage, learned in-band; -1 until the first report. */
69+
private int expectedSourcePartitionCount = -1;
70+
71+
/** Latest committed watermark per source partition (kept monotonic non-decreasing). */
72+
private final Map<Integer, Instant> committedWatermarkByPartition = new HashMap<>();
73+
74+
/** Last watermark {@link #advance()} emitted, to enforce a non-decreasing output. */
75+
private Instant lastEmitted = BoundedWindow.TIMESTAMP_MIN_VALUE;
76+
77+
/**
78+
* Record a committed watermark reported for one source partition, together with the total source
79+
* partition count carried in-band with the report.
80+
*
81+
* @param sourcePartition the source partition the report is for, in {@code [0,
82+
* totalSourcePartitions)}
83+
* @param committedWatermark the committed watermark for that partition
84+
* @param totalSourcePartitions the total number of upstream source partitions feeding this stage
85+
*/
86+
public void observe(int sourcePartition, Instant committedWatermark, int totalSourcePartitions) {
87+
if (committedWatermark == null) {
88+
throw new IllegalArgumentException("committedWatermark must not be null");
89+
}
90+
if (totalSourcePartitions <= 0) {
91+
throw new IllegalArgumentException(
92+
"totalSourcePartitions must be positive: " + totalSourcePartitions);
93+
}
94+
if (sourcePartition < 0 || sourcePartition >= totalSourcePartitions) {
95+
throw new IllegalArgumentException(
96+
"sourcePartition "
97+
+ sourcePartition
98+
+ " out of range for totalSourcePartitions "
99+
+ totalSourcePartitions);
100+
}
101+
if (totalSourcePartitions != expectedSourcePartitionCount) {
102+
// The source partition set changed (e.g. a repartition). The previous per-partition
103+
// watermarks describe a different partitioning, so drop them entirely and re-open the hold
104+
// until the new full set reports. The output watermark still cannot regress (lastEmitted is
105+
// retained).
106+
expectedSourcePartitionCount = totalSourcePartitions;
107+
committedWatermarkByPartition.clear();
108+
}
109+
// A source partition's watermark is monotonic non-decreasing; ignore an out-of-order lower
110+
// report.
111+
committedWatermarkByPartition.merge(
112+
sourcePartition, committedWatermark, (oldW, newW) -> newW.isAfter(oldW) ? newW : oldW);
113+
}
114+
115+
/** True once a committed watermark has been seen for every current source partition. */
116+
public boolean isReady() {
117+
return expectedSourcePartitionCount > 0
118+
&& committedWatermarkByPartition.size() == expectedSourcePartitionCount;
119+
}
120+
121+
/**
122+
* Advance and return the stage input watermark.
123+
*
124+
* <p>Returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} while the stage is still holding (not
125+
* every source partition has reported) — the caller emits nothing meaningful downstream in that
126+
* case. Once ready, returns {@code min()} over all source partitions, clamped to never regress
127+
* below the previously emitted value. The sequence of values returned across calls is
128+
* non-decreasing.
129+
*/
130+
public Instant advance() {
131+
if (!isReady()) {
132+
return BoundedWindow.TIMESTAMP_MIN_VALUE;
133+
}
134+
// isReady() guarantees the map is non-empty, so the seed is always replaced by a real value.
135+
Instant min = BoundedWindow.TIMESTAMP_MAX_VALUE;
136+
for (Instant w : committedWatermarkByPartition.values()) {
137+
if (w.isBefore(min)) {
138+
min = w;
139+
}
140+
}
141+
Instant emit = min.isAfter(lastEmitted) ? min : lastEmitted;
142+
lastEmitted = emit;
143+
return emit;
144+
}
145+
146+
/** The total source partition count learned in-band, or -1 if nothing reported yet. */
147+
@VisibleForTesting
148+
int expectedSourcePartitionCount() {
149+
return expectedSourcePartitionCount;
150+
}
151+
152+
/** How many distinct source partitions have reported so far. */
153+
@VisibleForTesting
154+
int reportedPartitionCount() {
155+
return committedWatermarkByPartition.size();
156+
}
157+
}

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

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.runners.kafka.streams.translation;
19+
20+
import static org.hamcrest.CoreMatchers.is;
21+
import static org.hamcrest.MatcherAssert.assertThat;
22+
import static org.junit.Assert.assertThrows;
23+
24+
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
25+
import org.joda.time.Instant;
26+
import org.junit.Test;
27+
28+
/** Tests for {@link WatermarkManager}. */
29+
public class WatermarkManagerTest {
30+
31+
private static Instant ts(long millis) {
32+
return new Instant(millis);
33+
}
34+
35+
@Test
36+
public void holdsBeforeAnyReport() {
37+
WatermarkManager manager = new WatermarkManager();
38+
assertThat(manager.isReady(), is(false));
39+
assertThat(manager.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE));
40+
}
41+
42+
@Test
43+
public void holdsUntilEverySourcePartitionReports() {
44+
WatermarkManager manager = new WatermarkManager();
45+
manager.observe(0, ts(100L), 4);
46+
manager.observe(1, ts(100L), 4);
47+
manager.observe(2, ts(100L), 4);
48+
// Three of four partitions reported — still holding.
49+
assertThat(manager.isReady(), is(false));
50+
assertThat(manager.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE));
51+
52+
manager.observe(3, ts(100L), 4);
53+
assertThat(manager.isReady(), is(true));
54+
assertThat(manager.advance(), is(ts(100L)));
55+
}
56+
57+
@Test
58+
public void emitsMinAcrossPartitionsOnceReady() {
59+
WatermarkManager manager = new WatermarkManager();
60+
manager.observe(0, ts(300L), 4);
61+
manager.observe(1, ts(100L), 4);
62+
manager.observe(2, ts(500L), 4);
63+
manager.observe(3, ts(200L), 4);
64+
// min(300, 100, 500, 200) = 100
65+
assertThat(manager.advance(), is(ts(100L)));
66+
}
67+
68+
@Test
69+
public void perPartitionWatermarkIsMonotonic() {
70+
WatermarkManager manager = new WatermarkManager();
71+
manager.observe(0, ts(100L), 1);
72+
assertThat(manager.advance(), is(ts(100L)));
73+
// A lower out-of-order report for the same partition is ignored.
74+
manager.observe(0, ts(50L), 1);
75+
assertThat(manager.advance(), is(ts(100L)));
76+
}
77+
78+
@Test
79+
public void partitionCountChangeClearsReportsAndReopensHold() {
80+
WatermarkManager manager = new WatermarkManager();
81+
manager.observe(0, ts(100L), 2);
82+
manager.observe(1, ts(100L), 2);
83+
assertThat(manager.advance(), is(ts(100L)));
84+
85+
// Source set grows to 4. The previous reports are dropped, so the stage holds again until all
86+
// four of the new partition set have reported.
87+
manager.observe(0, ts(200L), 4);
88+
assertThat(manager.reportedPartitionCount(), is(1));
89+
assertThat(manager.isReady(), is(false));
90+
assertThat(manager.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE));
91+
92+
manager.observe(1, ts(200L), 4);
93+
manager.observe(2, ts(200L), 4);
94+
manager.observe(3, ts(200L), 4);
95+
assertThat(manager.isReady(), is(true));
96+
assertThat(manager.advance(), is(ts(200L)));
97+
}
98+
99+
@Test
100+
public void emittedWatermarkDoesNotRegressAfterRepartition() {
101+
WatermarkManager manager = new WatermarkManager();
102+
manager.observe(0, ts(100L), 2);
103+
manager.observe(1, ts(100L), 2);
104+
assertThat(manager.advance(), is(ts(100L)));
105+
106+
// After a repartition to 3, the new partitions report older watermarks. Once ready again the
107+
// min is 50, but the emitted stage watermark must not go backwards below 100.
108+
manager.observe(0, ts(50L), 3);
109+
manager.observe(1, ts(50L), 3);
110+
manager.observe(2, ts(50L), 3);
111+
assertThat(manager.isReady(), is(true));
112+
assertThat(manager.advance(), is(ts(100L)));
113+
}
114+
115+
@Test
116+
public void partitionCountDecreaseClearsAndRecomputes() {
117+
WatermarkManager manager = new WatermarkManager();
118+
manager.observe(0, ts(100L), 4);
119+
manager.observe(1, ts(100L), 4);
120+
manager.observe(2, ts(100L), 4);
121+
manager.observe(3, ts(50L), 4);
122+
assertThat(manager.advance(), is(ts(50L)));
123+
124+
// Source set shrinks to 2. Reports are cleared; the stage holds until {0, 1} report again.
125+
manager.observe(0, ts(100L), 2);
126+
assertThat(manager.expectedSourcePartitionCount(), is(2));
127+
assertThat(manager.reportedPartitionCount(), is(1));
128+
assertThat(manager.isReady(), is(false));
129+
130+
manager.observe(1, ts(100L), 2);
131+
assertThat(manager.isReady(), is(true));
132+
// min is 100; the non-regression clamp keeps it >= the previously emitted 50.
133+
assertThat(manager.advance(), is(ts(100L)));
134+
}
135+
136+
@Test
137+
public void rejectsNullWatermark() {
138+
WatermarkManager manager = new WatermarkManager();
139+
assertThrows(IllegalArgumentException.class, () -> manager.observe(0, null, 4));
140+
}
141+
142+
@Test
143+
public void rejectsNonPositiveTotalSourcePartitions() {
144+
WatermarkManager manager = new WatermarkManager();
145+
assertThrows(IllegalArgumentException.class, () -> manager.observe(0, ts(100L), 0));
146+
assertThrows(IllegalArgumentException.class, () -> manager.observe(0, ts(100L), -1));
147+
}
148+
149+
@Test
150+
public void rejectsOutOfRangeSourcePartition() {
151+
WatermarkManager manager = new WatermarkManager();
152+
assertThrows(IllegalArgumentException.class, () -> manager.observe(-1, ts(100L), 4));
153+
assertThrows(IllegalArgumentException.class, () -> manager.observe(4, ts(100L), 4));
154+
}
155+
}

0 commit comments

Comments
 (0)