[FLINK-36949] Make getCurrentKey of async state operators return the …

…ground truth
Zakelly · Dec 23, 2024 · 15825fe · 15825fe
1 parent 1523f2c
commit 15825fe
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/...time/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java b/...time/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
@@ -243,6 +243,10 @@ public void setCurrentContext(RecordContext<K> switchingContext) {
         }
     }
 
+    public RecordContext<K> getCurrentContext() {
+        return currentContext;
+    }
+
     /**
      * Dispose a context.
      *

diff --git a/.../org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java b/.../org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
@@ -305,7 +305,7 @@ public void setKeyContextElement2(StreamRecord record) throws Exception {
 
     @Override
     public Object getCurrentKey() {
-        return currentProcessingContext.getKey();
+        return asyncExecutionController.getCurrentContext().getKey();
     }
 
     // ------------------------------------------------------------------------

diff --git a/...rg/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java b/...rg/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
@@ -154,7 +154,7 @@ public final <T> void setAsyncKeyedContextElement(
 
     @Override
     public Object getCurrentKey() {
-        return currentProcessingContext.getKey();
+        return asyncExecutionController.getCurrentContext().getKey();
     }
 
     @Override
-Original file line number
+Diff line change
@@ Expand Up @@
             }
         }
+        public RecordContext<K> getCurrentContext() {
+            return currentContext;
+        }
         /**
          * Dispose a context.
          *
@@ Expand Down @@