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

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

commit 1ff7433bcca8b3a2792df2e5a3b0b421bfe4ba4c
Author: Zdenek Tison <[email protected]>
AuthorDate: Tue Jan 28 17:20:04 2025 +0100

    [FLINK-37232][runtime] Fix for broken synchronization assumption on the 
AdaptiveScheduler's side introduced by FLIP-272
---
 .../scheduler/adaptive/AdaptiveScheduler.java      |  4 +-
 .../runtime/scheduler/adaptive/Executing.java      |  6 +-
 .../scheduler/adaptive/FailureResultUtil.java      |  2 +-
 .../runtime/scheduler/adaptive/Restarting.java     | 42 +++++++++---
 .../scheduler/adaptive/StateTransitions.java       | 11 ++--
 .../scheduler/adaptive/AdaptiveSchedulerTest.java  |  4 +-
 .../runtime/scheduler/adaptive/ExecutingTest.java  | 25 +++++---
 .../scheduler/adaptive/MockRestartingContext.java  | 14 ++++
 .../runtime/scheduler/adaptive/RestartingTest.java | 72 ++++++++++++++++-----
 .../scheduler/adaptive/StopWithSavepointTest.java  |  9 ++-
 .../UpdateJobResourceRequirementsITCase.java       | 74 +++++++++++++++++++++-
 11 files changed, 215 insertions(+), 48 deletions(-)

diff --git 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java
 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java
index 31505241fc1..e74fa69d4e4 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java
@@ -1246,7 +1246,7 @@ public class AdaptiveScheduler
             ExecutionGraphHandler executionGraphHandler,
             OperatorCoordinatorHandler operatorCoordinatorHandler,
             Duration backoffTime,
-            boolean forcedRestart,
+            @Nullable VertexParallelism restartWithParallelism,
             List<ExceptionHistoryEntry> failureCollection) {
 
         for (ExecutionVertex executionVertex : 
executionGraph.getAllExecutionVertices()) {
@@ -1267,7 +1267,7 @@ public class AdaptiveScheduler
                         operatorCoordinatorHandler,
                         LOG,
                         backoffTime,
-                        forcedRestart,
+                        restartWithParallelism,
                         userCodeClassLoader,
                         failureCollection));
 
diff --git 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Executing.java
 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Executing.java
index c4d23dc7be3..f699d13373e 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Executing.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Executing.java
@@ -154,7 +154,11 @@ class Executing extends StateWithExecutionGraph
                 getExecutionGraphHandler(),
                 getOperatorCoordinatorHandler(),
                 Duration.ofMillis(0L),
-                true,
+                context.getAvailableVertexParallelism()
+                        .orElseThrow(
+                                () ->
+                                        new IllegalStateException(
+                                                "Resources must be available 
when rescaling.")),
                 getFailures());
     }
 
diff --git 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/FailureResultUtil.java
 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/FailureResultUtil.java
index bc16cafbf14..83183ef7378 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/FailureResultUtil.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/FailureResultUtil.java
@@ -30,7 +30,7 @@ public class FailureResultUtil {
                     sweg.getExecutionGraphHandler(),
                     sweg.getOperatorCoordinatorHandler(),
                     failureResult.getBackoffTime(),
-                    false,
+                    null,
                     sweg.getFailures());
         } else {
             sweg.getLogger().info("Failing job.", 
failureResult.getFailureCause());
diff --git 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Restarting.java
 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Restarting.java
index 1dd3f29778f..7ce04dbcc37 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Restarting.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/Restarting.java
@@ -22,6 +22,7 @@ import org.apache.flink.api.common.JobStatus;
 import org.apache.flink.runtime.executiongraph.ExecutionGraph;
 import org.apache.flink.runtime.scheduler.ExecutionGraphHandler;
 import org.apache.flink.runtime.scheduler.OperatorCoordinatorHandler;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism;
 import 
org.apache.flink.runtime.scheduler.exceptionhistory.ExceptionHistoryEntry;
 import org.apache.flink.util.Preconditions;
 
@@ -32,6 +33,7 @@ import javax.annotation.Nullable;
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledFuture;
 
@@ -44,7 +46,7 @@ class Restarting extends StateWithExecutionGraph {
 
     @Nullable private ScheduledFuture<?> goToSubsequentStateFuture;
 
-    private final boolean forcedRestart;
+    private final @Nullable VertexParallelism restartWithParallelism;
 
     Restarting(
             Context context,
@@ -53,7 +55,7 @@ class Restarting extends StateWithExecutionGraph {
             OperatorCoordinatorHandler operatorCoordinatorHandler,
             Logger logger,
             Duration backoffTime,
-            boolean forcedRestart,
+            @Nullable VertexParallelism restartWithParallelism,
             ClassLoader userCodeClassLoader,
             List<ExceptionHistoryEntry> failureCollection) {
         super(
@@ -66,7 +68,7 @@ class Restarting extends StateWithExecutionGraph {
                 failureCollection);
         this.context = context;
         this.backoffTime = backoffTime;
-        this.forcedRestart = forcedRestart;
+        this.restartWithParallelism = restartWithParallelism;
 
         getExecutionGraph().cancel();
     }
@@ -112,13 +114,31 @@ class Restarting extends StateWithExecutionGraph {
     }
 
     private void goToSubsequentState() {
-        if (forcedRestart) {
+        if (availableParallelismNotChanged(restartWithParallelism)) {
             context.goToCreatingExecutionGraph(getExecutionGraph());
         } else {
             context.goToWaitingForResources(getExecutionGraph());
         }
     }
 
+    private boolean availableParallelismNotChanged(VertexParallelism 
restartWithParallelism) {
+        if (this.restartWithParallelism == null) {
+            return false;
+        }
+
+        return context.getAvailableVertexParallelism()
+                .map(
+                        vertexParallelism ->
+                                vertexParallelism.getVertices().stream()
+                                        .allMatch(
+                                                vertex ->
+                                                        
restartWithParallelism.getParallelism(
+                                                                        vertex)
+                                                                == 
vertexParallelism.getParallelism(
+                                                                        
vertex)))
+                .orElse(false);
+    }
+
     /** Context of the {@link Restarting} state. */
     interface Context
             extends StateWithExecutionGraph.Context,
@@ -137,6 +157,12 @@ class Restarting extends StateWithExecutionGraph {
          * @return a ScheduledFuture representing pending completion of the 
task
          */
         ScheduledFuture<?> runIfState(State expectedState, Runnable action, 
Duration delay);
+
+        /**
+         * Returns the {@link VertexParallelism} that can be provided by the 
currently available
+         * slots.
+         */
+        Optional<VertexParallelism> getAvailableVertexParallelism();
     }
 
     static class Factory implements StateFactory<Restarting> {
@@ -147,7 +173,7 @@ class Restarting extends StateWithExecutionGraph {
         private final ExecutionGraphHandler executionGraphHandler;
         private final OperatorCoordinatorHandler operatorCoordinatorHandler;
         private final Duration backoffTime;
-        private final boolean forcedRestart;
+        private final @Nullable VertexParallelism restartWithParallelism;
         private final ClassLoader userCodeClassLoader;
         private final List<ExceptionHistoryEntry> failureCollection;
 
@@ -158,7 +184,7 @@ class Restarting extends StateWithExecutionGraph {
                 OperatorCoordinatorHandler operatorCoordinatorHandler,
                 Logger log,
                 Duration backoffTime,
-                boolean forcedRestart,
+                @Nullable VertexParallelism restartWithParallelism,
                 ClassLoader userCodeClassLoader,
                 List<ExceptionHistoryEntry> failureCollection) {
             this.context = context;
@@ -167,7 +193,7 @@ class Restarting extends StateWithExecutionGraph {
             this.executionGraphHandler = executionGraphHandler;
             this.operatorCoordinatorHandler = operatorCoordinatorHandler;
             this.backoffTime = backoffTime;
-            this.forcedRestart = forcedRestart;
+            this.restartWithParallelism = restartWithParallelism;
             this.userCodeClassLoader = userCodeClassLoader;
             this.failureCollection = failureCollection;
         }
@@ -184,7 +210,7 @@ class Restarting extends StateWithExecutionGraph {
                     operatorCoordinatorHandler,
                     log,
                     backoffTime,
-                    forcedRestart,
+                    restartWithParallelism,
                     userCodeClassLoader,
                     failureCollection);
         }
diff --git 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/StateTransitions.java
 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/StateTransitions.java
index 7f949afeee1..be56a487c76 100644
--- 
a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/StateTransitions.java
+++ 
b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/StateTransitions.java
@@ -23,6 +23,7 @@ import 
org.apache.flink.runtime.executiongraph.ArchivedExecutionGraph;
 import org.apache.flink.runtime.executiongraph.ExecutionGraph;
 import org.apache.flink.runtime.scheduler.ExecutionGraphHandler;
 import org.apache.flink.runtime.scheduler.OperatorCoordinatorHandler;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism;
 import 
org.apache.flink.runtime.scheduler.exceptionhistory.ExceptionHistoryEntry;
 
 import javax.annotation.Nullable;
@@ -128,9 +129,11 @@ public interface StateTransitions {
          *     Restarting} state
          * @param backoffTime backoffTime to wait before transitioning to the 
{@link Restarting}
          *     state
-         * @param forcedRestart if the {@link WaitingForResources} state 
should be omitted and the
-         *     {@link CreatingExecutionGraph} state should be entered directly 
from the {@link
-         *     Restarting} state
+         * @param restartWithParallelism the {@link VertexParallelism} that 
triggered the
+         *     restarting. The {@code AdaptiveScheduler} should transition 
directly to {@link
+         *     CreatingExecutionGraph} if the available parallelism hasn't 
changed while cancelling
+         *     the job. If {@code null} is passed or the parallelism changed, 
{@link
+         *     WaitingForResources} state should be the subsequent state.
          * @param failureCollection collection of failures that are propagated
          */
         void goToRestarting(
@@ -138,7 +141,7 @@ public interface StateTransitions {
                 ExecutionGraphHandler executionGraphHandler,
                 OperatorCoordinatorHandler operatorCoordinatorHandler,
                 Duration backoffTime,
-                boolean forcedRestart,
+                @Nullable VertexParallelism restartWithParallelism,
                 List<ExceptionHistoryEntry> failureCollection);
     }
 
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 c3a77d26fcd..f87da1da8ad 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
@@ -102,6 +102,7 @@ import 
org.apache.flink.runtime.scheduler.VertexParallelismInformation;
 import org.apache.flink.runtime.scheduler.VertexParallelismStore;
 import org.apache.flink.runtime.scheduler.adaptive.allocator.TestingSlot;
 import 
org.apache.flink.runtime.scheduler.adaptive.allocator.TestingSlotAllocator;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism;
 import 
org.apache.flink.runtime.scheduler.exceptionhistory.ExceptionHistoryEntry;
 import 
org.apache.flink.runtime.scheduler.exceptionhistory.RootExceptionHistoryEntry;
 import org.apache.flink.runtime.slots.ResourceRequirement;
@@ -728,7 +729,8 @@ public class AdaptiveSchedulerTest {
                                     executionGraphHandler,
                                     operatorCoordinatorHandler,
                                     Duration.ZERO,
-                                    true,
+                                    new VertexParallelism(
+                                            
Collections.singletonMap(JOB_VERTEX.getID(), 1)),
                                     failureCollection));
         }
 
diff --git 
a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ExecutingTest.java
 
b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ExecutingTest.java
index 7fb25cc3818..509ea292018 100644
--- 
a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ExecutingTest.java
+++ 
b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ExecutingTest.java
@@ -353,7 +353,7 @@ class ExecutingTest {
             ctx.setExpectRestarting(
                     restartingArguments -> {
                         
assertThat(restartingArguments.getBackoffTime()).isEqualTo(duration);
-                        
assertThat(restartingArguments.isForcedRestart()).isFalse();
+                        
assertThat(restartingArguments.getRestartWithParallelism()).isEmpty();
                     });
             ctx.setHowToHandleFailure(f -> FailureResult.canRestart(f, 
duration));
             exec.handleGlobalFailure(
@@ -439,7 +439,7 @@ class ExecutingTest {
             ctx.setExpectRestarting(
                     restartingArguments -> {
                         assertThat(restartingArguments).isNotNull();
-                        
assertThat(restartingArguments.isForcedRestart()).isFalse();
+                        
assertThat(restartingArguments.getRestartWithParallelism()).isEmpty();
                     });
 
             Exception exception = new RuntimeException();
@@ -616,9 +616,14 @@ class ExecutingTest {
     public void testOmitsWaitingForResourcesStateWhenRestarting() throws 
Exception {
         try (MockExecutingContext ctx = new MockExecutingContext()) {
             final Executing testInstance = new 
ExecutingStateBuilder().build(ctx);
+            final VertexParallelism vertexParallelism =
+                    new VertexParallelism(Collections.singletonMap(new 
JobVertexID(), 2));
+            ctx.setVertexParallelism(vertexParallelism);
             ctx.setExpectRestarting(
                     restartingArguments ->
-                            
assertThat(restartingArguments.isForcedRestart()).isTrue());
+                            
assertThat(restartingArguments.getRestartWithParallelism())
+                                    .hasValue(vertexParallelism));
+
             testInstance.transitionToSubsequentState();
         }
     }
@@ -824,7 +829,7 @@ class ExecutingTest {
                 ExecutionGraphHandler executionGraphHandler,
                 OperatorCoordinatorHandler operatorCoordinatorHandler,
                 Duration backoffTime,
-                boolean forcedRestart,
+                @Nullable VertexParallelism restartWithParallelism,
                 List<ExceptionHistoryEntry> failureCollection) {
             restartingStateValidator.validateInput(
                     new RestartingArguments(
@@ -832,7 +837,7 @@ class ExecutingTest {
                             executionGraphHandler,
                             operatorCoordinatorHandler,
                             backoffTime,
-                            forcedRestart));
+                            restartWithParallelism));
             hadStateTransition = true;
         }
 
@@ -948,25 +953,25 @@ class ExecutingTest {
 
     static class RestartingArguments extends CancellingArguments {
         private final Duration backoffTime;
-        private final boolean forcedRestart;
+        private final @Nullable VertexParallelism restartWithParallelism;
 
         public RestartingArguments(
                 ExecutionGraph executionGraph,
                 ExecutionGraphHandler executionGraphHandler,
                 OperatorCoordinatorHandler operatorCoordinatorHandler,
                 Duration backoffTime,
-                boolean forcedRestart) {
+                @Nullable VertexParallelism restartWithParallelism) {
             super(executionGraph, executionGraphHandler, 
operatorCoordinatorHandler);
             this.backoffTime = backoffTime;
-            this.forcedRestart = forcedRestart;
+            this.restartWithParallelism = restartWithParallelism;
         }
 
         public Duration getBackoffTime() {
             return backoffTime;
         }
 
-        public boolean isForcedRestart() {
-            return forcedRestart;
+        public Optional<VertexParallelism> getRestartWithParallelism() {
+            return Optional.ofNullable(restartWithParallelism);
         }
     }
 
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 a507c9d0b63..c4e2023b4c3 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
@@ -22,6 +22,7 @@ import 
org.apache.flink.core.testutils.CompletedScheduledFuture;
 import org.apache.flink.runtime.executiongraph.ExecutionGraph;
 import org.apache.flink.runtime.scheduler.ExecutionGraphHandler;
 import org.apache.flink.runtime.scheduler.OperatorCoordinatorHandler;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism;
 import 
org.apache.flink.runtime.scheduler.exceptionhistory.ExceptionHistoryEntry;
 import 
org.apache.flink.runtime.scheduler.exceptionhistory.RootExceptionHistoryEntry;
 
@@ -29,6 +30,7 @@ import javax.annotation.Nullable;
 
 import java.time.Duration;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.ScheduledFuture;
 import java.util.function.Consumer;
 
@@ -47,6 +49,8 @@ class MockRestartingContext extends 
MockStateWithExecutionGraphContext
     private final StateValidator<ExecutionGraph> 
creatingExecutionGraphStateValidator =
             new StateValidator<>("CreatingExecutionGraph");
 
+    @Nullable private VertexParallelism availableVertexParallelism;
+
     public void 
setExpectCancelling(Consumer<ExecutingTest.CancellingArguments> asserter) {
         cancellingStateValidator.expectInput(asserter);
     }
@@ -59,6 +63,11 @@ class MockRestartingContext extends 
MockStateWithExecutionGraphContext
         creatingExecutionGraphStateValidator.expectInput(assertNonNull());
     }
 
+    public void setAvailableVertexParallelism(
+            @Nullable VertexParallelism availableVertexParallelism) {
+        this.availableVertexParallelism = availableVertexParallelism;
+    }
+
     @Override
     public void goToCanceling(
             ExecutionGraph executionGraph,
@@ -94,6 +103,11 @@ class MockRestartingContext extends 
MockStateWithExecutionGraphContext
         return CompletedScheduledFuture.create(null);
     }
 
+    @Override
+    public Optional<VertexParallelism> getAvailableVertexParallelism() {
+        return Optional.ofNullable(availableVertexParallelism);
+    }
+
     @Override
     public void close() throws Exception {
         super.close();
diff --git 
a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/RestartingTest.java
 
b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/RestartingTest.java
index 1686d90b0db..b987d43e154 100644
--- 
a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/RestartingTest.java
+++ 
b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/RestartingTest.java
@@ -21,18 +21,26 @@ package org.apache.flink.runtime.scheduler.adaptive;
 import org.apache.flink.api.common.JobStatus;
 import org.apache.flink.runtime.executiongraph.ExecutionGraph;
 import org.apache.flink.runtime.failure.FailureEnricherUtils;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
 import org.apache.flink.runtime.scheduler.ExecutionGraphHandler;
 import org.apache.flink.runtime.scheduler.OperatorCoordinatorHandler;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism;
 
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.ValueSource;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.annotation.Nullable;
+
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Optional;
+import java.util.stream.Stream;
 
+import static java.util.Collections.singletonMap;
 import static 
org.apache.flink.runtime.scheduler.adaptive.WaitingForResourcesTest.assertNonNull;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -52,12 +60,14 @@ class RestartingTest {
     }
 
     @ParameterizedTest
-    @ValueSource(booleans = {true, false})
-    public void 
testTransitionToSubsequentStateWhenCancellationComplete(boolean forcedRestart)
-            throws Exception {
+    @MethodSource("provideRestartWithParallelism")
+    public void testTransitionToSubsequentStateWhenCancellationComplete(
+            Optional<VertexParallelism> restartWithParallelism) throws 
Exception {
         try (MockRestartingContext ctx = new MockRestartingContext()) {
-            Restarting restarting = createRestartingState(ctx, forcedRestart);
-            if (forcedRestart) {
+            
restartWithParallelism.ifPresent(ctx::setAvailableVertexParallelism);
+            Restarting restarting = createRestartingState(ctx, 
restartWithParallelism.orElse(null));
+
+            if (restartWithParallelism.isPresent()) {
                 ctx.setExpectCreatingExecutionGraph();
             } else {
                 ctx.setExpectWaitingForResources();
@@ -66,6 +76,22 @@ class RestartingTest {
         }
     }
 
+    @Test
+    public void testTransitionToSubsequentStateWhenResourceChanged() throws 
Exception {
+        try (MockRestartingContext ctx = new MockRestartingContext()) {
+            JobVertexID jobVertexId = new JobVertexID();
+            VertexParallelism availableParallelism =
+                    new VertexParallelism(singletonMap(jobVertexId, 1));
+            VertexParallelism requiredParallelismForForcedRestart =
+                    new VertexParallelism(singletonMap(jobVertexId, 2));
+
+            ctx.setAvailableVertexParallelism(availableParallelism);
+            Restarting restarting = createRestartingState(ctx, 
requiredParallelismForForcedRestart);
+            ctx.setExpectWaitingForResources();
+            restarting.onGloballyTerminalState(JobStatus.CANCELED);
+        }
+    }
+
     @Test
     void testCancel() throws Exception {
         try (MockRestartingContext ctx = new MockRestartingContext()) {
@@ -115,16 +141,19 @@ class RestartingTest {
     }
 
     @ParameterizedTest
-    @ValueSource(booleans = {true, false})
-    public void testStateDoesNotExposeGloballyTerminalExecutionGraph(boolean 
forcedRestart)
-            throws Exception {
+    @MethodSource("provideRestartWithParallelism")
+    public void testStateDoesNotExposeGloballyTerminalExecutionGraph(
+            Optional<VertexParallelism> restartWithParallelism) throws 
Exception {
         try (MockRestartingContext ctx = new MockRestartingContext()) {
+            
restartWithParallelism.ifPresent(ctx::setAvailableVertexParallelism);
             StateTrackingMockExecutionGraph mockExecutionGraph =
                     new StateTrackingMockExecutionGraph();
-            Restarting restarting = createRestartingState(ctx, 
mockExecutionGraph, forcedRestart);
+            Restarting restarting =
+                    createRestartingState(
+                            ctx, mockExecutionGraph, 
restartWithParallelism.orElse(null));
 
             // ideally we'd just delay the state transitions, but the context 
does not support that
-            if (forcedRestart) {
+            if (restartWithParallelism.isPresent()) {
                 ctx.setExpectCreatingExecutionGraph();
             } else {
                 ctx.setExpectWaitingForResources();
@@ -141,17 +170,21 @@ class RestartingTest {
         }
     }
 
-    public Restarting createRestartingState(MockRestartingContext ctx, boolean 
forcedRestart) {
-        return createRestartingState(ctx, new 
StateTrackingMockExecutionGraph(), forcedRestart);
+    public Restarting createRestartingState(
+            MockRestartingContext ctx, @Nullable VertexParallelism 
restartWithParallelism) {
+        return createRestartingState(
+                ctx, new StateTrackingMockExecutionGraph(), 
restartWithParallelism);
     }
 
     public Restarting createRestartingState(
             MockRestartingContext ctx, ExecutionGraph executionGraph) {
-        return createRestartingState(ctx, executionGraph, false);
+        return createRestartingState(ctx, executionGraph, null);
     }
 
     public Restarting createRestartingState(
-            MockRestartingContext ctx, ExecutionGraph executionGraph, boolean 
forcedRestart) {
+            MockRestartingContext ctx,
+            ExecutionGraph executionGraph,
+            @Nullable VertexParallelism restartWithParallelism) {
         final ExecutionGraphHandler executionGraphHandler =
                 new ExecutionGraphHandler(
                         executionGraph,
@@ -168,7 +201,7 @@ class RestartingTest {
                 operatorCoordinatorHandler,
                 log,
                 Duration.ZERO,
-                forcedRestart,
+                restartWithParallelism,
                 ClassLoader.getSystemClassLoader(),
                 new ArrayList<>());
     }
@@ -176,4 +209,11 @@ class RestartingTest {
     public Restarting createRestartingState(MockRestartingContext ctx) {
         return createRestartingState(ctx, new 
StateTrackingMockExecutionGraph());
     }
+
+    private static Stream<Arguments> provideRestartWithParallelism() {
+        return Stream.of(
+                Arguments.of(Optional.empty()),
+                Arguments.of(
+                        Optional.of(new VertexParallelism(singletonMap(new 
JobVertexID(), 1)))));
+    }
 }
diff --git 
a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/StopWithSavepointTest.java
 
b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/StopWithSavepointTest.java
index 33744f6c50a..31a76e0b52e 100644
--- 
a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/StopWithSavepointTest.java
+++ 
b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/StopWithSavepointTest.java
@@ -28,6 +28,7 @@ import 
org.apache.flink.runtime.executiongraph.TaskExecutionStateTransition;
 import org.apache.flink.runtime.failure.FailureEnricherUtils;
 import org.apache.flink.runtime.scheduler.ExecutionGraphHandler;
 import org.apache.flink.runtime.scheduler.OperatorCoordinatorHandler;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism;
 import 
org.apache.flink.runtime.scheduler.exceptionhistory.ExceptionHistoryEntry;
 import 
org.apache.flink.runtime.scheduler.exceptionhistory.TestingAccessExecution;
 import 
org.apache.flink.runtime.scheduler.stopwithsavepoint.StopWithSavepointStoppingException;
@@ -37,6 +38,8 @@ import org.junit.jupiter.api.Test;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.annotation.Nullable;
+
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -267,7 +270,7 @@ class StopWithSavepointTest {
             ctx.setExpectRestarting(
                     (restartingArguments) -> {
                         assertThat(restartingArguments).isNotNull();
-                        
assertThat(restartingArguments.isForcedRestart()).isFalse();
+                        
assertThat(restartingArguments.getRestartWithParallelism()).isEmpty();
                     });
 
             Exception exception = new RuntimeException();
@@ -580,7 +583,7 @@ class StopWithSavepointTest {
                 ExecutionGraphHandler executionGraphHandler,
                 OperatorCoordinatorHandler operatorCoordinatorHandler,
                 Duration backoffTime,
-                boolean forcedRestart,
+                @Nullable VertexParallelism restartWithParallelism,
                 List<ExceptionHistoryEntry> failureCollection) {
             if (hadStateTransition) {
                 throw new IllegalStateException("Only one state transition is 
allowed.");
@@ -592,7 +595,7 @@ class StopWithSavepointTest {
                             executionGraphHandler,
                             operatorCoordinatorHandler,
                             backoffTime,
-                            forcedRestart));
+                            restartWithParallelism));
             hadStateTransition = true;
         }
 
diff --git 
a/flink-tests/src/test/java/org/apache/flink/test/scheduling/UpdateJobResourceRequirementsITCase.java
 
b/flink-tests/src/test/java/org/apache/flink/test/scheduling/UpdateJobResourceRequirementsITCase.java
index 84257ca2dd5..fc03c17b182 100644
--- 
a/flink-tests/src/test/java/org/apache/flink/test/scheduling/UpdateJobResourceRequirementsITCase.java
+++ 
b/flink-tests/src/test/java/org/apache/flink/test/scheduling/UpdateJobResourceRequirementsITCase.java
@@ -23,18 +23,21 @@ import 
org.apache.flink.client.program.rest.RestClusterClient;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.JobManagerOptions;
 import org.apache.flink.configuration.WebOptions;
+import org.apache.flink.runtime.execution.Environment;
 import org.apache.flink.runtime.execution.ExecutionState;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.jobgraph.JobGraphTestUtils;
 import org.apache.flink.runtime.jobgraph.JobResourceRequirements;
 import org.apache.flink.runtime.jobgraph.JobVertex;
 import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup;
+import org.apache.flink.runtime.minicluster.MiniCluster;
 import 
org.apache.flink.runtime.rest.handler.legacy.messages.ClusterOverviewWithVersion;
 import org.apache.flink.runtime.rest.messages.job.JobDetailsInfo;
 import org.apache.flink.runtime.testtasks.BlockingNoOpInvokable;
 import org.apache.flink.runtime.testutils.CommonTestUtils;
 import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
 import org.apache.flink.test.junit5.InjectClusterClient;
+import org.apache.flink.test.junit5.InjectMiniCluster;
 import org.apache.flink.test.junit5.MiniClusterExtension;
 import org.apache.flink.util.TestLoggerExtension;
 
@@ -44,6 +47,7 @@ import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 import java.time.Duration;
+import java.util.concurrent.CountDownLatch;
 
 /** Tests for the manual rescaling of Flink jobs using the REST API. */
 @ExtendWith(TestLoggerExtension.class)
@@ -68,15 +72,18 @@ public class UpdateJobResourceRequirementsITCase {
         // - lower slot idle timeout -> controls how fast we return idle slots 
to TM
         configuration.set(WebOptions.REFRESH_INTERVAL, Duration.ofMillis(50L));
         configuration.set(JobManagerOptions.SLOT_IDLE_TIMEOUT, 
Duration.ofMillis(50L));
-
         return configuration;
     }
 
     private RestClusterClient<?> restClusterClient;
+    private MiniCluster miniCluster;
 
     @BeforeEach
-    void beforeEach(@InjectClusterClient RestClusterClient<?> 
restClusterClient) {
+    void beforeEach(
+            @InjectClusterClient RestClusterClient<?> restClusterClient,
+            @InjectMiniCluster MiniCluster miniCluster) {
         this.restClusterClient = restClusterClient;
+        this.miniCluster = miniCluster;
     }
 
     @Test
@@ -182,6 +189,47 @@ public class UpdateJobResourceRequirementsITCase {
                 NUMBER_OF_SLOTS - runningTasksAfterRescale);
     }
 
+    @Test
+    void testResourcesUnavailableAfterRescale() throws Exception {
+        final int initialRunningTasks = 3;
+        final int runningTasksAfterRescale = 2;
+
+        final JobVertex jobVertex = new JobVertex("Operator");
+        jobVertex.setParallelism(initialRunningTasks);
+        jobVertex.setInvokableClass(BlockingCancelNoOpInvokable.class);
+        BlockingCancelNoOpInvokable.reset();
+
+        final JobGraph jobGraph = 
JobGraphTestUtils.streamingJobGraph(jobVertex);
+
+        try {
+            restClusterClient.submitJob(jobGraph).join();
+
+            final JobID jobId = jobGraph.getJobID();
+            waitForRunningTasks(restClusterClient, jobId, initialRunningTasks);
+
+            final JobResourceRequirements jobResourceRequirements =
+                    JobResourceRequirements.newBuilder()
+                            .setParallelismForJobVertex(
+                                    jobVertex.getID(), 1, 
runningTasksAfterRescale)
+                            .build();
+
+            // wait for the job to be cancelled before rescaling
+            restClusterClient.updateJobResourceRequirements(jobId, 
jobResourceRequirements).join();
+            BlockingCancelNoOpInvokable.cancelWasCalled.await();
+
+            // terminate a task manager to make resources unavailable
+            miniCluster.terminateTaskManager(0);
+            BlockingCancelNoOpInvokable.cancelBlocking.countDown();
+
+            // add a new task manager to make resources available again and 
verify that the job
+            // becomes running
+            miniCluster.startTaskManager();
+            waitForRunningTasks(restClusterClient, jobId, 
runningTasksAfterRescale);
+        } finally {
+            restClusterClient.cancel(jobGraph.getJobID()).join();
+        }
+    }
+
     private void runRescalingTest(
             JobGraph jobGraph,
             JobResourceRequirements newJobVertexParallelism,
@@ -234,4 +282,26 @@ public class UpdateJobResourceRequirementsITCase {
                     return clusterOverview.getNumSlotsAvailable() == 
desiredNumberOfAvailableSlots;
                 });
     }
+
+    public static class BlockingCancelNoOpInvokable extends 
BlockingNoOpInvokable {
+
+        private static CountDownLatch cancelWasCalled;
+        private static CountDownLatch cancelBlocking;
+
+        public BlockingCancelNoOpInvokable(Environment environment) {
+            super(environment);
+        }
+
+        @Override
+        public void cancel() throws Exception {
+            cancelWasCalled.countDown();
+            cancelBlocking.await();
+            super.cancel();
+        }
+
+        public static void reset() {
+            cancelWasCalled = new CountDownLatch(1);
+            cancelBlocking = new CountDownLatch(1);
+        }
+    }
 }


Reply via email to