This is an automated email from the ASF dual-hosted git repository.

danny0405 pushed a commit to branch release-1.1.2
in repository https://gitbox.apache.org/repos/asf/hudi.git

commit fcec6cc80a0c2e6baf78d11562e257ef62c2f067
Author: Danny Chan <[email protected]>
AuthorDate: Wed Sep 16 12:30:09 2026 +0800

    fix(flink): use restored checkpoint IDs for writers and RLI caches (#19952)
    
    * fix(flink): use restored checkpoint IDs when initializing writers
    
    (cherry picked from commit 43632ec5161fa672bbeb4e5d7e4688793792e80c)
---
 .../sink/common/AbstractStreamWriteFunction.java   |  46 +----
 .../sink/TestStreamWriteOperatorCoordinator.java   | 134 ++++++++++++--
 .../common/TestAbstractStreamWriteFunction.java    | 204 +++++++++++++++++++++
 3 files changed, 326 insertions(+), 58 deletions(-)

diff --git 
a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/common/AbstractStreamWriteFunction.java
 
b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/common/AbstractStreamWriteFunction.java
index deedfbba8b0e..0884ea47517a 100644
--- 
a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/common/AbstractStreamWriteFunction.java
+++ 
b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/common/AbstractStreamWriteFunction.java
@@ -31,7 +31,6 @@ import org.apache.hudi.util.FlinkWriteClients;
 import org.apache.hudi.util.StreamerUtil;
 import org.apache.hudi.utils.RuntimeContextUtils;
 
-import org.apache.flink.api.common.JobID;
 import org.apache.flink.api.common.state.ListState;
 import org.apache.flink.api.common.state.ListStateDescriptor;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
@@ -46,7 +45,6 @@ import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.stream.StreamSupport;
 
 /**
  * Base infrastructures for streaming writer function.
@@ -100,11 +98,6 @@ public abstract class AbstractStreamWriteFunction<I>
    */
   private transient ListState<WriteMetadataEvent> writeMetadataState;
 
-  /**
-   * List state of the JobID.
-   */
-  private transient ListState<JobID> jobIdState;
-
   /**
    * Write status list for the current checkpoint.
    */
@@ -141,15 +134,13 @@ public abstract class AbstractStreamWriteFunction<I>
             "write-metadata-state",
             TypeInformation.of(WriteMetadataEvent.class)
         ));
-    this.jobIdState = context.getOperatorStateStore().getListState(
-        new ListStateDescriptor<>(
-            "job-id-state",
-            TypeInformation.of(JobID.class)
-        ));
 
     int attemptId = RuntimeContextUtils.getAttemptNumber(getRuntimeContext());
     if (context.isRestored()) {
-      initCheckpointId(attemptId, 
context.getRestoredCheckpointId().orElse(-1L));
+      // sets up the known checkpoint id as the last successful checkpoint id 
for purposes of:
+      // 1). old events cleaning;
+      // 2). instant time request for current checkpoint.
+      this.checkpointId = context.getRestoredCheckpointId().orElse(-1L);
     }
     sendBootstrapEvent(attemptId, context.isRestored());
   }
@@ -162,8 +153,6 @@ public abstract class AbstractStreamWriteFunction<I>
     snapshotState();
     // Reload the snapshot state as the current state.
     reloadWriteMetaState();
-    // Reload the job ID state
-    reloadJobIdState();
     // Update checkpoint id
     this.checkpointId = functionSnapshotContext.getCheckpointId();
   }
@@ -199,24 +188,7 @@ public abstract class AbstractStreamWriteFunction<I>
   //  Utilities
   // -------------------------------------------------------------------------
 
-  private void initCheckpointId(int attemptId, long restoredCheckpointId) 
throws Exception {
-    if (attemptId <= 0) {
-      // returns early if the job/task is initially started.
-      return;
-    }
-    JobID currentJobId = RuntimeContextUtils.getJobId(getRuntimeContext());
-    if (StreamSupport.stream(this.jobIdState.get().spliterator(), false)
-        .noneMatch(currentJobId::equals)) {
-      // do not set up the checkpoint id if the state comes from the old job.
-      return;
-    }
-    // sets up the known checkpoint id as the last successful checkpoint id 
for purposes of:
-    // 1). old events cleaning;
-    // 2). instant time request for current checkpoint.
-    this.checkpointId = restoredCheckpointId;
-  }
-
-  private void sendBootstrapEvent(int attemptId, boolean isRestored) throws 
Exception {
+  protected void sendBootstrapEvent(int attemptId, boolean isRestored) throws 
Exception {
     if (attemptId <= 0) {
       if (isRestored) {
         HoodieTimeline pendingTimeline = 
this.metaClient.getActiveTimeline().filterPendingExcludingCompaction();
@@ -257,14 +229,6 @@ public abstract class AbstractStreamWriteFunction<I>
     writeStatuses.clear();
   }
 
-  /**
-   * Reload job id state as current job id.
-   */
-  private void reloadJobIdState() throws Exception {
-    this.jobIdState.clear();
-    this.jobIdState.add(RuntimeContextUtils.getJobId(getRuntimeContext()));
-  }
-
   public void handleOperatorEvent(OperatorEvent event) {
     ValidationUtils.checkArgument(event instanceof CommitAckEvent,
         "The write function can only handle CommitAckEvent");
diff --git 
a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/TestStreamWriteOperatorCoordinator.java
 
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/TestStreamWriteOperatorCoordinator.java
index 5ae64be04f20..c2e7a0da6ef8 100644
--- 
a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/TestStreamWriteOperatorCoordinator.java
+++ 
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/TestStreamWriteOperatorCoordinator.java
@@ -58,8 +58,11 @@ import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.Mockito;
 import org.slf4j.Logger;
 
@@ -136,36 +139,40 @@ public class TestStreamWriteOperatorCoordinator {
    * the event buffers for restoration during start. In case 2, global 
failover resets an already
    * started coordinator, so it recommits its live event buffers directly.
    */
-  @Test
-  public void testCheckpointAndRestore() throws Exception {
+  @ParameterizedTest
+  @ValueSource(booleans = {false, true})
+  public void testCheckpointAndRestore(boolean restartCoordinator) throws 
Exception {
     Configuration conf = 
TestConfigurations.getDefaultConf(tempFile.getAbsolutePath());
-    conf.set(FlinkOptions.TABLE_TYPE, HoodieTableType.MERGE_ON_READ.name());
-    coordinator = startCoordinator(conf, 2);
 
     requestInstantTime(-1);
     String instant = coordinator.getInstant();
     assertNotEquals("", instant);
 
+    // Writer checkpoint 1 finishes after the coordinator has taken its 
snapshot.
+    coordinator.checkpointCoordinator(1, new CompletableFuture<>());
     OperatorEvent event0 = createOperatorEvent(0, instant, "par1", true, 0.1);
     OperatorEvent event1 = createOperatorEvent(1, instant, "par2", true, 0.2);
     coordinator.handleEventFromOperator(0, event0);
     coordinator.handleEventFromOperator(1, event1);
 
     CompletableFuture<byte[]> future = new CompletableFuture<>();
-    coordinator.checkpointCoordinator(1, future);
-
-    // Case 1: job restart restores checkpoint data before the coordinator 
starts.
-    try (StreamWriteOperatorCoordinator restoredCoordinator = 
createCoordinator(conf, 2)) {
-      restoredCoordinator.resetToCheckpoint(1, future.get());
+    coordinator.checkpointCoordinator(2, future);
 
-      WriteMetadataEvent[] eventBuffer = 
restoredCoordinator.getEventBuffer(-1);
+    if (restartCoordinator) {
+      coordinator.close();
+      coordinator = createCoordinator(conf, 4);
+      coordinator.resetToCheckpoint(2, future.get());
+      WriteMetadataEvent[] eventBuffer = coordinator.getEventBuffer(-1);
       assertEquals(2, eventBuffer.length);
+      // Exercise recommit during start, not only checkpoint deserialization.
+      coordinator.start();
+    } else {
+      // Global failover recommits the live buffers of the already started 
coordinator.
+      coordinator.resetToCheckpoint(2, future.get());
     }
 
-    // Case 2: global failover recommits the live buffers of the already 
started coordinator.
-    coordinator.resetToCheckpoint(1, future.get());
-
-    assertNull(coordinator.getEventBuffer());
+    assertNull(coordinator.getEventBuffer(-1));
+    assertNull(((MockOperatorCoordinatorContext) 
coordinator.getContext()).getJobFailureReason());
     assertTrue(StreamerUtil.createMetaClient(conf).reloadActiveTimeline()
         .filterCompletedInstants().containsInstant(instant));
   }
@@ -185,11 +192,93 @@ public class TestStreamWriteOperatorCoordinator {
         "Receive an unexpected event for instant abc from task 0");
   }
 
+  @Test
+  void testDeferredRecommitAfterScaleUp() throws Exception {
+    Configuration conf = 
TestConfigurations.getDefaultConf(tempFile.getAbsolutePath());
+    String restoredInstant = restoreFirstCheckpointAfterScaleUp(conf);
+    String nextInstant = requestInstantTime(1);
+    assertNotEquals(restoredInstant, nextInstant);
+    assertEquals(nextInstant, coordinator.getInstant());
+    coordinator.checkpointCoordinator(2, new CompletableFuture<>());
+    sendCheckpointEvents(1, nextInstant, 4);
+
+    coordinator.notifyCheckpointComplete(2);
+
+    HoodieTimeline completed = StreamerUtil.createMetaClient(conf)
+        .reloadActiveTimeline().filterCompletedInstants();
+    assertTrue(completed.containsInstant(restoredInstant));
+    assertTrue(completed.containsInstant(nextInstant));
+    assertEquals(2, 
completed.readCommitMetadata(INSTANT_GENERATOR.createNewInstant(
+        HoodieInstant.State.COMPLETED, HoodieTimeline.DELTA_COMMIT_ACTION, 
restoredInstant))
+        .getPartitionToWriteStats().size(), "Both original writers must be 
included in the restored commit");
+    assertEquals(4, 
completed.readCommitMetadata(INSTANT_GENERATOR.createNewInstant(
+        HoodieInstant.State.COMPLETED, HoodieTimeline.DELTA_COMMIT_ACTION, 
nextInstant))
+        .getPartitionToWriteStats().size(), "All scaled-up writers must be 
included in the next commit");
+    assertNull(coordinator.getEventBuffer(-1));
+    assertNull(coordinator.getEventBuffer(1));
+    assertNull(((MockOperatorCoordinatorContext) 
coordinator.getContext()).getJobFailureReason());
+  }
+
+  @Disabled("https://github.com/apache/hudi/issues/19922: deferred bootstrap 
metadata is omitted from the next checkpoint")
+  @Test
+  void testDeferredRecommitSurvivesAnotherRestart() throws Exception {
+    Configuration conf = 
TestConfigurations.getDefaultConf(tempFile.getAbsolutePath());
+    String restoredInstant = restoreFirstCheckpointAfterScaleUp(conf);
+    String nextInstant = requestInstantTime(1);
+    CompletableFuture<byte[]> secondCheckpoint = new CompletableFuture<>();
+    coordinator.checkpointCoordinator(2, secondCheckpoint);
+    sendCheckpointEvents(1, nextInstant, 4);
+    coordinator.close();
+
+    // Checkpoint 2 succeeded, but the job stopped before 
notifyCheckpointComplete could commit it.
+    // The writers now restore only checkpoint 2's batch; the coordinator must 
preserve the older batch.
+    coordinator = createCoordinator(conf, 4);
+    coordinator.resetToCheckpoint(2, secondCheckpoint.get());
+    coordinator.start();
+    setSynchronousExecutors(coordinator);
+    for (int task = 0; task < 4; task++) {
+      coordinator.handleEventFromOperator(task, createBootstrapEvent(task, 1, 
nextInstant, "new" + task));
+    }
+    coordinator.notifyCheckpointComplete(3);
+
+    HoodieTimeline completed = 
StreamerUtil.createMetaClient(conf).reloadActiveTimeline().filterCompletedInstants();
+    assertTrue(completed.containsInstant(nextInstant));
+    assertTrue(completed.containsInstant(restoredInstant), "The first batch 
must survive a second restart before its deferred commit");
+  }
+
+  private String restoreFirstCheckpointAfterScaleUp(Configuration conf) throws 
Exception {
+    resetToMergeOnRead(conf);
+    String restoredInstant = requestInstantTime(-1);
+    CompletableFuture<byte[]> firstCheckpoint = new CompletableFuture<>();
+    // The first coordinator snapshot precedes the writer snapshots, so only 
writer state has this batch.
+    coordinator.checkpointCoordinator(1, firstCheckpoint);
+    coordinator.close();
+
+    coordinator = createCoordinator(conf, 4);
+    coordinator.resetToCheckpoint(1, firstCheckpoint.get());
+    coordinator.start();
+    setSynchronousExecutors(coordinator);
+    coordinator.handleEventFromOperator(0, createBootstrapEvent(0, -1, 
restoredInstant, "par1"));
+    coordinator.handleEventFromOperator(1, createBootstrapEvent(1, -1, 
restoredInstant, "par2"));
+    WriteMetadataEvent[] buffer = coordinator.getEventBuffer(-1);
+    assertFalse(java.util.Arrays.stream(buffer).allMatch(event -> event != 
null && event.isBootstrap()));
+    assertEquals("", coordinator.getInstant());
+    
assertFalse(StreamerUtil.createMetaClient(conf).reloadActiveTimeline().filterCompletedInstants().containsInstant(restoredInstant));
+    return restoredInstant;
+  }
+
+  private void sendCheckpointEvents(long checkpointId, String instant, int 
parallelism) {
+    for (int task = 0; task < parallelism; task++) {
+      coordinator.handleEventFromOperator(task,
+          createOperatorEvent(task, checkpointId, instant, "new" + task, 
false, true, 0.1));
+    }
+    assertNull(((MockOperatorCoordinatorContext) 
coordinator.getContext()).getJobFailureReason());
+  }
+
   @Test
   public void testEventReset() throws Exception {
     Configuration conf = 
TestConfigurations.getDefaultConf(tempFile.getAbsolutePath());
-    conf.set(FlinkOptions.TABLE_TYPE, HoodieTableType.MERGE_ON_READ.name());
-    coordinator = startCoordinator(conf, 2);
+    resetToMergeOnRead(conf);
     CompletableFuture<byte[]> future = new CompletableFuture<>();
     coordinator.checkpointCoordinator(1, future);
     String instant = requestInstantTime(0);
@@ -613,13 +702,24 @@ public class TestStreamWriteOperatorCoordinator {
     }
   }
 
+  private void resetToMergeOnRead(Configuration conf) throws Exception {
+    coordinator.close();
+    reset();
+    conf.set(FlinkOptions.TABLE_TYPE, HoodieTableType.MERGE_ON_READ.name());
+    coordinator = startCoordinator(conf, 2);
+  }
+
   private static StreamWriteOperatorCoordinator startCoordinator(Configuration 
conf, int subTasks) throws Exception {
     StreamWriteOperatorCoordinator coordinator = createCoordinator(conf, 
subTasks);
     coordinator.start();
+    setSynchronousExecutors(coordinator);
+    return coordinator;
+  }
+
+  private static void setSynchronousExecutors(StreamWriteOperatorCoordinator 
coordinator) throws Exception {
     MockOperatorCoordinatorContext coordinatorContext = 
(MockOperatorCoordinatorContext) coordinator.getContext();
     coordinator.setExecutor(new MockCoordinatorExecutor(coordinatorContext));
     coordinator.setInstantRequestExecutor(new 
MockCoordinatorExecutor(coordinatorContext));
-    return coordinator;
   }
 
   private static StreamWriteOperatorCoordinator 
createCoordinator(Configuration conf, int subTasks) {
diff --git 
a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/common/TestAbstractStreamWriteFunction.java
 
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/common/TestAbstractStreamWriteFunction.java
new file mode 100644
index 000000000000..b87ee0aa0a95
--- /dev/null
+++ 
b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/common/TestAbstractStreamWriteFunction.java
@@ -0,0 +1,204 @@
+/*
+ * 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.hudi.sink.common;
+
+import org.apache.hudi.client.HoodieFlinkWriteClient;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.common.table.timeline.HoodieTimeline;
+import org.apache.hudi.sink.event.Correspondent;
+import org.apache.hudi.sink.event.WriteMetadataEvent;
+import org.apache.hudi.sink.utils.MockOperatorStateStore;
+import org.apache.hudi.util.FlinkWriteClients;
+import org.apache.hudi.util.StreamerUtil;
+import org.apache.hudi.utils.RuntimeContextUtils;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.functions.RuntimeContext;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.state.FunctionInitializationContext;
+import 
org.apache.flink.streaming.api.operators.collect.utils.MockFunctionSnapshotContext;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.util.Collector;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.MockedStatic;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.OptionalLong;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests checkpoint identity and bootstrap events independently of the writer 
implementation.
+ */
+class TestAbstractStreamWriteFunction {
+  private final Configuration conf = new Configuration();
+  private final JobID jobId = new JobID();
+  private final List<OperatorEvent> events = new ArrayList<>();
+
+  private final MockOperatorStateStore stateStore = new 
MockOperatorStateStore();
+  private final HoodieTimeline pendingTimeline = mock(HoodieTimeline.class);
+  private final Correspondent correspondent = mock(Correspondent.class);
+  private final TestWriteFunction function = new TestWriteFunction(conf);
+
+  @AfterEach
+  void tearDown() throws Exception {
+    function.close();
+  }
+
+  @ParameterizedTest
+  @ValueSource(ints = {0, 1})
+  void testFreshStartUsesInitialCheckpointId(int attempt) throws Exception {
+    initialize(-1L, attempt);
+
+    assertEquals("002", function.instantToWrite(true));
+    verify(correspondent).requestInstantTime(-1L);
+    if (attempt == 0) {
+      assertTrue(events.isEmpty());
+    } else {
+      assertCleanupEvent(-1L);
+    }
+  }
+
+  @ParameterizedTest
+  @CsvSource({"0, SAME", "0, DIFFERENT", "0, MISSING", "1, SAME", "1, 
DIFFERENT", "1, MISSING"})
+  void testRestoredCheckpointId(int attempt, SavedJobState savedJob) throws 
Exception {
+    // Old savepoints may contain this state; new scale-up subtasks may have 
no operator state at all.
+    if (savedJob != SavedJobState.MISSING) {
+      stateStore.getListState(new ListStateDescriptor<>("job-id-state", 
TypeInformation.of(JobID.class)))
+          .add(savedJob == SavedJobState.SAME ? jobId : new JobID());
+    }
+    initialize(42L, attempt);
+
+    assertEquals("002", function.instantToWrite(true));
+    verify(correspondent).requestInstantTime(42L);
+    if (attempt == 0) {
+      assertTrue(events.isEmpty());
+    } else {
+      assertCleanupEvent(42L);
+    }
+
+    function.snapshotState(new MockFunctionSnapshotContext(43L));
+    WriteMetadataEvent snapshot = writeMetadataState().get().iterator().next();
+    assertEquals(42L, snapshot.getCheckpointId(), "Saved metadata belongs to 
the batch before checkpoint 43");
+    assertEquals("002", snapshot.getInstantTime());
+    assertTrue(snapshot.isBootstrap());
+    function.instantToWrite(true);
+    verify(correspondent).requestInstantTime(43L);
+  }
+
+  @Test
+  void testRestoredMetadataKeepsOriginalCheckpointId() throws Exception {
+    WriteMetadataEvent restored = WriteMetadataEvent.builder()
+        .taskID(0)
+        .checkpointId(41L)
+        .instantTime("001")
+        .writeStatus(Collections.emptyList())
+        .bootstrap(true)
+        .lastBatch(true)
+        .build();
+    writeMetadataState().add(restored);
+    when(pendingTimeline.containsInstant("001")).thenReturn(true);
+
+    initialize(42L, 0);
+
+    assertEquals(1, events.size());
+    WriteMetadataEvent bootstrap = (WriteMetadataEvent) events.get(0);
+    assertEquals(2, bootstrap.getTaskID(), "Restore must use the new subtask 
ID");
+    assertEquals(41L, bootstrap.getCheckpointId());
+    assertEquals("001", bootstrap.getInstantTime());
+    function.instantToWrite(true);
+    verify(correspondent).requestInstantTime(42L);
+  }
+
+  private void initialize(long checkpointId, int attempt) throws Exception {
+    RuntimeContext runtimeContext = mock(RuntimeContext.class);
+    FunctionInitializationContext context = 
mock(FunctionInitializationContext.class);
+    when(context.getOperatorStateStore()).thenReturn(stateStore);
+    when(context.isRestored()).thenReturn(checkpointId >= 0);
+    when(context.getRestoredCheckpointId()).thenReturn(checkpointId >= 0 ? 
OptionalLong.of(checkpointId) : OptionalLong.empty());
+    when(correspondent.requestInstantTime(anyLong())).thenReturn("002");
+    function.setRuntimeContext(runtimeContext);
+    function.setCorrespondent(correspondent);
+    function.setOperatorEventGateway(events::add);
+
+    try (MockedStatic<StreamerUtil> streamerUtil = 
mockStatic(StreamerUtil.class);
+         MockedStatic<FlinkWriteClients> writeClients = 
mockStatic(FlinkWriteClients.class);
+         MockedStatic<RuntimeContextUtils> runtimeContextUtils = 
mockStatic(RuntimeContextUtils.class)) {
+      HoodieTableMetaClient metaClient = mock(HoodieTableMetaClient.class, 
RETURNS_DEEP_STUBS);
+      
when(metaClient.getActiveTimeline().filterPendingExcludingCompaction()).thenReturn(pendingTimeline);
+      streamerUtil.when(() -> 
StreamerUtil.createMetaClient(conf)).thenReturn(metaClient);
+      writeClients.when(() -> FlinkWriteClients.createWriteClient(conf, 
runtimeContext)).thenReturn(mock(HoodieFlinkWriteClient.class));
+      runtimeContextUtils.when(() -> 
RuntimeContextUtils.getJobId(runtimeContext)).thenReturn(jobId);
+      runtimeContextUtils.when(() -> 
RuntimeContextUtils.getIndexOfThisSubtask(runtimeContext)).thenReturn(2);
+      runtimeContextUtils.when(() -> 
RuntimeContextUtils.getAttemptNumber(runtimeContext)).thenReturn(attempt);
+      function.initializeState(context);
+    }
+  }
+
+  private void assertCleanupEvent(long checkpointId) {
+    assertEquals(1, events.size());
+    WriteMetadataEvent bootstrap = (WriteMetadataEvent) events.get(0);
+    assertEquals(2, bootstrap.getTaskID());
+    assertEquals(checkpointId, bootstrap.getCheckpointId());
+    assertEquals(WriteMetadataEvent.BOOTSTRAP_INSTANT, 
bootstrap.getInstantTime());
+    assertTrue(bootstrap.isBootstrap());
+    assertTrue(bootstrap.getWriteStatuses().isEmpty());
+  }
+
+  private ListState<WriteMetadataEvent> writeMetadataState() throws Exception {
+    return stateStore.getListState(new 
ListStateDescriptor<>("write-metadata-state", 
TypeInformation.of(WriteMetadataEvent.class)));
+  }
+
+  private enum SavedJobState {
+    SAME, DIFFERENT, MISSING
+  }
+
+  private static class TestWriteFunction extends 
AbstractStreamWriteFunction<RowData> {
+    TestWriteFunction(Configuration conf) {
+      super(conf);
+    }
+
+    @Override
+    public void snapshotState() {
+      currentInstant = instantToWrite(true);
+    }
+
+    @Override
+    public void processElement(RowData value, Context ctx, Collector<RowData> 
out) {
+      // No record processing is needed to exercise the base checkpoint 
lifecycle.
+    }
+  }
+}

Reply via email to