You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by ro...@apache.org on 2025/07/23 11:55:58 UTC
(flink) branch master updated: [FLINK-38114] Allow asynchronous offloading of TaskRestore
This is an automated email from the ASF dual-hosted git repository.
roman pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 6b0426c7c08 [FLINK-38114] Allow asynchronous offloading of TaskRestore
6b0426c7c08 is described below
commit 6b0426c7c08e9f5d10adf0cd2b6573e5ea37c6eb
Author: Aleksandr Iushmanov <ai...@confluent.io>
AuthorDate: Tue Jul 22 16:29:52 2025 +0100
[FLINK-38114] Allow asynchronous offloading of TaskRestore
---
.../org/apache/flink/testutils/TestingUtils.java | 29 +++
.../deployment/TaskDeploymentDescriptor.java | 40 ++-
.../TaskDeploymentDescriptorFactory.java | 11 +-
.../flink/runtime/executiongraph/Execution.java | 284 +++++++++++++++++----
.../runtime/executiongraph/ExecutionVertex.java | 5 +
.../flink/runtime/taskexecutor/TaskExecutor.java | 4 +-
.../TaskDeploymentDescriptorBuilder.java | 12 +-
.../deployment/TaskDeploymentDescriptorTest.java | 6 +-
...hDeploymentWithSmallBlobCacheSizeLimitTest.java | 12 +
.../executiongraph/ExecutionGraphSuspendTest.java | 205 +++++++++------
.../runtime/executiongraph/ExecutionTest.java | 116 ++++++++-
.../InteractionsCountingTaskManagerGateway.java | 11 +-
.../executiongraph/utils/ExecutionUtils.java | 41 +++
.../scheduler/UpdatePartitionConsumersTest.java | 43 +++-
.../runtime/scheduler/SchedulerTestingUtils.java | 80 ++++--
15 files changed, 698 insertions(+), 201 deletions(-)
diff --git a/flink-core/src/test/java/org/apache/flink/testutils/TestingUtils.java b/flink-core/src/test/java/org/apache/flink/testutils/TestingUtils.java
index 2bceea09fda..d64790dffee 100644
--- a/flink-core/src/test/java/org/apache/flink/testutils/TestingUtils.java
+++ b/flink-core/src/test/java/org/apache/flink/testutils/TestingUtils.java
@@ -31,6 +31,7 @@ import java.util.Collection;
import java.util.UUID;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
import java.util.function.Predicate;
import static org.assertj.core.api.Assertions.assertThat;
@@ -53,10 +54,38 @@ public class TestingUtils {
return Duration.ofDays(365L);
}
+ // To make debugging logs easier, we use a custom thread factory that names the thread.
+ static class JmMainSingleThreadPoolFactory implements ThreadFactory {
+ public Thread newThread(final Runnable r) {
+ return new Thread(r, "jm-main-thread");
+ }
+ }
+
+ // To make debugging logs easier, we use a custom thread factory that names the thread.
+ static class AsyncSingleThreadPoolFactory implements ThreadFactory {
+ public Thread newThread(final Runnable r) {
+ return new Thread(r, "async-single-thread-pool");
+ }
+ }
+
public static TestExecutorExtension<ScheduledExecutorService> defaultExecutorExtension() {
return new TestExecutorExtension<>(Executors::newSingleThreadScheduledExecutor);
}
+ public static TestExecutorExtension<ScheduledExecutorService> jmMainThreadExecutorExtension() {
+ return new TestExecutorExtension<>(
+ () ->
+ Executors.newSingleThreadScheduledExecutor(
+ new JmMainSingleThreadPoolFactory()));
+ }
+
+ public static TestExecutorExtension<ScheduledExecutorService> jmAsyncThreadExecutorExtension() {
+ return new TestExecutorExtension<>(
+ () ->
+ Executors.newSingleThreadScheduledExecutor(
+ new AsyncSingleThreadPoolFactory()));
+ }
+
public static TestExecutorResource<ScheduledExecutorService> defaultExecutorResource() {
return new TestExecutorResource<>(Executors::newSingleThreadScheduledExecutor);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
index 5684066735f..5e8aaca91c7 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
@@ -118,6 +118,12 @@ public final class TaskDeploymentDescriptor implements Serializable {
*/
private transient TaskInformation taskInformation;
+ /** Information to restore the task. This can be null if there is no state to restore. */
+ @Nullable private final MaybeOffloaded<JobManagerTaskRestore> serializedTaskRestore;
+
+ /** Information to restore the task. This can be null if there is no state to restore. */
+ @Nullable private transient JobManagerTaskRestore taskRestore;
+
/**
* The ID referencing the job this task belongs to.
*
@@ -138,16 +144,13 @@ public final class TaskDeploymentDescriptor implements Serializable {
/** The list of consumed intermediate result partitions. */
private final List<InputGateDeploymentDescriptor> inputGates;
- /** Information to restore the task. This can be null if there is no state to restore. */
- @Nullable private final JobManagerTaskRestore taskRestore;
-
public TaskDeploymentDescriptor(
JobID jobId,
MaybeOffloaded<JobInformation> serializedJobInformation,
MaybeOffloaded<TaskInformation> serializedTaskInformation,
ExecutionAttemptID executionAttemptId,
AllocationID allocationId,
- @Nullable JobManagerTaskRestore taskRestore,
+ @Nullable MaybeOffloaded<JobManagerTaskRestore> serializedTaskRestore,
List<ResultPartitionDeploymentDescriptor> resultPartitionDeploymentDescriptors,
List<InputGateDeploymentDescriptor> inputGateDeploymentDescriptors) {
@@ -159,7 +162,7 @@ public final class TaskDeploymentDescriptor implements Serializable {
this.executionId = Preconditions.checkNotNull(executionAttemptId);
this.allocationId = Preconditions.checkNotNull(allocationId);
- this.taskRestore = taskRestore;
+ this.serializedTaskRestore = serializedTaskRestore;
this.producedPartitions = Preconditions.checkNotNull(resultPartitionDeploymentDescriptors);
this.inputGates = Preconditions.checkNotNull(inputGateDeploymentDescriptors);
@@ -241,8 +244,16 @@ public final class TaskDeploymentDescriptor implements Serializable {
}
@Nullable
- public JobManagerTaskRestore getTaskRestore() {
- return taskRestore;
+ public JobManagerTaskRestore getTaskRestore() throws IOException, ClassNotFoundException {
+ if (taskRestore != null || serializedTaskRestore == null) {
+ return taskRestore;
+ }
+ if (serializedTaskRestore instanceof NonOffloaded) {
+ NonOffloaded<JobManagerTaskRestore> taskRestore =
+ (NonOffloaded<JobManagerTaskRestore>) serializedTaskRestore;
+ return taskRestore.serializedValue.deserializeValue(getClass().getClassLoader());
+ }
+ throw new IllegalStateException("Trying to work with offloaded serialized task restore.");
}
public AllocationID getAllocationId() {
@@ -311,6 +322,21 @@ public final class TaskDeploymentDescriptor implements Serializable {
this.taskInformation = taskInformation.deepCopy();
}
+ if (serializedTaskRestore == null) {
+ this.taskRestore = null;
+ } else if (serializedTaskRestore instanceof Offloaded) {
+ final PermanentBlobKey blobKey =
+ ((Offloaded<JobManagerTaskRestore>) serializedTaskRestore).serializedValueKey;
+
+ Preconditions.checkNotNull(blobService);
+
+ final File dataFile = blobService.getFile(jobId, blobKey);
+ taskRestore =
+ InstantiationUtil.deserializeObject(
+ new BufferedInputStream(Files.newInputStream(dataFile.toPath())),
+ getClass().getClassLoader());
+ }
+
for (InputGateDeploymentDescriptor inputGate : inputGates) {
inputGate.tryLoadAndDeserializeShuffleDescriptors(
blobService, jobId, shuffleDescriptorsCache);
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorFactory.java
index 23d7ed20df0..1936777797a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorFactory.java
@@ -112,7 +112,7 @@ public class TaskDeploymentDescriptorFactory {
public TaskDeploymentDescriptor createDeploymentDescriptor(
Execution execution,
AllocationID allocationID,
- @Nullable JobManagerTaskRestore taskRestore,
+ @Nullable Either<SerializedValue<JobManagerTaskRestore>, PermanentBlobKey> taskRestore,
Collection<ResultPartitionDeploymentDescriptor> producedPartitions)
throws IOException, ClusterDatasetCorruptedException {
final ExecutionVertex executionVertex = execution.getVertex();
@@ -124,7 +124,7 @@ public class TaskDeploymentDescriptorFactory {
executionVertex.getJobVertex().getTaskInformationOrBlobKey()),
execution.getAttemptId(),
allocationID,
- taskRestore,
+ taskRestore != null ? getSerializedTaskRestore(taskRestore) : null,
new ArrayList<>(producedPartitions),
createInputGateDeploymentDescriptors(executionVertex));
}
@@ -295,6 +295,13 @@ public class TaskDeploymentDescriptorFactory {
: new TaskDeploymentDescriptor.Offloaded<>(taskInfo.right());
}
+ private static MaybeOffloaded<JobManagerTaskRestore> getSerializedTaskRestore(
+ Either<SerializedValue<JobManagerTaskRestore>, PermanentBlobKey> either) {
+ return either.isLeft()
+ ? new TaskDeploymentDescriptor.NonOffloaded<>(either.left())
+ : new TaskDeploymentDescriptor.Offloaded<>(either.right());
+ }
+
public static ShuffleDescriptor getConsumedPartitionShuffleDescriptor(
IntermediateResultPartition consumedPartition,
PartitionLocationConstraint partitionDeploymentConstraint,
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
index d2683b3373c..a1301f604dc 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
@@ -24,6 +24,8 @@ import org.apache.flink.api.common.accumulators.Accumulator;
import org.apache.flink.core.io.InputSplit;
import org.apache.flink.runtime.JobException;
import org.apache.flink.runtime.accumulators.StringifiedAccumulatorResult;
+import org.apache.flink.runtime.blob.BlobWriter;
+import org.apache.flink.runtime.blob.PermanentBlobKey;
import org.apache.flink.runtime.checkpoint.CheckpointOptions;
import org.apache.flink.runtime.checkpoint.JobManagerTaskRestore;
import org.apache.flink.runtime.clusterframework.types.AllocationID;
@@ -43,6 +45,7 @@ import org.apache.flink.runtime.jobmaster.LogicalSlot;
import org.apache.flink.runtime.messages.Acknowledge;
import org.apache.flink.runtime.operators.coordination.OperatorEvent;
import org.apache.flink.runtime.operators.coordination.TaskNotRunningException;
+import org.apache.flink.runtime.scheduler.ClusterDatasetCorruptedException;
import org.apache.flink.runtime.scheduler.strategy.ConsumerVertexGroup;
import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
import org.apache.flink.runtime.shuffle.PartitionDescriptor;
@@ -51,17 +54,20 @@ import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
import org.apache.flink.runtime.shuffle.ShuffleMaster;
import org.apache.flink.runtime.taskexecutor.TaskExecutorOperatorEventGateway;
import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.types.Either;
import org.apache.flink.util.CollectionUtil;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.FlinkException;
import org.apache.flink.util.OptionalFailure;
import org.apache.flink.util.SerializedValue;
import org.apache.flink.util.concurrent.FutureUtils;
+import org.apache.flink.util.function.FunctionUtils;
import org.slf4j.Logger;
import javax.annotation.Nullable;
+import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Collection;
@@ -74,7 +80,9 @@ import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
+import java.util.function.Supplier;
import java.util.stream.Collectors;
import static org.apache.flink.runtime.deployment.TaskDeploymentDescriptorFactory.getConsumedPartitionShuffleDescriptor;
@@ -165,6 +173,13 @@ public class Execution
private Optional<ErrorInfo> failureCause =
Optional.empty(); // once an ErrorInfo is set, never changes
+ /**
+ * Future that indicates {@link TaskDeploymentDescriptor} being created after the most recent
+ * deploy. This is only used for testing purposes to handle race conditions caused by thread
+ * switching between main and async pool in #deploy().
+ */
+ @Nullable private volatile CompletableFuture<Void> tddCreatedDuringDeployFuture;
+
/**
* Information to restore the task on recovery, such as checkpoint id and task state snapshot.
*/
@@ -595,14 +610,8 @@ public class Execution
try {
// race double check, did we fail/cancel and do we need to release the slot?
- if (this.state != DEPLOYING) {
- slot.releaseSlot(
- new FlinkException(
- "Actual state of execution "
- + this
- + " ("
- + state
- + ") does not match expected state DEPLOYING."));
+ if (state != DEPLOYING) {
+ releaseSlotWhenNotInDeployingState(slot);
return;
}
@@ -615,18 +624,6 @@ public class Execution
getAssignedResourceLocation(),
slot.getAllocationId());
- final TaskDeploymentDescriptor deployment =
- vertex.getExecutionGraphAccessor()
- .getTaskDeploymentDescriptorFactory()
- .createDeploymentDescriptor(
- this,
- slot.getAllocationId(),
- taskRestore,
- producedPartitions.values());
-
- // null taskRestore to let it be GC'ed
- taskRestore = null;
-
final TaskManagerGateway taskManagerGateway = slot.getTaskManagerGateway();
final ComponentMainThreadExecutor jobMasterMainThreadExecutor =
@@ -636,38 +633,32 @@ public class Execution
// We run the submission in the future executor so that the serialization of large TDDs
// does not block
// the main thread and sync back to the main thread once submission is completed.
- CompletableFuture.supplyAsync(
- () -> taskManagerGateway.submitTask(deployment, rpcTimeout), executor)
- .thenCompose(Function.identity())
+ AtomicReference<Either<SerializedValue<JobManagerTaskRestore>, PermanentBlobKey>>
+ maybeOffloadedTaskRestoreCleanupRef = new AtomicReference<>();
+ final CompletableFuture<TaskDeploymentDescriptor> taskDeploymentDescriptorFuture =
+ CompletableFuture.supplyAsync(
+ initOffloadedTaskRestoreRef(
+ // Passing a copy of the task restore from the
+ // main thread to the I/O executor in order to
+ // avoid synchronization issues
+ taskRestore, maybeOffloadedTaskRestoreCleanupRef),
+ executor)
+ // back to main thread because this accesses execution graph
+ // internals
+ .thenComposeAsync(
+ tryGetTaskDeploymentDescriptorForSlot(slot),
+ jobMasterMainThreadExecutor);
+ this.tddCreatedDuringDeployFuture = taskDeploymentDescriptorFuture.thenRun(() -> {});
+ // back to io executor for TDD serialization
+ taskDeploymentDescriptorFuture
+ .thenComposeAsync(
+ deploymentDescriptor ->
+ taskManagerGateway.submitTask(deploymentDescriptor, rpcTimeout),
+ executor)
.whenCompleteAsync(
- (ack, failure) -> {
- if (failure == null) {
- vertex.notifyCompletedDeployment(this);
- } else {
- final Throwable actualFailure =
- ExceptionUtils.stripCompletionException(failure);
-
- if (actualFailure instanceof TimeoutException) {
- String taskname =
- vertex.getTaskNameWithSubtaskIndex()
- + " ("
- + attemptId
- + ')';
-
- markFailed(
- new Exception(
- "Cannot deploy task "
- + taskname
- + " - TaskManager ("
- + getAssignedResourceLocation()
- + ") not responding after a rpcTimeout of "
- + rpcTimeout,
- actualFailure));
- } else {
- markFailed(actualFailure);
- }
- }
- },
+ (ack, failure) ->
+ handleDeploymentCompletionAndCleanup(
+ maybeOffloadedTaskRestoreCleanupRef, ack, failure),
jobMasterMainThreadExecutor);
} catch (Throwable t) {
@@ -676,6 +667,185 @@ public class Execution
}
}
+ private void handleDeploymentCompletionAndCleanup(
+ AtomicReference<Either<SerializedValue<JobManagerTaskRestore>, PermanentBlobKey>>
+ cleanupRef,
+ @Nullable final Acknowledge ack,
+ @Nullable final Throwable failure) {
+
+ // Running in async thread to avoid blocking the main thread on cleanup.
+ CompletableFuture.runAsync(() -> cleanUpOffloadedTaskRestore(cleanupRef), executor)
+ .exceptionally(
+ cleanupError -> {
+ LOG.warn(
+ "Failed to cleanup offloaded task restore for "
+ + "{} (attempt #{}) with attempt "
+ + "id {} and vertex id {} to {}",
+ vertex.getTaskNameWithSubtaskIndex(),
+ getAttemptNumber(),
+ attemptId,
+ vertex.getID(),
+ getAssignedResourceLocation(),
+ cleanupError);
+ return null;
+ });
+
+ finalizeDeploymentAndCleanUpInMainThread(ack, failure);
+ }
+
+ private Supplier<Either<SerializedValue<JobManagerTaskRestore>, PermanentBlobKey>>
+ initOffloadedTaskRestoreRef(
+ final JobManagerTaskRestore taskRestoreSnapshot,
+ final AtomicReference<
+ Either<
+ SerializedValue<JobManagerTaskRestore>,
+ PermanentBlobKey>>
+ maybeOffloadedTaskRestoreCleanupRef) {
+ return FunctionUtils.uncheckedSupplier(
+ () -> {
+ maybeOffloadedTaskRestoreCleanupRef.set(
+ getMaybeOffloadedTaskRestore(taskRestoreSnapshot));
+ return maybeOffloadedTaskRestoreCleanupRef.get();
+ });
+ }
+
+ private void finalizeDeploymentAndCleanUpInMainThread(
+ @Nullable final Acknowledge ack, @Nullable final Throwable failure) {
+ // Verify that we are back in the job master main thread
+ assertRunningInJobMasterMainThread();
+ // null taskRestore to let it be GC'ed
+ taskRestore = null;
+
+ if (failure == null) {
+ vertex.notifyCompletedDeployment(this);
+ } else {
+ final Throwable actualFailure = ExceptionUtils.stripCompletionException(failure);
+
+ if (actualFailure instanceof TimeoutException) {
+ String taskname = vertex.getTaskNameWithSubtaskIndex() + " (" + attemptId + ')';
+
+ markFailed(
+ new Exception(
+ "Cannot deploy task "
+ + taskname
+ + " - TaskManager ("
+ + getAssignedResourceLocation()
+ + ") not responding after a rpcTimeout of "
+ + rpcTimeout,
+ actualFailure));
+ } else {
+ markFailed(actualFailure);
+ }
+ }
+ }
+
+ private Function<
+ Either<SerializedValue<JobManagerTaskRestore>, PermanentBlobKey>,
+ CompletableFuture<TaskDeploymentDescriptor>>
+ tryGetTaskDeploymentDescriptorForSlot(final LogicalSlot slot) {
+ return FunctionUtils.uncheckedFunction(
+ (maybeOffloadedTaskRestore) -> {
+ // Check that we are in jm main thread while creating
+ // task deployment descriptor.
+ assertRunningInJobMasterMainThread();
+ if (state != DEPLOYING) {
+ return FutureUtils.completedExceptionally(
+ new IllegalStateException(
+ String.format(
+ "Cannot deploy %s (attempt #%s) with attempt "
+ + "id %s and vertex id %s to %s "
+ + "with allocation id %s "
+ + "because execution state has switched to %s "
+ + "during task restore offload",
+ vertex.getTaskNameWithSubtaskIndex(),
+ getAttemptNumber(),
+ attemptId,
+ vertex.getID(),
+ getAssignedResourceLocation(),
+ slot.getAllocationId(),
+ state)));
+ }
+ return CompletableFuture.completedFuture(
+ getDeploymentDescriptor(maybeOffloadedTaskRestore, slot));
+ });
+ }
+
+ /**
+ * Helper method to release the slot when Execution is not in DEPLOYING state.
+ *
+ * @param slot The slot to release.
+ */
+ private void releaseSlotWhenNotInDeployingState(final LogicalSlot slot) {
+ slot.releaseSlot(
+ new FlinkException(
+ "Actual state of execution "
+ + this
+ + " ("
+ + state
+ + ") does not match expected state DEPLOYING."));
+ }
+
+ /**
+ * Cleans up the offloaded task restore if it was offloaded to the blob store.
+ *
+ * @param maybeOffloadedTaskRestoreCleanupRef The reference to the offloaded task restore
+ */
+ private void cleanUpOffloadedTaskRestore(
+ final AtomicReference<Either<SerializedValue<JobManagerTaskRestore>, PermanentBlobKey>>
+ maybeOffloadedTaskRestoreCleanupRef) {
+ if (maybeOffloadedTaskRestoreCleanupRef.get() != null
+ && maybeOffloadedTaskRestoreCleanupRef.get().isRight()) {
+ vertex.getExecutionGraphAccessor()
+ .deleteBlobs(
+ Collections.singletonList(
+ maybeOffloadedTaskRestoreCleanupRef.get().right()));
+ maybeOffloadedTaskRestoreCleanupRef.set(null);
+ }
+ }
+
+ /**
+ * Creates the task deployment descriptor for the task.
+ *
+ * @param maybeOffloadedTaskRestore The serialized task restore information or the blob key if
+ * it was offloaded
+ * @param slot The slot to which the task is deployed
+ * @return The task deployment descriptor
+ * @throws IOException If an I/O error occurs
+ * @throws ClusterDatasetCorruptedException If the cluster dataset is corrupted
+ */
+ private TaskDeploymentDescriptor getDeploymentDescriptor(
+ final Either<SerializedValue<JobManagerTaskRestore>, PermanentBlobKey>
+ maybeOffloadedTaskRestore,
+ final LogicalSlot slot)
+ throws IOException, ClusterDatasetCorruptedException {
+ return vertex.getExecutionGraphAccessor()
+ .getTaskDeploymentDescriptorFactory()
+ .createDeploymentDescriptor(
+ this,
+ slot.getAllocationId(),
+ maybeOffloadedTaskRestore,
+ producedPartitions.values());
+ }
+
+ /**
+ * Returns the serialized task restore information or the blob key if it was offloaded.
+ *
+ * @param taskRestoreSnapshot Task restore snapshot value at the time of deployment.
+ * @return The serialized task restore information or the blob key if it was offloaded
+ * @throws IOException If an I/O error occurs
+ */
+ @VisibleForTesting
+ protected Either<SerializedValue<JobManagerTaskRestore>, PermanentBlobKey>
+ getMaybeOffloadedTaskRestore(final JobManagerTaskRestore taskRestoreSnapshot)
+ throws IOException {
+ return taskRestoreSnapshot == null
+ ? null
+ : BlobWriter.serializeAndTryOffload(
+ taskRestoreSnapshot,
+ vertex.getJobId(),
+ vertex.getExecutionGraphAccessor().getBlobWriter());
+ }
+
public void cancel() {
// depending on the previous state, we go directly to cancelled (no cancel call necessary)
// -- or to canceling (cancel call needs to be sent to the task manager)
@@ -1637,6 +1807,18 @@ public class Execution
return new ArchivedExecution(this);
}
+ @VisibleForTesting
+ public CompletableFuture<Void> getTddCreationDuringDeployFuture() {
+ // Defensive copy in case future is set to null between check and join.
+ final CompletableFuture<Void> future = tddCreatedDuringDeployFuture;
+ if (future == null) {
+ throw new IllegalStateException(
+ "Task deployment descriptor creation future is null, "
+ + "please ensure that this method is called after `Execution.deploy()`.");
+ }
+ return future;
+ }
+
private void assertRunningInJobMasterMainThread() {
vertex.getExecutionGraphAccessor()
.getJobMasterMainThreadExecutor()
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
index 13e3fb3564a..81d70ec6202 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
@@ -576,6 +576,11 @@ public class ExecutionVertex
return currentExecution == execution;
}
+ @VisibleForTesting
+ public CompletableFuture<Void> getTddCreationDuringDeployFuture() {
+ return currentExecution.getTddCreationDuringDeployFuture();
+ }
+
// --------------------------------------------------------------------------------------------
// Utilities
// --------------------------------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java
index f713162dfda..17b2adea2e8 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskExecutor.java
@@ -717,9 +717,11 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
// deserialize the pre-serialized information
final JobInformation jobInformation;
final TaskInformation taskInformation;
+ final JobManagerTaskRestore taskRestore;
try {
jobInformation = tdd.getJobInformation();
taskInformation = tdd.getTaskInformation();
+ taskRestore = tdd.getTaskRestore();
} catch (IOException | ClassNotFoundException e) {
throw new TaskSubmissionException(
"Could not deserialize the job or task information.", e);
@@ -812,8 +814,6 @@ public class TaskExecutor extends RpcEndpoint implements TaskExecutorGateway {
throw new TaskSubmissionException(e);
}
- final JobManagerTaskRestore taskRestore = tdd.getTaskRestore();
-
final TaskStateManager taskStateManager =
new TaskStateManagerImpl(
jobId,
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorBuilder.java
index 6a7edeab017..b7f660db02f 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorBuilder.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorBuilder.java
@@ -49,7 +49,7 @@ public class TaskDeploymentDescriptorBuilder {
private List<ResultPartitionDeploymentDescriptor> producedPartitions;
private List<InputGateDeploymentDescriptor> inputGates;
- @Nullable private JobManagerTaskRestore taskRestore;
+ @Nullable private MaybeOffloaded<JobManagerTaskRestore> serializedTaskRestore;
private TaskDeploymentDescriptorBuilder(JobID jobId, String invokableClassName)
throws IOException {
@@ -71,7 +71,7 @@ public class TaskDeploymentDescriptorBuilder {
this.allocationId = new AllocationID();
this.producedPartitions = Collections.emptyList();
this.inputGates = Collections.emptyList();
- this.taskRestore = null;
+ this.serializedTaskRestore = null;
}
public TaskDeploymentDescriptorBuilder setSerializedJobInformation(
@@ -113,9 +113,9 @@ public class TaskDeploymentDescriptorBuilder {
return this;
}
- public TaskDeploymentDescriptorBuilder setTaskRestore(
- @Nullable JobManagerTaskRestore taskRestore) {
- this.taskRestore = taskRestore;
+ public TaskDeploymentDescriptorBuilder setSerializedTaskRestore(
+ @Nullable MaybeOffloaded<JobManagerTaskRestore> serializedTaskRestore) {
+ this.serializedTaskRestore = serializedTaskRestore;
return this;
}
@@ -126,7 +126,7 @@ public class TaskDeploymentDescriptorBuilder {
serializedTaskInformation,
executionId,
allocationId,
- taskRestore,
+ serializedTaskRestore,
producedPartitions,
inputGates);
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorTest.java
index edae266fd4c..be47f6cd578 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorTest.java
@@ -82,8 +82,8 @@ class TaskDeploymentDescriptorTest {
private static final List<PermanentBlobKey> requiredJars = new ArrayList<>(0);
private static final List<URL> requiredClasspaths = new ArrayList<>(0);
private static final TaskStateSnapshot taskStateHandles = new TaskStateSnapshot();
- private static final JobManagerTaskRestore taskRestore =
- new JobManagerTaskRestore(1L, taskStateHandles);
+ private final SerializedValue<JobManagerTaskRestore> taskRestore =
+ new SerializedValue<>(new JobManagerTaskRestore(1L, taskStateHandles));
private final SerializedValue<ExecutionConfig> executionConfig =
new SerializedValue<>(new ExecutionConfig());
@@ -237,7 +237,7 @@ class TaskDeploymentDescriptorTest {
taskInformation,
execId,
allocationId,
- taskRestore,
+ new TaskDeploymentDescriptor.NonOffloaded<>(taskRestore),
producedResults,
inputGates);
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraphDeploymentWithSmallBlobCacheSizeLimitTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraphDeploymentWithSmallBlobCacheSizeLimitTest.java
index 9fd9d0cddf9..c6cef169d96 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraphDeploymentWithSmallBlobCacheSizeLimitTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/DefaultExecutionGraphDeploymentWithSmallBlobCacheSizeLimitTest.java
@@ -25,6 +25,8 @@ import org.apache.flink.runtime.blob.BlobCacheSizeTracker;
import org.apache.flink.runtime.blob.BlobServer;
import org.apache.flink.runtime.blob.PermanentBlobCache;
import org.apache.flink.runtime.blob.VoidBlobStore;
+import org.apache.flink.runtime.checkpoint.JobManagerTaskRestore;
+import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
import org.apache.flink.runtime.client.JobExecutionException;
import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutorServiceAdapter;
import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
@@ -59,6 +61,7 @@ import java.util.concurrent.BlockingQueue;
import static org.apache.flink.runtime.util.JobVertexConnectionUtils.connectNewDataSetAsInput;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/**
* Tests {@link ExecutionGraph} deployment when job and task information are offloaded into the BLOB
@@ -121,6 +124,8 @@ class DefaultExecutionGraphDeploymentWithSmallBlobCacheSizeLimitTest
taskManagerGateway.setSubmitConsumer(
FunctionUtils.uncheckedConsumer(
taskDeploymentDescriptor -> {
+ assertThatThrownBy(taskDeploymentDescriptor::getTaskRestore)
+ .isInstanceOf(IllegalStateException.class);
taskDeploymentDescriptor.loadBigData(
blobCache,
new NoOpGroupCache<>(),
@@ -138,9 +143,12 @@ class DefaultExecutionGraphDeploymentWithSmallBlobCacheSizeLimitTest
new TestingLogicalSlotBuilder()
.setTaskManagerGateway(taskManagerGateway)
.createTestingLogicalSlot();
+ JobManagerTaskRestore jobManagerTaskRestore =
+ new JobManagerTaskRestore(0, new TaskStateSnapshot());
final Execution execution = ev.getCurrentExecutionAttempt();
execution.transitionState(ExecutionState.SCHEDULED);
execution.registerProducedPartitions(slot.getTaskManagerLocation()).get();
+ execution.setInitialState(jobManagerTaskRestore);
ev.deployToSlot(slot);
assertThat(ev.getExecutionState()).isEqualTo(ExecutionState.DEPLOYING);
@@ -150,6 +158,10 @@ class DefaultExecutionGraphDeploymentWithSmallBlobCacheSizeLimitTest
List<InputGateDeploymentDescriptor> igdds = tdd.getInputGates();
assertThat(igdds).hasSize(ev.getAllConsumedPartitionGroups().size());
+ assertThat(tdd.getTaskRestore()).isNotNull();
+ assertThat(tdd.getTaskRestore().getRestoreCheckpointId())
+ .isEqualTo(jobManagerTaskRestore.getRestoreCheckpointId());
+
if (igdds.size() > 0) {
checkShuffleDescriptors(igdds.get(0), ev.getConsumedPartitionGroup(0));
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphSuspendTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphSuspendTest.java
index 822c4e467d8..7353518cc07 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphSuspendTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionGraphSuspendTest.java
@@ -38,6 +38,7 @@ import org.apache.flink.util.concurrent.ManuallyTriggeredScheduledExecutor;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import static org.assertj.core.api.Assertions.assertThat;
@@ -47,7 +48,11 @@ class ExecutionGraphSuspendTest {
@RegisterExtension
static final TestExecutorExtension<ScheduledExecutorService> EXECUTOR_RESOURCE =
- TestingUtils.defaultExecutorExtension();
+ TestingUtils.jmAsyncThreadExecutorExtension();
+
+ @RegisterExtension
+ static final TestExecutorExtension<ScheduledExecutorService> JM_MAIN_THREAD_EXECUTOR_RESOURCE =
+ TestingUtils.jmMainThreadExecutorExtension();
/**
* Going into SUSPENDED out of CREATED should immediately cancel everything and not send out RPC
@@ -59,17 +64,19 @@ class ExecutionGraphSuspendTest {
new InteractionsCountingTaskManagerGateway();
final int parallelism = 10;
final SchedulerBase scheduler = createScheduler(gateway, parallelism);
- final ExecutionGraph eg = scheduler.getExecutionGraph();
+ runInMainThread(
+ () -> {
+ final ExecutionGraph eg = scheduler.getExecutionGraph();
- assertThat(eg.getState()).isEqualTo(JobStatus.CREATED);
+ assertThat(eg.getState()).isEqualTo(JobStatus.CREATED);
- // suspend
+ // suspend
+ scheduler.closeAsync();
- scheduler.closeAsync();
-
- assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
- validateAllVerticesInState(eg, ExecutionState.CANCELED);
- validateCancelRpcCalls(gateway, 0);
+ assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ validateAllVerticesInState(eg, ExecutionState.CANCELED);
+ validateCancelRpcCalls(gateway, 0);
+ });
ensureCannotLeaveSuspendedState(scheduler, gateway);
}
@@ -83,18 +90,22 @@ class ExecutionGraphSuspendTest {
final int parallelism = 10;
final InteractionsCountingTaskManagerGateway gateway =
new InteractionsCountingTaskManagerGateway(parallelism);
+
final SchedulerBase scheduler = createScheduler(gateway, parallelism);
- final ExecutionGraph eg = scheduler.getExecutionGraph();
+ runInMainThread(
+ () -> {
+ final ExecutionGraph eg = scheduler.getExecutionGraph();
- scheduler.startScheduling();
- assertThat(eg.getState()).isEqualTo(JobStatus.RUNNING);
- validateAllVerticesInState(eg, ExecutionState.DEPLOYING);
+ scheduler.startScheduling();
+ assertThat(eg.getState()).isEqualTo(JobStatus.RUNNING);
+ validateAllVerticesInState(eg, ExecutionState.DEPLOYING);
- // suspend
- scheduler.closeAsync();
+ // suspend
+ scheduler.closeAsync();
- assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
- validateCancelRpcCalls(gateway, parallelism);
+ assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ validateCancelRpcCalls(gateway, parallelism);
+ });
ensureCannotLeaveSuspendedState(scheduler, gateway);
}
@@ -108,19 +119,22 @@ class ExecutionGraphSuspendTest {
final InteractionsCountingTaskManagerGateway gateway =
new InteractionsCountingTaskManagerGateway(parallelism);
final SchedulerBase scheduler = createScheduler(gateway, parallelism);
- final ExecutionGraph eg = scheduler.getExecutionGraph();
+ runInMainThread(
+ () -> {
+ final ExecutionGraph eg = scheduler.getExecutionGraph();
- scheduler.startScheduling();
- ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
+ scheduler.startScheduling();
+ ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
- assertThat(eg.getState()).isEqualTo(JobStatus.RUNNING);
- validateAllVerticesInState(eg, ExecutionState.RUNNING);
+ assertThat(eg.getState()).isEqualTo(JobStatus.RUNNING);
+ validateAllVerticesInState(eg, ExecutionState.RUNNING);
- // suspend
- scheduler.closeAsync();
+ // suspend
+ scheduler.closeAsync();
- assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
- validateCancelRpcCalls(gateway, parallelism);
+ assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ validateCancelRpcCalls(gateway, parallelism);
+ });
ensureCannotLeaveSuspendedState(scheduler, gateway);
}
@@ -132,20 +146,23 @@ class ExecutionGraphSuspendTest {
final InteractionsCountingTaskManagerGateway gateway =
new InteractionsCountingTaskManagerGateway(parallelism);
final SchedulerBase scheduler = createScheduler(gateway, parallelism);
- final ExecutionGraph eg = scheduler.getExecutionGraph();
+ runInMainThread(
+ () -> {
+ final ExecutionGraph eg = scheduler.getExecutionGraph();
- scheduler.startScheduling();
- ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
+ scheduler.startScheduling();
+ ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
- scheduler.handleGlobalFailure(new Exception("fail global"));
+ scheduler.handleGlobalFailure(new Exception("fail global"));
- assertThat(eg.getState()).isEqualTo(JobStatus.FAILING);
- validateCancelRpcCalls(gateway, parallelism);
+ assertThat(eg.getState()).isEqualTo(JobStatus.FAILING);
+ validateCancelRpcCalls(gateway, parallelism);
- // suspend
- scheduler.closeAsync();
+ // suspend
+ scheduler.closeAsync();
- assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ });
ensureCannotLeaveSuspendedState(scheduler, gateway);
}
@@ -156,24 +173,27 @@ class ExecutionGraphSuspendTest {
new InteractionsCountingTaskManagerGateway();
final int parallelism = 10;
final SchedulerBase scheduler = createScheduler(gateway, parallelism);
- final ExecutionGraph eg = scheduler.getExecutionGraph();
+ runInMainThread(
+ () -> {
+ final ExecutionGraph eg = scheduler.getExecutionGraph();
- scheduler.startScheduling();
- ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
+ scheduler.startScheduling();
+ ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
- scheduler.handleGlobalFailure(new Exception("fail global"));
+ scheduler.handleGlobalFailure(new Exception("fail global"));
- assertThat(eg.getState()).isEqualTo(JobStatus.FAILING);
- validateCancelRpcCalls(gateway, parallelism);
+ assertThat(eg.getState()).isEqualTo(JobStatus.FAILING);
+ validateCancelRpcCalls(gateway, parallelism);
- ExecutionGraphTestUtils.completeCancellingForAllVertices(eg);
- assertThat(eg.getState()).isEqualTo(JobStatus.FAILED);
+ ExecutionGraphTestUtils.completeCancellingForAllVertices(eg);
+ assertThat(eg.getState()).isEqualTo(JobStatus.FAILED);
- // suspend
- scheduler.closeAsync();
+ // suspend
+ scheduler.closeAsync();
- // still in failed state
- assertThat(eg.getState()).isEqualTo(JobStatus.FAILED);
+ // still in failed state
+ assertThat(eg.getState()).isEqualTo(JobStatus.FAILED);
+ });
validateCancelRpcCalls(gateway, parallelism);
}
@@ -184,20 +204,23 @@ class ExecutionGraphSuspendTest {
final InteractionsCountingTaskManagerGateway gateway =
new InteractionsCountingTaskManagerGateway(parallelism);
final SchedulerBase scheduler = createScheduler(gateway, parallelism);
- final ExecutionGraph eg = scheduler.getExecutionGraph();
+ runInMainThread(
+ () -> {
+ final ExecutionGraph eg = scheduler.getExecutionGraph();
- scheduler.startScheduling();
- ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
+ scheduler.startScheduling();
+ ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
- scheduler.cancel();
+ scheduler.cancel();
- assertThat(eg.getState()).isEqualTo(JobStatus.CANCELLING);
- validateCancelRpcCalls(gateway, parallelism);
+ assertThat(eg.getState()).isEqualTo(JobStatus.CANCELLING);
+ validateCancelRpcCalls(gateway, parallelism);
- // suspend
- scheduler.closeAsync();
+ // suspend
+ scheduler.closeAsync();
- assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ });
ensureCannotLeaveSuspendedState(scheduler, gateway);
}
@@ -209,27 +232,30 @@ class ExecutionGraphSuspendTest {
new InteractionsCountingTaskManagerGateway();
final int parallelism = 10;
final SchedulerBase scheduler = createScheduler(gateway, parallelism);
- final ExecutionGraph eg = scheduler.getExecutionGraph();
+ runInMainThread(
+ () -> {
+ final ExecutionGraph eg = scheduler.getExecutionGraph();
- scheduler.startScheduling();
- ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
+ scheduler.startScheduling();
+ ExecutionGraphTestUtils.switchAllVerticesToRunning(eg);
- scheduler.cancel();
+ scheduler.cancel();
- assertThat(eg.getState()).isEqualTo(JobStatus.CANCELLING);
- validateCancelRpcCalls(gateway, parallelism);
+ assertThat(eg.getState()).isEqualTo(JobStatus.CANCELLING);
+ validateCancelRpcCalls(gateway, parallelism);
- ExecutionGraphTestUtils.completeCancellingForAllVertices(eg);
+ ExecutionGraphTestUtils.completeCancellingForAllVertices(eg);
- FlinkAssertions.assertThatFuture(eg.getTerminationFuture())
- .eventuallySucceeds()
- .isEqualTo(JobStatus.CANCELED);
+ FlinkAssertions.assertThatFuture(eg.getTerminationFuture())
+ .eventuallySucceeds()
+ .isEqualTo(JobStatus.CANCELED);
- // suspend
- scheduler.closeAsync();
+ // suspend
+ scheduler.closeAsync();
- // still in failed state
- assertThat(eg.getState()).isEqualTo(JobStatus.CANCELED);
+ // still in failed state
+ assertThat(eg.getState()).isEqualTo(JobStatus.CANCELED);
+ });
validateCancelRpcCalls(gateway, parallelism);
}
@@ -278,24 +304,28 @@ class ExecutionGraphSuspendTest {
final ExecutionGraph eg = scheduler.getExecutionGraph();
gateway.waitUntilAllTasksAreSubmitted();
- assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
- gateway.resetCounts();
- scheduler.handleGlobalFailure(new Exception("fail"));
- assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
- validateNoInteractions(gateway);
+ runInMainThread(
+ () -> {
+ assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ gateway.resetCounts();
- scheduler.cancel();
- assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
- validateNoInteractions(gateway);
+ scheduler.handleGlobalFailure(new Exception("fail"));
+ assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ validateNoInteractions(gateway);
- scheduler.closeAsync();
- assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
- validateNoInteractions(gateway);
+ scheduler.cancel();
+ assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ validateNoInteractions(gateway);
- for (ExecutionVertex ev : eg.getAllExecutionVertices()) {
- assertThat(ev.getCurrentExecutionAttempt().getAttemptNumber()).isZero();
- }
+ scheduler.closeAsync();
+ assertThat(eg.getState()).isEqualTo(JobStatus.SUSPENDED);
+ validateNoInteractions(gateway);
+
+ for (ExecutionVertex ev : eg.getAllExecutionVertices()) {
+ assertThat(ev.getCurrentExecutionAttempt().getAttemptNumber()).isZero();
+ }
+ });
}
private static void validateNoInteractions(InteractionsCountingTaskManagerGateway gateway) {
@@ -322,7 +352,8 @@ class ExecutionGraphSuspendTest {
final SchedulerBase scheduler =
new DefaultSchedulerBuilder(
JobGraphTestUtils.streamingJobGraph(vertex),
- ComponentMainThreadExecutorServiceAdapter.forMainThread(),
+ ComponentMainThreadExecutorServiceAdapter.forSingleThreadExecutor(
+ JM_MAIN_THREAD_EXECUTOR_RESOURCE.getExecutor()),
EXECUTOR_RESOURCE.getExecutor())
.setExecutionSlotAllocatorFactory(
SchedulerTestingUtils.newSlotSharingExecutionSlotAllocatorFactory(
@@ -332,4 +363,8 @@ class ExecutionGraphSuspendTest {
.build();
return scheduler;
}
+
+ private static void runInMainThread(final Runnable runnable) {
+ CompletableFuture.runAsync(runnable, JM_MAIN_THREAD_EXECUTOR_RESOURCE.getExecutor()).join();
+ }
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionTest.java
index 1e574b7160e..426914b4c21 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/ExecutionTest.java
@@ -18,13 +18,19 @@
package org.apache.flink.runtime.executiongraph;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.runtime.blob.PermanentBlobKey;
+import org.apache.flink.runtime.blob.TestingBlobWriter;
import org.apache.flink.runtime.checkpoint.JobManagerTaskRestore;
import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutorServiceAdapter;
+import org.apache.flink.runtime.execution.ExecutionState;
import org.apache.flink.runtime.executiongraph.utils.SimpleAckingTaskManagerGateway;
+import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobgraph.JobGraphTestUtils;
import org.apache.flink.runtime.jobgraph.JobVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.jobmaster.TestingLogicalSlotBuilder;
import org.apache.flink.runtime.jobmaster.slotpool.PhysicalSlot;
import org.apache.flink.runtime.scheduler.DefaultSchedulerBuilder;
import org.apache.flink.runtime.scheduler.SchedulerBase;
@@ -41,7 +47,10 @@ import org.junit.jupiter.api.extension.RegisterExtension;
import javax.annotation.Nonnull;
+import java.io.IOException;
+import java.io.InputStream;
import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ScheduledExecutorService;
import static org.assertj.core.api.Assertions.assertThat;
@@ -121,7 +130,7 @@ class ExecutionTest {
final SchedulerBase scheduler =
new DefaultSchedulerBuilder(
JobGraphTestUtils.streamingJobGraph(jobVertex),
- ComponentMainThreadExecutorServiceAdapter.forMainThread(),
+ testMainThreadUtil.getMainThreadExecutor(),
EXECUTOR_RESOURCE.getExecutor())
.setExecutionSlotAllocatorFactory(
SchedulerTestingUtils.newSlotSharingExecutionSlotAllocatorFactory(
@@ -142,9 +151,14 @@ class ExecutionTest {
assertThat(execution.getTaskRestore()).isNotNull();
// schedule the execution vertex and wait for its deployment
- scheduler.startScheduling();
+ testMainThreadUtil.execute(scheduler::startScheduling);
- assertThat(execution.getTaskRestore()).isNull();
+ // After the execution has been deployed, the task restore state should be nulled.
+ while (execution.getTaskRestore() != null) {
+ // Using busy waiting because there is no `future` that would indicate the completion
+ // of the deployment and I want to keep production code clean.
+ Thread.sleep(10);
+ }
}
@Test
@@ -235,6 +249,102 @@ class ExecutionTest {
});
}
+ @Test
+ void testExecutionCancelledDuringTaskRestoreOffload() {
+ final CountDownLatch offloadLatch = new CountDownLatch(1);
+
+ final Execution cancelledExecution =
+ testMainThreadUtil.execute(
+ () -> {
+ final ConditionalLatchBlockingBlobWriter blobWriter =
+ new ConditionalLatchBlockingBlobWriter(offloadLatch);
+
+ final JobVertex jobVertex = createNoOpJobVertex();
+ final JobGraph jobGraph =
+ JobGraphTestUtils.streamingJobGraph(jobVertex);
+ final ExecutionGraph executionGraph =
+ TestingDefaultExecutionGraphBuilder.newBuilder()
+ .setJobGraph(jobGraph)
+ .setBlobWriter(blobWriter)
+ .build(EXECUTOR_RESOURCE.getExecutor());
+ executionGraph.start(testMainThreadUtil.getMainThreadExecutor());
+
+ final ExecutionJobVertex executionJobVertex =
+ executionGraph.getJobVertex(jobVertex.getID());
+ final ExecutionVertex executionVertex =
+ executionJobVertex.getTaskVertices()[0];
+ final Execution execution =
+ executionVertex.getCurrentExecutionAttempt();
+
+ final JobManagerTaskRestore taskRestoreState =
+ new JobManagerTaskRestore(1L, new TaskStateSnapshot());
+ execution.setInitialState(taskRestoreState);
+ execution.tryAssignResource(
+ new TestingLogicalSlotBuilder().createTestingLogicalSlot());
+ execution.transitionState(ExecutionState.SCHEDULED);
+
+ blobWriter.enableBlocking();
+ execution.deploy();
+ execution.cancel();
+
+ return execution;
+ });
+
+ offloadLatch.countDown();
+
+ cancelledExecution
+ .getTddCreationDuringDeployFuture()
+ .handle(
+ (result, exception) -> {
+ assertThat(exception)
+ .isNotNull()
+ .describedAs("Expected IllegalStateException to be thrown");
+
+ final Throwable rootCause = exception.getCause();
+ assertThat(rootCause).isInstanceOf(IllegalStateException.class);
+ assertThat(rootCause.getMessage())
+ .contains("execution state has switched");
+ assertThat(rootCause.getMessage())
+ .contains("during task restore offload");
+ return null;
+ })
+ .join();
+ }
+
+ /**
+ * Custom BlobWriter that conditionally blocks putPermanent calls on a latch to simulate slow
+ * offloading.
+ */
+ private static class ConditionalLatchBlockingBlobWriter extends TestingBlobWriter {
+ private final CountDownLatch offloadLatch;
+ private volatile boolean blockingEnabled = false;
+
+ public ConditionalLatchBlockingBlobWriter(CountDownLatch offloadLatch) {
+ super();
+ this.offloadLatch = offloadLatch;
+ }
+
+ public void enableBlocking() {
+ this.blockingEnabled = true;
+ }
+
+ @Override
+ public PermanentBlobKey putPermanent(JobID jobId, InputStream inputStream)
+ throws IOException {
+ if (blockingEnabled) {
+ try {
+ // Block until the latch is released
+ offloadLatch.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException("Interrupted while waiting for offload latch", e);
+ }
+ }
+ // Delegate to parent implementation
+ return super.putPermanent(jobId, inputStream);
+ }
+ }
+
@Nonnull
private JobVertex createNoOpJobVertex() {
final JobVertex jobVertex = new JobVertex("Test vertex", new JobVertexID());
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/InteractionsCountingTaskManagerGateway.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/InteractionsCountingTaskManagerGateway.java
index abc4b534401..afddf41263c 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/InteractionsCountingTaskManagerGateway.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/InteractionsCountingTaskManagerGateway.java
@@ -33,20 +33,21 @@ class InteractionsCountingTaskManagerGateway extends SimpleAckingTaskManagerGate
private final AtomicInteger submitTaskCount = new AtomicInteger(0);
- private CountDownLatch submitLatch;
+ private final CountDownLatch submitOrCancelLatch;
public InteractionsCountingTaskManagerGateway() {
- submitLatch = new CountDownLatch(0);
+ submitOrCancelLatch = new CountDownLatch(0);
}
public InteractionsCountingTaskManagerGateway(final int expectedSubmitCount) {
- this.submitLatch = new CountDownLatch(expectedSubmitCount);
+ this.submitOrCancelLatch = new CountDownLatch(expectedSubmitCount);
}
@Override
public CompletableFuture<Acknowledge> cancelTask(
ExecutionAttemptID executionAttemptID, Duration timeout) {
cancelTaskCount.incrementAndGet();
+ submitOrCancelLatch.countDown();
return CompletableFuture.completedFuture(Acknowledge.get());
}
@@ -54,7 +55,7 @@ class InteractionsCountingTaskManagerGateway extends SimpleAckingTaskManagerGate
public CompletableFuture<Acknowledge> submitTask(
TaskDeploymentDescriptor tdd, Duration timeout) {
submitTaskCount.incrementAndGet();
- submitLatch.countDown();
+ submitOrCancelLatch.countDown();
return CompletableFuture.completedFuture(Acknowledge.get());
}
@@ -77,7 +78,7 @@ class InteractionsCountingTaskManagerGateway extends SimpleAckingTaskManagerGate
void waitUntilAllTasksAreSubmitted() {
try {
- submitLatch.await();
+ submitOrCancelLatch.await();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/utils/ExecutionUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/utils/ExecutionUtils.java
new file mode 100644
index 00000000000..fc3f31b881f
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/executiongraph/utils/ExecutionUtils.java
@@ -0,0 +1,41 @@
+/*
+ * 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.executiongraph.utils;
+
+import org.apache.flink.runtime.executiongraph.ExecutionVertex;
+
+public final class ExecutionUtils {
+
+ private ExecutionUtils() {
+ // Private constructor for utility class.
+ }
+
+ /**
+ * Waits for the task deployment descriptors to be created for the given execution vertices.
+ *
+ * @param executionVertices the execution vertices to wait for
+ */
+ public static void waitForTaskDeploymentDescriptorsCreation(
+ final ExecutionVertex... executionVertices) {
+ for (ExecutionVertex ev : executionVertices) {
+ ev.getTddCreationDuringDeployFuture().join();
+ }
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/UpdatePartitionConsumersTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/UpdatePartitionConsumersTest.java
index 569d6be7856..fdf76e59672 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/UpdatePartitionConsumersTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/UpdatePartitionConsumersTest.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.jobmanager.scheduler;
+import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutor;
import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutorServiceAdapter;
import org.apache.flink.runtime.deployment.InputGateDeploymentDescriptor;
import org.apache.flink.runtime.deployment.TaskDeploymentDescriptor;
@@ -49,10 +50,12 @@ import org.junit.Test;
import java.util.List;
import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
+import static org.apache.flink.runtime.executiongraph.utils.ExecutionUtils.waitForTaskDeploymentDescriptorsCreation;
import static org.apache.flink.runtime.util.JobVertexConnectionUtils.connectNewDataSetAsInput;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.instanceOf;
@@ -132,10 +135,17 @@ public class UpdatePartitionConsumersTest extends TestLogger {
final SimpleAckingTaskManagerGateway taskManagerGateway =
new SimpleAckingTaskManagerGateway();
+ final ScheduledExecutorService mainThreadExecutorService =
+ Executors.newSingleThreadScheduledExecutor();
+
+ final ComponentMainThreadExecutor singleThreadMainThreadExecutor =
+ ComponentMainThreadExecutorServiceAdapter.forSingleThreadExecutor(
+ mainThreadExecutorService);
+
final SchedulerBase scheduler =
new DefaultSchedulerBuilder(
jobGraph,
- ComponentMainThreadExecutorServiceAdapter.forMainThread(),
+ singleThreadMainThreadExecutor,
EXECUTOR_RESOURCE.getExecutor())
.setExecutionSlotAllocatorFactory(
new TestExecutionSlotAllocatorFactory(taskManagerGateway))
@@ -159,21 +169,27 @@ public class UpdatePartitionConsumersTest extends TestLogger {
}
});
- scheduler.startScheduling();
+ CompletableFuture.runAsync(scheduler::startScheduling, singleThreadMainThreadExecutor)
+ .join();
+
+ waitForTaskDeploymentDescriptorsCreation(ev1, ev2, ev3, ev4);
assertThat(ev1.getExecutionState(), is(ExecutionState.DEPLOYING));
assertThat(ev2.getExecutionState(), is(ExecutionState.DEPLOYING));
assertThat(ev3.getExecutionState(), is(ExecutionState.DEPLOYING));
assertThat(ev4.getExecutionState(), is(ExecutionState.DEPLOYING));
- updateState(scheduler, ev1, ExecutionState.INITIALIZING);
- updateState(scheduler, ev1, ExecutionState.RUNNING);
- updateState(scheduler, ev2, ExecutionState.INITIALIZING);
- updateState(scheduler, ev2, ExecutionState.RUNNING);
- updateState(scheduler, ev3, ExecutionState.INITIALIZING);
- updateState(scheduler, ev3, ExecutionState.RUNNING);
- updateState(scheduler, ev4, ExecutionState.INITIALIZING);
- updateState(scheduler, ev4, ExecutionState.RUNNING);
+ mainThreadExecutorService.execute(
+ () -> {
+ updateState(scheduler, ev1, ExecutionState.INITIALIZING);
+ updateState(scheduler, ev1, ExecutionState.RUNNING);
+ updateState(scheduler, ev2, ExecutionState.INITIALIZING);
+ updateState(scheduler, ev2, ExecutionState.RUNNING);
+ updateState(scheduler, ev3, ExecutionState.INITIALIZING);
+ updateState(scheduler, ev3, ExecutionState.RUNNING);
+ updateState(scheduler, ev4, ExecutionState.INITIALIZING);
+ updateState(scheduler, ev4, ExecutionState.RUNNING);
+ });
final InputGateDeploymentDescriptor ev4Igdd2 =
ev4TddFuture.get(TIMEOUT, TimeUnit.MILLISECONDS).getInputGates().get(1);
@@ -196,8 +212,11 @@ public class UpdatePartitionConsumersTest extends TestLogger {
updatePartitionFuture.complete(null);
});
- updateState(scheduler, ev1, ExecutionState.FINISHED);
- updateState(scheduler, ev3, ExecutionState.FINISHED);
+ mainThreadExecutorService.execute(
+ () -> {
+ updateState(scheduler, ev1, ExecutionState.FINISHED);
+ updateState(scheduler, ev3, ExecutionState.FINISHED);
+ });
updatePartitionFuture.get(TIMEOUT, TimeUnit.MILLISECONDS);
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/SchedulerTestingUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/SchedulerTestingUtils.java
index 135e0737fa3..0b53800e159 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/SchedulerTestingUtils.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/SchedulerTestingUtils.java
@@ -63,6 +63,7 @@ import org.apache.flink.runtime.state.StateBackend;
import org.apache.flink.runtime.taskmanager.TaskExecutionState;
import org.apache.flink.util.SerializedValue;
import org.apache.flink.util.TernaryBoolean;
+import org.apache.flink.util.function.RunnableWithException;
import javax.annotation.Nullable;
@@ -81,6 +82,7 @@ import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import static org.apache.flink.runtime.executiongraph.ExecutionGraphTestUtils.finishJobVertex;
+import static org.apache.flink.runtime.executiongraph.utils.ExecutionUtils.waitForTaskDeploymentDescriptorsCreation;
import static org.apache.flink.runtime.testutils.CommonTestUtils.waitUntilCondition;
import static org.apache.flink.runtime.util.JobVertexConnectionUtils.connectNewDataSetAsInput;
import static org.assertj.core.api.Assertions.assertThat;
@@ -415,35 +417,44 @@ public class SchedulerTestingUtils {
final ExecutionGraph executionGraph = scheduler.getExecutionGraph();
final TestingLogicalSlotBuilder slotBuilder = new TestingLogicalSlotBuilder();
- CompletableFuture.runAsync(
- () -> {
- try {
- if (isAdaptive) {
- initializeExecutionJobVertex(producer.getID(), executionGraph);
- }
- // Deploy upstream source vertices
- deployTasks(executionGraph, producer.getID(), slotBuilder);
- // Transition upstream vertices into FINISHED
- finishJobVertex(executionGraph, producer.getID());
- // Deploy downstream sink vertices
- for (JobVertex consumer : consumers) {
- if (isAdaptive) {
- initializeExecutionJobVertex(
- consumer.getID(), executionGraph);
- }
- deployTasks(executionGraph, consumer.getID(), slotBuilder);
- }
- } catch (Exception e) {
- throw new RuntimeException("Exceptions shouldn't happen here.", e);
- }
- },
- mainThreadExecutor)
- .join();
+ runUnsafe(
+ () -> {
+ initializeExecutionJobVertex(producer.getID(), executionGraph, isAdaptive);
+ deployTasks(executionGraph, producer.getID(), slotBuilder);
+ },
+ mainThreadExecutor);
+
+ waitForTaskDeploymentDescriptorsCreation(
+ Objects.requireNonNull(executionGraph.getJobVertex(producer.getID()))
+ .getTaskVertices());
+
+ // Transition upstream vertices into FINISHED
+ runUnsafe(() -> finishJobVertex(executionGraph, producer.getID()), mainThreadExecutor);
+
+ // Deploy downstream sink vertices
+ for (JobVertex consumer : consumers) {
+ runUnsafe(
+ () -> {
+ initializeExecutionJobVertex(consumer.getID(), executionGraph, isAdaptive);
+ deployTasks(executionGraph, consumer.getID(), slotBuilder);
+ },
+ mainThreadExecutor);
+ waitForTaskDeploymentDescriptorsCreation(
+ Objects.requireNonNull(executionGraph.getJobVertex(consumer.getID()))
+ .getTaskVertices());
+ }
+
return scheduler;
}
private static void initializeExecutionJobVertex(
- JobVertexID jobVertex, ExecutionGraph executionGraph) {
+ JobVertexID jobVertex,
+ ExecutionGraph executionGraph,
+ final boolean adaptiveSchedulerEnabled) {
+ if (!adaptiveSchedulerEnabled) {
+ // This method call only needed for adaptive scheduler, no-op otherwise.
+ return;
+ }
try {
executionGraph.initializeJobVertex(
executionGraph.getJobVertex(jobVertex), System.currentTimeMillis());
@@ -530,4 +541,23 @@ public class SchedulerTestingUtils {
ExecutionAttemptID attemptId) {
return new TaskExecutionState(attemptId, ExecutionState.CANCELED);
}
+
+ private static void runUnsafe(
+ final RunnableWithException callback, final ComponentMainThreadExecutor executor) {
+ try {
+ CompletableFuture.runAsync(
+ () -> {
+ try {
+ callback.run();
+ } catch (Throwable e) {
+ throw new RuntimeException(
+ "Exception shouldn't happen here.", e);
+ }
+ },
+ executor)
+ .join();
+ } catch (Exception e) {
+ throw new RuntimeException("Exception shouldn't happen here.", e);
+ }
+ }
}