[FLINK-37208][Runtime] Properly notify a new key is selected for asyn…

…c state operators (#26063)
apache · Jan 24, 2025 · 9912774 · 9912774
1 parent b615a7e
commit 9912774
Show file tree

Hide file tree

Showing 16 changed files with 172 additions and 108 deletions.
diff --git a/...am/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedProcessOperator.java b/...am/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedProcessOperator.java
@@ -26,7 +26,6 @@
 import org.apache.flink.datastream.impl.common.OutputCollector;
 import org.apache.flink.datastream.impl.common.TimestampCollector;
 import org.apache.flink.datastream.impl.context.DefaultNonPartitionedContext;
-import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 
 import javax.annotation.Nullable;
 
@@ -83,10 +82,8 @@ protected NonPartitionedContext<OUT> getNonPartitionedContext() {
     }
 
     @Override
-    @SuppressWarnings({"rawtypes"})
-    public void setKeyContextElement1(StreamRecord record) throws Exception {
-        super.setKeyContextElement1(record);
-        keySet.add(getCurrentKey());
+    public void newKeySelected(Object newKey) {
+        keySet.add(newKey);
     }
 
     @Override

diff --git a/.../apache/flink/datastream/impl/operators/BaseKeyedTwoInputNonBroadcastProcessOperator.java b/.../apache/flink/datastream/impl/operators/BaseKeyedTwoInputNonBroadcastProcessOperator.java
@@ -26,7 +26,6 @@
 import org.apache.flink.datastream.impl.common.OutputCollector;
 import org.apache.flink.datastream.impl.common.TimestampCollector;
 import org.apache.flink.datastream.impl.context.DefaultNonPartitionedContext;
-import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 
 import javax.annotation.Nullable;
 
@@ -89,17 +88,8 @@ protected NonPartitionedContext<OUT> getNonPartitionedContext() {
     }
 
     @Override
-    @SuppressWarnings({"rawtypes"})
-    public void setKeyContextElement1(StreamRecord record) throws Exception {
-        super.setKeyContextElement1(record);
-        keySet.add(getCurrentKey());
-    }
-
-    @Override
-    @SuppressWarnings({"rawtypes"})
-    public void setKeyContextElement2(StreamRecord record) throws Exception {
-        super.setKeyContextElement2(record);
-        keySet.add(getCurrentKey());
+    public void newKeySelected(Object newKey) {
+        keySet.add(newKey);
     }
 
     @Override

diff --git a/...in/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoOutputProcessOperator.java b/...in/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoOutputProcessOperator.java
@@ -26,7 +26,6 @@
 import org.apache.flink.datastream.impl.common.OutputCollector;
 import org.apache.flink.datastream.impl.common.TimestampCollector;
 import org.apache.flink.datastream.impl.context.DefaultTwoOutputNonPartitionedContext;
-import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.util.OutputTag;
 import org.apache.flink.util.Preconditions;
 
@@ -105,10 +104,8 @@ protected TwoOutputNonPartitionedContext<OUT_MAIN, OUT_SIDE> getNonPartitionedCo
     }
 
     @Override
-    @SuppressWarnings({"rawtypes"})
-    public void setKeyContextElement1(StreamRecord record) throws Exception {
-        super.setKeyContextElement1(record);
-        keySet.add(getCurrentKey());
+    public void newKeySelected(Object newKey) {
+        keySet.add(newKey);
     }
 
     @Override

diff --git a/...ava/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java b/...ava/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java
@@ -34,7 +34,6 @@
 import org.apache.flink.streaming.api.operators.InternalTimer;
 import org.apache.flink.streaming.api.operators.InternalTimerService;
 import org.apache.flink.streaming.api.operators.Triggerable;
-import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 
 import javax.annotation.Nullable;
 
@@ -117,11 +116,8 @@ protected NonPartitionedContext<OUT> getNonPartitionedContext() {
     }
 
     @Override
-    @SuppressWarnings({"rawtypes"})
-    // Only element from input1 should be considered as the other side is broadcast input.
-    public void setKeyContextElement1(StreamRecord record) throws Exception {
-        super.setKeyContextElement1(record);
-        keySet.add(getCurrentKey());
+    public void newKeySelected(Object newKey) {
+        keySet.add(newKey);
     }
 
     @Override

diff --git a/...tream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java b/...tream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java
@@ -30,8 +30,8 @@
 import org.apache.flink.datastream.impl.operators.MockSumAggregateProcessFunction;
 import org.apache.flink.streaming.api.operators.collect.utils.MockOperatorStateStore;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
 import org.apache.flink.streaming.util.MockStreamingRuntimeContext;
+import org.apache.flink.streaming.util.asyncprocessing.AsyncKeyedOneInputStreamOperatorTestHarness;
 
 import org.junit.jupiter.api.Test;
 
@@ -98,8 +98,8 @@ void testListState() throws Exception {
                 new KeyedProcessOperator<>(
                         function, (KeySelector<Integer, Integer>) value -> value);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -120,8 +120,8 @@ void testAggState() throws Exception {
         KeyedProcessOperator<Integer, Integer, Integer> processOperator =
                 new KeyedProcessOperator<>(function);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -148,8 +148,8 @@ void testValueState() throws Exception {
         KeyedProcessOperator<Integer, Integer, Integer> processOperator =
                 new KeyedProcessOperator<>(function);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -175,8 +175,8 @@ void testMapState() throws Exception {
         KeyedProcessOperator<Integer, Integer, Integer> processOperator =
                 new KeyedProcessOperator<>(function);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -203,8 +203,8 @@ void testReducingState() throws Exception {
         KeyedProcessOperator<Integer, Integer, Integer> processOperator =
                 new KeyedProcessOperator<>(function);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -231,8 +231,8 @@ void testBroadcastMapState() throws Exception {
         KeyedProcessOperator<Integer, Integer, Integer> processOperator =
                 new KeyedProcessOperator<>(function);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -260,8 +260,8 @@ void testBroadcastListState() throws Exception {
                 new KeyedProcessOperator<>(
                         function, (KeySelector<Integer, Integer>) value -> value);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {

diff --git a/...am/src/test/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperatorTest.java b/...am/src/test/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperatorTest.java
@@ -25,7 +25,7 @@
 import org.apache.flink.datastream.api.context.PartitionedContext;
 import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.apache.flink.streaming.util.asyncprocessing.AsyncKeyedOneInputStreamOperatorTestHarness;
 
 import org.junit.jupiter.api.Test;
 
@@ -51,8 +51,8 @@ public void processRecord(
                             }
                         });
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -98,8 +98,8 @@ public void endInput(NonPartitionedContext<Integer> ctx) {
                             }
                         });
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -133,8 +133,8 @@ public void processRecord(
                         // -1 is an invalid key in this suite.
                         (ignore) -> -1);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {

diff --git a/...org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperatorTest.java b/...org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperatorTest.java
@@ -25,7 +25,7 @@
 import org.apache.flink.datastream.api.context.PartitionedContext;
 import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness;
+import org.apache.flink.streaming.util.asyncprocessing.AsyncKeyedTwoInputStreamOperatorTestHarness;
 
 import org.junit.jupiter.api.Test;
 
@@ -61,8 +61,8 @@ public void processRecordFromBroadcastInput(
                             }
                         });
 
-        try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
-                new KeyedTwoInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
+                AsyncKeyedTwoInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Long>) (data) -> (long) (data + 1),
                         (KeySelector<Long, Long>) value -> value + 1,
@@ -130,11 +130,11 @@ public void endBroadcastInput(NonPartitionedContext<Long> ctx) {
                             }
                         });
 
-        try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
-                new KeyedTwoInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
+                AsyncKeyedTwoInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Long>) Long::valueOf,
-                        (KeySelector<Long, Long>) value -> value,
+                        null,
                         Types.LONG)) {
             testHarness.open();
             testHarness.processElement1(new StreamRecord<>(1)); // key is 1L
@@ -175,8 +175,8 @@ public void processRecordFromBroadcastInput(
                         // -1 is an invalid key in this suite.
                         (out) -> -1L);
 
-        try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
-                new KeyedTwoInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
+                AsyncKeyedTwoInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Long>) Long::valueOf,
                         (KeySelector<Long, Long>) value -> value,

diff --git a/.../apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperatorTest.java b/.../apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperatorTest.java
@@ -25,7 +25,7 @@
 import org.apache.flink.datastream.api.context.PartitionedContext;
 import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness;
+import org.apache.flink.streaming.util.asyncprocessing.AsyncKeyedTwoInputStreamOperatorTestHarness;
 
 import org.junit.jupiter.api.Test;
 
@@ -59,8 +59,8 @@ public void processRecordFromSecondInput(
                             }
                         });
 
-        try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
-                new KeyedTwoInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
+                AsyncKeyedTwoInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Long>) (data) -> (long) (data + 1),
                         (KeySelector<Long, Long>) value -> value + 1,
@@ -134,8 +134,8 @@ public void endSecondInput(NonPartitionedContext<Long> ctx) {
                             }
                         });
 
-        try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
-                new KeyedTwoInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
+                AsyncKeyedTwoInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Long>) Long::valueOf,
                         (KeySelector<Long, Long>) value -> value,
@@ -183,15 +183,23 @@ public void processRecordFromSecondInput(
                         // -1 is an invalid key in this suite.
                         (out) -> -1L);
 
-        try (KeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
-                new KeyedTwoInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
+                AsyncKeyedTwoInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Long>) Long::valueOf,
                         (KeySelector<Long, Long>) value -> value,
                         Types.LONG)) {
             testHarness.open();
             assertThatThrownBy(() -> testHarness.processElement1(new StreamRecord<>(1)))
                     .isInstanceOf(IllegalStateException.class);
+        }
+        try (AsyncKeyedTwoInputStreamOperatorTestHarness<Long, Integer, Long, Long> testHarness =
+                AsyncKeyedTwoInputStreamOperatorTestHarness.create(
+                        processOperator,
+                        (KeySelector<Integer, Long>) Long::valueOf,
+                        (KeySelector<Long, Long>) value -> value,
+                        Types.LONG)) {
+            testHarness.open();
             assertThatThrownBy(() -> testHarness.processElement2(new StreamRecord<>(1L)))
                     .isInstanceOf(IllegalStateException.class);
         }

diff --git a/...st/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperatorTest.java b/...st/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperatorTest.java
@@ -25,7 +25,7 @@
 import org.apache.flink.datastream.api.context.TwoOutputPartitionedContext;
 import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.apache.flink.streaming.util.asyncprocessing.AsyncKeyedOneInputStreamOperatorTestHarness;
 import org.apache.flink.util.OutputTag;
 
 import org.junit.jupiter.api.Test;
@@ -59,8 +59,8 @@ public void processRecord(
                         },
                         sideOutputTag);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -116,8 +116,8 @@ public void endInput(
                         },
                         sideOutputTag);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
@@ -161,14 +161,21 @@ public void processRecord(
                         // -1 is an invalid key in this suite.
                         (KeySelector<Long, Integer>) value -> -1);
 
-        try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
-                new KeyedOneInputStreamOperatorTestHarness<>(
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
                         processOperator,
                         (KeySelector<Integer, Integer>) value -> value,
                         Types.INT)) {
             testHarness.open();
             assertThatThrownBy(() -> testHarness.processElement(new StreamRecord<>(1)))
                     .isInstanceOf(IllegalStateException.class);
+        }
+        try (AsyncKeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
+                AsyncKeyedOneInputStreamOperatorTestHarness.create(
+                        processOperator,
+                        (KeySelector<Integer, Integer>) value -> value,
+                        Types.INT)) {
+            testHarness.open();
             emitToFirstOutput.set(false);
             assertThatThrownBy(() -> testHarness.processElement(new StreamRecord<>(1)))
                     .isInstanceOf(IllegalStateException.class);