[KafkaIO] Advance idle partition watermark without requiring a record (#40176)

Co-authored-by: Udaya Chathuranga <udaya.chathuranga@trivago.com>

Author: udayaw

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/CustomTimestampPolicyWithLimitedDelay.java

@@ -31,11 +31,18 @@
  * watermark at any time is '({@code Min(now(), Max(event timestamp so far)) - max delay})'.
  * However, watermark is never set in future and capped to 'now - max delay'. In addition, watermark
  * advanced to 'now - max delay' when a partition is idle.
+ *
+ * <p>By default the idle advance requires at least one record to have been read, so a partition
+ * which is caught up but has delivered nothing since the job started holds its watermark at {@link
+ * BoundedWindow#TIMESTAMP_MIN_VALUE} until its first record arrives. See {@link
+ * #CustomTimestampPolicyWithLimitedDelay(SerializableFunction, Duration, Optional, boolean)} to opt
+ * out of that.
  */
 public class CustomTimestampPolicyWithLimitedDelay<K, V> extends TimestampPolicy<K, V> {
 
   private final Duration maxDelay;
   private final SerializableFunction<KafkaRecord<K, V>, Instant> timestampFunction;
+  private final boolean advanceWatermarkBeforeFirstRecord;
   private Instant maxEventTimestamp;
 
   /**
@@ -53,8 +60,28 @@ public CustomTimestampPolicyWithLimitedDelay(
       SerializableFunction<KafkaRecord<K, V>, Instant> timestampFunction,
       Duration maxDelay,
       Optional<Instant> previousWatermark) {
+    this(timestampFunction, maxDelay, previousWatermark, false);
+  }
+
+  /**
+   * As {@link #CustomTimestampPolicyWithLimitedDelay(SerializableFunction, Duration, Optional)},
+   * with control over whether an idle partition may advance its watermark before it has delivered
+   * any record.
+   *
+   * @param advanceWatermarkBeforeFirstRecord When true, a partition with a backlog of zero advances
+   *     its watermark whether or not a record has been read. A zero backlog means the reader is at
+   *     the log end, so no unread record can arrive late. Leaving this false keeps such a partition
+   *     pinned at {@link BoundedWindow#TIMESTAMP_MIN_VALUE} until its first record arrives, which
+   *     holds back the watermark of every stage reading it.
+   */
+  public CustomTimestampPolicyWithLimitedDelay(
+      SerializableFunction<KafkaRecord<K, V>, Instant> timestampFunction,
+      Duration maxDelay,
+      Optional<Instant> previousWatermark,
+      boolean advanceWatermarkBeforeFirstRecord) {
     this.maxDelay = maxDelay;
     this.timestampFunction = timestampFunction;
+    this.advanceWatermarkBeforeFirstRecord = advanceWatermarkBeforeFirstRecord;
 
     // 'previousWatermark' is not the same as maxEventTimestamp (e.g. it could have been in future).
     // Initialize it such that watermark before reading any event same as previousWatermark.
@@ -90,7 +117,8 @@ Instant getWatermark(PartitionContext ctx, Instant now) {
       return now.minus(maxDelay); // (a) above.
     } else if (ctx.getMessageBacklog() == 0
         && ctx.getBacklogCheckTime().minus(maxDelay).isAfter(maxEventTimestamp) // Idle
-        && maxEventTimestamp.getMillis() > 0) { // Read at least one record with positive timestamp.
+        // Read at least one record with positive timestamp, unless opted out of.
+        && (maxEventTimestamp.getMillis() > 0 || advanceWatermarkBeforeFirstRecord)) {
       return ctx.getBacklogCheckTime().minus(maxDelay);
     } else {
       return maxEventTimestamp.minus(maxDelay);

sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/CustomTimestampPolicyWithLimitedDelayTest.java

@@ -115,4 +115,40 @@ record -> new Instant(record.getTimestamp()), maxDelay, Optional.empty());
 
     assertThat(policy.getWatermark(ctx, now), is(backlogCheckTime.minus(maxDelay)));
   }
+
+  private static CustomTimestampPolicyWithLimitedDelay<String, String> policy(
+      Duration maxDelay, boolean advanceWatermarkBeforeFirstRecord) {
+    return new CustomTimestampPolicyWithLimitedDelay<>(
+        record -> new Instant(record.getTimestamp()),
+        maxDelay,
+        Optional.empty(),
+        advanceWatermarkBeforeFirstRecord);
+  }
+
+  @Test
+  public void testIdleWatermarkIsPinnedBeforeFirstRecordByDefault() {
+    // Default behaviour: a caught up partition which has delivered nothing stays at the floor.
+    Duration maxDelay = Duration.standardSeconds(60);
+    CustomTimestampPolicyWithLimitedDelay<String, String> policy = policy(maxDelay, false);
+
+    Instant now = Instant.now();
+    TimestampPolicy.PartitionContext ctx = mock(TimestampPolicy.PartitionContext.class);
+    when(ctx.getMessageBacklog()).thenReturn(0L);
+    when(ctx.getBacklogCheckTime()).thenReturn(now);
+
+    assertThat(policy.getWatermark(ctx, now), is(BoundedWindow.TIMESTAMP_MIN_VALUE));
+  }
+
+  @Test
+  public void testIdleWatermarkAdvancesBeforeFirstRecordWhenEnabled() {
+    Duration maxDelay = Duration.standardSeconds(60);
+    CustomTimestampPolicyWithLimitedDelay<String, String> policy = policy(maxDelay, true);
+
+    Instant now = Instant.now();
+    TimestampPolicy.PartitionContext ctx = mock(TimestampPolicy.PartitionContext.class);
+    when(ctx.getMessageBacklog()).thenReturn(0L);
+    when(ctx.getBacklogCheckTime()).thenReturn(now);
+
+    assertThat(policy.getWatermark(ctx, now), is(now.minus(maxDelay)));
+  }
 }