From fd451b9e231c7528ea9a3768aefe008f3ce64e9a Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Sun, 14 Jun 2026 12:30:28 +0500 Subject: [PATCH 1/2] Add in-memory WatermarkManager core (per-source-partition tracking) First slice of the Kafka Streams runner WatermarkManager, decoupled from the Kafka wiring so it can be unit-tested in isolation. A stage's input watermark is min() over its upstream source partitions' committed watermarks. Tracking is keyed on source partitions rather than producer instances: the total source-partition count travels in-band with every report, and a partition is always owned by exactly one live instance (reassigned to a new owner on failure), so a killed instance never leaves the reader stuck and no instance-liveness tracking is needed. This was validated in a standalone Kafka Streams PoC before implementation. WatermarkManager holds at BoundedWindow.TIMESTAMP_MIN_VALUE until every source partition has reported, then emits min(); the emitted watermark is clamped to be non-decreasing, and a change in totalSourcePartitions re-opens the hold (the "revert" case, without an explicit epoch). Wiring into ExecutableStageProcessor (flush coupled to the offset commit, fan-out to all downstream partitions, real-Kafka integration tests) is a follow-up. Refs #18479 --- .../streams/translation/WatermarkManager.java | 148 ++++++++++++++++++ .../translation/WatermarkManagerTest.java | 136 ++++++++++++++++ 2 files changed, 284 insertions(+) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java new file mode 100644 index 000000000000..fa9f924225c6 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java @@ -0,0 +1,148 @@ +/* + * 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.HashMap; +import java.util.Map; +import java.util.OptionalLong; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; + +/** + * In-memory tracker of a single fused stage's input watermark, computed from the committed + * watermarks reported by its upstream source partitions. + * + *

This is the core of the Kafka Streams runner's watermark propagation, decoupled from the Kafka + * wiring so it can be unit-tested in isolation. The wiring that produces the reports (flushing + * {@code (sourcePartition, committedWatermark, totalSourcePartitions)} atomically with each offset + * commit and fanning it out to every downstream partition) and consumes them lands in a follow-up. + * + *

Why source partitions, not producer instances

+ * + *

The question a stage has to answer is "have I received the watermark from every upstream + * producer, so that {@code min()} across them is meaningful?". Counting producer instances + * is hard: an instance can be killed without notice, leaving stale state, and the number changes on + * every rebalance. Counting source partitions is robust instead, because the partition count + * is fixed and known: it travels in-band with every report ({@code totalSourcePartitions}), a + * partition is always owned by exactly one live instance, and when an instance dies its partitions + * are reassigned and the new owner keeps reporting. So the manager only ever reasons about + * partitions, never about instances. (Design agreed with the mentor; see the watermark + * coordination-channel PoC findings.) + * + *

Holding until complete

+ * + *

Until a committed watermark has been seen for every source partition, the stage's input + * watermark is undefined and must be held at {@link BoundedWindow#TIMESTAMP_MIN_VALUE} — i.e. the + * stage emits no watermark downstream. {@link #advance()} returns {@link OptionalLong#empty()} in + * that case. A change in {@code totalSourcePartitions} (e.g. a repartition) re-opens this hold + * until the new full set has reported, which subsumes the "new epoch / revert" rule without an + * explicit epoch. + * + *

Monotonicity

+ * + *

Beam watermarks must be non-decreasing. Each source partition's watermark is held monotonic (a + * lower report is ignored), and the emitted stage watermark is additionally clamped so it never + * regresses below the previously emitted value — relevant if a newly appeared partition reports an + * older watermark after the stage had already advanced. + * + *

Not thread-safe; the caller (a single Kafka Streams processor thread) serializes access. + */ +public final class WatermarkManager { + + /** Millis the stage holds at while it cannot yet emit a watermark (the Beam minimum). */ + public static final long HOLD_MILLIS = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(); + + /** Total source partitions feeding this stage, learned in-band; -1 until the first report. */ + private int expectedSourcePartitionCount = -1; + + /** Latest committed watermark per source partition (kept monotonic non-decreasing). */ + private final Map committedWatermarkByPartition = new HashMap<>(); + + /** Last watermark {@link #advance()} emitted, to enforce a non-decreasing output. */ + private long lastEmittedMillis = HOLD_MILLIS; + + /** + * Record a committed watermark reported for one source partition, together with the total source + * partition count carried in-band with the report. + * + * @param sourcePartition the source partition the report is for, in {@code [0, + * totalSourcePartitions)} + * @param committedWatermarkMillis the committed watermark for that partition, in event-time + * millis + * @param totalSourcePartitions the total number of source partitions feeding this stage + */ + public void observe( + int sourcePartition, long committedWatermarkMillis, int totalSourcePartitions) { + if (totalSourcePartitions <= 0) { + throw new IllegalArgumentException( + "totalSourcePartitions must be positive: " + totalSourcePartitions); + } + if (sourcePartition < 0 || sourcePartition >= totalSourcePartitions) { + throw new IllegalArgumentException( + "sourcePartition " + + sourcePartition + + " out of range for totalSourcePartitions " + + totalSourcePartitions); + } + if (totalSourcePartitions != expectedSourcePartitionCount) { + expectedSourcePartitionCount = totalSourcePartitions; + // On a partition-count decrease, drop reports for partitions that no longer exist so + // completeness and min() are computed over the current partition set only. + committedWatermarkByPartition.keySet().removeIf(p -> p >= totalSourcePartitions); + } + // A source partition's watermark is monotonic non-decreasing; ignore an out-of-order lower + // report. + committedWatermarkByPartition.merge(sourcePartition, committedWatermarkMillis, Math::max); + } + + /** True once a committed watermark has been seen for every current source partition. */ + public boolean isComplete() { + return expectedSourcePartitionCount > 0 + && committedWatermarkByPartition.size() == expectedSourcePartitionCount; + } + + /** + * Advance and return the stage input watermark. + * + *

Returns {@link OptionalLong#empty()} while the stage is still holding (not every source + * partition has reported) — the caller must emit nothing downstream in that case. Once complete, + * returns {@code min()} over all source partitions, clamped to never regress below the previously + * emitted value. The sequence of present values returned across calls is non-decreasing. + */ + public OptionalLong advance() { + if (!isComplete()) { + return OptionalLong.empty(); + } + long min = Long.MAX_VALUE; + for (long w : committedWatermarkByPartition.values()) { + min = Math.min(min, w); + } + long emit = Math.max(min, lastEmittedMillis); + lastEmittedMillis = emit; + return OptionalLong.of(emit); + } + + /** The total source partition count learned in-band, or -1 if nothing reported yet. */ + public int expectedSourcePartitionCount() { + return expectedSourcePartitionCount; + } + + /** How many distinct source partitions have reported so far. */ + public int reportedPartitionCount() { + return committedWatermarkByPartition.size(); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java new file mode 100644 index 000000000000..cf1564592581 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java @@ -0,0 +1,136 @@ +/* + * 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 static org.junit.Assert.assertThrows; + +import java.util.OptionalLong; +import org.junit.Test; + +/** Tests for {@link WatermarkManager}. */ +public class WatermarkManagerTest { + + @Test + public void holdsBeforeAnyReport() { + WatermarkManager manager = new WatermarkManager(); + assertThat(manager.isComplete(), is(false)); + assertThat(manager.advance().isPresent(), is(false)); + } + + @Test + public void holdsUntilEverySourcePartitionReports() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, 100L, 4); + manager.observe(1, 100L, 4); + manager.observe(2, 100L, 4); + // Three of four partitions reported — still holding. + assertThat(manager.isComplete(), is(false)); + assertThat(manager.advance().isPresent(), is(false)); + + manager.observe(3, 100L, 4); + assertThat(manager.isComplete(), is(true)); + assertThat(manager.advance(), is(OptionalLong.of(100L))); + } + + @Test + public void emitsMinAcrossPartitionsOnceComplete() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, 300L, 4); + manager.observe(1, 100L, 4); + manager.observe(2, 500L, 4); + manager.observe(3, 200L, 4); + // min(300, 100, 500, 200) = 100 + assertThat(manager.advance(), is(OptionalLong.of(100L))); + } + + @Test + public void perPartitionWatermarkIsMonotonic() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, 100L, 1); + assertThat(manager.advance(), is(OptionalLong.of(100L))); + // A lower out-of-order report for the same partition is ignored. + manager.observe(0, 50L, 1); + assertThat(manager.advance(), is(OptionalLong.of(100L))); + } + + @Test + public void emittedWatermarkDoesNotRegressWhenNewPartitionReportsOlderWatermark() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, 100L, 2); + manager.observe(1, 100L, 2); + assertThat(manager.advance(), is(OptionalLong.of(100L))); + + // A repartition grows the source set to 3; a freshly appeared partition reports an older + // watermark. The stage watermark must not go backwards. + manager.observe(2, 50L, 3); + assertThat(manager.isComplete(), is(true)); + assertThat(manager.advance(), is(OptionalLong.of(100L))); + } + + @Test + public void partitionCountIncreaseReopensHold() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, 100L, 2); + manager.observe(1, 100L, 2); + assertThat(manager.advance(), is(OptionalLong.of(100L))); + + // Source set grows to 4; until partitions 2 and 3 report, the stage holds again. + manager.observe(2, 200L, 4); + assertThat(manager.isComplete(), is(false)); + assertThat(manager.advance().isPresent(), is(false)); + + manager.observe(3, 200L, 4); + assertThat(manager.isComplete(), is(true)); + // min(100, 100, 200, 200) = 100, which also satisfies the non-regression clamp. + assertThat(manager.advance(), is(OptionalLong.of(100L))); + } + + @Test + public void partitionCountDecreasePrunesStalePartitions() { + WatermarkManager manager = new WatermarkManager(); + manager.observe(0, 100L, 4); + manager.observe(1, 100L, 4); + manager.observe(2, 100L, 4); + manager.observe(3, 50L, 4); + assertThat(manager.advance(), is(OptionalLong.of(50L))); + + // Source set shrinks to 2; partitions 2 and 3 no longer exist and are dropped. The remaining + // set {0, 1} is complete and its min is 100, but the output must not regress below 50. + manager.observe(0, 100L, 2); + assertThat(manager.expectedSourcePartitionCount(), is(2)); + assertThat(manager.reportedPartitionCount(), is(2)); + assertThat(manager.isComplete(), is(true)); + assertThat(manager.advance(), is(OptionalLong.of(100L))); + } + + @Test + public void rejectsNonPositiveTotalSourcePartitions() { + WatermarkManager manager = new WatermarkManager(); + assertThrows(IllegalArgumentException.class, () -> manager.observe(0, 100L, 0)); + assertThrows(IllegalArgumentException.class, () -> manager.observe(0, 100L, -1)); + } + + @Test + public void rejectsOutOfRangeSourcePartition() { + WatermarkManager manager = new WatermarkManager(); + assertThrows(IllegalArgumentException.class, () -> manager.observe(-1, 100L, 4)); + assertThrows(IllegalArgumentException.class, () -> manager.observe(4, 100L, 4)); + } +} From 61b0b64cf22e9021e29d3f13b4cf60c656ae0e9f Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Mon, 15 Jun 2026 18:49:02 +0500 Subject: [PATCH 2/2] Address review: use Instant, isReady, clear-on-repartition - Switch from long event-time millis to org.joda.time.Instant throughout, matching the Beam codebase convention; advance() returns an Instant and a holding stage returns BoundedWindow.TIMESTAMP_MIN_VALUE instead of an empty OptionalLong. - Drop the HOLD_MILLIS constant; "watermark hold" is a distinct Beam concept and this was only the minimum timestamp. - On a source-partition-count change, clear the whole per-partition map rather than pruning only out-of-range partitions, so completeness and min() are always computed over a single consistent partitioning. - Rename isComplete() to isReady() ("complete" could be read as the watermark having reached TIMESTAMP_MAX_VALUE). - Make the inspection getters package-private and @VisibleForTesting. - Clarify in the docs that the tracked partitions are the upstream/parent stage's output (repartition-topic) partitions that feed this stage. --- .../streams/translation/WatermarkManager.java | 81 +++++----- .../translation/WatermarkManagerTest.java | 143 ++++++++++-------- 2 files changed, 126 insertions(+), 98 deletions(-) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java index fa9f924225c6..cd0fca3654ae 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManager.java @@ -19,12 +19,14 @@ import java.util.HashMap; import java.util.Map; -import java.util.OptionalLong; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.joda.time.Instant; /** * In-memory tracker of a single fused stage's input watermark, computed from the committed - * watermarks reported by its upstream source partitions. + * watermarks reported by the upstream source partitions that feed it (the output / + * repartition-topic partitions of the parent stage). * *

This is the core of the Kafka Streams runner's watermark propagation, decoupled from the Kafka * wiring so it can be unit-tested in isolation. The wiring that produces the reports (flushing @@ -43,12 +45,12 @@ * partitions, never about instances. (Design agreed with the mentor; see the watermark * coordination-channel PoC findings.) * - *

Holding until complete

+ *

Holding until ready

* *

Until a committed watermark has been seen for every source partition, the stage's input - * watermark is undefined and must be held at {@link BoundedWindow#TIMESTAMP_MIN_VALUE} — i.e. the - * stage emits no watermark downstream. {@link #advance()} returns {@link OptionalLong#empty()} in - * that case. A change in {@code totalSourcePartitions} (e.g. a repartition) re-opens this hold + * watermark is undefined and {@link #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} — + * i.e. the stage emits no meaningful watermark downstream. A change in {@code + * totalSourcePartitions} (e.g. a repartition) clears the accumulated reports and re-opens this hold * until the new full set has reported, which subsumes the "new epoch / revert" rule without an * explicit epoch. * @@ -63,17 +65,14 @@ */ public final class WatermarkManager { - /** Millis the stage holds at while it cannot yet emit a watermark (the Beam minimum). */ - public static final long HOLD_MILLIS = BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(); - /** Total source partitions feeding this stage, learned in-band; -1 until the first report. */ private int expectedSourcePartitionCount = -1; /** Latest committed watermark per source partition (kept monotonic non-decreasing). */ - private final Map committedWatermarkByPartition = new HashMap<>(); + private final Map committedWatermarkByPartition = new HashMap<>(); /** Last watermark {@link #advance()} emitted, to enforce a non-decreasing output. */ - private long lastEmittedMillis = HOLD_MILLIS; + private Instant lastEmitted = BoundedWindow.TIMESTAMP_MIN_VALUE; /** * Record a committed watermark reported for one source partition, together with the total source @@ -81,12 +80,13 @@ public final class WatermarkManager { * * @param sourcePartition the source partition the report is for, in {@code [0, * totalSourcePartitions)} - * @param committedWatermarkMillis the committed watermark for that partition, in event-time - * millis - * @param totalSourcePartitions the total number of source partitions feeding this stage + * @param committedWatermark the committed watermark for that partition + * @param totalSourcePartitions the total number of upstream source partitions feeding this stage */ - public void observe( - int sourcePartition, long committedWatermarkMillis, int totalSourcePartitions) { + public void observe(int sourcePartition, Instant committedWatermark, int totalSourcePartitions) { + if (committedWatermark == null) { + throw new IllegalArgumentException("committedWatermark must not be null"); + } if (totalSourcePartitions <= 0) { throw new IllegalArgumentException( "totalSourcePartitions must be positive: " + totalSourcePartitions); @@ -99,18 +99,21 @@ public void observe( + totalSourcePartitions); } if (totalSourcePartitions != expectedSourcePartitionCount) { + // The source partition set changed (e.g. a repartition). The previous per-partition + // watermarks describe a different partitioning, so drop them entirely and re-open the hold + // until the new full set reports. The output watermark still cannot regress (lastEmitted is + // retained). expectedSourcePartitionCount = totalSourcePartitions; - // On a partition-count decrease, drop reports for partitions that no longer exist so - // completeness and min() are computed over the current partition set only. - committedWatermarkByPartition.keySet().removeIf(p -> p >= totalSourcePartitions); + committedWatermarkByPartition.clear(); } // A source partition's watermark is monotonic non-decreasing; ignore an out-of-order lower // report. - committedWatermarkByPartition.merge(sourcePartition, committedWatermarkMillis, Math::max); + committedWatermarkByPartition.merge( + sourcePartition, committedWatermark, (oldW, newW) -> newW.isAfter(oldW) ? newW : oldW); } /** True once a committed watermark has been seen for every current source partition. */ - public boolean isComplete() { + public boolean isReady() { return expectedSourcePartitionCount > 0 && committedWatermarkByPartition.size() == expectedSourcePartitionCount; } @@ -118,31 +121,37 @@ public boolean isComplete() { /** * Advance and return the stage input watermark. * - *

Returns {@link OptionalLong#empty()} while the stage is still holding (not every source - * partition has reported) — the caller must emit nothing downstream in that case. Once complete, - * returns {@code min()} over all source partitions, clamped to never regress below the previously - * emitted value. The sequence of present values returned across calls is non-decreasing. + *

Returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} while the stage is still holding (not + * every source partition has reported) — the caller emits nothing meaningful downstream in that + * case. Once ready, returns {@code min()} over all source partitions, clamped to never regress + * below the previously emitted value. The sequence of values returned across calls is + * non-decreasing. */ - public OptionalLong advance() { - if (!isComplete()) { - return OptionalLong.empty(); + public Instant advance() { + if (!isReady()) { + return BoundedWindow.TIMESTAMP_MIN_VALUE; } - long min = Long.MAX_VALUE; - for (long w : committedWatermarkByPartition.values()) { - min = Math.min(min, w); + // isReady() guarantees the map is non-empty, so the seed is always replaced by a real value. + Instant min = BoundedWindow.TIMESTAMP_MAX_VALUE; + for (Instant w : committedWatermarkByPartition.values()) { + if (w.isBefore(min)) { + min = w; + } } - long emit = Math.max(min, lastEmittedMillis); - lastEmittedMillis = emit; - return OptionalLong.of(emit); + Instant emit = min.isAfter(lastEmitted) ? min : lastEmitted; + lastEmitted = emit; + return emit; } /** The total source partition count learned in-band, or -1 if nothing reported yet. */ - public int expectedSourcePartitionCount() { + @VisibleForTesting + int expectedSourcePartitionCount() { return expectedSourcePartitionCount; } /** How many distinct source partitions have reported so far. */ - public int reportedPartitionCount() { + @VisibleForTesting + int reportedPartitionCount() { return committedWatermarkByPartition.size(); } } diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java index cf1564592581..1d642e5ea1a7 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkManagerTest.java @@ -21,116 +21,135 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertThrows; -import java.util.OptionalLong; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.joda.time.Instant; import org.junit.Test; /** Tests for {@link WatermarkManager}. */ public class WatermarkManagerTest { + private static Instant ts(long millis) { + return new Instant(millis); + } + @Test public void holdsBeforeAnyReport() { WatermarkManager manager = new WatermarkManager(); - assertThat(manager.isComplete(), is(false)); - assertThat(manager.advance().isPresent(), is(false)); + assertThat(manager.isReady(), is(false)); + assertThat(manager.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); } @Test public void holdsUntilEverySourcePartitionReports() { WatermarkManager manager = new WatermarkManager(); - manager.observe(0, 100L, 4); - manager.observe(1, 100L, 4); - manager.observe(2, 100L, 4); + manager.observe(0, ts(100L), 4); + manager.observe(1, ts(100L), 4); + manager.observe(2, ts(100L), 4); // Three of four partitions reported — still holding. - assertThat(manager.isComplete(), is(false)); - assertThat(manager.advance().isPresent(), is(false)); + assertThat(manager.isReady(), is(false)); + assertThat(manager.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); - manager.observe(3, 100L, 4); - assertThat(manager.isComplete(), is(true)); - assertThat(manager.advance(), is(OptionalLong.of(100L))); + manager.observe(3, ts(100L), 4); + assertThat(manager.isReady(), is(true)); + assertThat(manager.advance(), is(ts(100L))); } @Test - public void emitsMinAcrossPartitionsOnceComplete() { + public void emitsMinAcrossPartitionsOnceReady() { WatermarkManager manager = new WatermarkManager(); - manager.observe(0, 300L, 4); - manager.observe(1, 100L, 4); - manager.observe(2, 500L, 4); - manager.observe(3, 200L, 4); + manager.observe(0, ts(300L), 4); + manager.observe(1, ts(100L), 4); + manager.observe(2, ts(500L), 4); + manager.observe(3, ts(200L), 4); // min(300, 100, 500, 200) = 100 - assertThat(manager.advance(), is(OptionalLong.of(100L))); + assertThat(manager.advance(), is(ts(100L))); } @Test public void perPartitionWatermarkIsMonotonic() { WatermarkManager manager = new WatermarkManager(); - manager.observe(0, 100L, 1); - assertThat(manager.advance(), is(OptionalLong.of(100L))); + manager.observe(0, ts(100L), 1); + assertThat(manager.advance(), is(ts(100L))); // A lower out-of-order report for the same partition is ignored. - manager.observe(0, 50L, 1); - assertThat(manager.advance(), is(OptionalLong.of(100L))); + manager.observe(0, ts(50L), 1); + assertThat(manager.advance(), is(ts(100L))); } @Test - public void emittedWatermarkDoesNotRegressWhenNewPartitionReportsOlderWatermark() { + public void partitionCountChangeClearsReportsAndReopensHold() { WatermarkManager manager = new WatermarkManager(); - manager.observe(0, 100L, 2); - manager.observe(1, 100L, 2); - assertThat(manager.advance(), is(OptionalLong.of(100L))); - - // A repartition grows the source set to 3; a freshly appeared partition reports an older - // watermark. The stage watermark must not go backwards. - manager.observe(2, 50L, 3); - assertThat(manager.isComplete(), is(true)); - assertThat(manager.advance(), is(OptionalLong.of(100L))); + manager.observe(0, ts(100L), 2); + manager.observe(1, ts(100L), 2); + assertThat(manager.advance(), is(ts(100L))); + + // Source set grows to 4. The previous reports are dropped, so the stage holds again until all + // four of the new partition set have reported. + manager.observe(0, ts(200L), 4); + assertThat(manager.reportedPartitionCount(), is(1)); + assertThat(manager.isReady(), is(false)); + assertThat(manager.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE)); + + manager.observe(1, ts(200L), 4); + manager.observe(2, ts(200L), 4); + manager.observe(3, ts(200L), 4); + assertThat(manager.isReady(), is(true)); + assertThat(manager.advance(), is(ts(200L))); } @Test - public void partitionCountIncreaseReopensHold() { + public void emittedWatermarkDoesNotRegressAfterRepartition() { WatermarkManager manager = new WatermarkManager(); - manager.observe(0, 100L, 2); - manager.observe(1, 100L, 2); - assertThat(manager.advance(), is(OptionalLong.of(100L))); - - // Source set grows to 4; until partitions 2 and 3 report, the stage holds again. - manager.observe(2, 200L, 4); - assertThat(manager.isComplete(), is(false)); - assertThat(manager.advance().isPresent(), is(false)); - - manager.observe(3, 200L, 4); - assertThat(manager.isComplete(), is(true)); - // min(100, 100, 200, 200) = 100, which also satisfies the non-regression clamp. - assertThat(manager.advance(), is(OptionalLong.of(100L))); + manager.observe(0, ts(100L), 2); + manager.observe(1, ts(100L), 2); + assertThat(manager.advance(), is(ts(100L))); + + // After a repartition to 3, the new partitions report older watermarks. Once ready again the + // min is 50, but the emitted stage watermark must not go backwards below 100. + manager.observe(0, ts(50L), 3); + manager.observe(1, ts(50L), 3); + manager.observe(2, ts(50L), 3); + assertThat(manager.isReady(), is(true)); + assertThat(manager.advance(), is(ts(100L))); } @Test - public void partitionCountDecreasePrunesStalePartitions() { + public void partitionCountDecreaseClearsAndRecomputes() { WatermarkManager manager = new WatermarkManager(); - manager.observe(0, 100L, 4); - manager.observe(1, 100L, 4); - manager.observe(2, 100L, 4); - manager.observe(3, 50L, 4); - assertThat(manager.advance(), is(OptionalLong.of(50L))); - - // Source set shrinks to 2; partitions 2 and 3 no longer exist and are dropped. The remaining - // set {0, 1} is complete and its min is 100, but the output must not regress below 50. - manager.observe(0, 100L, 2); + manager.observe(0, ts(100L), 4); + manager.observe(1, ts(100L), 4); + manager.observe(2, ts(100L), 4); + manager.observe(3, ts(50L), 4); + assertThat(manager.advance(), is(ts(50L))); + + // Source set shrinks to 2. Reports are cleared; the stage holds until {0, 1} report again. + manager.observe(0, ts(100L), 2); assertThat(manager.expectedSourcePartitionCount(), is(2)); - assertThat(manager.reportedPartitionCount(), is(2)); - assertThat(manager.isComplete(), is(true)); - assertThat(manager.advance(), is(OptionalLong.of(100L))); + assertThat(manager.reportedPartitionCount(), is(1)); + assertThat(manager.isReady(), is(false)); + + manager.observe(1, ts(100L), 2); + assertThat(manager.isReady(), is(true)); + // min is 100; the non-regression clamp keeps it >= the previously emitted 50. + assertThat(manager.advance(), is(ts(100L))); + } + + @Test + public void rejectsNullWatermark() { + WatermarkManager manager = new WatermarkManager(); + assertThrows(IllegalArgumentException.class, () -> manager.observe(0, null, 4)); } @Test public void rejectsNonPositiveTotalSourcePartitions() { WatermarkManager manager = new WatermarkManager(); - assertThrows(IllegalArgumentException.class, () -> manager.observe(0, 100L, 0)); - assertThrows(IllegalArgumentException.class, () -> manager.observe(0, 100L, -1)); + assertThrows(IllegalArgumentException.class, () -> manager.observe(0, ts(100L), 0)); + assertThrows(IllegalArgumentException.class, () -> manager.observe(0, ts(100L), -1)); } @Test public void rejectsOutOfRangeSourcePartition() { WatermarkManager manager = new WatermarkManager(); - assertThrows(IllegalArgumentException.class, () -> manager.observe(-1, 100L, 4)); - assertThrows(IllegalArgumentException.class, () -> manager.observe(4, 100L, 4)); + assertThrows(IllegalArgumentException.class, () -> manager.observe(-1, ts(100L), 4)); + assertThrows(IllegalArgumentException.class, () -> manager.observe(4, ts(100L), 4)); } }