You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by le...@apache.org on 2024/04/17 09:34:53 UTC

(flink) branch master updated (31ea1a9358f -> fe8dde4ea48)

This is an automated email from the ASF dual-hosted git repository.

leiyanfei pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git


    from 31ea1a9358f [FLINK-34936][Checkpointing] Register reused shared state handle to FileMergingSnapshotManager
     new c7be45d0ba9 [FLINK-35025][Runtime/State] Erase the type parameter of record within AEC and related components.
     new e1edb24ecc2 [FLINK-35025][Runtime/State] Abstract stream operators for async state processing
     new fe8dde4ea48 [FLINK-35025][Runtime/State] Introduce element order for async state processing

The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.


Summary of changes:
 .../asyncprocessing/AsyncExecutionController.java  |  44 +++-
 .../asyncprocessing/ContextStateFutureImpl.java    |   4 +-
 .../runtime/asyncprocessing/KeyAccountingUnit.java |   9 +-
 .../runtime/asyncprocessing/RecordContext.java     |  14 +-
 .../asyncprocessing/StateFutureFactory.java        |   9 +-
 .../runtime/asyncprocessing/StateRequest.java      |   6 +-
 .../asyncprocessing/StateRequestBuffer.java        |  29 ++-
 .../AsyncExecutionControllerTest.java              |  68 +++++-
 .../ContextStateFutureImplTest.java                |  10 +-
 .../asyncprocessing/KeyAccountingUnitTest.java     |   2 +-
 .../streaming/api/operators/AbstractInput.java     |  23 +-
 .../api/operators/AbstractStreamOperator.java      |   4 +-
 .../api/operators/AbstractStreamOperatorV2.java    |   2 +-
 .../streaming/runtime/io/RecordProcessorUtils.java |  10 +
 .../AbstractAsyncStateStreamOperator.java          | 145 ++++++++++++
 .../AbstractAsyncStateStreamOperatorV2.java        | 125 +++++++++++
 .../asyncprocessing/AsyncStateProcessing.java      |  78 +++++++
 .../AsyncStateProcessingOperator.java              |  56 +++++
 .../operators/asyncprocessing/ElementOrder.java    |  42 ++++
 .../AbstractAsyncStateStreamOperatorTest.java      | 196 ++++++++++++++++
 .../AbstractAsyncStateStreamOperatorV2Test.java    | 248 +++++++++++++++++++++
 .../multipleinput/input/FirstInputOfTwoInput.java  |  18 +-
 .../operators/multipleinput/input/OneInput.java    |  18 +-
 .../multipleinput/input/SecondInputOfTwoInput.java |  18 +-
 24 files changed, 1108 insertions(+), 70 deletions(-)
 create mode 100644 flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
 create mode 100644 flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
 create mode 100644 flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
 create mode 100644 flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java
 create mode 100644 flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/ElementOrder.java
 create mode 100644 flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorTest.java
 create mode 100644 flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2Test.java


(flink) 01/03: [FLINK-35025][Runtime/State] Erase the type parameter of record within AEC and related components.

Posted by le...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

leiyanfei pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit c7be45d0ba9c4ab6c2ce7757b6c8cd8e1f0fe539
Author: Zakelly <za...@gmail.com>
AuthorDate: Fri Apr 12 18:34:07 2024 +0800

    [FLINK-35025][Runtime/State] Erase the type parameter of record within AEC and related components.
---
 .../asyncprocessing/AsyncExecutionController.java  | 23 ++++++++++-----------
 .../asyncprocessing/ContextStateFutureImpl.java    |  4 ++--
 .../runtime/asyncprocessing/KeyAccountingUnit.java |  9 ++++----
 .../runtime/asyncprocessing/RecordContext.java     | 14 ++++++-------
 .../asyncprocessing/StateFutureFactory.java        |  9 ++++----
 .../runtime/asyncprocessing/StateRequest.java      |  6 +++---
 .../asyncprocessing/StateRequestBuffer.java        | 21 +++++++++++--------
 .../AsyncExecutionControllerTest.java              | 24 +++++++++++-----------
 .../ContextStateFutureImplTest.java                | 10 ++++-----
 .../asyncprocessing/KeyAccountingUnitTest.java     |  2 +-
 10 files changed, 61 insertions(+), 61 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
index cf8304a71ea..8dd8d16e721 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
@@ -42,10 +42,9 @@ import java.util.concurrent.atomic.AtomicInteger;
  *     (back-pressure). It invokes {@link MailboxExecutor#yield()} to pause current operations,
  *     allowing for the execution of callbacks (mails in Mailbox).
  *
- * @param <R> the type of the record
  * @param <K> the type of the key
  */
-public class AsyncExecutionController<R, K> {
+public class AsyncExecutionController<K> {
 
     private static final Logger LOG = LoggerFactory.getLogger(AsyncExecutionController.class);
 
@@ -70,22 +69,22 @@ public class AsyncExecutionController<R, K> {
     private final MailboxExecutor mailboxExecutor;
 
     /** The key accounting unit which is used to detect the key conflict. */
-    final KeyAccountingUnit<R, K> keyAccountingUnit;
+    final KeyAccountingUnit<K> keyAccountingUnit;
 
     /**
      * A factory to build {@link org.apache.flink.core.state.InternalStateFuture}, this will auto
      * wire the created future with mailbox executor. Also conducting the context switch.
      */
-    private final StateFutureFactory<R, K> stateFutureFactory;
+    private final StateFutureFactory<K> stateFutureFactory;
 
     /** The state executor where the {@link StateRequest} is actually executed. */
     final StateExecutor stateExecutor;
 
     /** The corresponding context that currently runs in task thread. */
-    RecordContext<R, K> currentContext;
+    RecordContext<K> currentContext;
 
     /** The buffer to store the state requests to execute in batch. */
-    StateRequestBuffer<R, K> stateRequestsBuffer;
+    StateRequestBuffer<K> stateRequestsBuffer;
 
     /**
      * The number of in-flight records. Including the records in active buffer and blocking buffer.
@@ -123,7 +122,7 @@ public class AsyncExecutionController<R, K> {
      * @param key the given key.
      * @return the built record context.
      */
-    public RecordContext<R, K> buildContext(R record, K key) {
+    public RecordContext<K> buildContext(Object record, K key) {
         return new RecordContext<>(record, key, this::disposeContext);
     }
 
@@ -133,7 +132,7 @@ public class AsyncExecutionController<R, K> {
      *
      * @param switchingContext the context to switch.
      */
-    public void setCurrentContext(RecordContext<R, K> switchingContext) {
+    public void setCurrentContext(RecordContext<K> switchingContext) {
         currentContext = switchingContext;
     }
 
@@ -142,10 +141,10 @@ public class AsyncExecutionController<R, K> {
      *
      * @param toDispose the context to dispose.
      */
-    public void disposeContext(RecordContext<R, K> toDispose) {
+    public void disposeContext(RecordContext<K> toDispose) {
         keyAccountingUnit.release(toDispose.getRecord(), toDispose.getKey());
         inFlightRecordNum.decrementAndGet();
-        RecordContext<R, K> nextRecordCtx =
+        RecordContext<K> nextRecordCtx =
                 stateRequestsBuffer.tryActivateOneByKey(toDispose.getKey());
         if (nextRecordCtx != null) {
             Preconditions.checkState(
@@ -160,7 +159,7 @@ public class AsyncExecutionController<R, K> {
      * @param recordContext the given context.
      * @return true if occupy succeed or the key has already occupied by this context.
      */
-    boolean tryOccupyKey(RecordContext<R, K> recordContext) {
+    boolean tryOccupyKey(RecordContext<K> recordContext) {
         boolean occupied = recordContext.isKeyOccupied();
         if (!occupied
                 && keyAccountingUnit.occupy(recordContext.getRecord(), recordContext.getKey())) {
@@ -230,7 +229,7 @@ public class AsyncExecutionController<R, K> {
         if (currentContext.isKeyOccupied()) {
             return;
         }
-        RecordContext<R, K> storedContext = currentContext;
+        RecordContext<K> storedContext = currentContext;
         // 2. If the state request is for a newly entered record, the in-flight record number should
         // be less than the max in-flight record number.
         // Note: the currentContext may be updated by {@code StateFutureFactory#build}.
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/ContextStateFutureImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/ContextStateFutureImpl.java
index 31e89676523..9355a43a795 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/ContextStateFutureImpl.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/ContextStateFutureImpl.java
@@ -37,9 +37,9 @@ import org.apache.flink.core.state.StateFutureImpl;
  */
 public class ContextStateFutureImpl<T> extends StateFutureImpl<T> {
 
-    private final RecordContext<?, ?> recordContext;
+    private final RecordContext<?> recordContext;
 
-    ContextStateFutureImpl(CallbackRunner callbackRunner, RecordContext<?, ?> recordContext) {
+    ContextStateFutureImpl(CallbackRunner callbackRunner, RecordContext<?> recordContext) {
         super(callbackRunner);
         this.recordContext = recordContext;
         // When state request submitted, ref count +1, as described in FLIP-425:
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/KeyAccountingUnit.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/KeyAccountingUnit.java
index 3d8a0caa052..fc3ccb3e468 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/KeyAccountingUnit.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/KeyAccountingUnit.java
@@ -28,13 +28,12 @@ import java.util.concurrent.ConcurrentHashMap;
  * which is used to preserve the ordering of independent chained {@link
  * org.apache.flink.api.common.state.v2.StateFuture}.
  *
- * @param <R> the type of record
  * @param <K> the type of key
  */
-public class KeyAccountingUnit<R, K> {
+public class KeyAccountingUnit<K> {
 
     /** The in-flight records that are being processed, their keys are different from each other. */
-    private final Map<K, R> noConflictInFlightRecords;
+    private final Map<K, Object> noConflictInFlightRecords;
 
     public KeyAccountingUnit(int initCapacity) {
         this.noConflictInFlightRecords = new ConcurrentHashMap<>(initCapacity);
@@ -46,12 +45,12 @@ public class KeyAccountingUnit<R, K> {
      *
      * @return true if no one is occupying this key, and this record succeeds to take it.
      */
-    public boolean occupy(R record, K key) {
+    public boolean occupy(Object record, K key) {
         return noConflictInFlightRecords.putIfAbsent(key, record) == null;
     }
 
     /** Release a key, which is invoked when a {@link RecordContext} is released. */
-    public void release(R record, K key) {
+    public void release(Object record, K key) {
         if (noConflictInFlightRecords.remove(key) != record) {
             throw new IllegalStateException(
                     String.format(
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/RecordContext.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/RecordContext.java
index 3436a0063e8..f91eb2f20da 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/RecordContext.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/RecordContext.java
@@ -27,14 +27,12 @@ import java.util.function.Consumer;
  *
  * <p>Reference counting mechanism, please refer to {@link ContextStateFutureImpl}.
  *
- * @param <R> The type of the record that extends {@code
- *     org.apache.flink.streaming.runtime.streamrecord.StreamElement}.
  * @param <K> The type of the key inside the record.
  */
-public class RecordContext<R, K> extends ReferenceCounted {
+public class RecordContext<K> extends ReferenceCounted {
 
     /** The record to be processed. */
-    private final R record;
+    private final Object record;
 
     /** The key inside the record. */
     private final K key;
@@ -47,9 +45,9 @@ public class RecordContext<R, K> extends ReferenceCounted {
      * #referenceCountReachedZero()}, which may be called once the ref count reaches zero in any
      * thread.
      */
-    private final Consumer<RecordContext<R, K>> disposer;
+    private final Consumer<RecordContext<K>> disposer;
 
-    RecordContext(R record, K key, Consumer<RecordContext<R, K>> disposer) {
+    RecordContext(Object record, K key, Consumer<RecordContext<K>> disposer) {
         super(0);
         this.record = record;
         this.key = key;
@@ -57,7 +55,7 @@ public class RecordContext<R, K> extends ReferenceCounted {
         this.disposer = disposer;
     }
 
-    public R getRecord() {
+    public Object getRecord() {
         return record;
     }
 
@@ -96,7 +94,7 @@ public class RecordContext<R, K> extends ReferenceCounted {
         if (o == null || getClass() != o.getClass()) {
             return false;
         }
-        RecordContext<?, ?> that = (RecordContext<?, ?>) o;
+        RecordContext<?> that = (RecordContext<?>) o;
         if (!Objects.equals(record, that.record)) {
             return false;
         }
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateFutureFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateFutureFactory.java
index cbb0c122589..0ce6c89087a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateFutureFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateFutureFactory.java
@@ -25,19 +25,18 @@ import org.apache.flink.core.state.InternalStateFuture;
  * An internal factory for {@link InternalStateFuture} that build future with necessary context
  * switch and wired with mailbox executor.
  */
-public class StateFutureFactory<R, K> {
+public class StateFutureFactory<K> {
 
-    private final AsyncExecutionController<R, K> asyncExecutionController;
+    private final AsyncExecutionController<K> asyncExecutionController;
     private final MailboxExecutor mailboxExecutor;
 
     StateFutureFactory(
-            AsyncExecutionController<R, K> asyncExecutionController,
-            MailboxExecutor mailboxExecutor) {
+            AsyncExecutionController<K> asyncExecutionController, MailboxExecutor mailboxExecutor) {
         this.asyncExecutionController = asyncExecutionController;
         this.mailboxExecutor = mailboxExecutor;
     }
 
-    public <OUT> InternalStateFuture<OUT> create(RecordContext<R, K> context) {
+    public <OUT> InternalStateFuture<OUT> create(RecordContext<K> context) {
         return new ContextStateFutureImpl<>(
                 (runnable) ->
                         mailboxExecutor.submit(
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequest.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequest.java
index be4164df22d..aa050cf6557 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequest.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequest.java
@@ -49,14 +49,14 @@ public class StateRequest<K, IN, OUT> implements Serializable {
     private final InternalStateFuture<OUT> stateFuture;
 
     /** The record context of this request. */
-    private final RecordContext<?, K> context;
+    private final RecordContext<K> context;
 
     StateRequest(
             @Nullable State state,
             StateRequestType type,
             @Nullable IN payload,
             InternalStateFuture<OUT> stateFuture,
-            RecordContext<?, K> context) {
+            RecordContext<K> context) {
         this.state = state;
         this.type = type;
         this.payload = payload;
@@ -82,7 +82,7 @@ public class StateRequest<K, IN, OUT> implements Serializable {
         return stateFuture;
     }
 
-    RecordContext<?, K> getRecordContext() {
+    RecordContext<K> getRecordContext() {
         return context;
     }
 }
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestBuffer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestBuffer.java
index 9a0cde21936..0eea7bbcacc 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestBuffer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestBuffer.java
@@ -21,6 +21,8 @@ package org.apache.flink.runtime.asyncprocessing;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.NotThreadSafe;
 
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Deque;
 import java.util.HashMap;
 import java.util.LinkedList;
@@ -31,11 +33,10 @@ import java.util.Map;
  * A buffer to hold state requests to execute state requests in batch, which can only be manipulated
  * within task thread.
  *
- * @param <R> the type of the record
  * @param <K> the type of the key
  */
 @NotThreadSafe
-public class StateRequestBuffer<R, K> {
+public class StateRequestBuffer<K> {
     /**
      * The state requests in this buffer could be executed when the buffer is full or configured
      * batch size is reached. All operations on this buffer must be invoked in task thread.
@@ -76,8 +77,7 @@ public class StateRequestBuffer<R, K> {
      * @return The first record context with the same key in blocking queue, null if no such record.
      */
     @Nullable
-    @SuppressWarnings("rawtypes")
-    RecordContext<R, K> tryActivateOneByKey(K key) {
+    RecordContext<K> tryActivateOneByKey(K key) {
         if (!blockingQueue.containsKey(key)) {
             return null;
         }
@@ -88,7 +88,7 @@ public class StateRequestBuffer<R, K> {
             blockingQueue.remove(key);
         }
         blockingQueueSize--;
-        return (RecordContext<R, K>) stateRequest.getRecordContext();
+        return stateRequest.getRecordContext();
     }
 
     /**
@@ -117,9 +117,14 @@ public class StateRequestBuffer<R, K> {
      * @return A list of state requests.
      */
     List<StateRequest<?, ?, ?>> popActive(int n) {
-        LinkedList<StateRequest<?, ?, ?>> ret =
-                new LinkedList<>(activeQueue.subList(0, Math.min(activeQueue.size(), n)));
-        activeQueue.removeAll(ret);
+        final int count = Math.min(n, activeQueue.size());
+        if (count <= 0) {
+            return Collections.emptyList();
+        }
+        ArrayList<StateRequest<?, ?, ?>> ret = new ArrayList<>(count);
+        for (int i = 0; i < count; i++) {
+            ret.add(activeQueue.pop());
+        }
         return ret;
     }
 }
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
index bff46cbcde3..5d661958fc4 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
@@ -74,7 +74,7 @@ class AsyncExecutionControllerTest {
         String key1 = "key1";
         // Simulate the wrapping in {@link RecordProcessorUtils#getRecordProcessor()}, wrapping the
         // record and key with RecordContext.
-        RecordContext<String, String> recordContext1 = aec.buildContext(record1, key1);
+        RecordContext<String> recordContext1 = aec.buildContext(record1, key1);
         aec.setCurrentContext(recordContext1);
         // execute user code
         userCode.run();
@@ -99,14 +99,14 @@ class AsyncExecutionControllerTest {
         // ============================ element 2 & 3 ============================
         String record2 = "key1-r2";
         String key2 = "key1";
-        RecordContext<String, String> recordContext2 = aec.buildContext(record2, key2);
+        RecordContext<String> recordContext2 = aec.buildContext(record2, key2);
         aec.setCurrentContext(recordContext2);
         // execute user code
         userCode.run();
 
         String record3 = "key1-r3";
         String key3 = "key1";
-        RecordContext<String, String> recordContext3 = aec.buildContext(record3, key3);
+        RecordContext<String> recordContext3 = aec.buildContext(record3, key3);
         aec.setCurrentContext(recordContext3);
         // execute user code
         userCode.run();
@@ -151,7 +151,7 @@ class AsyncExecutionControllerTest {
         // ============================ element4 ============================
         String record4 = "key3-r3";
         String key4 = "key3";
-        RecordContext<String, String> recordContext4 = aec.buildContext(record4, key4);
+        RecordContext<String> recordContext4 = aec.buildContext(record4, key4);
         aec.setCurrentContext(recordContext4);
         // execute user code
         userCode.run();
@@ -182,21 +182,21 @@ class AsyncExecutionControllerTest {
 
         String record1 = "key1-r1";
         String key1 = "key1";
-        RecordContext<String, String> recordContext1 = aec.buildContext(record1, key1);
+        RecordContext<String> recordContext1 = aec.buildContext(record1, key1);
         aec.setCurrentContext(recordContext1);
         // execute user code
         userCode.run();
 
         String record2 = "key2-r1";
         String key2 = "key2";
-        RecordContext<String, String> recordContext2 = aec.buildContext(record2, key2);
+        RecordContext<String> recordContext2 = aec.buildContext(record2, key2);
         aec.setCurrentContext(recordContext2);
         // execute user code
         userCode.run();
 
         String record3 = "key1-r2";
         String key3 = "key1";
-        RecordContext<String, String> recordContext3 = aec.buildContext(record3, key3);
+        RecordContext<String> recordContext3 = aec.buildContext(record3, key3);
         aec.setCurrentContext(recordContext3);
         // execute user code
         userCode.run();
@@ -266,7 +266,7 @@ class AsyncExecutionControllerTest {
                 String record =
                         String.format("key%d-r%d", round * batchSize + i, round * batchSize + i);
                 String key = String.format("key%d", round * batchSize + i);
-                RecordContext<String, String> recordContext = aec.buildContext(record, key);
+                RecordContext<String> recordContext = aec.buildContext(record, key);
                 aec.setCurrentContext(recordContext);
                 userCode.run();
             }
@@ -279,7 +279,7 @@ class AsyncExecutionControllerTest {
         for (int i = 0; i < maxInFlight; i++) {
             String record = String.format("sameKey-r%d", i, i);
             String key = "sameKey";
-            RecordContext<String, String> recordContext = aec.buildContext(record, key);
+            RecordContext<String> recordContext = aec.buildContext(record, key);
             aec.setCurrentContext(recordContext);
             userCode.run();
         }
@@ -291,7 +291,7 @@ class AsyncExecutionControllerTest {
         for (int i = maxInFlight; i < 10 * maxInFlight; i++) {
             String record = String.format("sameKey-r%d", i, i);
             String key = "sameKey";
-            RecordContext<String, String> recordContext = aec.buildContext(record, key);
+            RecordContext<String> recordContext = aec.buildContext(record, key);
             aec.setCurrentContext(recordContext);
             userCode.run();
             assertThat(aec.inFlightRecordNum.get()).isEqualTo(maxInFlight + 1);
@@ -326,12 +326,12 @@ class AsyncExecutionControllerTest {
 
     static class TestValueState implements ValueState<Integer> {
 
-        private final AsyncExecutionController<String, String> asyncExecutionController;
+        private final AsyncExecutionController<String> asyncExecutionController;
 
         private final TestUnderlyingState underlyingState;
 
         public TestValueState(
-                AsyncExecutionController<String, String> aec, TestUnderlyingState underlyingState) {
+                AsyncExecutionController<String> aec, TestUnderlyingState underlyingState) {
             this.asyncExecutionController = aec;
             this.underlyingState = underlyingState;
         }
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/ContextStateFutureImplTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/ContextStateFutureImplTest.java
index d0a169f6444..0a2a40418ba 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/ContextStateFutureImplTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/ContextStateFutureImplTest.java
@@ -34,7 +34,7 @@ public class ContextStateFutureImplTest {
     @Test
     public void testThenApply() {
         SingleStepRunner runner = new SingleStepRunner();
-        RecordContext<String, String> recordContext = new RecordContext<>("a", "b", (e) -> {});
+        RecordContext<String> recordContext = new RecordContext<>("a", "b", (e) -> {});
 
         // validate
         ContextStateFutureImpl<Void> future =
@@ -58,7 +58,7 @@ public class ContextStateFutureImplTest {
     @Test
     public void testThenAccept() {
         SingleStepRunner runner = new SingleStepRunner();
-        RecordContext<String, String> recordContext = new RecordContext<>("a", "b", (e) -> {});
+        RecordContext<String> recordContext = new RecordContext<>("a", "b", (e) -> {});
 
         // validate
         ContextStateFutureImpl<Void> future =
@@ -82,7 +82,7 @@ public class ContextStateFutureImplTest {
     @Test
     public void testThenCompose() {
         SingleStepRunner runner = new SingleStepRunner();
-        RecordContext<String, String> recordContext = new RecordContext<>("a", "b", (e) -> {});
+        RecordContext<String> recordContext = new RecordContext<>("a", "b", (e) -> {});
 
         // validate
         ContextStateFutureImpl<Void> future =
@@ -106,7 +106,7 @@ public class ContextStateFutureImplTest {
     @Test
     public void testThenCombine() {
         SingleStepRunner runner = new SingleStepRunner();
-        RecordContext<String, String> recordContext = new RecordContext<>("a", "b", (e) -> {});
+        RecordContext<String> recordContext = new RecordContext<>("a", "b", (e) -> {});
 
         // validate
         ContextStateFutureImpl<Void> future1 =
@@ -156,7 +156,7 @@ public class ContextStateFutureImplTest {
     @Test
     public void testComplex() {
         SingleStepRunner runner = new SingleStepRunner();
-        RecordContext<String, String> recordContext = new RecordContext<>("a", "b", (e) -> {});
+        RecordContext<String> recordContext = new RecordContext<>("a", "b", (e) -> {});
 
         for (int i = 0; i < 32; i++) { // 2^5 for completion status combination
             // Each bit of 'i' represents the complete status when the user code executes.
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/KeyAccountingUnitTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/KeyAccountingUnitTest.java
index a2e2c60369a..ad047945a6c 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/KeyAccountingUnitTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/KeyAccountingUnitTest.java
@@ -28,7 +28,7 @@ class KeyAccountingUnitTest {
 
     @Test
     void testBasic() {
-        KeyAccountingUnit<String, Integer> keyAccountingUnit = new KeyAccountingUnit<>(10);
+        KeyAccountingUnit<Integer> keyAccountingUnit = new KeyAccountingUnit<>(10);
         assertThat(keyAccountingUnit.occupy("record1", 1)).isTrue();
         assertThat(keyAccountingUnit.occupy("record1", 1)).isFalse();
         assertThat(keyAccountingUnit.occupy("record2", 2)).isTrue();


(flink) 03/03: [FLINK-35025][Runtime/State] Introduce element order for async state processing

Posted by le...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

leiyanfei pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit fe8dde4ea489f466982ba896c1e2b22a7466e719
Author: Zakelly <za...@gmail.com>
AuthorDate: Mon Apr 15 20:53:47 2024 +0800

    [FLINK-35025][Runtime/State] Introduce element order for async state processing
---
 .../asyncprocessing/AsyncExecutionController.java  |  21 ++++
 .../asyncprocessing/StateRequestBuffer.java        |   8 +-
 .../AsyncExecutionControllerTest.java              |  44 +++++++
 .../AbstractAsyncStateStreamOperator.java          |  42 ++++---
 .../AbstractAsyncStateStreamOperatorV2.java        |  17 +++
 .../asyncprocessing/AsyncStateProcessing.java      |  23 +++-
 .../AsyncStateProcessingOperator.java              |  12 ++
 ...teProcessingOperator.java => ElementOrder.java} |  28 +++--
 .../AbstractAsyncStateStreamOperatorTest.java      | 108 ++++++++++++++++-
 .../AbstractAsyncStateStreamOperatorV2Test.java    | 130 +++++++++++++++++++--
 10 files changed, 380 insertions(+), 53 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
index 8dd8d16e721..1964ec3f378 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionController.java
@@ -21,7 +21,9 @@ package org.apache.flink.runtime.asyncprocessing;
 import org.apache.flink.api.common.operators.MailboxExecutor;
 import org.apache.flink.api.common.state.v2.State;
 import org.apache.flink.core.state.InternalStateFuture;
+import org.apache.flink.util.FlinkRuntimeException;
 import org.apache.flink.util.Preconditions;
+import org.apache.flink.util.function.ThrowingRunnable;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -248,4 +250,23 @@ public class AsyncExecutionController<K> {
         setCurrentContext(storedContext);
         inFlightRecordNum.incrementAndGet();
     }
+
+    /**
+     * A helper to request a {@link StateRequestType#SYNC_POINT} and run a callback if it finishes
+     * (once the record is not blocked).
+     *
+     * @param callback the callback to run if it finishes (once the record is not blocked).
+     */
+    public void syncPointRequestWithCallback(ThrowingRunnable<Exception> callback) {
+        handleRequest(null, StateRequestType.SYNC_POINT, null)
+                .thenAccept(
+                        v -> {
+                            try {
+                                callback.run();
+                            } catch (Exception e) {
+                                // TODO: Properly handle the exception and fail the entire job.
+                                throw new FlinkRuntimeException("Unexpected runtime exception", e);
+                            }
+                        });
+    }
 }
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestBuffer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestBuffer.java
index 0eea7bbcacc..9fa10bfee5b 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestBuffer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/StateRequestBuffer.java
@@ -60,7 +60,11 @@ public class StateRequestBuffer<K> {
     }
 
     void enqueueToActive(StateRequest<K, ?, ?> request) {
-        activeQueue.add(request);
+        if (request.getRequestType() == StateRequestType.SYNC_POINT) {
+            request.getFuture().complete(null);
+        } else {
+            activeQueue.add(request);
+        }
     }
 
     void enqueueToBlocking(StateRequest<K, ?, ?> request) {
@@ -83,7 +87,7 @@ public class StateRequestBuffer<K> {
         }
 
         StateRequest<K, ?, ?> stateRequest = blockingQueue.get(key).removeFirst();
-        activeQueue.add(stateRequest);
+        enqueueToActive(stateRequest);
         if (blockingQueue.get(key).isEmpty()) {
             blockingQueue.remove(key);
         }
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
index 5d661958fc4..d36a5e53ae6 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/asyncprocessing/AsyncExecutionControllerTest.java
@@ -300,6 +300,50 @@ class AsyncExecutionControllerTest {
         }
     }
 
+    @Test
+    public void testSyncPoint() {
+        AtomicInteger counter = new AtomicInteger(0);
+
+        // Test the sync point processing without a key occupied.
+        RecordContext<String> recordContext = aec.buildContext("record", "key");
+        aec.setCurrentContext(recordContext);
+        recordContext.retain();
+        aec.syncPointRequestWithCallback(counter::incrementAndGet);
+        assertThat(counter.get()).isEqualTo(1);
+        assertThat(recordContext.getReferenceCount()).isEqualTo(1);
+        assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(0);
+        assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(0);
+        assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(1);
+        recordContext.release();
+        assertThat(aec.keyAccountingUnit.occupiedCount()).isEqualTo(0);
+
+        counter.set(0);
+        // Test the sync point processing with a key occupied.
+        RecordContext<String> recordContext1 = aec.buildContext("record1", "occupied");
+        aec.setCurrentContext(recordContext1);
+        userCode.run();
+
+        RecordContext<String> recordContext2 = aec.buildContext("record2", "occupied");
+        aec.setCurrentContext(recordContext2);
+        aec.syncPointRequestWithCallback(counter::incrementAndGet);
+        recordContext2.retain();
+        assertThat(counter.get()).isEqualTo(0);
+        assertThat(recordContext2.getReferenceCount()).isGreaterThan(1);
+        assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
+        assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(1);
+        aec.triggerIfNeeded(true);
+        assertThat(counter.get()).isEqualTo(0);
+        assertThat(recordContext2.getReferenceCount()).isGreaterThan(1);
+        assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(1);
+        assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(1);
+        aec.triggerIfNeeded(true);
+        assertThat(counter.get()).isEqualTo(1);
+        assertThat(recordContext2.getReferenceCount()).isEqualTo(1);
+        assertThat(aec.stateRequestsBuffer.activeQueueSize()).isEqualTo(0);
+        assertThat(aec.stateRequestsBuffer.blockingQueueSize()).isEqualTo(0);
+        recordContext2.release();
+    }
+
     private StateExecutor createStateExecutor() {
         TestAsyncStateBackend testAsyncStateBackend = new TestAsyncStateBackend();
         assertThat(testAsyncStateBackend.supportsAsyncKeyedStateBackend()).isTrue();
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
index ff02953cebd..e57f8b66781 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
@@ -32,6 +32,7 @@ import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.runtime.tasks.StreamTask;
 import org.apache.flink.util.function.ThrowingConsumer;
+import org.apache.flink.util.function.ThrowingRunnable;
 
 /**
  * This operator is an abstract class that give the {@link AbstractStreamOperator} the ability to
@@ -65,6 +66,11 @@ public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStre
         return true;
     }
 
+    @Override
+    public ElementOrder getElementOrder() {
+        return ElementOrder.RECORD_ORDER;
+    }
+
     @Override
     @SuppressWarnings("unchecked")
     public final <T> void setAsyncKeyedContextElement(
@@ -91,6 +97,12 @@ public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStre
         currentProcessingContext.release();
     }
 
+    @Override
+    @SuppressWarnings("unchecked")
+    public final void preserveRecordOrderAndProcess(ThrowingRunnable<Exception> processing) {
+        asyncExecutionController.syncPointRequestWithCallback(processing);
+    }
+
     @Override
     @SuppressWarnings("unchecked")
     public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId) {
@@ -99,26 +111,21 @@ public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStre
         if (this instanceof TwoInputStreamOperator) {
             switch (inputId) {
                 case 1:
-                    return record -> {
-                        setAsyncKeyedContextElement(record, (KeySelector<T, ?>) stateKeySelector1);
-                        ((TwoInputStreamOperator) this).processElement1(record);
-                        postProcessElement();
-                    };
+                    return AsyncStateProcessing.<T>makeRecordProcessor(
+                            this,
+                            (KeySelector) stateKeySelector1,
+                            ((TwoInputStreamOperator) this)::processElement1);
                 case 2:
-                    return record -> {
-                        setAsyncKeyedContextElement(record, (KeySelector<T, ?>) stateKeySelector2);
-                        ((TwoInputStreamOperator) this).processElement2(record);
-                        postProcessElement();
-                    };
+                    return AsyncStateProcessing.<T>makeRecordProcessor(
+                            this,
+                            (KeySelector) stateKeySelector2,
+                            ((TwoInputStreamOperator) this)::processElement2);
                 default:
                     break;
             }
         } else if (this instanceof Input && inputId == 1) {
-            return record -> {
-                setAsyncKeyedContextElement(record, (KeySelector<T, ?>) stateKeySelector1);
-                ((Input) this).processElement(record);
-                postProcessElement();
-            };
+            return AsyncStateProcessing.<T>makeRecordProcessor(
+                    this, (KeySelector) stateKeySelector1, ((Input) this)::processElement);
         }
         throw new IllegalArgumentException(
                 String.format(
@@ -130,4 +137,9 @@ public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStre
     AsyncExecutionController<?> getAsyncExecutionController() {
         return asyncExecutionController;
     }
+
+    @VisibleForTesting
+    RecordContext getCurrentProcessingContext() {
+        return currentProcessingContext;
+    }
 }
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
index 3e7254b7567..000a63aeea1 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
@@ -29,6 +29,7 @@ import org.apache.flink.streaming.api.operators.StreamOperatorParameters;
 import org.apache.flink.streaming.api.operators.StreamTaskStateInitializer;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.util.function.ThrowingConsumer;
+import org.apache.flink.util.function.ThrowingRunnable;
 
 /**
  * This operator is an abstract class that give the {@link AbstractStreamOperatorV2} the ability to
@@ -67,6 +68,11 @@ public abstract class AbstractAsyncStateStreamOperatorV2<OUT> extends AbstractSt
         return true;
     }
 
+    @Override
+    public ElementOrder getElementOrder() {
+        return ElementOrder.RECORD_ORDER;
+    }
+
     @Override
     @SuppressWarnings("unchecked")
     public final <T> void setAsyncKeyedContextElement(
@@ -93,6 +99,12 @@ public abstract class AbstractAsyncStateStreamOperatorV2<OUT> extends AbstractSt
         currentProcessingContext.release();
     }
 
+    @Override
+    @SuppressWarnings("unchecked")
+    public final void preserveRecordOrderAndProcess(ThrowingRunnable<Exception> processing) {
+        asyncExecutionController.syncPointRequestWithCallback(processing);
+    }
+
     @Override
     public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId) {
         // The real logic should be in First/SecondInputOfTwoInput#getRecordProcessor.
@@ -105,4 +117,9 @@ public abstract class AbstractAsyncStateStreamOperatorV2<OUT> extends AbstractSt
     AsyncExecutionController<?> getAsyncExecutionController() {
         return asyncExecutionController;
     }
+
+    @VisibleForTesting
+    RecordContext getCurrentProcessingContext() {
+        return currentProcessingContext;
+    }
 }
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
index ec033b3906f..563f3fd1ad8 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
@@ -56,10 +56,23 @@ public interface AsyncStateProcessing {
             AsyncStateProcessingOperator asyncOperator,
             KeySelector<T, ?> keySelector,
             ThrowingConsumer<StreamRecord<T>, Exception> processor) {
-        return (record) -> {
-            asyncOperator.setAsyncKeyedContextElement(record, keySelector);
-            processor.accept(record);
-            asyncOperator.postProcessElement();
-        };
+        switch (asyncOperator.getElementOrder()) {
+            case RECORD_ORDER:
+                return (record) -> {
+                    asyncOperator.setAsyncKeyedContextElement(record, keySelector);
+                    asyncOperator.preserveRecordOrderAndProcess(() -> processor.accept(record));
+                    asyncOperator.postProcessElement();
+                };
+            case FIRST_STATE_ORDER:
+                return (record) -> {
+                    asyncOperator.setAsyncKeyedContextElement(record, keySelector);
+                    processor.accept(record);
+                    asyncOperator.postProcessElement();
+                };
+            default:
+                throw new UnsupportedOperationException(
+                        "Unknown element order for async processing:"
+                                + asyncOperator.getElementOrder());
+        }
     }
 }
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java
index d8b52b6dd4f..e24344fece5 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java
@@ -21,6 +21,7 @@ package org.apache.flink.streaming.runtime.operators.asyncprocessing;
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.api.java.functions.KeySelector;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.util.function.ThrowingRunnable;
 
 /**
  * A more detailed interface based on {@link AsyncStateProcessing}, which gives the essential
@@ -29,6 +30,9 @@ import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 @Internal
 public interface AsyncStateProcessingOperator extends AsyncStateProcessing {
 
+    /** Get the {@link ElementOrder} of this operator. */
+    ElementOrder getElementOrder();
+
     /**
      * Set key context for async state processing.
      *
@@ -41,4 +45,12 @@ public interface AsyncStateProcessingOperator extends AsyncStateProcessing {
 
     /** A callback that will be triggered after an element finishes {@code processElement}. */
     void postProcessElement();
+
+    /**
+     * Check the order of same-key record, and then process the record. Mainly used when the {@link
+     * #getElementOrder()} returns {@link ElementOrder#RECORD_ORDER}.
+     *
+     * @param processing the record processing logic.
+     */
+    void preserveRecordOrderAndProcess(ThrowingRunnable<Exception> processing);
 }
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/ElementOrder.java
similarity index 54%
copy from flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java
copy to flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/ElementOrder.java
index d8b52b6dd4f..2d7b2a63c12 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/ElementOrder.java
@@ -19,26 +19,24 @@
 package org.apache.flink.streaming.runtime.operators.asyncprocessing;
 
 import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.java.functions.KeySelector;
-import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 
 /**
- * A more detailed interface based on {@link AsyncStateProcessing}, which gives the essential
- * methods for an operator to perform async state processing.
+ * This enum defines the element order of being processed. Only the elements with the same key
+ * should be considered here. We should keep this internal and away from API module for now, until
+ * we could see the concrete need for {@link #FIRST_STATE_ORDER} from average users.
  */
 @Internal
-public interface AsyncStateProcessingOperator extends AsyncStateProcessing {
-
+public enum ElementOrder {
     /**
-     * Set key context for async state processing.
-     *
-     * @param record the record.
-     * @param keySelector the key selector to select a key from record.
-     * @param <T> the type of the record.
+     * Treat the record processing as a whole, meaning that any {@code processElement} call for the
+     * elements with same key should follow the order of record arrival AND no parallel run is
+     * allowed.
      */
-    <T> void setAsyncKeyedContextElement(StreamRecord<T> record, KeySelector<T, ?> keySelector)
-            throws Exception;
+    RECORD_ORDER,
 
-    /** A callback that will be triggered after an element finishes {@code processElement}. */
-    void postProcessElement();
+    /**
+     * The {@code processElement} call will be invoked on record arrival, but may be blocked at the
+     * first state accessing if there is a preceding same-key record under processing.
+     */
+    FIRST_STATE_ORDER,
 }
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorTest.java
index 20c01c8bc4f..a82a0d2968b 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorTest.java
@@ -25,20 +25,27 @@ import org.apache.flink.runtime.state.VoidNamespace;
 import org.apache.flink.streaming.api.operators.InternalTimer;
 import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
 import org.apache.flink.streaming.api.operators.Triggerable;
+import org.apache.flink.streaming.runtime.io.RecordProcessorUtils;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.apache.flink.util.function.ThrowingConsumer;
 
 import org.junit.jupiter.api.Test;
 
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicInteger;
+
 import static org.assertj.core.api.Assertions.assertThat;
 
 /** Basic tests for {@link AbstractAsyncStateStreamOperator}. */
 public class AbstractAsyncStateStreamOperatorTest {
 
     protected KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
-            createTestHarness(int maxParalelism, int numSubtasks, int subtaskIndex)
+            createTestHarness(
+                    int maxParalelism, int numSubtasks, int subtaskIndex, ElementOrder elementOrder)
                     throws Exception {
-        TestOperator testOperator = new TestOperator();
+        TestOperator testOperator = new TestOperator(elementOrder);
         return new KeyedOneInputStreamOperatorTestHarness<>(
                 testOperator,
                 new TestKeySelector(),
@@ -51,7 +58,7 @@ public class AbstractAsyncStateStreamOperatorTest {
     @Test
     public void testCreateAsyncExecutionController() throws Exception {
         try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
-                testHarness = createTestHarness(128, 1, 0)) {
+                testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
             testHarness.open();
             assertThat(testHarness.getOperator())
                     .isInstanceOf(AbstractAsyncStateStreamOperator.class);
@@ -62,6 +69,68 @@ public class AbstractAsyncStateStreamOperatorTest {
         }
     }
 
+    @Test
+    public void testRecordProcessorWithFirstStateOrder() throws Exception {
+        try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
+                testHarness = createTestHarness(128, 1, 0, ElementOrder.FIRST_STATE_ORDER)) {
+            testHarness.open();
+            TestOperator testOperator = (TestOperator) testHarness.getOperator();
+            ThrowingConsumer<StreamRecord<Tuple2<Integer, String>>, Exception> processor =
+                    RecordProcessorUtils.getRecordProcessor(testOperator);
+            ExecutorService anotherThread = Executors.newSingleThreadExecutor();
+            // Trigger the processor
+            anotherThread.execute(
+                    () -> {
+                        try {
+                            processor.accept(new StreamRecord<>(Tuple2.of(5, "5")));
+                        } catch (Exception e) {
+                        }
+                    });
+
+            Thread.sleep(1000);
+            assertThat(testOperator.getProcessed()).isEqualTo(1);
+            assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(1);
+
+            // Proceed processing
+            testOperator.proceed();
+            anotherThread.shutdown();
+            Thread.sleep(1000);
+            assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(0);
+        }
+    }
+
+    @Test
+    public void testRecordProcessorWithRecordOrder() throws Exception {
+        try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
+                testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
+            testHarness.open();
+            TestOperator testOperator = (TestOperator) testHarness.getOperator();
+            ThrowingConsumer<StreamRecord<Tuple2<Integer, String>>, Exception> processor =
+                    RecordProcessorUtils.getRecordProcessor(testOperator);
+            ExecutorService anotherThread = Executors.newSingleThreadExecutor();
+            // Trigger the processor
+            anotherThread.execute(
+                    () -> {
+                        try {
+                            processor.accept(new StreamRecord<>(Tuple2.of(5, "5")));
+                        } catch (Exception e) {
+                        }
+                    });
+
+            Thread.sleep(1000);
+            assertThat(testOperator.getProcessed()).isEqualTo(1);
+            // Why greater than 1:  +1 when enter the processor; +1 when handle the SYNC_POINT
+            assertThat(testOperator.getCurrentProcessingContext().getReferenceCount())
+                    .isGreaterThan(1);
+
+            // Proceed processing
+            testOperator.proceed();
+            anotherThread.shutdown();
+            Thread.sleep(1000);
+            assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(0);
+        }
+    }
+
     /** A simple testing operator. */
     private static class TestOperator extends AbstractAsyncStateStreamOperator<String>
             implements OneInputStreamOperator<Tuple2<Integer, String>, String>,
@@ -69,14 +138,33 @@ public class AbstractAsyncStateStreamOperatorTest {
 
         private static final long serialVersionUID = 1L;
 
+        private final ElementOrder elementOrder;
+
+        final AtomicInteger processed = new AtomicInteger(0);
+
+        final Object objectToWait = new Object();
+
+        TestOperator(ElementOrder elementOrder) {
+            this.elementOrder = elementOrder;
+        }
+
         @Override
         public void open() throws Exception {
             super.open();
         }
 
         @Override
-        public void processElement(StreamRecord<Tuple2<Integer, String>> element)
-                throws Exception {}
+        public ElementOrder getElementOrder() {
+            return elementOrder;
+        }
+
+        @Override
+        public void processElement(StreamRecord<Tuple2<Integer, String>> element) throws Exception {
+            processed.incrementAndGet();
+            synchronized (objectToWait) {
+                objectToWait.wait();
+            }
+        }
 
         @Override
         public void onEventTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {}
@@ -84,6 +172,16 @@ public class AbstractAsyncStateStreamOperatorTest {
         @Override
         public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer)
                 throws Exception {}
+
+        public int getProcessed() {
+            return processed.get();
+        }
+
+        public void proceed() {
+            synchronized (objectToWait) {
+                objectToWait.notify();
+            }
+        }
     }
 
     /** {@link KeySelector} for tests. */
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2Test.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2Test.java
index 088729b5067..5cbb5d71c3e 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2Test.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2Test.java
@@ -32,13 +32,18 @@ import org.apache.flink.streaming.api.operators.StreamOperator;
 import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
 import org.apache.flink.streaming.api.operators.StreamOperatorParameters;
 import org.apache.flink.streaming.api.operators.Triggerable;
+import org.apache.flink.streaming.runtime.io.RecordProcessorUtils;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.apache.flink.util.function.ThrowingConsumer;
 
 import org.junit.jupiter.api.Test;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -46,10 +51,11 @@ import static org.assertj.core.api.Assertions.assertThat;
 public class AbstractAsyncStateStreamOperatorV2Test {
 
     protected KeyedOneInputStreamOperatorV2TestHarness<Integer, Tuple2<Integer, String>, String>
-            createTestHarness(int maxParalelism, int numSubtasks, int subtaskIndex)
+            createTestHarness(
+                    int maxParalelism, int numSubtasks, int subtaskIndex, ElementOrder elementOrder)
                     throws Exception {
         return new KeyedOneInputStreamOperatorV2TestHarness<>(
-                new TestOperatorFactory(),
+                new TestOperatorFactory(elementOrder),
                 new AbstractAsyncStateStreamOperatorTest.TestKeySelector(),
                 BasicTypeInfo.INT_TYPE_INFO,
                 maxParalelism,
@@ -60,7 +66,7 @@ public class AbstractAsyncStateStreamOperatorV2Test {
     @Test
     public void testCreateAsyncExecutionController() throws Exception {
         try (KeyedOneInputStreamOperatorV2TestHarness<Integer, Tuple2<Integer, String>, String>
-                testHarness = createTestHarness(128, 1, 0)) {
+                testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
             testHarness.open();
             assertThat(testHarness.getBaseOperator())
                     .isInstanceOf(AbstractAsyncStateStreamOperatorV2.class);
@@ -71,6 +77,70 @@ public class AbstractAsyncStateStreamOperatorV2Test {
         }
     }
 
+    @Test
+    public void testRecordProcessorWithFirstStateOrder() throws Exception {
+        try (KeyedOneInputStreamOperatorV2TestHarness<Integer, Tuple2<Integer, String>, String>
+                testHarness = createTestHarness(128, 1, 0, ElementOrder.FIRST_STATE_ORDER)) {
+            testHarness.open();
+            SingleInputTestOperator testOperator =
+                    (SingleInputTestOperator) testHarness.getBaseOperator();
+            ThrowingConsumer<StreamRecord<Tuple2<Integer, String>>, Exception> processor =
+                    RecordProcessorUtils.getRecordProcessor(testOperator.getInputs().get(0));
+            ExecutorService anotherThread = Executors.newSingleThreadExecutor();
+            // Trigger the processor
+            anotherThread.execute(
+                    () -> {
+                        try {
+                            processor.accept(new StreamRecord<>(Tuple2.of(5, "5")));
+                        } catch (Exception e) {
+                        }
+                    });
+
+            Thread.sleep(1000);
+            assertThat(testOperator.getProcessed()).isEqualTo(1);
+            assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(1);
+
+            // Proceed processing
+            testOperator.proceed();
+            anotherThread.shutdown();
+            Thread.sleep(1000);
+            assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(0);
+        }
+    }
+
+    @Test
+    public void testRecordProcessorWithRecordOrder() throws Exception {
+        try (KeyedOneInputStreamOperatorV2TestHarness<Integer, Tuple2<Integer, String>, String>
+                testHarness = createTestHarness(128, 1, 0, ElementOrder.RECORD_ORDER)) {
+            testHarness.open();
+            SingleInputTestOperator testOperator =
+                    (SingleInputTestOperator) testHarness.getBaseOperator();
+            ThrowingConsumer<StreamRecord<Tuple2<Integer, String>>, Exception> processor =
+                    RecordProcessorUtils.getRecordProcessor(testOperator.getInputs().get(0));
+            ExecutorService anotherThread = Executors.newSingleThreadExecutor();
+            // Trigger the processor
+            anotherThread.execute(
+                    () -> {
+                        try {
+                            processor.accept(new StreamRecord<>(Tuple2.of(5, "5")));
+                        } catch (Exception e) {
+                        }
+                    });
+
+            Thread.sleep(1000);
+            assertThat(testOperator.getProcessed()).isEqualTo(1);
+            // Why greater than 1:  +1 when enter the processor; +1 when handle the SYNC_POINT
+            assertThat(testOperator.getCurrentProcessingContext().getReferenceCount())
+                    .isGreaterThan(1);
+
+            // Proceed processing
+            testOperator.proceed();
+            anotherThread.shutdown();
+            Thread.sleep(1000);
+            assertThat(testOperator.getCurrentProcessingContext().getReferenceCount()).isEqualTo(0);
+        }
+    }
+
     private static class KeyedOneInputStreamOperatorV2TestHarness<K, IN, OUT>
             extends KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> {
         public KeyedOneInputStreamOperatorV2TestHarness(
@@ -90,10 +160,17 @@ public class AbstractAsyncStateStreamOperatorV2Test {
     }
 
     private static class TestOperatorFactory extends AbstractStreamOperatorFactory<String> {
+
+        private final ElementOrder elementOrder;
+
+        TestOperatorFactory(ElementOrder elementOrder) {
+            this.elementOrder = elementOrder;
+        }
+
         @Override
         public <T extends StreamOperator<String>> T createStreamOperator(
                 StreamOperatorParameters<String> parameters) {
-            return (T) new SingleInputTestOperator(parameters);
+            return (T) new SingleInputTestOperator(parameters, elementOrder);
         }
 
         @Override
@@ -111,8 +188,29 @@ public class AbstractAsyncStateStreamOperatorV2Test {
 
         private static final long serialVersionUID = 1L;
 
-        public SingleInputTestOperator(StreamOperatorParameters<String> parameters) {
+        final AtomicInteger processed = new AtomicInteger(0);
+
+        private final ElementOrder elementOrder;
+
+        final Object objectToWait = new Object();
+
+        final Input input;
+
+        public SingleInputTestOperator(
+                StreamOperatorParameters<String> parameters, ElementOrder elementOrder) {
             super(parameters, 1);
+            this.elementOrder = elementOrder;
+            input =
+                    new AbstractInput<Tuple2<Integer, String>, String>(this, 1) {
+                        @Override
+                        public void processElement(StreamRecord<Tuple2<Integer, String>> element)
+                                throws Exception {
+                            processed.incrementAndGet();
+                            synchronized (objectToWait) {
+                                objectToWait.wait();
+                            }
+                        }
+                    };
         }
 
         @Override
@@ -122,12 +220,12 @@ public class AbstractAsyncStateStreamOperatorV2Test {
 
         @Override
         public List<Input> getInputs() {
-            return Collections.singletonList(
-                    new AbstractInput<Tuple2<Integer, String>, String>(this, 1) {
-                        @Override
-                        public void processElement(StreamRecord<Tuple2<Integer, String>> element)
-                                throws Exception {}
-                    });
+            return Collections.singletonList(input);
+        }
+
+        @Override
+        public ElementOrder getElementOrder() {
+            return elementOrder;
         }
 
         @Override
@@ -136,5 +234,15 @@ public class AbstractAsyncStateStreamOperatorV2Test {
         @Override
         public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer)
                 throws Exception {}
+
+        public int getProcessed() {
+            return processed.get();
+        }
+
+        public void proceed() {
+            synchronized (objectToWait) {
+                objectToWait.notify();
+            }
+        }
     }
 }


(flink) 02/03: [FLINK-35025][Runtime/State] Abstract stream operators for async state processing

Posted by le...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

leiyanfei pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit e1edb24ecc2e28987803ded3b460323c0179a052
Author: Zakelly <za...@gmail.com>
AuthorDate: Fri Apr 12 21:52:29 2024 +0800

    [FLINK-35025][Runtime/State] Abstract stream operators for async state processing
---
 .../streaming/api/operators/AbstractInput.java     |  23 +++-
 .../api/operators/AbstractStreamOperator.java      |   4 +-
 .../api/operators/AbstractStreamOperatorV2.java    |   2 +-
 .../streaming/runtime/io/RecordProcessorUtils.java |  10 ++
 .../AbstractAsyncStateStreamOperator.java          | 133 ++++++++++++++++++++
 .../AbstractAsyncStateStreamOperatorV2.java        | 108 ++++++++++++++++
 .../asyncprocessing/AsyncStateProcessing.java      |  65 ++++++++++
 .../AsyncStateProcessingOperator.java              |  44 +++++++
 .../AbstractAsyncStateStreamOperatorTest.java      |  98 +++++++++++++++
 .../AbstractAsyncStateStreamOperatorV2Test.java    | 140 +++++++++++++++++++++
 .../multipleinput/input/FirstInputOfTwoInput.java  |  18 ++-
 .../operators/multipleinput/input/OneInput.java    |  18 ++-
 .../multipleinput/input/SecondInputOfTwoInput.java |  18 ++-
 13 files changed, 674 insertions(+), 7 deletions(-)

diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractInput.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractInput.java
index 9982d12c8e8..20c3006d838 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractInput.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractInput.java
@@ -19,12 +19,16 @@
 package org.apache.flink.streaming.api.operators;
 
 import org.apache.flink.annotation.Experimental;
+import org.apache.flink.annotation.Internal;
 import org.apache.flink.api.java.functions.KeySelector;
 import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.operators.asyncprocessing.AsyncStateProcessing;
+import org.apache.flink.streaming.runtime.operators.asyncprocessing.AsyncStateProcessingOperator;
 import org.apache.flink.streaming.runtime.streamrecord.LatencyMarker;
 import org.apache.flink.streaming.runtime.streamrecord.RecordAttributes;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
+import org.apache.flink.util.function.ThrowingConsumer;
 
 import javax.annotation.Nullable;
 
@@ -35,7 +39,8 @@ import static org.apache.flink.util.Preconditions.checkArgument;
  * AbstractStreamOperatorV2}.
  */
 @Experimental
-public abstract class AbstractInput<IN, OUT> implements Input<IN>, KeyContextHandler {
+public abstract class AbstractInput<IN, OUT>
+        implements Input<IN>, KeyContextHandler, AsyncStateProcessing {
     /**
      * {@code KeySelector} for extracting a key from an element being processed. This is used to
      * scope keyed state to a key. This is null if the operator is not a keyed operator.
@@ -86,4 +91,20 @@ public abstract class AbstractInput<IN, OUT> implements Input<IN>, KeyContextHan
     public boolean hasKeyContext() {
         return stateKeySelector != null;
     }
+
+    @Internal
+    @Override
+    public final boolean isAsyncStateProcessingEnabled() {
+        return (owner instanceof AsyncStateProcessingOperator)
+                && ((AsyncStateProcessingOperator) owner).isAsyncStateProcessingEnabled();
+    }
+
+    @Internal
+    @Override
+    public final ThrowingConsumer<StreamRecord<IN>, Exception> getRecordProcessor(int inputId) {
+        return AsyncStateProcessing.makeRecordProcessor(
+                (AsyncStateProcessingOperator) owner,
+                (KeySelector) stateKeySelector,
+                this::processElement);
+    }
 }
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
index 632b89b16ae..76fef0c26cf 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
@@ -128,7 +128,7 @@ public abstract class AbstractStreamOperator<OUT>
      *
      * <p>This is for elements from the first input.
      */
-    private transient KeySelector<?, ?> stateKeySelector1;
+    protected transient KeySelector<?, ?> stateKeySelector1;
 
     /**
      * {@code KeySelector} for extracting a key from an element being processed. This is used to
@@ -136,7 +136,7 @@ public abstract class AbstractStreamOperator<OUT>
      *
      * <p>This is for elements from the second input.
      */
-    private transient KeySelector<?, ?> stateKeySelector2;
+    protected transient KeySelector<?, ?> stateKeySelector2;
 
     private transient StreamOperatorStateHandler stateHandler;
 
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperatorV2.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperatorV2.java
index c16e5f85d4e..aed9892c8d5 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperatorV2.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperatorV2.java
@@ -190,7 +190,7 @@ public abstract class AbstractStreamOperatorV2<OUT>
     }
 
     @Override
-    public final void initializeState(StreamTaskStateInitializer streamTaskStateManager)
+    public void initializeState(StreamTaskStateInitializer streamTaskStateManager)
             throws Exception {
         final TypeSerializer<?> keySerializer =
                 config.getStateKeySerializer(getUserCodeClassloader());
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/io/RecordProcessorUtils.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/io/RecordProcessorUtils.java
index f1636f3285d..d6fbfb3636f 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/io/RecordProcessorUtils.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/io/RecordProcessorUtils.java
@@ -23,6 +23,7 @@ import org.apache.flink.streaming.api.operators.KeyContextHandler;
 import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
 import org.apache.flink.streaming.api.operators.StreamOperator;
 import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
+import org.apache.flink.streaming.runtime.operators.asyncprocessing.AsyncStateProcessing;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.util.FlinkRuntimeException;
 import org.apache.flink.util.function.ThrowingConsumer;
@@ -54,6 +55,9 @@ public class RecordProcessorUtils {
 
         if (canOmitSetKeyContext) {
             return input::processElement;
+        } else if (input instanceof AsyncStateProcessing
+                && ((AsyncStateProcessing) input).isAsyncStateProcessingEnabled()) {
+            return ((AsyncStateProcessing) input).getRecordProcessor(1);
         } else {
             return record -> {
                 input.setKeyContextElement(record);
@@ -82,6 +86,9 @@ public class RecordProcessorUtils {
 
         if (canOmitSetKeyContext) {
             return operator::processElement1;
+        } else if (operator instanceof AsyncStateProcessing
+                && ((AsyncStateProcessing) operator).isAsyncStateProcessingEnabled()) {
+            return ((AsyncStateProcessing) operator).getRecordProcessor(1);
         } else {
             return record -> {
                 operator.setKeyContextElement1(record);
@@ -110,6 +117,9 @@ public class RecordProcessorUtils {
 
         if (canOmitSetKeyContext) {
             return operator::processElement2;
+        } else if (operator instanceof AsyncStateProcessing
+                && ((AsyncStateProcessing) operator).isAsyncStateProcessingEnabled()) {
+            return ((AsyncStateProcessing) operator).getRecordProcessor(2);
         } else {
             return record -> {
                 operator.setKeyContextElement2(record);
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
new file mode 100644
index 00000000000..ff02953cebd
--- /dev/null
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperator.java
@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.operators.asyncprocessing;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.common.operators.MailboxExecutor;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.asyncprocessing.AsyncExecutionController;
+import org.apache.flink.runtime.asyncprocessing.RecordContext;
+import org.apache.flink.streaming.api.graph.StreamConfig;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.Input;
+import org.apache.flink.streaming.api.operators.Output;
+import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
+import org.apache.flink.util.function.ThrowingConsumer;
+
+/**
+ * This operator is an abstract class that give the {@link AbstractStreamOperator} the ability to
+ * perform {@link AsyncStateProcessing}. The aim is to make any subclass of {@link
+ * AbstractStreamOperator} could manipulate async state with only a change of base class.
+ */
+@Internal
+@SuppressWarnings("rawtypes")
+public abstract class AbstractAsyncStateStreamOperator<OUT> extends AbstractStreamOperator<OUT>
+        implements AsyncStateProcessingOperator {
+
+    private AsyncExecutionController asyncExecutionController;
+
+    private RecordContext currentProcessingContext;
+
+    /** Initialize necessary state components for {@link AbstractStreamOperator}. */
+    @Override
+    public void setup(
+            StreamTask<?, ?> containingTask,
+            StreamConfig config,
+            Output<StreamRecord<OUT>> output) {
+        super.setup(containingTask, config, output);
+        // TODO: properly read config and setup
+        final MailboxExecutor mailboxExecutor =
+                containingTask.getEnvironment().getMainMailboxExecutor();
+        this.asyncExecutionController = new AsyncExecutionController(mailboxExecutor, null);
+    }
+
+    @Override
+    public boolean isAsyncStateProcessingEnabled() {
+        return true;
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public final <T> void setAsyncKeyedContextElement(
+            StreamRecord<T> record, KeySelector<T, ?> keySelector) throws Exception {
+        currentProcessingContext =
+                asyncExecutionController.buildContext(
+                        record.getValue(), keySelector.getKey(record.getValue()));
+        // The processElement will be treated as a callback for dummy request. So reference
+        // counting should be maintained.
+        // When state request submitted, ref count +1, as described in FLIP-425:
+        // To cover the statements without a callback, in addition to the reference count marked
+        // in Fig.5, each state request itself is also protected by a paired reference count.
+        currentProcessingContext.retain();
+        asyncExecutionController.setCurrentContext(currentProcessingContext);
+    }
+
+    @Override
+    public final void postProcessElement() {
+        // The processElement will be treated as a callback for dummy request. So reference
+        // counting should be maintained.
+        // When a state request completes, ref count -1, as described in FLIP-425:
+        // To cover the statements without a callback, in addition to the reference count marked
+        // in Fig.5, each state request itself is also protected by a paired reference count.
+        currentProcessingContext.release();
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId) {
+        // Ideally, only TwoStreamInputOperator/OneInputStreamOperator(Input) will invoke here.
+        // Only those operators have the definition of processElement(1/2).
+        if (this instanceof TwoInputStreamOperator) {
+            switch (inputId) {
+                case 1:
+                    return record -> {
+                        setAsyncKeyedContextElement(record, (KeySelector<T, ?>) stateKeySelector1);
+                        ((TwoInputStreamOperator) this).processElement1(record);
+                        postProcessElement();
+                    };
+                case 2:
+                    return record -> {
+                        setAsyncKeyedContextElement(record, (KeySelector<T, ?>) stateKeySelector2);
+                        ((TwoInputStreamOperator) this).processElement2(record);
+                        postProcessElement();
+                    };
+                default:
+                    break;
+            }
+        } else if (this instanceof Input && inputId == 1) {
+            return record -> {
+                setAsyncKeyedContextElement(record, (KeySelector<T, ?>) stateKeySelector1);
+                ((Input) this).processElement(record);
+                postProcessElement();
+            };
+        }
+        throw new IllegalArgumentException(
+                String.format(
+                        "Unsupported operator type %s with input id %d",
+                        getClass().getName(), inputId));
+    }
+
+    @VisibleForTesting
+    AsyncExecutionController<?> getAsyncExecutionController() {
+        return asyncExecutionController;
+    }
+}
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
new file mode 100644
index 00000000000..3e7254b7567
--- /dev/null
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2.java
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.operators.asyncprocessing;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.api.common.operators.MailboxExecutor;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.asyncprocessing.AsyncExecutionController;
+import org.apache.flink.runtime.asyncprocessing.RecordContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperatorV2;
+import org.apache.flink.streaming.api.operators.StreamOperatorParameters;
+import org.apache.flink.streaming.api.operators.StreamTaskStateInitializer;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.util.function.ThrowingConsumer;
+
+/**
+ * This operator is an abstract class that give the {@link AbstractStreamOperatorV2} the ability to
+ * perform {@link AsyncStateProcessing}. The aim is to make any subclass of {@link
+ * AbstractStreamOperatorV2} could manipulate async state with only a change of base class.
+ */
+@Internal
+@SuppressWarnings("rawtypes")
+public abstract class AbstractAsyncStateStreamOperatorV2<OUT> extends AbstractStreamOperatorV2<OUT>
+        implements AsyncStateProcessingOperator {
+
+    private final MailboxExecutor mailboxExecutor;
+
+    private AsyncExecutionController asyncExecutionController;
+
+    private RecordContext currentProcessingContext;
+
+    public AbstractAsyncStateStreamOperatorV2(
+            StreamOperatorParameters<OUT> parameters, int numberOfInputs) {
+        super(parameters, numberOfInputs);
+        this.mailboxExecutor =
+                parameters.getContainingTask().getEnvironment().getMainMailboxExecutor();
+    }
+
+    /** Initialize necessary state components for {@link AbstractStreamOperatorV2}. */
+    @Override
+    public final void initializeState(StreamTaskStateInitializer streamTaskStateManager)
+            throws Exception {
+        super.initializeState(streamTaskStateManager);
+        // TODO: Read config and properly set.
+        this.asyncExecutionController = new AsyncExecutionController(mailboxExecutor, null);
+    }
+
+    @Override
+    public boolean isAsyncStateProcessingEnabled() {
+        return true;
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public final <T> void setAsyncKeyedContextElement(
+            StreamRecord<T> record, KeySelector<T, ?> keySelector) throws Exception {
+        currentProcessingContext =
+                asyncExecutionController.buildContext(
+                        record.getValue(), keySelector.getKey(record.getValue()));
+        // The processElement will be treated as a callback for dummy request. So reference
+        // counting should be maintained.
+        // When state request submitted, ref count +1, as described in FLIP-425:
+        // To cover the statements without a callback, in addition to the reference count marked
+        // in Fig.5, each state request itself is also protected by a paired reference count.
+        currentProcessingContext.retain();
+        asyncExecutionController.setCurrentContext(currentProcessingContext);
+    }
+
+    @Override
+    public final void postProcessElement() {
+        // The processElement will be treated as a callback for dummy request. So reference
+        // counting should be maintained.
+        // When a state request completes, ref count -1, as described in FLIP-425:
+        // To cover the statements without a callback, in addition to the reference count marked
+        // in Fig.5, each state request itself is also protected by a paired reference count.
+        currentProcessingContext.release();
+    }
+
+    @Override
+    public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId) {
+        // The real logic should be in First/SecondInputOfTwoInput#getRecordProcessor.
+        throw new UnsupportedOperationException(
+                "Never getRecordProcessor from AbstractAsyncStateStreamOperatorV2,"
+                        + " since this part is handled by the Input.");
+    }
+
+    @VisibleForTesting
+    AsyncExecutionController<?> getAsyncExecutionController() {
+        return asyncExecutionController;
+    }
+}
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
new file mode 100644
index 00000000000..ec033b3906f
--- /dev/null
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessing.java
@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.operators.asyncprocessing;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.util.function.ThrowingConsumer;
+
+/** This class defines the basic interfaces to process a state in operator/input layer. */
+@Internal
+public interface AsyncStateProcessing {
+
+    /**
+     * Get if the async state processing is enabled for this input/operator.
+     *
+     * @return ture if async state processing is enabled.
+     */
+    boolean isAsyncStateProcessingEnabled();
+
+    /**
+     * Get the record processor that could process record from input, which is the only entry for
+     * async processing.
+     *
+     * @param inputId the input identifier, start from 1. Borrow the design from {@code
+     *     org.apache.flink.streaming.api.operators.AbstractInput#inputId}. This is only relevant if
+     *     there is multiple inputs for the instance.
+     */
+    <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId);
+
+    /**
+     * Static method helper to make a record processor with given infos.
+     *
+     * @param asyncOperator the operator that can process state asynchronously.
+     * @param keySelector the key selector.
+     * @param processor the record processing logic.
+     * @return the built record processor that can returned by {@link #getRecordProcessor(int)}.
+     */
+    static <T> ThrowingConsumer<StreamRecord<T>, Exception> makeRecordProcessor(
+            AsyncStateProcessingOperator asyncOperator,
+            KeySelector<T, ?> keySelector,
+            ThrowingConsumer<StreamRecord<T>, Exception> processor) {
+        return (record) -> {
+            asyncOperator.setAsyncKeyedContextElement(record, keySelector);
+            processor.accept(record);
+            asyncOperator.postProcessElement();
+        };
+    }
+}
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java
new file mode 100644
index 00000000000..d8b52b6dd4f
--- /dev/null
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AsyncStateProcessingOperator.java
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.operators.asyncprocessing;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+
+/**
+ * A more detailed interface based on {@link AsyncStateProcessing}, which gives the essential
+ * methods for an operator to perform async state processing.
+ */
+@Internal
+public interface AsyncStateProcessingOperator extends AsyncStateProcessing {
+
+    /**
+     * Set key context for async state processing.
+     *
+     * @param record the record.
+     * @param keySelector the key selector to select a key from record.
+     * @param <T> the type of the record.
+     */
+    <T> void setAsyncKeyedContextElement(StreamRecord<T> record, KeySelector<T, ?> keySelector)
+            throws Exception;
+
+    /** A callback that will be triggered after an element finishes {@code processElement}. */
+    void postProcessElement();
+}
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorTest.java
new file mode 100644
index 00000000000..20c01c8bc4f
--- /dev/null
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorTest.java
@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.operators.asyncprocessing;
+
+import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.state.VoidNamespace;
+import org.apache.flink.streaming.api.operators.InternalTimer;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.operators.Triggerable;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Basic tests for {@link AbstractAsyncStateStreamOperator}. */
+public class AbstractAsyncStateStreamOperatorTest {
+
+    protected KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
+            createTestHarness(int maxParalelism, int numSubtasks, int subtaskIndex)
+                    throws Exception {
+        TestOperator testOperator = new TestOperator();
+        return new KeyedOneInputStreamOperatorTestHarness<>(
+                testOperator,
+                new TestKeySelector(),
+                BasicTypeInfo.INT_TYPE_INFO,
+                maxParalelism,
+                numSubtasks,
+                subtaskIndex);
+    }
+
+    @Test
+    public void testCreateAsyncExecutionController() throws Exception {
+        try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String>
+                testHarness = createTestHarness(128, 1, 0)) {
+            testHarness.open();
+            assertThat(testHarness.getOperator())
+                    .isInstanceOf(AbstractAsyncStateStreamOperator.class);
+            assertThat(
+                            ((AbstractAsyncStateStreamOperator) testHarness.getOperator())
+                                    .getAsyncExecutionController())
+                    .isNotNull();
+        }
+    }
+
+    /** A simple testing operator. */
+    private static class TestOperator extends AbstractAsyncStateStreamOperator<String>
+            implements OneInputStreamOperator<Tuple2<Integer, String>, String>,
+                    Triggerable<Integer, VoidNamespace> {
+
+        private static final long serialVersionUID = 1L;
+
+        @Override
+        public void open() throws Exception {
+            super.open();
+        }
+
+        @Override
+        public void processElement(StreamRecord<Tuple2<Integer, String>> element)
+                throws Exception {}
+
+        @Override
+        public void onEventTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {}
+
+        @Override
+        public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer)
+                throws Exception {}
+    }
+
+    /** {@link KeySelector} for tests. */
+    static class TestKeySelector implements KeySelector<Tuple2<Integer, String>, Integer> {
+        private static final long serialVersionUID = 1L;
+
+        @Override
+        public Integer getKey(Tuple2<Integer, String> value) {
+            return value.f0;
+        }
+    }
+}
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2Test.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2Test.java
new file mode 100644
index 00000000000..088729b5067
--- /dev/null
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/asyncprocessing/AbstractAsyncStateStreamOperatorV2Test.java
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.runtime.operators.asyncprocessing;
+
+import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.state.VoidNamespace;
+import org.apache.flink.streaming.api.operators.AbstractInput;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperatorFactory;
+import org.apache.flink.streaming.api.operators.Input;
+import org.apache.flink.streaming.api.operators.InternalTimer;
+import org.apache.flink.streaming.api.operators.MultipleInputStreamOperator;
+import org.apache.flink.streaming.api.operators.StreamOperator;
+import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
+import org.apache.flink.streaming.api.operators.StreamOperatorParameters;
+import org.apache.flink.streaming.api.operators.Triggerable;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Basic tests for {@link AbstractAsyncStateStreamOperatorV2}. */
+public class AbstractAsyncStateStreamOperatorV2Test {
+
+    protected KeyedOneInputStreamOperatorV2TestHarness<Integer, Tuple2<Integer, String>, String>
+            createTestHarness(int maxParalelism, int numSubtasks, int subtaskIndex)
+                    throws Exception {
+        return new KeyedOneInputStreamOperatorV2TestHarness<>(
+                new TestOperatorFactory(),
+                new AbstractAsyncStateStreamOperatorTest.TestKeySelector(),
+                BasicTypeInfo.INT_TYPE_INFO,
+                maxParalelism,
+                numSubtasks,
+                subtaskIndex);
+    }
+
+    @Test
+    public void testCreateAsyncExecutionController() throws Exception {
+        try (KeyedOneInputStreamOperatorV2TestHarness<Integer, Tuple2<Integer, String>, String>
+                testHarness = createTestHarness(128, 1, 0)) {
+            testHarness.open();
+            assertThat(testHarness.getBaseOperator())
+                    .isInstanceOf(AbstractAsyncStateStreamOperatorV2.class);
+            assertThat(
+                            ((AbstractAsyncStateStreamOperatorV2) testHarness.getBaseOperator())
+                                    .getAsyncExecutionController())
+                    .isNotNull();
+        }
+    }
+
+    private static class KeyedOneInputStreamOperatorV2TestHarness<K, IN, OUT>
+            extends KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> {
+        public KeyedOneInputStreamOperatorV2TestHarness(
+                StreamOperatorFactory<OUT> operatorFactory,
+                final KeySelector<IN, K> keySelector,
+                TypeInformation<K> keyType,
+                int maxParallelism,
+                int numSubtasks,
+                int subtaskIndex)
+                throws Exception {
+            super(operatorFactory, keySelector, keyType, maxParallelism, numSubtasks, subtaskIndex);
+        }
+
+        public StreamOperator<OUT> getBaseOperator() {
+            return operator;
+        }
+    }
+
+    private static class TestOperatorFactory extends AbstractStreamOperatorFactory<String> {
+        @Override
+        public <T extends StreamOperator<String>> T createStreamOperator(
+                StreamOperatorParameters<String> parameters) {
+            return (T) new SingleInputTestOperator(parameters);
+        }
+
+        @Override
+        public Class<? extends StreamOperator> getStreamOperatorClass(ClassLoader classLoader) {
+            return SingleInputTestOperator.class;
+        }
+    }
+
+    /**
+     * Testing operator that can respond to commands by either setting/deleting state, emitting
+     * state or setting timers.
+     */
+    private static class SingleInputTestOperator extends AbstractAsyncStateStreamOperatorV2<String>
+            implements MultipleInputStreamOperator<String>, Triggerable<Integer, VoidNamespace> {
+
+        private static final long serialVersionUID = 1L;
+
+        public SingleInputTestOperator(StreamOperatorParameters<String> parameters) {
+            super(parameters, 1);
+        }
+
+        @Override
+        public void open() throws Exception {
+            super.open();
+        }
+
+        @Override
+        public List<Input> getInputs() {
+            return Collections.singletonList(
+                    new AbstractInput<Tuple2<Integer, String>, String>(this, 1) {
+                        @Override
+                        public void processElement(StreamRecord<Tuple2<Integer, String>> element)
+                                throws Exception {}
+                    });
+        }
+
+        @Override
+        public void onEventTime(InternalTimer<Integer, VoidNamespace> timer) throws Exception {}
+
+        @Override
+        public void onProcessingTime(InternalTimer<Integer, VoidNamespace> timer)
+                throws Exception {}
+    }
+}
diff --git a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/FirstInputOfTwoInput.java b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/FirstInputOfTwoInput.java
index 077744b4c42..4f090c29008 100644
--- a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/FirstInputOfTwoInput.java
+++ b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/FirstInputOfTwoInput.java
@@ -18,16 +18,19 @@
 
 package org.apache.flink.table.runtime.operators.multipleinput.input;
 
+import org.apache.flink.annotation.Internal;
 import org.apache.flink.streaming.api.operators.Input;
 import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
 import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.operators.asyncprocessing.AsyncStateProcessing;
 import org.apache.flink.streaming.runtime.streamrecord.LatencyMarker;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
 import org.apache.flink.table.data.RowData;
+import org.apache.flink.util.function.ThrowingConsumer;
 
 /** {@link Input} for the first input of {@link TwoInputStreamOperator}. */
-public class FirstInputOfTwoInput extends InputBase {
+public class FirstInputOfTwoInput extends InputBase implements AsyncStateProcessing {
 
     private final TwoInputStreamOperator<RowData, RowData, RowData> operator;
 
@@ -54,4 +57,17 @@ public class FirstInputOfTwoInput extends InputBase {
     public void processWatermarkStatus(WatermarkStatus watermarkStatus) throws Exception {
         operator.processWatermarkStatus1(watermarkStatus);
     }
+
+    @Internal
+    @Override
+    public final boolean isAsyncStateProcessingEnabled() {
+        return (operator instanceof AsyncStateProcessing)
+                && ((AsyncStateProcessing) operator).isAsyncStateProcessingEnabled();
+    }
+
+    @Internal
+    @Override
+    public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId) {
+        return ((AsyncStateProcessing) operator).getRecordProcessor(1);
+    }
 }
diff --git a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/OneInput.java b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/OneInput.java
index 764f8aded7d..f4992249e7d 100644
--- a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/OneInput.java
+++ b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/OneInput.java
@@ -18,16 +18,19 @@
 
 package org.apache.flink.table.runtime.operators.multipleinput.input;
 
+import org.apache.flink.annotation.Internal;
 import org.apache.flink.streaming.api.operators.Input;
 import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
 import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.operators.asyncprocessing.AsyncStateProcessing;
 import org.apache.flink.streaming.runtime.streamrecord.LatencyMarker;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
 import org.apache.flink.table.data.RowData;
+import org.apache.flink.util.function.ThrowingConsumer;
 
 /** {@link Input} for {@link OneInputStreamOperator}. */
-public class OneInput extends InputBase {
+public class OneInput extends InputBase implements AsyncStateProcessing {
 
     private final OneInputStreamOperator<RowData, RowData> operator;
 
@@ -54,4 +57,17 @@ public class OneInput extends InputBase {
     public void processWatermarkStatus(WatermarkStatus watermarkStatus) throws Exception {
         operator.processWatermarkStatus(watermarkStatus);
     }
+
+    @Internal
+    @Override
+    public final boolean isAsyncStateProcessingEnabled() {
+        return (operator instanceof AsyncStateProcessing)
+                && ((AsyncStateProcessing) operator).isAsyncStateProcessingEnabled();
+    }
+
+    @Internal
+    @Override
+    public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId) {
+        return ((AsyncStateProcessing) operator).getRecordProcessor(1);
+    }
 }
diff --git a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/SecondInputOfTwoInput.java b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/SecondInputOfTwoInput.java
index cd364ab7108..7158bd11527 100644
--- a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/SecondInputOfTwoInput.java
+++ b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/multipleinput/input/SecondInputOfTwoInput.java
@@ -18,16 +18,19 @@
 
 package org.apache.flink.table.runtime.operators.multipleinput.input;
 
+import org.apache.flink.annotation.Internal;
 import org.apache.flink.streaming.api.operators.Input;
 import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
 import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.operators.asyncprocessing.AsyncStateProcessing;
 import org.apache.flink.streaming.runtime.streamrecord.LatencyMarker;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
 import org.apache.flink.table.data.RowData;
+import org.apache.flink.util.function.ThrowingConsumer;
 
 /** {@link Input} for the second input of {@link TwoInputStreamOperator}. */
-public class SecondInputOfTwoInput extends InputBase {
+public class SecondInputOfTwoInput extends InputBase implements AsyncStateProcessing {
 
     private final TwoInputStreamOperator<RowData, RowData, RowData> operator;
 
@@ -54,4 +57,17 @@ public class SecondInputOfTwoInput extends InputBase {
     public void processWatermarkStatus(WatermarkStatus watermarkStatus) throws Exception {
         operator.processWatermarkStatus2(watermarkStatus);
     }
+
+    @Internal
+    @Override
+    public final boolean isAsyncStateProcessingEnabled() {
+        return (operator instanceof AsyncStateProcessing)
+                && ((AsyncStateProcessing) operator).isAsyncStateProcessingEnabled();
+    }
+
+    @Internal
+    @Override
+    public final <T> ThrowingConsumer<StreamRecord<T>, Exception> getRecordProcessor(int inputId) {
+        return ((AsyncStateProcessing) operator).getRecordProcessor(2);
+    }
 }