[FLINK-36963][Runtime] Fix wrong context maintain around `asyncProces…

…sWithKey`
Zakelly · Dec 25, 2024 · 5da1693 · 5da1693
1 parent 56ba88b
commit 5da1693
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 16 deletions.
diff --git a/.../org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java b/.../org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
@@ -77,6 +77,7 @@ public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStre
 
     private AsyncExecutionController asyncExecutionController;
 
+    /** Act as a cache for {@link #setAsyncKeyedContextElement} and {@link #postProcessElement}. */
     private RecordContext currentProcessingContext;
 
     private Environment environment;
@@ -175,20 +176,18 @@ public final void preserveRecordOrderAndProcess(ThrowingRunnable<Exception> proc
     @Override
     @SuppressWarnings("unchecked")
     public <K> void asyncProcessWithKey(K key, ThrowingRunnable<Exception> processing) {
-        RecordContext<K> previousContext = currentProcessingContext;
-
+        RecordContext<K> oldContext = asyncExecutionController.getCurrentContext();
         // build a context and switch to the new context
-        currentProcessingContext = asyncExecutionController.buildContext(null, key, true);
-        currentProcessingContext.retain();
-        asyncExecutionController.setCurrentContext(currentProcessingContext);
+        RecordContext<K> newContext = asyncExecutionController.buildContext(null, key, true);
+        newContext.retain();
+        asyncExecutionController.setCurrentContext(newContext);
         // Same logic as RECORD_ORDER, since FIRST_STATE_ORDER is problematic when the call's key
         // pass the same key in.
         preserveRecordOrderAndProcess(processing);
-        postProcessElement();
+        newContext.release();
 
         // switch to original context
-        asyncExecutionController.setCurrentContext(previousContext);
-        currentProcessingContext = previousContext;
+        asyncExecutionController.setCurrentContext(oldContext);
     }
 
     @Override

diff --git a/...rg/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java b/...rg/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
@@ -73,6 +73,7 @@ public abstract class AbstractAsyncStateStreamOperatorV2<OUT> extends AbstractSt
     private final Environment environment;
     private AsyncExecutionController asyncExecutionController;
 
+    /** Act as a cache for {@link #setAsyncKeyedContextElement} and {@link #postProcessElement}. */
     private RecordContext currentProcessingContext;
 
     public AbstractAsyncStateStreamOperatorV2(
@@ -182,20 +183,18 @@ public final void preserveRecordOrderAndProcess(ThrowingRunnable<Exception> proc
     @Override
     @SuppressWarnings("unchecked")
     public <K> void asyncProcessWithKey(K key, ThrowingRunnable<Exception> processing) {
-        RecordContext<K> previousContext = currentProcessingContext;
-
+        RecordContext<K> oldContext = asyncExecutionController.getCurrentContext();
         // build a context and switch to the new context
-        currentProcessingContext = asyncExecutionController.buildContext(null, key, true);
-        currentProcessingContext.retain();
-        asyncExecutionController.setCurrentContext(currentProcessingContext);
+        RecordContext<K> newContext = asyncExecutionController.buildContext(null, key, true);
+        newContext.retain();
+        asyncExecutionController.setCurrentContext(newContext);
         // Same logic as RECORD_ORDER, since FIRST_STATE_ORDER is problematic when the call's key
         // pass the same key in.
         preserveRecordOrderAndProcess(processing);
-        postProcessElement();
+        newContext.release();
 
         // switch to original context
-        asyncExecutionController.setCurrentContext(previousContext);
-        currentProcessingContext = previousContext;
+        asyncExecutionController.setCurrentContext(oldContext);
     }
 
     @Override