[FLINK-38114] Allow asynchronous offloading of TaskRestore

Author: Izeren

flink-core/src/test/java/org/apache/flink/testutils/TestingUtils.java

@@ -31,6 +31,7 @@
 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 static Duration infiniteDuration() {
         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);
     }

flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java

@@ -118,6 +118,12 @@ public Offloaded(PermanentBlobKey serializedValueKey) {
      */
     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 Offloaded(PermanentBlobKey serializedValueKey) {
     /** 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 TaskDeploymentDescriptor(
         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 List<InputGateDeploymentDescriptor> getInputGates() {
     }
 
     @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 void loadBigData(
             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);

flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorFactory.java

@@ -112,7 +112,7 @@ public MaybeOffloaded<JobInformation> getSerializedJobInformation() {
     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 TaskDeploymentDescriptor createDeploymentDescriptor(
                         executionVertex.getJobVertex().getTaskInformationOrBlobKey()),
                 execution.getAttemptId(),
                 allocationID,
-                taskRestore,
+                taskRestore != null ? getSerializedTaskRestore(taskRestore) : null,
                 new ArrayList<>(producedPartitions),
                 createInputGateDeploymentDescriptors(executionVertex));
     }
@@ -295,6 +295,13 @@ private static MaybeOffloaded<TaskInformation> getSerializedTaskInformation(
                 : 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,