feat: Option to skip replaying events after failure recovery in metronome source (#337)

Author: ferenc-csaky

connectors/datagen-connectors/src/main/java/com/datasqrl/flinkrunner/connector/datagen/metronome/MetronomeDynamicTableSourceFactory.java

@@ -18,6 +18,7 @@
 import com.google.auto.service.AutoService;
 import java.util.Set;
 import org.apache.flink.configuration.ConfigOption;
+import org.apache.flink.configuration.ConfigOptions;
 import org.apache.flink.table.api.ValidationException;
 import org.apache.flink.table.connector.source.DynamicTableSource;
 import org.apache.flink.table.factories.DynamicTableSourceFactory;
@@ -32,14 +33,22 @@ public class MetronomeDynamicTableSourceFactory implements DynamicTableSourceFac
 
   public static final String IDENTIFIER = "metronome";
 
+  public static final ConfigOption<Boolean> REPLAY_ON_FAILURE =
+      ConfigOptions.key("replay-on-failure")
+          .booleanType()
+          .defaultValue(true)
+          .withDescription(
+              "Whether to emit sequence numbers missed while recovering from a checkpoint.");
+
   @Override
   public DynamicTableSource createDynamicTableSource(Context context) {
-    FactoryUtil.createTableFactoryHelper(this, context).validate();
+    var helper = FactoryUtil.createTableFactoryHelper(this, context);
+    helper.validate();
 
     var rowDataType = context.getPhysicalRowDataType();
     validateSchema(rowDataType);
 
-    return new MetronomeTableSource(rowDataType);
+    return new MetronomeTableSource(rowDataType, helper.getOptions().get(REPLAY_ON_FAILURE));
   }
 
   @Override
@@ -54,7 +63,7 @@ public Set<ConfigOption<?>> requiredOptions() {
 
   @Override
   public Set<ConfigOption<?>> optionalOptions() {
-    return Set.of();
+    return Set.of(REPLAY_ON_FAILURE);
   }
 
   private static void validateSchema(DataType rowDataType) {

connectors/datagen-connectors/src/main/java/com/datasqrl/flinkrunner/connector/datagen/metronome/MetronomeReader.java

@@ -38,26 +38,34 @@ public class MetronomeReader implements SourceReader<RowData, MetronomeSplit> {
   public static final long UNINITIALIZED = Long.MIN_VALUE;
 
   private final SourceReaderContext readerContext;
+  private final boolean replayOnFailure;
   @Nullable private final Long numberOfRows;
   private final ScheduledExecutorService availabilityExecutor;
 
   private CompletableFuture<Void> availability;
   private boolean noMoreSplits;
   private boolean closed;
   private boolean splitAssigned;
+  private boolean skipReplayOnRecovery;
   private long lastEmittedNumber;
   private long startTimestampSec;
 
-  public MetronomeReader(SourceReaderContext readerContext, @Nullable Long numberOfRows) {
+  public MetronomeReader(
+      SourceReaderContext readerContext, boolean replayOnFailure, @Nullable Long numberOfRows) {
     this.readerContext = readerContext;
     this.numberOfRows = numberOfRows;
+    this.replayOnFailure = replayOnFailure;
     this.availability = new CompletableFuture<>();
     this.availabilityExecutor =
         Executors.newSingleThreadScheduledExecutor(new ExecutorThreadFactory("metronome-source"));
     this.lastEmittedNumber = 0L;
     this.startTimestampSec = UNINITIALIZED;
   }
 
+  public MetronomeReader(SourceReaderContext readerContext, @Nullable Long numberOfRows) {
+    this(readerContext, true, numberOfRows);
+  }
+
   @Override
   public void start() {
     readerContext.sendSplitRequest();
@@ -68,9 +76,10 @@ public void start() {
    * second boundary.
    *
    * <p>The emitted sequence number is derived from wall-clock epoch seconds relative to the first
-   * observed start second. If the reader wakes up late or recovers from a checkpoint, repeated
-   * calls emit all missing sequence numbers with the current wall-clock timestamp until the source
-   * catches up.
+   * observed start second. If the reader wakes up late, repeated calls emit all missing sequence
+   * numbers with the current wall-clock timestamp until the source catches up. If checkpoint
+   * recovery replay is disabled, only the first poll after restoring checkpointed progress skips
+   * the recovery backlog.
    */
   @Override
   public InputStatus pollNext(ReaderOutput<RowData> output) {
@@ -92,16 +101,27 @@ public InputStatus pollNext(ReaderOutput<RowData> output) {
       targetNumber = Math.min(targetNumber, numberOfRows);
     }
 
+    // One-shot checkpoint recovery skip; normal late wakeups still replay.
+    boolean skipReplay = skipReplayOnRecovery;
+    skipReplayOnRecovery = false;
+
     if (lastEmittedNumber < targetNumber) {
       long nextNumber = lastEmittedNumber + 1;
       long eventTimestampMillis = currentTimestampSec * 1000L;
       var row =
           GenericRowData.of(
               nextNumber, TimestampData.fromInstant(Instant.ofEpochSecond(currentTimestampSec)));
       lastEmittedNumber = nextNumber;
+      if (skipReplay) {
+        // Realign time so the restored backlog is dropped after this record.
+        startTimestampSec = currentTimestampSec - lastEmittedNumber;
+      }
       output.collect(row, eventTimestampMillis);
 
-      return lastEmittedNumber < targetNumber ? InputStatus.MORE_AVAILABLE : nextStatus();
+      // Skip waits for the next tick; otherwise drain due records until caught up.
+      return !skipReplay && lastEmittedNumber < targetNumber
+          ? InputStatus.MORE_AVAILABLE
+          : nextStatus();
     }
 
     return nextStatus();
@@ -111,7 +131,7 @@ public InputStatus pollNext(ReaderOutput<RowData> output) {
    * Snapshots the assigned split as the reader's progress state.
    *
    * <p>The split carries the only source progress that must survive failover: the last emitted
-   * sequence number and the source's fixed start epoch second.
+   * sequence number and the source's effective start epoch second.
    */
   @Override
   public List<MetronomeSplit> snapshotState(long checkpointId) {
@@ -134,6 +154,9 @@ public void addSplits(List<MetronomeSplit> splits) {
     for (var split : splits) {
       lastEmittedNumber = split.lastEmittedNumber();
       startTimestampSec = split.startTimestampSec();
+      // Non-initial split means checkpoint/savepoint recovery.
+      skipReplayOnRecovery =
+          !replayOnFailure && (startTimestampSec != UNINITIALIZED || lastEmittedNumber > 0L);
       splitAssigned = true;
       break;
     }

connectors/datagen-connectors/src/main/java/com/datasqrl/flinkrunner/connector/datagen/metronome/MetronomeSource.java

@@ -38,10 +38,15 @@ public class MetronomeSource implements Source<RowData, MetronomeSplit, Metronom
 
   @Serial private static final long serialVersionUID = 1L;
 
+  private final boolean replayOnFailure;
   @Nullable private final Long numberOfRows;
 
   public MetronomeSource() {
-    this(null);
+    this(true, null);
+  }
+
+  public MetronomeSource(@Nullable Long numberOfRows) {
+    this(true, numberOfRows);
   }
 
   @Override
@@ -51,7 +56,7 @@ public Boundedness getBoundedness() {
 
   @Override
   public SourceReader<RowData, MetronomeSplit> createReader(SourceReaderContext readerContext) {
-    return new MetronomeReader(readerContext, numberOfRows);
+    return new MetronomeReader(readerContext, replayOnFailure, numberOfRows);
   }
 
   @Override

connectors/datagen-connectors/src/main/java/com/datasqrl/flinkrunner/connector/datagen/metronome/MetronomeTableSource.java

@@ -32,11 +32,16 @@ public class MetronomeTableSource implements ScanTableSource, SupportsLimitPushD
   private static final int SOURCE_PARALLELISM = 1;
 
   private final DataType rowDataType;
+  private final boolean replayOnFailure;
 
   @Nullable private Long numberOfRows;
 
   public MetronomeTableSource(DataType rowDataType) {
-    this(rowDataType, null);
+    this(rowDataType, true, null);
+  }
+
+  public MetronomeTableSource(DataType rowDataType, boolean replayOnFailure) {
+    this(rowDataType, replayOnFailure, null);
   }
 
   @Override
@@ -46,12 +51,13 @@ public ChangelogMode getChangelogMode() {
 
   @Override
   public ScanRuntimeProvider getScanRuntimeProvider(ScanContext scanContext) {
-    return SourceProvider.of(new MetronomeSource(numberOfRows), SOURCE_PARALLELISM);
+    return SourceProvider.of(
+        new MetronomeSource(replayOnFailure, numberOfRows), SOURCE_PARALLELISM);
   }
 
   @Override
   public DynamicTableSource copy() {
-    return new MetronomeTableSource(rowDataType, numberOfRows);
+    return new MetronomeTableSource(rowDataType, replayOnFailure, numberOfRows);
   }
 
   @Override

connectors/datagen-connectors/src/test/java/com/datasqrl/flinkrunner/connector/datagen/metronome/MetronomeSourceIT.java

@@ -21,6 +21,7 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 import org.apache.flink.api.common.eventtime.Watermark;
 import org.apache.flink.api.connector.source.ReaderOutput;
 import org.apache.flink.api.connector.source.SourceEvent;
@@ -72,10 +73,28 @@ ts TIMESTAMP_LTZ(3),
         .containsExactly(Row.of(1L), Row.of(2L), Row.of(3L));
   }
 
+  @Test
+  @Timeout(30)
+  void acceptsDisabledReplayOnFailureOptionInSqlDdl() throws Exception {
+    tEnv.executeSql(
+        """
+            CREATE TABLE metronome_source_no_replay (
+              num BIGINT,
+              ts TIMESTAMP_LTZ(3),
+              WATERMARK FOR ts AS ts
+            ) WITH (
+              'connector' = 'metronome',
+              'replay-on-failure' = 'false'
+            )""");
+
+    assertThat(collectRows("SELECT num FROM metronome_source_no_replay", 1))
+        .containsExactly(Row.of(1L));
+  }
+
   @Test
   void restoresReaderProgressWithCurrentTimestampFromCheckpointedSplit() {
     long startTimestampSec = Instant.now().getEpochSecond() - 3L;
-    var reader = new MetronomeReader(unusedReaderContext(), 3L);
+    var reader = new MetronomeReader(unusedReaderContext(), true, 3L);
     var output = new CollectingReaderOutput();
 
     try {
@@ -95,6 +114,55 @@ void restoresReaderProgressWithCurrentTimestampFromCheckpointedSplit() {
     }
   }
 
+  @Test
+  void skipsMissedEventsOnCheckpointRecoveryWhenReplayOnFailureIsDisabled() {
+    long startTimestampSec = Instant.now().getEpochSecond() - 3L;
+    var reader = new MetronomeReader(unusedReaderContext(), false, 5L);
+    var output = new CollectingReaderOutput();
+
+    try {
+      reader.addSplits(List.of(new MetronomeSplit(1L, startTimestampSec)));
+
+      long beforePollSec = Instant.now().getEpochSecond();
+      assertThat(reader.pollNext(output)).isEqualTo(InputStatus.NOTHING_AVAILABLE);
+      long afterPollSec = Instant.now().getEpochSecond();
+
+      assertThat(output.numbers()).containsExactly(2L);
+      assertThat(output.eventTimestampSecs().get(0)).isBetween(beforePollSec, afterPollSec);
+      assertThat(output.rowTimestampSecs().get(0)).isBetween(beforePollSec, afterPollSec);
+
+      var snapshot = reader.snapshotState(1L);
+      assertThat(snapshot).hasSize(1);
+      assertThat(snapshot.get(0).lastEmittedNumber()).isEqualTo(2L);
+      assertThat(snapshot.get(0).startTimestampSec())
+          .isBetween(beforePollSec - 2L, afterPollSec - 2L);
+    } finally {
+      reader.close();
+    }
+  }
+
+  @Test
+  void replaysNormalWakeUpLatenessWhenReplayOnFailureIsDisabled() throws Exception {
+    long startTimestampSec = Instant.now().getEpochSecond() - 3L;
+    var reader = new MetronomeReader(unusedReaderContext(), false, 5L);
+    var output = new CollectingReaderOutput();
+
+    try {
+      reader.addSplits(List.of(new MetronomeSplit(1L, startTimestampSec)));
+
+      assertThat(reader.pollNext(output)).isEqualTo(InputStatus.NOTHING_AVAILABLE);
+      var snapshot = reader.snapshotState(1L);
+      assertThat(snapshot).hasSize(1);
+
+      sleepUntilAtLeastSecond(snapshot.get(0).startTimestampSec() + 4L);
+
+      assertThat(reader.pollNext(output)).isEqualTo(InputStatus.MORE_AVAILABLE);
+      assertThat(output.numbers()).containsExactly(2L, 3L);
+    } finally {
+      reader.close();
+    }
+  }
+
   private List<Row> collectRows(String sql, int rowCount) throws Exception {
     List<Row> rows = new ArrayList<>();
     try (CloseableIterator<Row> iterator = tEnv.executeSql(sql).collect()) {
@@ -140,6 +208,12 @@ public UserCodeClassLoader getUserCodeClassLoader() {
     };
   }
 
+  private static void sleepUntilAtLeastSecond(long epochSecond) throws InterruptedException {
+    while (Instant.now().getEpochSecond() < epochSecond) {
+      TimeUnit.MILLISECONDS.sleep(10L);
+    }
+  }
+
   private static final class CollectingReaderOutput implements ReaderOutput<RowData> {
 
     private final List<RowData> rows = new ArrayList<>();