You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by za...@apache.org on 2025/01/24 14:29:22 UTC
(flink) branch release-2.0 updated: [FLINK-37208][Runtime] Properly notify a new key is selected for async state operators (#26068)
This is an automated email from the ASF dual-hosted git repository.
zakelly pushed a commit to branch release-2.0
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/release-2.0 by this push:
new 6a5fab8036f [FLINK-37208][Runtime] Properly notify a new key is selected for async state operators (#26068)
6a5fab8036f is described below
commit 6a5fab8036f534409d774e8724f2cd01ca5d24f8
Author: Zakelly <za...@gmail.com>
AuthorDate: Fri Jan 24 21:58:19 2025 +0800
[FLINK-37208][Runtime] Properly notify a new key is selected for async state operators (#26068)
---
.../impl/operators/BaseKeyedProcessOperator.java | 7 ++--
...seKeyedTwoInputNonBroadcastProcessOperator.java | 14 ++------
.../BaseKeyedTwoOutputProcessOperator.java | 7 ++--
.../KeyedTwoInputBroadcastProcessOperator.java | 8 ++---
.../impl/context/DefaultStateManagerTest.java | 30 ++++++++---------
.../impl/operators/KeyedProcessOperatorTest.java | 14 ++++----
.../KeyedTwoInputBroadcastProcessOperatorTest.java | 16 ++++-----
...yedTwoInputNonBroadcastProcessOperatorTest.java | 22 +++++++++----
.../KeyedTwoOutputProcessOperatorTest.java | 21 ++++++++----
.../state/api/input/KeyedStateInputFormatTest.java | 27 ++++++++++-----
.../AbstractAsyncStateStreamOperator.java | 36 ++++++++++++++------
.../AbstractAsyncStateStreamOperatorV2.java | 38 ++++++++++++++++++----
.../asyncprocessing/AsyncStateProcessing.java | 5 +++
.../operators/AsyncIntervalJoinOperatorTest.java | 12 +++----
...syncKeyedTwoInputStreamOperatorTestHarness.java | 10 ++++++
.../join/stream/StreamingJoinOperatorTestBase.java | 13 +++++---
16 files changed, 172 insertions(+), 108 deletions(-)
diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedProcessOperator.java
index d47c3a0d231..2028b5ee969 100644
--- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedProcessOperator.java
+++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedProcessOperator.java
@@ -26,7 +26,6 @@ import org.apache.flink.datastream.impl.common.KeyCheckedOutputCollector;
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 @@ public class BaseKeyedProcessOperator<KEY, IN, OUT> extends ProcessOperator<IN,
}
@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/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoInputNonBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoInputNonBroadcastProcessOperator.java
index 0861189a2ab..8d1f821beb9 100644
--- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoInputNonBroadcastProcessOperator.java
+++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoInputNonBroadcastProcessOperator.java
@@ -26,7 +26,6 @@ import org.apache.flink.datastream.impl.common.KeyCheckedOutputCollector;
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 @@ public class BaseKeyedTwoInputNonBroadcastProcessOperator<KEY, IN1, IN2, OUT>
}
@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/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoOutputProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoOutputProcessOperator.java
index cc3ba7b6363..6648388cb9b 100644
--- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoOutputProcessOperator.java
+++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/BaseKeyedTwoOutputProcessOperator.java
@@ -26,7 +26,6 @@ import org.apache.flink.datastream.impl.common.KeyCheckedOutputCollector;
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 @@ public class BaseKeyedTwoOutputProcessOperator<KEY, IN, OUT_MAIN, OUT_SIDE>
}
@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/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java
index bb985643e76..86b334c1f1e 100644
--- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java
+++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java
@@ -34,7 +34,6 @@ import org.apache.flink.runtime.state.VoidNamespaceSerializer;
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 @@ public class KeyedTwoInputBroadcastProcessOperator<KEY, IN1, IN2, OUT>
}
@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/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java
index 5366fca7962..2d60e31dd9d 100644
--- a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java
+++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java
@@ -30,8 +30,8 @@ import org.apache.flink.datastream.impl.operators.MockRecudingMultiplierProcessF
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 @@ class DefaultStateManagerTest {
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 @@ class DefaultStateManagerTest {
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 @@ class DefaultStateManagerTest {
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 @@ class DefaultStateManagerTest {
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 @@ class DefaultStateManagerTest {
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 @@ class DefaultStateManagerTest {
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 @@ class DefaultStateManagerTest {
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/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperatorTest.java
index f3cd604c445..49c4d5f1b7e 100644
--- a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperatorTest.java
+++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperatorTest.java
@@ -25,7 +25,7 @@ import org.apache.flink.datastream.api.context.NonPartitionedContext;
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 @@ class KeyedProcessOperatorTest {
}
});
- 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 @@ class KeyedProcessOperatorTest {
}
});
- 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 @@ class KeyedProcessOperatorTest {
// -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/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperatorTest.java
index c9e4b65276c..9684d2decaa 100644
--- a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperatorTest.java
+++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperatorTest.java
@@ -25,7 +25,7 @@ import org.apache.flink.datastream.api.context.NonPartitionedContext;
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 @@ class KeyedTwoInputBroadcastProcessOperatorTest {
}
});
- 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 @@ class KeyedTwoInputBroadcastProcessOperatorTest {
}
});
- 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 @@ class KeyedTwoInputBroadcastProcessOperatorTest {
// -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/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperatorTest.java
index 4ff5501768f..96107c7b6f1 100644
--- a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperatorTest.java
+++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperatorTest.java
@@ -25,7 +25,7 @@ import org.apache.flink.datastream.api.context.NonPartitionedContext;
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 @@ class KeyedTwoInputNonBroadcastProcessOperatorTest {
}
});
- 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 @@ class KeyedTwoInputNonBroadcastProcessOperatorTest {
}
});
- 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,8 +183,8 @@ class KeyedTwoInputNonBroadcastProcessOperatorTest {
// -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,
@@ -192,6 +192,14 @@ class KeyedTwoInputNonBroadcastProcessOperatorTest {
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/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperatorTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperatorTest.java
index 7d0680325b9..f494f7ac633 100644
--- a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperatorTest.java
+++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperatorTest.java
@@ -25,7 +25,7 @@ import org.apache.flink.datastream.api.context.TwoOutputNonPartitionedContext;
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 @@ class KeyedTwoOutputProcessOperatorTest {
},
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 @@ class KeyedTwoOutputProcessOperatorTest {
},
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 @@ class KeyedTwoOutputProcessOperatorTest {
// -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);
diff --git a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/KeyedStateInputFormatTest.java b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/KeyedStateInputFormatTest.java
index c2cb6aa597e..fa653e71f5f 100644
--- a/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/KeyedStateInputFormatTest.java
+++ b/flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/KeyedStateInputFormatTest.java
@@ -43,6 +43,7 @@ import org.apache.flink.streaming.api.operators.StreamFlatMap;
import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.MockStreamingRuntimeContext;
+import org.apache.flink.streaming.util.asyncprocessing.AsyncKeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.util.Collector;
import org.junit.Assert;
@@ -74,7 +75,8 @@ class KeyedStateInputFormatTest {
void testCreatePartitionedInputSplits(boolean asyncState) throws Exception {
OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
- OperatorSubtaskState state = createOperatorSubtaskState(createFlatMap(asyncState));
+ OperatorSubtaskState state =
+ createOperatorSubtaskState(createFlatMap(asyncState), asyncState);
OperatorState operatorState = new OperatorState(null, null, operatorID, 1, 128);
operatorState.putState(0, state);
@@ -95,7 +97,8 @@ class KeyedStateInputFormatTest {
void testMaxParallelismRespected(boolean asyncState) throws Exception {
OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
- OperatorSubtaskState state = createOperatorSubtaskState(createFlatMap(asyncState));
+ OperatorSubtaskState state =
+ createOperatorSubtaskState(createFlatMap(asyncState), asyncState);
OperatorState operatorState = new OperatorState(null, null, operatorID, 1, 128);
operatorState.putState(0, state);
@@ -118,7 +121,8 @@ class KeyedStateInputFormatTest {
void testReadState(boolean asyncState) throws Exception {
OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
- OperatorSubtaskState state = createOperatorSubtaskState(createFlatMap(asyncState));
+ OperatorSubtaskState state =
+ createOperatorSubtaskState(createFlatMap(asyncState), asyncState);
OperatorState operatorState = new OperatorState(null, null, operatorID, 1, 128);
operatorState.putState(0, state);
@@ -143,7 +147,8 @@ class KeyedStateInputFormatTest {
void testReadMultipleOutputPerKey(boolean asyncState) throws Exception {
OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
- OperatorSubtaskState state = createOperatorSubtaskState(createFlatMap(asyncState));
+ OperatorSubtaskState state =
+ createOperatorSubtaskState(createFlatMap(asyncState), asyncState);
OperatorState operatorState = new OperatorState(null, null, operatorID, 1, 128);
operatorState.putState(0, state);
@@ -169,7 +174,8 @@ class KeyedStateInputFormatTest {
void testInvalidProcessReaderFunctionFails(boolean asyncState) throws Exception {
OperatorID operatorID = OperatorIDGenerator.fromUid("uid");
- OperatorSubtaskState state = createOperatorSubtaskState(createFlatMap(asyncState));
+ OperatorSubtaskState state =
+ createOperatorSubtaskState(createFlatMap(asyncState), asyncState);
OperatorState operatorState = new OperatorState(null, null, operatorID, 1, 128);
operatorState.putState(0, state);
@@ -194,7 +200,7 @@ class KeyedStateInputFormatTest {
OperatorSubtaskState state =
createOperatorSubtaskState(
- new KeyedProcessOperator<>(new StatefulFunctionWithTime()));
+ new KeyedProcessOperator<>(new StatefulFunctionWithTime()), false);
OperatorState operatorState = new OperatorState(null, null, operatorID, 1, 128);
operatorState.putState(0, state);
@@ -252,10 +258,13 @@ class KeyedStateInputFormatTest {
}
private OperatorSubtaskState createOperatorSubtaskState(
- OneInputStreamOperator<Integer, Void> operator) throws Exception {
+ OneInputStreamOperator<Integer, Void> operator, boolean async) throws Exception {
try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Void> testHarness =
- new KeyedOneInputStreamOperatorTestHarness<>(
- operator, id -> id, Types.INT, 128, 1, 0)) {
+ async
+ ? AsyncKeyedOneInputStreamOperatorTestHarness.create(
+ operator, id -> id, Types.INT, 128, 1, 0)
+ : new KeyedOneInputStreamOperatorTestHarness<>(
+ operator, id -> id, Types.INT, 128, 1, 0)) {
testHarness.setup(VoidSerializer.INSTANCE);
testHarness.open();
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
index 32f88ff7759..15191551fb3 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
@@ -166,6 +166,18 @@ public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStre
// in Fig.5, each state request itself is also protected by a paired reference count.
currentProcessingContext.retain();
asyncExecutionController.setCurrentContext(currentProcessingContext);
+ newKeySelected(currentProcessingContext.getKey());
+ }
+
+ /**
+ * A hook that will be invoked after a new key is selected. It is not recommended to perform
+ * async state here. Only some synchronous logic is suggested.
+ *
+ * @param newKey the new key selected.
+ */
+ public void newKeySelected(Object newKey) {
+ // By default, do nothing.
+ // Subclass could override this method to do some operations when a new key is selected.
}
@Override
@@ -300,32 +312,36 @@ public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStre
}
@Override
- @SuppressWarnings("unchecked")
public void setKeyContextElement1(StreamRecord record) throws Exception {
+ // This method is invoked only when isAsyncStateProcessingEnabled() is false.
super.setKeyContextElement1(record);
if (stateKeySelector1 != null) {
- setAsyncKeyedContextElement(record, stateKeySelector1);
+ newKeySelected(getCurrentKey());
}
}
@Override
- @SuppressWarnings("unchecked")
public void setKeyContextElement2(StreamRecord record) throws Exception {
+ // This method is invoked only when isAsyncStateProcessingEnabled() is false.
super.setKeyContextElement2(record);
if (stateKeySelector2 != null) {
- setAsyncKeyedContextElement(record, stateKeySelector2);
+ newKeySelected(getCurrentKey());
}
}
@Override
public Object getCurrentKey() {
- 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.");
+ if (isAsyncStateProcessingEnabled()) {
+ 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();
+ } else {
+ return super.getCurrentKey();
}
- return currentContext.getKey();
}
// ------------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
index 955faa770b3..407dd177817 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
@@ -163,17 +163,43 @@ public abstract class AbstractAsyncStateStreamOperatorV2<OUT> extends AbstractSt
// in Fig.5, each state request itself is also protected by a paired reference count.
currentProcessingContext.retain();
asyncExecutionController.setCurrentContext(currentProcessingContext);
+ newKeySelected(currentProcessingContext.getKey());
+ }
+
+ /**
+ * A hook that will be invoked after a new key is selected. It is not recommended to perform
+ * async state here. Only some synchronous logic is suggested.
+ *
+ * @param newKey the new key selected.
+ */
+ public void newKeySelected(Object newKey) {
+ // By default, do nothing.
+ // Subclass could override this method to do some operations when a new key is selected.
+ }
+
+ @Override
+ protected <T> void internalSetKeyContextElement(
+ StreamRecord<T> record, KeySelector<T, ?> selector) throws Exception {
+ // This method is invoked only when isAsyncStateProcessingEnabled() is false.
+ super.internalSetKeyContextElement(record, selector);
+ if (selector != null) {
+ newKeySelected(getCurrentKey());
+ }
}
@Override
public Object getCurrentKey() {
- 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.");
+ if (isAsyncStateProcessingEnabled()) {
+ 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();
+ } else {
+ return super.getCurrentKey();
}
- return currentContext.getKey();
}
@Override
diff --git a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
index 563f3fd1ad8..6e151b52ebc 100644
--- a/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
+++ b/flink-runtime/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
@@ -56,6 +56,11 @@ public interface AsyncStateProcessing {
AsyncStateProcessingOperator asyncOperator,
KeySelector<T, ?> keySelector,
ThrowingConsumer<StreamRecord<T>, Exception> processor) {
+ if (keySelector == null) {
+ // A non-keyed input does not need to set the key context and perform async context
+ // switches.
+ return processor;
+ }
switch (asyncOperator.getElementOrder()) {
case RECORD_ORDER:
return (record) -> {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/operators/AsyncIntervalJoinOperatorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/operators/AsyncIntervalJoinOperatorTest.java
index ab1999e2484..80bb3450cc1 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/operators/AsyncIntervalJoinOperatorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/operators/AsyncIntervalJoinOperatorTest.java
@@ -412,7 +412,7 @@ class AsyncIntervalJoinOperatorTest {
});
try (TestHarness testHarness =
- TestHarness.create(
+ TestHarness.createOne(
op,
(elem) -> elem.key,
(elem) -> elem.key,
@@ -454,7 +454,7 @@ class AsyncIntervalJoinOperatorTest {
});
try (TestHarness testHarness =
- TestHarness.create(
+ TestHarness.createOne(
op,
(elem) -> elem.key,
(elem) -> elem.key,
@@ -493,7 +493,7 @@ class AsyncIntervalJoinOperatorTest {
});
try (TestHarness testHarness =
- TestHarness.create(
+ TestHarness.createOne(
op,
(elem) -> elem.key,
(elem) -> elem.key,
@@ -661,7 +661,7 @@ class AsyncIntervalJoinOperatorTest {
TestElem.serializer(),
new PassthroughFunction());
- return TestHarness.create(
+ return TestHarness.createOne(
operator,
(elem) -> elem.key, // key
(elem) -> elem.key, // key
@@ -690,7 +690,7 @@ class AsyncIntervalJoinOperatorTest {
new PassthroughFunction());
TestHarness t =
- TestHarness.create(
+ TestHarness.createOne(
operator,
(elem) -> elem.key, // key
(elem) -> elem.key, // key
@@ -966,7 +966,7 @@ class AsyncIntervalJoinOperatorTest {
super(executor, operator, keySelector1, keySelector2, keyType, 1, 1, 0);
}
- public static TestHarness create(
+ public static TestHarness createOne(
TwoInputStreamOperator<TestElem, TestElem, Tuple2<TestElem, TestElem>> operator,
KeySelector<TestElem, String> keySelector1,
KeySelector<TestElem, String> keySelector2,
diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/util/asyncprocessing/AsyncKeyedTwoInputStreamOperatorTestHarness.java b/flink-runtime/src/test/java/org/apache/flink/streaming/util/asyncprocessing/AsyncKeyedTwoInputStreamOperatorTestHarness.java
index 17eacf983dd..ae666e3f541 100644
--- a/flink-runtime/src/test/java/org/apache/flink/streaming/util/asyncprocessing/AsyncKeyedTwoInputStreamOperatorTestHarness.java
+++ b/flink-runtime/src/test/java/org/apache/flink/streaming/util/asyncprocessing/AsyncKeyedTwoInputStreamOperatorTestHarness.java
@@ -88,6 +88,16 @@ public class AsyncKeyedTwoInputStreamOperatorTestHarness<K, IN1, IN2, OUT>
return future.get();
}
+ public static <K, IN1, IN2, OUT>
+ AsyncKeyedTwoInputStreamOperatorTestHarness<K, IN1, IN2, OUT> create(
+ TwoInputStreamOperator<IN1, IN2, OUT> operator,
+ KeySelector<IN1, K> keySelector1,
+ KeySelector<IN2, K> keySelector2,
+ TypeInformation<K> keyType)
+ throws Exception {
+ return create(operator, keySelector1, keySelector2, keyType, 1, 1, 0);
+ }
+
public static <K, IN1, IN2, OUT>
AsyncKeyedTwoInputStreamOperatorTestHarness<K, IN1, IN2, OUT> create(
TwoInputStreamOperator<IN1, IN2, OUT> operator,
diff --git a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/join/stream/StreamingJoinOperatorTestBase.java b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/join/stream/StreamingJoinOperatorTestBase.java
index 4ba36ab0d08..851d35939e4 100644
--- a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/join/stream/StreamingJoinOperatorTestBase.java
+++ b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/join/stream/StreamingJoinOperatorTestBase.java
@@ -19,7 +19,9 @@
package org.apache.flink.table.runtime.operators.join.stream;
import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
+import org.apache.flink.streaming.runtime.operators.asyncprocessing.AsyncStateProcessingOperator;
import org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness;
+import org.apache.flink.streaming.util.asyncprocessing.AsyncKeyedTwoInputStreamOperatorTestHarness;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.runtime.generated.GeneratedJoinCondition;
import org.apache.flink.table.runtime.keyselector.RowDataKeySelector;
@@ -102,12 +104,13 @@ public abstract class StreamingJoinOperatorTestBase {
@BeforeEach
void beforeEach(TestInfo testInfo) throws Exception {
+ TwoInputStreamOperator<RowData, RowData, RowData> operator = createJoinOperator(testInfo);
testHarness =
- new KeyedTwoInputStreamOperatorTestHarness<>(
- createJoinOperator(testInfo),
- leftKeySelector,
- rightKeySelector,
- joinKeyTypeInfo);
+ operator instanceof AsyncStateProcessingOperator
+ ? AsyncKeyedTwoInputStreamOperatorTestHarness.create(
+ operator, leftKeySelector, rightKeySelector, joinKeyTypeInfo)
+ : new KeyedTwoInputStreamOperatorTestHarness<>(
+ operator, leftKeySelector, rightKeySelector, joinKeyTypeInfo);
testHarness.open();
// extend for mini-batch join test
assertor =