task.cancellation.timeout, the TaskManager providing that slot is marked failed and has to be reprovisioned before the slot is returned to the pool. This value should be configured high enough to cover that delay across all restarted vertices, to avoid restarting with fewer resources than were available right before the restart.task.cancellation.timeout, the TaskManager providing that slot is marked failed and has to be reprovisioned before the slot is returned to the pool. This value should be configured high enough to cover that delay across all restarted vertices, to avoid restarting with fewer resources than were available right before the restart.task.cancellation.timeout, the TaskManager providing that slot is marked failed and has to be reprovisioned before the slot is returned to the pool. This value should be configured high enough to cover that delay across all restarted vertices, to avoid restarting with fewer resources than were available right before the restart.This is kept in its own file rather than {@link AdaptiveSchedulerTest} to avoid growing that + * file past the checkstyle {@code FileLength} limit, following the precedent set by extracting + * {@code LocalRecoveryTest} and other companion test files out of {@code AdaptiveSchedulerTest}. + */ +class AdaptiveSchedulerFreeSlotVertexParallelismTest extends AdaptiveSchedulerTestBase { + + @Test + void testFreeSlotVertexParallelismExcludesReservedSlots() throws Exception { + final JobGraph jobGraph = createJobGraph(); + final DefaultDeclarativeSlotPool declarativeSlotPool = + createDeclarativeSlotPool(jobGraph.getJobID(), singleThreadMainThreadExecutor); + + scheduler = prepareSchedulerWithNoTimeouts(jobGraph, declarativeSlotPool).build(); + + final SubmissionBufferingTaskManagerGateway taskManagerGateway = + new SubmissionBufferingTaskManagerGateway(1); + + startTestInstanceInMainThread(); + + // one slot is offered and immediately claimed by the running (parallelism-1) job. + runInMainThread( + () -> + declarativeSlotPool.offerSlots( + createSlotOffersForResourceRequirements( + ResourceCounter.withResource(ResourceProfile.UNKNOWN, 1)), + new LocalTaskManagerLocation(), + taskManagerGateway, + System.currentTimeMillis())); + taskManagerGateway.waitForSubmissions(1); + + // a second slot joins, but stays genuinely free: nothing has requested it yet. + runInMainThread( + () -> + declarativeSlotPool.offerSlots( + createSlotOffersForResourceRequirements( + ResourceCounter.withResource(ResourceProfile.UNKNOWN, 1)), + new LocalTaskManagerLocation(), + taskManagerGateway, + System.currentTimeMillis())); + + runInMainThread( + () -> { + // over-counted view: 1 reserved (in use by the running task) + 1 free = 2. + // This is correct for Executing's use (predicting a restart's target). + assertThat(scheduler.getAvailableVertexParallelism()) + .hasValueSatisfying( + parallelism -> + assertThat( + parallelism.getParallelism( + JOB_VERTEX.getID())) + .isEqualTo(2)); + + // free-slots-based view: only the genuinely free slot counts = 1. + // This is what a restart can *actually* reserve right now. + assertThat(scheduler.getFreeSlotVertexParallelism()) + .hasValueSatisfying( + parallelism -> + assertThat( + parallelism.getParallelism( + JOB_VERTEX.getID())) + .isEqualTo(1)); + }); + } + + private static JobGraph createJobGraph() { + return streamingJobGraph(JOB_VERTEX); + } + + private static DefaultDeclarativeSlotPool createDeclarativeSlotPool( + JobID jobId, ComponentMainThreadExecutor mainThreadExecutor) { + return new DefaultDeclarativeSlotPool( + jobId, + new DefaultAllocatedSlotPool(), + ignored -> {}, + DEFAULT_TIMEOUT, + DEFAULT_TIMEOUT, + Duration.ZERO, + mainThreadExecutor); + } + + private AdaptiveSchedulerBuilder prepareSchedulerWithNoTimeouts( + JobGraph jobGraph, DefaultDeclarativeSlotPool declarativeSlotPool) { + return new AdaptiveSchedulerBuilder( + jobGraph, singleThreadMainThreadExecutor, EXECUTOR_RESOURCE.getExecutor()) + .setDeclarativeSlotPool(declarativeSlotPool) + .setJobMasterConfiguration(createConfigurationWithNoTimeouts()); + } + + private static Configuration createConfigurationWithNoTimeouts() { + return new Configuration() + .set( + JobManagerOptions.SCHEDULER_SUBMISSION_RESOURCE_WAIT_TIMEOUT, + Duration.ofMillis(-1L)) + .set( + JobManagerOptions.SCHEDULER_RESCALE_RESOURCE_STABILIZATION_TIMEOUT, + Duration.ZERO) + .set( + JobManagerOptions.SCHEDULER_SUBMISSION_RESOURCE_STABILIZATION_TIMEOUT, + Duration.ofMillis(1L)); + } +} diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerRescaleRestartTimingTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerRescaleRestartTimingTest.java new file mode 100644 index 0000000000000..0c93fcb061e3f --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerRescaleRestartTimingTest.java @@ -0,0 +1,200 @@ +/* + * 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.runtime.scheduler.adaptive; + +import org.apache.flink.api.common.JobID; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.JobManagerOptions; +import org.apache.flink.runtime.clusterframework.types.ResourceProfile; +import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutor; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobmaster.slotpool.DefaultAllocatedSlotPool; +import org.apache.flink.runtime.jobmaster.slotpool.DefaultDeclarativeSlotPool; +import org.apache.flink.runtime.scheduler.adaptive.AdaptiveSchedulerTest.SubmissionBufferingTaskManagerGateway; +import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism; +import org.apache.flink.runtime.taskmanager.LocalTaskManagerLocation; +import org.apache.flink.runtime.testutils.CommonTestUtils; +import org.apache.flink.runtime.util.ResourceCounter; + +import org.junit.jupiter.api.Test; + +import java.time.Duration; +import java.util.Collections; + +import static org.apache.flink.runtime.jobgraph.JobGraphTestUtils.streamingJobGraph; +import static org.apache.flink.runtime.jobmaster.slotpool.SlotPoolTestUtils.createSlotOffersForResourceRequirements; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration test proving that {@link WaitingForResources}, driven by the real {@link + * DefaultStateTransitionManager} (not a no-op test double), correctly gates a rescale-triggered + * restart on genuinely free slots reaching the pre-restart target, and falls back once the + * rescale resource-stabilization timeout elapses. + * + *
Kept in its own file, following the precedent set by {@link
+ * AdaptiveSchedulerFreeSlotVertexParallelismTest} and other companion test files extracted out of
+ * {@link AdaptiveSchedulerTest}.
+ */
+class AdaptiveSchedulerRescaleRestartTimingTest extends AdaptiveSchedulerTestBase {
+
+ private static final int RETRY_INTERVAL_MILLIS = 20;
+ private static final int RETRY_ATTEMPTS = 250;
+
+ @Test
+ void testWaitingForResourcesDoesNotTransitionUntilFreeSlotsReachRescaleTarget()
+ throws Exception {
+ final JobGraph jobGraph = createJobGraph();
+ final DefaultDeclarativeSlotPool declarativeSlotPool =
+ createDeclarativeSlotPool(jobGraph.getJobID(), singleThreadMainThreadExecutor);
+
+ // long enough that the stabilization timeout cannot fire during this test.
+ scheduler =
+ prepareScheduler(jobGraph, declarativeSlotPool, Duration.ofSeconds(10)).build();
+
+ final SubmissionBufferingTaskManagerGateway taskManagerGateway =
+ new SubmissionBufferingTaskManagerGateway(2);
+
+ // go straight to the restart-triggered WaitingForResources from the initial Created
+ // state, the same way Restarting#goToSubsequentState does - never through the plain
+ // submission path (startScheduling()), which would race a second, unrelated
+ // WaitingForResources transition using the submission timeout config.
+ final VertexParallelism restartTarget = vertexParallelism(2);
+ runInMainThread(
+ () ->
+ scheduler.goToWaitingForResources(
+ new StateTrackingMockExecutionGraph(), restartTarget));
+
+ assertThat(scheduler.getState()).isInstanceOf(WaitingForResources.class);
+
+ offerSlots(declarativeSlotPool, taskManagerGateway, 1);
+
+ // only 1 of the 2 targeted slots is free: must not have shortcut to
+ // CreatingExecutionGraph, even though 1 slot is already "sufficient" to run the job at a
+ // lower parallelism. No sleep is needed here: offerSlots() runs synchronously on the main
+ // thread executor, and no stabilization work gets scheduled while desired resources
+ // (gated on the restart target) aren't met, so the state is already final by the time it
+ // returns.
+ assertThat(scheduler.getState()).isInstanceOf(WaitingForResources.class);
+
+ offerSlots(declarativeSlotPool, taskManagerGateway, 1);
+
+ CommonTestUtils.waitUntilCondition(
+ () -> scheduler.getState() instanceof CreatingExecutionGraph,
+ RETRY_INTERVAL_MILLIS,
+ RETRY_ATTEMPTS);
+ }
+
+ @Test
+ void testWaitingForResourcesFallsBackAfterRescaleResourceStabilizationTimeoutElapses()
+ throws Exception {
+ final JobGraph jobGraph = createJobGraph();
+ final DefaultDeclarativeSlotPool declarativeSlotPool =
+ createDeclarativeSlotPool(jobGraph.getJobID(), singleThreadMainThreadExecutor);
+
+ // short enough to keep the test fast, but comfortably longer than the offerSlots() call
+ // below so the fallback can only be triggered by the timeout, not by a race with it.
+ final Duration rescaleResourceStabilizationTimeout = Duration.ofMillis(300);
+ scheduler =
+ prepareScheduler(jobGraph, declarativeSlotPool, rescaleResourceStabilizationTimeout)
+ .build();
+
+ final int requiredParallelism = 2;
+ final SubmissionBufferingTaskManagerGateway taskManagerGateway =
+ new SubmissionBufferingTaskManagerGateway(requiredParallelism - 1);
+
+ // go straight to the restart-triggered WaitingForResources from the initial Created
+ // state, as in the test above - never through the plain submission path.
+ final VertexParallelism restartTarget = vertexParallelism(requiredParallelism);
+ runInMainThread(
+ () ->
+ scheduler.goToWaitingForResources(
+ new StateTrackingMockExecutionGraph(), restartTarget));
+
+ assertThat(scheduler.getState()).isInstanceOf(WaitingForResources.class);
+
+ // only 1 of the 2 targeted slots is ever offered.
+ offerSlots(declarativeSlotPool, taskManagerGateway, requiredParallelism - 1);
+
+ // the target is never reached, but the stabilization timeout must still force the
+ // transition once it elapses, rather than waiting forever.
+ CommonTestUtils.waitUntilCondition(
+ () -> scheduler.getState() instanceof CreatingExecutionGraph,
+ RETRY_INTERVAL_MILLIS,
+ RETRY_ATTEMPTS);
+ }
+
+ private VertexParallelism vertexParallelism(int parallelism) {
+ return new VertexParallelism(Collections.singletonMap(JOB_VERTEX.getID(), parallelism));
+ }
+
+ private void offerSlots(
+ DefaultDeclarativeSlotPool declarativeSlotPool,
+ SubmissionBufferingTaskManagerGateway taskManagerGateway,
+ int numSlots) {
+ runInMainThread(
+ () ->
+ declarativeSlotPool.offerSlots(
+ createSlotOffersForResourceRequirements(
+ ResourceCounter.withResource(
+ ResourceProfile.UNKNOWN, numSlots)),
+ new LocalTaskManagerLocation(),
+ taskManagerGateway,
+ System.currentTimeMillis()));
+ }
+
+ private static JobGraph createJobGraph() {
+ return streamingJobGraph(JOB_VERTEX);
+ }
+
+ private static DefaultDeclarativeSlotPool createDeclarativeSlotPool(
+ JobID jobId, ComponentMainThreadExecutor mainThreadExecutor) {
+ return new DefaultDeclarativeSlotPool(
+ jobId,
+ new DefaultAllocatedSlotPool(),
+ ignored -> {},
+ DEFAULT_TIMEOUT,
+ DEFAULT_TIMEOUT,
+ Duration.ZERO,
+ mainThreadExecutor);
+ }
+
+ private AdaptiveSchedulerBuilder prepareScheduler(
+ JobGraph jobGraph,
+ DefaultDeclarativeSlotPool declarativeSlotPool,
+ Duration rescaleResourceStabilizationTimeout) {
+ return new AdaptiveSchedulerBuilder(
+ jobGraph, singleThreadMainThreadExecutor, EXECUTOR_RESOURCE.getExecutor())
+ .setDeclarativeSlotPool(declarativeSlotPool)
+ .setJobMasterConfiguration(
+ createConfiguration(rescaleResourceStabilizationTimeout));
+ }
+
+ private static Configuration createConfiguration(
+ Duration rescaleResourceStabilizationTimeout) {
+ return new Configuration()
+ .set(
+ JobManagerOptions.SCHEDULER_SUBMISSION_RESOURCE_WAIT_TIMEOUT,
+ Duration.ofMillis(-1L))
+ .set(
+ JobManagerOptions.SCHEDULER_RESCALE_RESOURCE_STABILIZATION_TIMEOUT,
+ rescaleResourceStabilizationTimeout)
+ .set(
+ JobManagerOptions.SCHEDULER_SUBMISSION_RESOURCE_STABILIZATION_TIMEOUT,
+ Duration.ofMillis(1L));
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerTest.java
index 4feef4ca0f54f..ec7e1cc535764 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerTest.java
@@ -58,6 +58,7 @@
import org.apache.flink.runtime.executiongraph.ArchivedExecutionJobVertex;
import org.apache.flink.runtime.executiongraph.ArchivedExecutionVertex;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
+import org.apache.flink.runtime.executiongraph.ExecutionGraph;
import org.apache.flink.runtime.executiongraph.TaskExecutionStateTransition;
import org.apache.flink.runtime.executiongraph.failover.FixedDelayRestartBackoffTimeStrategy;
import org.apache.flink.runtime.executiongraph.failover.NoRestartBackoffTimeStrategy;
@@ -70,6 +71,7 @@
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobgraph.JobResourceRequirements;
import org.apache.flink.runtime.jobgraph.JobVertex;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.jobgraph.JobVertexResourceRequirements;
import org.apache.flink.runtime.jobgraph.OperatorID;
import org.apache.flink.runtime.jobgraph.tasks.CheckpointCoordinatorConfiguration;
@@ -1446,6 +1448,9 @@ private static Configuration createConfigurationWithNoTimeouts() {
.set(
JobManagerOptions.SCHEDULER_SUBMISSION_RESOURCE_STABILIZATION_TIMEOUT,
Duration.ofMillis(1L))
+ .set(
+ JobManagerOptions.SCHEDULER_RESCALE_RESOURCE_STABILIZATION_TIMEOUT,
+ Duration.ofMillis(1L))
.set(
JobManagerOptions.SCHEDULER_EXECUTING_COOLDOWN_AFTER_RESCALING,
Duration.ofMillis(1L))
@@ -2297,6 +2302,90 @@ void testScalingIntervalConfigurationIsRespected() throws ConfigurationException
.isEqualTo(scalingStabilizationTimeout);
}
+ @Test
+ void testRescaleResourceStabilizationTimeoutConfigurationIsRespected()
+ throws ConfigurationException {
+ final Duration rescaleResourceStabilizationTimeout = Duration.ofMillis(4242);
+ final Configuration configuration = createConfigurationWithNoTimeouts();
+ configuration.set(
+ JobManagerOptions.SCHEDULER_RESCALE_RESOURCE_STABILIZATION_TIMEOUT,
+ rescaleResourceStabilizationTimeout);
+
+ final AdaptiveScheduler.Settings settings = AdaptiveScheduler.Settings.of(configuration);
+ assertThat(settings.getRescaleResourceStabilizationTimeout())
+ .isEqualTo(rescaleResourceStabilizationTimeout);
+ }
+
+ @Test
+ void testRescaleResourceStabilizationTimeoutDefault() throws ConfigurationException {
+ final AdaptiveScheduler.Settings settings =
+ AdaptiveScheduler.Settings.of(new Configuration());
+
+ assertThat(settings.getRescaleResourceStabilizationTimeout())
+ .isEqualTo(
+ JobManagerOptions.SCHEDULER_RESCALE_RESOURCE_STABILIZATION_TIMEOUT
+ .defaultValue());
+ }
+
+ @Test
+ void testRescaleResourceStabilizationTimeoutIsDisabledInReactiveMode()
+ throws ConfigurationException {
+ final AdaptiveScheduler.Settings settings =
+ AdaptiveScheduler.Settings.of(
+ new Configuration()
+ .set(
+ JobManagerOptions.SCHEDULER_MODE,
+ SchedulerExecutionMode.REACTIVE));
+
+ assertThat(settings.getRescaleResourceStabilizationTimeout()).isEqualTo(Duration.ZERO);
+ }
+
+ @Test
+ void testGoToWaitingForResourcesForRestartConfiguresStateTransitionManagerFactory()
+ throws Exception {
+ final TestingStateTransitionManagerFactory factory =
+ new TestingStateTransitionManagerFactory(
+ ctx -> TestingStateTransitionManager.withNoOp());
+
+ final Duration rescaleResourceStabilizationTimeout = Duration.ofMillis(1234);
+ final Configuration configuration = createConfigurationWithNoTimeouts();
+ configuration.set(
+ JobManagerOptions.SCHEDULER_RESCALE_RESOURCE_STABILIZATION_TIMEOUT,
+ rescaleResourceStabilizationTimeout);
+
+ scheduler =
+ new AdaptiveSchedulerBuilder(
+ createJobGraph(),
+ singleThreadMainThreadExecutor,
+ EXECUTOR_RESOURCE.getExecutor())
+ .setStateTransitionManagerFactory(factory)
+ .setJobMasterConfiguration(configuration)
+ .build();
+
+ final JobVertexID jobVertexId = new JobVertexID();
+ final VertexParallelism restartWithParallelism =
+ new VertexParallelism(Collections.singletonMap(jobVertexId, 2));
+ final ExecutionGraph mockExecutionGraph = new StateTrackingMockExecutionGraph();
+
+ final OneShotLatch latch = new OneShotLatch();
+ singleThreadMainThreadExecutor.execute(
+ () -> {
+ scheduler.goToWaitingForResources(mockExecutionGraph, restartWithParallelism);
+ latch.trigger();
+ });
+ latch.await();
+
+ assertThat(scheduler.getState()).isInstanceOf(WaitingForResources.class);
+ assertThat(factory.getCooldownTimeout()).isEqualTo(Duration.ZERO);
+ assertThat(factory.getResourceStabilizationTimeout())
+ .isEqualTo(rescaleResourceStabilizationTimeout);
+ // Same shape as submission: the eager, zero-delay evaluation is a fast path for "target
+ // already reached", not the mechanism enforcing the wait. The wait itself is enforced by
+ // hasDesiredResources() (gated on the restart target) versus hasSufficientResources()
+ // (the plain, restart-agnostic check) inside the stabilization phase machine.
+ assertThat(factory.getMaximumDelayForTrigger()).isEqualTo(Duration.ZERO);
+ }
+
@Test
void testOnCompletedCheckpointIsHandledInMainThread() throws Exception {
testCheckpointStatsEventBeingExecutedInTheMainThread(
@@ -2482,6 +2571,18 @@ public StateTransitionManager create(
return stateTransitionManagerCreator.apply(context);
}
+
+ public Duration getCooldownTimeout() {
+ return cooldownTimeout;
+ }
+
+ public Duration getResourceStabilizationTimeout() {
+ return resourceStabilizationTimeout;
+ }
+
+ public Duration getMaximumDelayForTrigger() {
+ return maximumDelayForTrigger;
+ }
}
private AdaptiveScheduler createSchedulerThatReachesExecutingState(
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/CreatedTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/CreatedTest.java
index 0617c83131d00..df106171bbd66 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/CreatedTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/CreatedTest.java
@@ -21,6 +21,7 @@
import org.apache.flink.api.common.JobStatus;
import org.apache.flink.runtime.executiongraph.ArchivedExecutionGraph;
import org.apache.flink.runtime.executiongraph.ExecutionGraph;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism;
import org.apache.flink.runtime.scheduler.adaptive.timeline.RescaleTimeline;
import org.junit.jupiter.api.Test;
@@ -66,7 +67,9 @@ public void setExpectWaitingForResources() {
}
@Override
- public void goToWaitingForResources(@Nullable ExecutionGraph previousExecutionGraph) {
+ public void goToWaitingForResources(
+ @Nullable ExecutionGraph previousExecutionGraph,
+ @Nullable VertexParallelism targetVertexParallelism) {
waitingForResourcesStateValidator.validateInput(null);
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/CreatingExecutionGraphTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/CreatingExecutionGraphTest.java
index 4e6f4f6d3cf09..9766ea9e25a76 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/CreatingExecutionGraphTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/CreatingExecutionGraphTest.java
@@ -28,6 +28,7 @@
import org.apache.flink.runtime.scheduler.ExecutionGraphHandler;
import org.apache.flink.runtime.scheduler.GlobalFailureHandler;
import org.apache.flink.runtime.scheduler.OperatorCoordinatorHandler;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism;
import org.apache.flink.runtime.scheduler.adaptive.timeline.RescaleTimeline;
import org.apache.flink.runtime.scheduler.exceptionhistory.ExceptionHistoryEntry;
import org.apache.flink.util.FlinkException;
@@ -233,7 +234,9 @@ public CreatingExecutionGraph.AssignmentResult tryToAssignSlots(
}
@Override
- public void goToWaitingForResources(@Nullable ExecutionGraph previousExecutionGraph) {
+ public void goToWaitingForResources(
+ @Nullable ExecutionGraph previousExecutionGraph,
+ @Nullable VertexParallelism targetVertexParallelism) {
waitingForResourcesStateValidator.validateInput(null);
registerStateTransition();
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/MockRestartingContext.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/MockRestartingContext.java
index d0a5db815fe01..39efe906515db 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/MockRestartingContext.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/MockRestartingContext.java
@@ -36,6 +36,7 @@
import java.util.function.Consumer;
import static org.apache.flink.runtime.scheduler.adaptive.WaitingForResourcesTest.assertNonNull;
+import static org.assertj.core.api.Assertions.assertThat;
/** Mock the {@link StateWithExecutionGraph.Context} for restarting state. */
class MockRestartingContext extends MockStateWithExecutionGraphContext
@@ -44,13 +45,13 @@ class MockRestartingContext extends MockStateWithExecutionGraphContext
private final StateValidator