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

zakelly pushed a commit to branch release-2.0
in repository https://gitbox.apache.org/repos/asf/flink.git


The following commit(s) were added to refs/heads/release-2.0 by this push:
     new 3b2f94ecf2b [FLINK-37206][Runtime] Fix initialization of batching 
timer service in async state operators (#26092)
3b2f94ecf2b is described below

commit 3b2f94ecf2bbebe81af3d5e65d803bd5473a1374
Author: Zakelly <[email protected]>
AuthorDate: Wed Feb 5 10:17:26 2025 +0800

    [FLINK-37206][Runtime] Fix initialization of batching timer service in 
async state operators (#26092)
---
 .../AbstractAsyncStateStreamOperator.java          |   8 +-
 .../AbstractAsyncStateStreamOperatorV2.java        |   8 +-
 .../v2/adaptor/AsyncKeyedStateBackendAdaptor.java  |   2 -
 .../state/BatchExecutionInternalTimeService.java   |  14 +-
 .../BatchExecutionInternalTimeServiceManager.java  |  42 ++-
 ...ExecutionInternalTimeServiceWithAsyncState.java |  87 +++++
 ...utionInternalTimeServiceWithAsyncStateTest.java | 369 +++++++++++++++++++++
 7 files changed, 509 insertions(+), 21 deletions(-)

diff --git 
a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
 
b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
index 15191551fb3..94dc17a9016 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperator.java
@@ -43,6 +43,7 @@ import 
org.apache.flink.streaming.api.operators.InternalTimerService;
 import org.apache.flink.streaming.api.operators.InternalTimerServiceAsyncImpl;
 import org.apache.flink.streaming.api.operators.Triggerable;
 import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
+import 
org.apache.flink.streaming.api.operators.sorted.state.BatchExecutionInternalTimeServiceWithAsyncState;
 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;
@@ -307,7 +308,12 @@ public abstract class 
AbstractAsyncStateStreamOperator<OUT> extends AbstractStre
         InternalTimerService<N> service =
                 keyedTimeServiceHandler.getInternalTimerService(
                         name, keySerializer, namespaceSerializer, triggerable);
-        ((InternalTimerServiceAsyncImpl<K, N>) 
service).setup(asyncExecutionController);
+        if (service instanceof InternalTimerServiceAsyncImpl) {
+            ((InternalTimerServiceAsyncImpl<K, N>) 
service).setup(asyncExecutionController);
+        } else if (service instanceof 
BatchExecutionInternalTimeServiceWithAsyncState) {
+            ((BatchExecutionInternalTimeServiceWithAsyncState<K, N>) service)
+                    .setup(asyncExecutionController);
+        }
         return service;
     }
 
diff --git 
a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
 
b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
index 407dd177817..0036055eb6a 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/runtime/asyncprocessing/operators/AbstractAsyncStateStreamOperatorV2.java
@@ -39,6 +39,7 @@ import 
org.apache.flink.streaming.api.operators.InternalTimerService;
 import org.apache.flink.streaming.api.operators.InternalTimerServiceAsyncImpl;
 import org.apache.flink.streaming.api.operators.StreamOperatorParameters;
 import org.apache.flink.streaming.api.operators.Triggerable;
+import 
org.apache.flink.streaming.api.operators.sorted.state.BatchExecutionInternalTimeServiceWithAsyncState;
 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;
@@ -286,7 +287,12 @@ public abstract class 
AbstractAsyncStateStreamOperatorV2<OUT> extends AbstractSt
         InternalTimerService<N> service =
                 keyedTimeServiceHandler.getInternalTimerService(
                         name, keySerializer, namespaceSerializer, triggerable);
-        ((InternalTimerServiceAsyncImpl<K, N>) 
service).setup(asyncExecutionController);
+        if (service instanceof InternalTimerServiceAsyncImpl) {
+            ((InternalTimerServiceAsyncImpl<K, N>) 
service).setup(asyncExecutionController);
+        } else if (service instanceof 
BatchExecutionInternalTimeServiceWithAsyncState) {
+            ((BatchExecutionInternalTimeServiceWithAsyncState<K, N>) service)
+                    .setup(asyncExecutionController);
+        }
         return service;
     }
 
diff --git 
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/adaptor/AsyncKeyedStateBackendAdaptor.java
 
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/adaptor/AsyncKeyedStateBackendAdaptor.java
index 87cffb259f7..9604b0656a9 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/adaptor/AsyncKeyedStateBackendAdaptor.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/adaptor/AsyncKeyedStateBackendAdaptor.java
@@ -18,7 +18,6 @@
 
 package org.apache.flink.runtime.state.v2.adaptor;
 
-import org.apache.flink.annotation.VisibleForTesting;
 import org.apache.flink.api.common.state.CheckpointListener;
 import org.apache.flink.api.common.state.InternalCheckpointListener;
 import org.apache.flink.api.common.state.v2.State;
@@ -199,7 +198,6 @@ public class AsyncKeyedStateBackendAdaptor<K> implements 
AsyncKeyedStateBackend<
         return keyedStateBackend.isSafeToReuseKVState();
     }
 
-    @VisibleForTesting
     public CheckpointableKeyedStateBackend<K> getKeyedStateBackend() {
         return keyedStateBackend;
     }
diff --git 
a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeService.java
 
b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeService.java
index de590ffc4e2..771afecd2ba 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeService.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeService.java
@@ -40,25 +40,23 @@ public class BatchExecutionInternalTimeService<K, N> 
implements InternalTimerSer
     private static final Logger LOG =
             LoggerFactory.getLogger(BatchExecutionInternalTimeService.class);
 
-    private final ProcessingTimeService processingTimeService;
+    final ProcessingTimeService processingTimeService;
 
     /** Processing time timers that are currently in-flight. */
-    private final KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<K, N>>
-            processingTimeTimersQueue;
+    final KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<K, N>> 
processingTimeTimersQueue;
 
     /** Event time timers that are currently in-flight. */
-    private final KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<K, N>>
-            eventTimeTimersQueue;
+    final KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<K, N>> 
eventTimeTimersQueue;
 
     /**
      * The local event time, as denoted by the last received {@link
      * org.apache.flink.streaming.api.watermark.Watermark Watermark}.
      */
-    private long currentWatermark = Long.MIN_VALUE;
+    long currentWatermark = Long.MIN_VALUE;
 
-    private final Triggerable<K, N> triggerTarget;
+    final Triggerable<K, N> triggerTarget;
 
-    private K currentKey;
+    K currentKey;
 
     BatchExecutionInternalTimeService(
             ProcessingTimeService processingTimeService, Triggerable<K, N> 
triggerTarget) {
diff --git 
a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceManager.java
 
b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceManager.java
index 5d027671b43..22817331631 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceManager.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceManager.java
@@ -25,6 +25,7 @@ import 
org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider;
 import org.apache.flink.runtime.state.KeyedStateBackend;
 import org.apache.flink.runtime.state.KeyedStateCheckpointOutputStream;
 import org.apache.flink.runtime.state.PriorityQueueSetFactory;
+import org.apache.flink.runtime.state.v2.adaptor.AsyncKeyedStateBackendAdaptor;
 import org.apache.flink.streaming.api.operators.InternalTimeServiceManager;
 import org.apache.flink.streaming.api.operators.InternalTimerService;
 import org.apache.flink.streaming.api.operators.KeyContext;
@@ -38,7 +39,6 @@ import java.util.HashMap;
 import java.util.Map;
 
 import static org.apache.flink.util.Preconditions.checkNotNull;
-import static org.apache.flink.util.Preconditions.checkState;
 
 /**
  * An implementation of a {@link InternalTimeServiceManager} that manages 
timers with a single
@@ -51,8 +51,14 @@ public class BatchExecutionInternalTimeServiceManager<K>
     private final Map<String, BatchExecutionInternalTimeService<K, ?>> 
timerServices =
             new HashMap<>();
 
-    public BatchExecutionInternalTimeServiceManager(ProcessingTimeService 
processingTimeService) {
+    // In batch mode, there is a chance that the operator is {@link 
AsyncStateProcessing} and we
+    // should perform correctly when the timer fires.
+    private final boolean asyncStateProcessingMode;
+
+    public BatchExecutionInternalTimeServiceManager(
+            ProcessingTimeService processingTimeService, boolean 
asyncStateProcessingMode) {
         this.processingTimeService = checkNotNull(processingTimeService);
+        this.asyncStateProcessingMode = asyncStateProcessingMode;
     }
 
     @Override
@@ -66,7 +72,11 @@ public class BatchExecutionInternalTimeServiceManager<K>
                 (BatchExecutionInternalTimeService<K, N>) 
timerServices.get(name);
         if (timerService == null) {
             timerService =
-                    new 
BatchExecutionInternalTimeService<>(processingTimeService, triggerable);
+                    asyncStateProcessingMode
+                            ? new 
BatchExecutionInternalTimeServiceWithAsyncState<>(
+                                    processingTimeService, triggerable)
+                            : new BatchExecutionInternalTimeService<>(
+                                    processingTimeService, triggerable);
             timerServices.put(name, timerService);
         }
 
@@ -93,6 +103,7 @@ public class BatchExecutionInternalTimeServiceManager<K>
         throw new UnsupportedOperationException("Checkpoints are not supported 
in BATCH execution");
     }
 
+    @SuppressWarnings("unchecked")
     public static <K> InternalTimeServiceManager<K> create(
             TaskIOMetricGroup taskIOMetricGroup,
             PriorityQueueSetFactory factory,
@@ -102,14 +113,27 @@ public class BatchExecutionInternalTimeServiceManager<K>
             ProcessingTimeService processingTimeService,
             Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates,
             StreamTaskCancellationContext cancellationContext) {
-        checkState(
-                factory instanceof BatchExecutionKeyedStateBackend,
-                "Batch execution specific time service can work only with 
BatchExecutionKeyedStateBackend");
+        BatchExecutionKeyedStateBackend<K> theFactory = null;
+        boolean asyncStateProcessingMode = false;
+        if (factory instanceof BatchExecutionKeyedStateBackend) {
+            theFactory = (BatchExecutionKeyedStateBackend<K>) factory;
+        } else if (factory instanceof AsyncKeyedStateBackendAdaptor) {
+            KeyedStateBackend<K> keyedStateBackend =
+                    ((AsyncKeyedStateBackendAdaptor<K>) 
factory).getKeyedStateBackend();
+            if (keyedStateBackend instanceof BatchExecutionKeyedStateBackend) {
+                theFactory = (BatchExecutionKeyedStateBackend<K>) 
keyedStateBackend;
+                asyncStateProcessingMode = true;
+            }
+        }
+        if (theFactory == null) {
+            throw new IllegalStateException(
+                    "Batch execution specific time service can work only with 
BatchExecutionKeyedStateBackend");
+        }
 
         BatchExecutionInternalTimeServiceManager<K> timeServiceManager =
-                new 
BatchExecutionInternalTimeServiceManager<>(processingTimeService);
-        ((BatchExecutionKeyedStateBackend) factory)
-                .registerKeySelectionListener(timeServiceManager);
+                new BatchExecutionInternalTimeServiceManager<>(
+                        processingTimeService, asyncStateProcessingMode);
+        theFactory.registerKeySelectionListener(timeServiceManager);
         return timeServiceManager;
     }
 
diff --git 
a/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceWithAsyncState.java
 
b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceWithAsyncState.java
new file mode 100644
index 00000000000..da24a9ed804
--- /dev/null
+++ 
b/flink-runtime/src/main/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceWithAsyncState.java
@@ -0,0 +1,87 @@
+/*
+ * 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.api.operators.sorted.state;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.runtime.asyncprocessing.AsyncExecutionController;
+import org.apache.flink.runtime.asyncprocessing.RecordContext;
+import org.apache.flink.streaming.api.operators.InternalTimer;
+import org.apache.flink.streaming.api.operators.InternalTimerService;
+import org.apache.flink.streaming.api.operators.Triggerable;
+import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService;
+import org.apache.flink.util.function.ThrowingRunnable;
+
+/**
+ * An implementation of a {@link InternalTimerService} that manages timers 
with a single active key
+ * at a time. Can be used in a BATCH execution mode cooperating with async 
state operators.
+ */
+@Internal
+public class BatchExecutionInternalTimeServiceWithAsyncState<K, N>
+        extends BatchExecutionInternalTimeService<K, N> {
+
+    private AsyncExecutionController<K> asyncExecutionController;
+
+    BatchExecutionInternalTimeServiceWithAsyncState(
+            ProcessingTimeService processingTimeService, Triggerable<K, N> 
triggerTarget) {
+        super(processingTimeService, triggerTarget);
+    }
+
+    /** Set up the async execution controller. */
+    public void setup(AsyncExecutionController<K> asyncExecutionController) {
+        if (asyncExecutionController != null) {
+            this.asyncExecutionController = asyncExecutionController;
+        }
+    }
+
+    /**
+     * Sets the current key. Timers that are due to be fired are collected and 
will be triggered.
+     */
+    @Override
+    public void setCurrentKey(K currentKey) throws Exception {
+        if (currentKey != null && currentKey.equals(this.currentKey)) {
+            return;
+        }
+        currentWatermark = Long.MAX_VALUE;
+        InternalTimer<K, N> timer;
+        while ((timer = eventTimeTimersQueue.poll()) != null) {
+            final InternalTimer<K, N> timerToTrigger = timer;
+            maintainContextAndProcess(
+                    timerToTrigger, () -> 
triggerTarget.onEventTime(timerToTrigger));
+        }
+        while ((timer = processingTimeTimersQueue.poll()) != null) {
+            final InternalTimer<K, N> timerToTrigger = timer;
+            maintainContextAndProcess(
+                    timerToTrigger, () -> 
triggerTarget.onProcessingTime(timerToTrigger));
+        }
+        currentWatermark = Long.MIN_VALUE;
+        this.currentKey = currentKey;
+    }
+
+    private void maintainContextAndProcess(
+            InternalTimer<K, N> timer, ThrowingRunnable<Exception> runnable) {
+        // Since we are in middle of processing a record, we need to maintain 
the context.
+        final RecordContext<K> previousContext = 
asyncExecutionController.getCurrentContext();
+        RecordContext<K> recordCtx = 
asyncExecutionController.buildContext(timer, timer.getKey());
+        recordCtx.retain();
+        asyncExecutionController.setCurrentContext(recordCtx);
+        asyncExecutionController.syncPointRequestWithCallback(runnable, true);
+        recordCtx.release();
+        asyncExecutionController.setCurrentContext(previousContext);
+    }
+}
diff --git 
a/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceWithAsyncStateTest.java
 
b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceWithAsyncStateTest.java
new file mode 100644
index 00000000000..18e94316cfc
--- /dev/null
+++ 
b/flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/sorted/state/BatchExecutionInternalTimeServiceWithAsyncStateTest.java
@@ -0,0 +1,369 @@
+/*
+ * 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.api.operators.sorted.state;
+
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.typeutils.base.IntSerializer;
+import org.apache.flink.runtime.asyncprocessing.AsyncExecutionController;
+import org.apache.flink.runtime.asyncprocessing.MockStateExecutor;
+import org.apache.flink.runtime.asyncprocessing.declare.DeclarationManager;
+import org.apache.flink.runtime.mailbox.SyncMailboxExecutor;
+import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.VoidNamespace;
+import org.apache.flink.runtime.state.VoidNamespaceSerializer;
+import org.apache.flink.runtime.state.v2.adaptor.AsyncKeyedStateBackendAdaptor;
+import org.apache.flink.streaming.api.operators.InternalTimeServiceManager;
+import org.apache.flink.streaming.api.operators.InternalTimer;
+import org.apache.flink.streaming.api.operators.InternalTimerService;
+import org.apache.flink.streaming.api.operators.KeyContext;
+import org.apache.flink.streaming.api.operators.Triggerable;
+import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.tasks.StreamTaskCancellationContext;
+import org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.Assertions.fail;
+
+/**
+ * Tests for {@link BatchExecutionInternalTimeServiceManager} and {@link
+ * BatchExecutionInternalTimeServiceWithAsyncState}.
+ */
+class BatchExecutionInternalTimeServiceWithAsyncStateTest {
+    public static final IntSerializer KEY_SERIALIZER = new IntSerializer();
+
+    BatchExecutionKeyedStateBackend<Integer> keyedStatedBackend;
+    InternalTimeServiceManager<Integer> timeServiceManager;
+    TestProcessingTimeService processingTimeService;
+    AsyncExecutionController<Integer> aec;
+
+    @BeforeEach
+    public void setup() {
+        keyedStatedBackend =
+                new BatchExecutionKeyedStateBackend<>(
+                        KEY_SERIALIZER, new KeyGroupRange(0, 1), new 
ExecutionConfig());
+        processingTimeService = new TestProcessingTimeService();
+        aec =
+                new AsyncExecutionController<>(
+                        new SyncMailboxExecutor(),
+                        (a, b) -> {},
+                        new MockStateExecutor(),
+                        new DeclarationManager(),
+                        1,
+                        100,
+                        1000,
+                        1,
+                        null,
+                        null);
+        timeServiceManager =
+                BatchExecutionInternalTimeServiceManager.create(
+                        
UnregisteredMetricGroups.createUnregisteredTaskMetricGroup()
+                                .getIOMetricGroup(),
+                        new 
AsyncKeyedStateBackendAdaptor<>(keyedStatedBackend),
+                        null,
+                        this.getClass().getClassLoader(),
+                        new DummyKeyContext(),
+                        processingTimeService,
+                        Collections.emptyList(),
+                        StreamTaskCancellationContext.alwaysRunning());
+    }
+
+    @Test
+    void testForEachEventTimeTimerUnsupported() {
+        BatchExecutionInternalTimeServiceWithAsyncState<Object, Object> 
timeService =
+                new BatchExecutionInternalTimeServiceWithAsyncState<>(
+                        new TestProcessingTimeService(),
+                        LambdaTrigger.eventTimeTrigger(timer -> {}));
+
+        assertThatThrownBy(
+                        () ->
+                                timeService.forEachEventTimeTimer(
+                                        (o, aLong) ->
+                                                fail(
+                                                        "The 
forEachEventTimeTimer() should not be supported")))
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessageContaining(
+                        "The BatchExecutionInternalTimeService should not be 
used in State Processor API");
+    }
+
+    @Test
+    void testForEachProcessingTimeTimerUnsupported() {
+        BatchExecutionInternalTimeServiceWithAsyncState<Object, Object> 
timeService =
+                new BatchExecutionInternalTimeServiceWithAsyncState<>(
+                        new TestProcessingTimeService(),
+                        LambdaTrigger.eventTimeTrigger(timer -> {}));
+
+        assertThatThrownBy(
+                        () ->
+                                timeService.forEachEventTimeTimer(
+                                        (o, aLong) ->
+                                                fail(
+                                                        "The 
forEachProcessingTimeTimer() should not be supported")))
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessageContaining(
+                        "The BatchExecutionInternalTimeService should not be 
used in State Processor API");
+    }
+
+    @Test
+    void testFiringEventTimeTimers() throws Exception {
+        List<Long> timers = new ArrayList<>();
+        InternalTimerService<VoidNamespace> timerService =
+                buildTimerService(
+                        LambdaTrigger.eventTimeTrigger(timer -> 
timers.add(timer.getTimestamp())));
+
+        keyedStatedBackend.setCurrentKey(1);
+        timerService.registerEventTimeTimer(VoidNamespace.INSTANCE, 123);
+
+        // advancing the watermark should not fire timers
+        timeServiceManager.advanceWatermark(new Watermark(1000));
+        timerService.deleteEventTimeTimer(VoidNamespace.INSTANCE, 123);
+        timerService.registerEventTimeTimer(VoidNamespace.INSTANCE, 150);
+
+        // changing the current key fires all timers
+        keyedStatedBackend.setCurrentKey(2);
+
+        assertThat(timers).containsExactly(150L);
+    }
+
+    @Test
+    void testSettingSameKeyDoesNotFireTimers() {
+        List<Long> timers = new ArrayList<>();
+        InternalTimerService<VoidNamespace> timerService =
+                buildTimerService(
+                        LambdaTrigger.eventTimeTrigger(timer -> 
timers.add(timer.getTimestamp())));
+
+        keyedStatedBackend.setCurrentKey(1);
+        timerService.registerEventTimeTimer(VoidNamespace.INSTANCE, 123);
+        keyedStatedBackend.setCurrentKey(1);
+
+        assertThat(timers).isEmpty();
+    }
+
+    @Test
+    void testCurrentWatermark() throws Exception {
+        List<Long> timers = new ArrayList<>();
+        TriggerWithTimerServiceAccess<Integer, VoidNamespace> eventTimeTrigger 
=
+                TriggerWithTimerServiceAccess.eventTimeTrigger(
+                        (timer, timerService) -> {
+                            
assertThat(timerService.currentWatermark()).isEqualTo(Long.MAX_VALUE);
+                            timers.add(timer.getTimestamp());
+                        });
+        InternalTimerService<VoidNamespace> timerService = 
buildTimerService(eventTimeTrigger);
+        eventTimeTrigger.setTimerService(timerService);
+
+        assertThat(timerService.currentWatermark()).isEqualTo(Long.MIN_VALUE);
+        keyedStatedBackend.setCurrentKey(1);
+        timerService.registerEventTimeTimer(VoidNamespace.INSTANCE, 123);
+        assertThat(timerService.currentWatermark()).isEqualTo(Long.MIN_VALUE);
+
+        // advancing the watermark to a value different than Long.MAX_VALUE 
should have no effect
+        timeServiceManager.advanceWatermark(new Watermark(1000));
+        assertThat(timerService.currentWatermark()).isEqualTo(Long.MIN_VALUE);
+
+        // changing the current key fires all timers
+        keyedStatedBackend.setCurrentKey(2);
+        assertThat(timerService.currentWatermark()).isEqualTo(Long.MIN_VALUE);
+        timerService.registerEventTimeTimer(VoidNamespace.INSTANCE, 124);
+
+        // advancing the watermark to Long.MAX_VALUE should fire remaining key
+        timeServiceManager.advanceWatermark(Watermark.MAX_WATERMARK);
+
+        assertThat(timers).containsExactly(123L, 124L);
+    }
+
+    @Test
+    void testProcessingTimeTimers() {
+        List<Long> timers = new ArrayList<>();
+        InternalTimerService<VoidNamespace> timerService =
+                buildTimerService(
+                        LambdaTrigger.processingTimeTrigger(
+                                timer -> timers.add(timer.getTimestamp())));
+
+        keyedStatedBackend.setCurrentKey(1);
+        timerService.registerProcessingTimeTimer(VoidNamespace.INSTANCE, 150);
+
+        // we should never register physical timers
+        assertThat(processingTimeService.getNumActiveTimers()).isZero();
+        // changing the current key fires all timers
+        keyedStatedBackend.setCurrentKey(2);
+
+        assertThat(timers).containsExactly(150L);
+    }
+
+    @Test
+    void testIgnoringEventTimeTimersFromWithinCallback() {
+        List<Long> timers = new ArrayList<>();
+        TriggerWithTimerServiceAccess<Integer, VoidNamespace> trigger =
+                TriggerWithTimerServiceAccess.eventTimeTrigger(
+                        (timer, ts) -> {
+                            timers.add(timer.getTimestamp());
+                            ts.registerEventTimeTimer(
+                                    VoidNamespace.INSTANCE, 
timer.getTimestamp() + 20);
+                        });
+        InternalTimerService<VoidNamespace> timerService = 
buildTimerService(trigger);
+        trigger.setTimerService(timerService);
+
+        keyedStatedBackend.setCurrentKey(1);
+        timerService.registerEventTimeTimer(VoidNamespace.INSTANCE, 150);
+
+        // we should never register physical timers
+        assertThat(processingTimeService.getNumActiveTimers()).isZero();
+        // changing the current key fires all timers
+        keyedStatedBackend.setCurrentKey(2);
+
+        // We check that the timer from the callback is ignored
+        assertThat(timers).containsExactly(150L);
+    }
+
+    @Test
+    void testIgnoringProcessingTimeTimersFromWithinCallback() {
+        List<Long> timers = new ArrayList<>();
+        TriggerWithTimerServiceAccess<Integer, VoidNamespace> trigger =
+                TriggerWithTimerServiceAccess.processingTimeTrigger(
+                        (timer, ts) -> {
+                            timers.add(timer.getTimestamp());
+                            ts.registerProcessingTimeTimer(
+                                    VoidNamespace.INSTANCE, 
timer.getTimestamp() + 20);
+                        });
+        InternalTimerService<VoidNamespace> timerService = 
buildTimerService(trigger);
+        trigger.setTimerService(timerService);
+
+        keyedStatedBackend.setCurrentKey(1);
+        timerService.registerProcessingTimeTimer(VoidNamespace.INSTANCE, 150);
+
+        // we should never register physical timers
+        assertThat(processingTimeService.getNumActiveTimers()).isZero();
+        // changing the current key fires all timers
+        keyedStatedBackend.setCurrentKey(2);
+
+        // We check that the timer from the callback is ignored
+        assertThat(timers).containsExactly(150L);
+    }
+
+    private InternalTimerService<VoidNamespace> buildTimerService(
+            Triggerable<Integer, VoidNamespace> trigger) {
+        InternalTimerService<VoidNamespace> timerService =
+                timeServiceManager.getInternalTimerService(
+                        "test", KEY_SERIALIZER, new VoidNamespaceSerializer(), 
trigger);
+        ((BatchExecutionInternalTimeServiceWithAsyncState<Integer, 
VoidNamespace>) timerService)
+                .setup(aec);
+        return timerService;
+    }
+
+    private static class TriggerWithTimerServiceAccess<K, N> implements 
Triggerable<K, N> {
+
+        private InternalTimerService<N> timerService;
+        private final BiConsumer<InternalTimer<K, N>, InternalTimerService<N>> 
eventTimeHandler;
+        private final BiConsumer<InternalTimer<K, N>, InternalTimerService<N>>
+                processingTimeHandler;
+
+        private TriggerWithTimerServiceAccess(
+                BiConsumer<InternalTimer<K, N>, InternalTimerService<N>> 
eventTimeHandler,
+                BiConsumer<InternalTimer<K, N>, InternalTimerService<N>> 
processingTimeHandler) {
+            this.eventTimeHandler = eventTimeHandler;
+            this.processingTimeHandler = processingTimeHandler;
+        }
+
+        public static <K, N> TriggerWithTimerServiceAccess<K, N> 
eventTimeTrigger(
+                BiConsumer<InternalTimer<K, N>, InternalTimerService<N>> 
eventTimeHandler) {
+            return new TriggerWithTimerServiceAccess<>(
+                    eventTimeHandler,
+                    (timer, timeService) ->
+                            fail("We did not expect processing timer to be 
triggered."));
+        }
+
+        public static <K, N> TriggerWithTimerServiceAccess<K, N> 
processingTimeTrigger(
+                BiConsumer<InternalTimer<K, N>, InternalTimerService<N>> 
processingTimeHandler) {
+            return new TriggerWithTimerServiceAccess<>(
+                    (timer, timeService) -> fail("We did not expect event 
timer to be triggered."),
+                    processingTimeHandler);
+        }
+
+        public void setTimerService(InternalTimerService<N> timerService) {
+            this.timerService = timerService;
+        }
+
+        @Override
+        public void onEventTime(InternalTimer<K, N> timer) throws Exception {
+            this.eventTimeHandler.accept(timer, timerService);
+        }
+
+        @Override
+        public void onProcessingTime(InternalTimer<K, N> timer) throws 
Exception {
+            this.processingTimeHandler.accept(timer, timerService);
+        }
+    }
+
+    private static class LambdaTrigger<K, N> implements Triggerable<K, N> {
+
+        private final Consumer<InternalTimer<K, N>> eventTimeHandler;
+        private final Consumer<InternalTimer<K, N>> processingTimeHandler;
+
+        public static <K, N> LambdaTrigger<K, N> eventTimeTrigger(
+                Consumer<InternalTimer<K, N>> eventTimeHandler) {
+            return new LambdaTrigger<>(
+                    eventTimeHandler,
+                    timer -> fail("We did not expect processing timer to be 
triggered."));
+        }
+
+        public static <K, N> LambdaTrigger<K, N> processingTimeTrigger(
+                Consumer<InternalTimer<K, N>> processingTimeHandler) {
+            return new LambdaTrigger<>(
+                    timer -> fail("We did not expect event timer to be 
triggered."),
+                    processingTimeHandler);
+        }
+
+        private LambdaTrigger(
+                Consumer<InternalTimer<K, N>> eventTimeHandler,
+                Consumer<InternalTimer<K, N>> processingTimeHandler) {
+            this.eventTimeHandler = eventTimeHandler;
+            this.processingTimeHandler = processingTimeHandler;
+        }
+
+        @Override
+        public void onEventTime(InternalTimer<K, N> timer) throws Exception {
+            this.eventTimeHandler.accept(timer);
+        }
+
+        @Override
+        public void onProcessingTime(InternalTimer<K, N> timer) throws 
Exception {
+            this.processingTimeHandler.accept(timer);
+        }
+    }
+
+    private static class DummyKeyContext implements KeyContext {
+        @Override
+        public void setCurrentKey(Object key) {}
+
+        @Override
+        public Object getCurrentKey() {
+            return null;
+        }
+    }
+}


Reply via email to