[FLINK-36949][Runtime] Make getCurrentKey of async state operators re…

…turn the ground truth (#25841)
apache · Dec 23, 2024 · f8014b6 · f8014b6
1 parent 5b65539
commit f8014b6
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 6 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,13 @@ public void setKeyContextElement2(StreamRecord record) throws Exception {
 
     @Override
     public Object getCurrentKey() {
-        return currentProcessingContext.getKey();
+        RecordContext currentContext = asyncExecutionController.getCurrentContext();
+        if (currentContext == null) {
+            throw new UnsupportedOperationException(
+                    "Have not set the current key yet, this may because the operator has not "
+                            + "started to run, or you are invoking this under a non-keyed context.");
+        }
+        return currentContext.getKey();
     }
 
     // ------------------------------------------------------------------------

diff --git a/...rg/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java b/...rg/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
@@ -154,7 +154,13 @@ public final <T> void setAsyncKeyedContextElement(
 
     @Override
     public Object getCurrentKey() {
-        return currentProcessingContext.getKey();
+        RecordContext currentContext = asyncExecutionController.getCurrentContext();
+        if (currentContext == null) {
+            throw new UnsupportedOperationException(
+                    "Have not set the current key yet, this may because the operator has not "
+                            + "started to run, or you are invoking this under a non-keyed context.");
+        }
+        return currentContext.getKey();
     }
 
     @Override

diff --git a/.../apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorTest.java b/.../apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorTest.java
@@ -293,6 +293,7 @@ void testWatermarkHooks() throws Exception {
                     testOperator.asyncProcessWithKey(
                             1L,
                             () -> {
+                                assertThat(testOperator.getCurrentKey()).isEqualTo(1L);
                                 testOperator.output(watermark.getTimestamp() + 1000L);
                             });
                     if (counter.incrementAndGet() % 2 == 0) {
@@ -484,13 +485,15 @@ protected void processWatermarkStatus(WatermarkStatus watermarkStatus, int index
 
         @Override
         public void onEventTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {
+            assertThat(getCurrentKey()).isEqualTo(timer.getKey());
             output.collect(
                     new StreamRecord<>(
                             "EventTimer-" + timer.getKey() + "-" + timer.getTimestamp()));
         }
 
         @Override
         public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {
+            assertThat(getCurrentKey()).isEqualTo(timer.getKey());
             output.collect(
                     new StreamRecord<>(
                             "ProcessingTimer-" + timer.getKey() + "-" + timer.getTimestamp()));
@@ -601,12 +604,14 @@ public void postProcessWatermark(Watermark watermark) throws Exception {
 
         @Override
         public void onEventTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {
+            assertThat(getCurrentKey()).isEqualTo(timer.getKey());
             output.collect(new StreamRecord<>(timer.getTimestamp()));
         }
 
         @Override
-        public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer)
-                throws Exception {}
+        public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {
+            assertThat(getCurrentKey()).isEqualTo(timer.getKey());
+        }
 
         @Override
         public void processElement1(StreamRecord<Long> element) throws Exception {

diff --git a/...java/org/apache/flink/streaming/api/operators/AbstractAsyncStateStreamOperatorV2Test.java b/...java/org/apache/flink/streaming/api/operators/AbstractAsyncStateStreamOperatorV2Test.java
@@ -290,6 +290,7 @@ void testWatermarkHooks() throws Exception {
                     operator.asyncProcessWithKey(
                             1L,
                             () -> {
+                                assertThat(operator.getCurrentKey()).isEqualTo(1L);
                                 operator.output(watermark.getTimestamp() + 1000L);
                             });
                     if (counter.incrementAndGet() % 2 == 0) {
@@ -620,13 +621,15 @@ public ElementOrder getElementOrder() {
 
         @Override
         public void onEventTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {
+            assertThat(getCurrentKey()).isEqualTo(timer.getKey());
             output.collect(
                     new StreamRecord<>(
                             "EventTimer-" + timer.getKey() + "-" + timer.getTimestamp()));
         }
 
         @Override
         public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {
+            assertThat(getCurrentKey()).isEqualTo(timer.getKey());
             output.collect(
                     new StreamRecord<>(
                             "ProcessingTimer-" + timer.getKey() + "-" + timer.getTimestamp()));
@@ -842,12 +845,14 @@ public void postProcessWatermark(Watermark watermark) throws Exception {
 
         @Override
         public void onEventTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {
+            assertThat(getCurrentKey()).isEqualTo(timer.getKey());
             output.collect(new StreamRecord<>(timer.getTimestamp()));
         }
 
         @Override
-        public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer)
-                throws Exception {}
+        public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {
+            assertThat(getCurrentKey()).isEqualTo(timer.getKey());
+        }
 
         private Input<Long> createInput(int idx) {
             return new AbstractInput<Long, Long>(this, idx) {
-Original file line number
+Diff line change
@@ Expand Up @@
             }
         }
+        public RecordContext<K> getCurrentContext() {
+            return currentContext;
+        }
         /**
          * Dispose a context.
          *
@@ Expand Down @@