[FLINK-36823][Checkpointing] Drain state requests after the user func…

…tion perform snapshotState
Zakelly · Dec 2, 2024 · 97f4ac4 · 97f4ac4
1 parent ea74da0
commit 97f4ac4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/.../org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java b/.../org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
@@ -405,6 +405,10 @@ RecordContext getCurrentProcessingContext() {
         return currentProcessingContext;
     }
 
+    public void drainStateRequests() {
+        asyncExecutionController.drainInflightRecords(0);
+    }
+
     @Override
     public void finish() throws Exception {
         super.finish();

diff --git a/...g/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateUdfStreamOperator.java b/...g/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateUdfStreamOperator.java
@@ -93,6 +93,8 @@ public void snapshotState(StateSnapshotContext context) throws Exception {
         super.snapshotState(context);
         StreamingFunctionUtils.snapshotFunctionState(
                 context, getOperatorStateBackend(), userFunction);
+        // Drain state requests in case the user function modifies the state.
+        drainStateRequests();
     }
 
     @Override