diff --git a/docs/layouts/shortcodes/generated/all_jobmanager_section.html b/docs/layouts/shortcodes/generated/all_jobmanager_section.html index 42eb43eea22dd6..961dcbedeffba8 100644 --- a/docs/layouts/shortcodes/generated/all_jobmanager_section.html +++ b/docs/layouts/shortcodes/generated/all_jobmanager_section.html @@ -126,7 +126,7 @@
jobmanager.io-pool.size
(none) Integer - The size of the IO thread pool to run blocking operations for all spawned JobMasters. This includes recovery and completion of checkpoints. Increase this value if you experience slow checkpoint operations when running many jobs. If no value is specified, then Flink defaults to the number of available CPU cores. + The size of the IO thread pool to run blocking operations for all spawned JobMasters. This includes recovery and completion of checkpoints. The same size is used for a separate adaptive graph-recovery pool, shared across jobs, so the two pools together can use twice this number of threads. Asynchronous checkpoint-trigger initialization and cleanup continue to share the original IO pool. Synchronous storage operations during graph construction still use recovery workers. Increase this value if you experience slow checkpoint operations when running many jobs. If no value is specified, then Flink defaults to the number of available CPU cores.
jobmanager.partition.hybrid.partition-data-consume-constraint
diff --git a/docs/layouts/shortcodes/generated/expert_jobmanager_section.html b/docs/layouts/shortcodes/generated/expert_jobmanager_section.html index c882966a31687e..14cd90b1925f63 100644 --- a/docs/layouts/shortcodes/generated/expert_jobmanager_section.html +++ b/docs/layouts/shortcodes/generated/expert_jobmanager_section.html @@ -18,7 +18,7 @@
jobmanager.io-pool.size
(none) Integer - The size of the IO thread pool to run blocking operations for all spawned JobMasters. This includes recovery and completion of checkpoints. Increase this value if you experience slow checkpoint operations when running many jobs. If no value is specified, then Flink defaults to the number of available CPU cores. + The size of the IO thread pool to run blocking operations for all spawned JobMasters. This includes recovery and completion of checkpoints. The same size is used for a separate adaptive graph-recovery pool, shared across jobs, so the two pools together can use twice this number of threads. Asynchronous checkpoint-trigger initialization and cleanup continue to share the original IO pool. Synchronous storage operations during graph construction still use recovery workers. Increase this value if you experience slow checkpoint operations when running many jobs. If no value is specified, then Flink defaults to the number of available CPU cores. diff --git a/docs/layouts/shortcodes/generated/job_manager_configuration.html b/docs/layouts/shortcodes/generated/job_manager_configuration.html index 1023058a36b599..1f0e275cd97e2d 100644 --- a/docs/layouts/shortcodes/generated/job_manager_configuration.html +++ b/docs/layouts/shortcodes/generated/job_manager_configuration.html @@ -126,7 +126,7 @@
jobmanager.io-pool.size
(none) Integer - The size of the IO thread pool to run blocking operations for all spawned JobMasters. This includes recovery and completion of checkpoints. Increase this value if you experience slow checkpoint operations when running many jobs. If no value is specified, then Flink defaults to the number of available CPU cores. + The size of the IO thread pool to run blocking operations for all spawned JobMasters. This includes recovery and completion of checkpoints. The same size is used for a separate adaptive graph-recovery pool, shared across jobs, so the two pools together can use twice this number of threads. Asynchronous checkpoint-trigger initialization and cleanup continue to share the original IO pool. Synchronous storage operations during graph construction still use recovery workers. Increase this value if you experience slow checkpoint operations when running many jobs. If no value is specified, then Flink defaults to the number of available CPU cores.
jobmanager.memory.enable-jvm-direct-memory-limit
diff --git a/flink-core/src/main/java/org/apache/flink/configuration/JobManagerOptions.java b/flink-core/src/main/java/org/apache/flink/configuration/JobManagerOptions.java index 847678a9e45dd4..c2801aeeccb7c0 100644 --- a/flink-core/src/main/java/org/apache/flink/configuration/JobManagerOptions.java +++ b/flink-core/src/main/java/org/apache/flink/configuration/JobManagerOptions.java @@ -504,6 +504,10 @@ public enum ArchivedApplicationStoreType { .withDescription( "The size of the IO thread pool to run blocking operations for all spawned JobMasters. " + "This includes recovery and completion of checkpoints. " + + "The same size is used for a separate adaptive graph-recovery pool, " + + "shared across jobs, so the two pools together can use twice this number of threads. " + + "Asynchronous checkpoint-trigger initialization and cleanup continue to share the original IO pool. " + + "Synchronous storage operations during graph construction still use recovery workers. " + "Increase this value if you experience slow checkpoint operations when running many jobs. " + "If no value is specified, then Flink defaults to the number of available CPU cores."); diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/DefaultSlotPoolServiceSchedulerFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/DefaultSlotPoolServiceSchedulerFactory.java index a49942b52c15ac..6cffad2545a1d4 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/DefaultSlotPoolServiceSchedulerFactory.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/DefaultSlotPoolServiceSchedulerFactory.java @@ -111,6 +111,7 @@ public SchedulerNG createScheduler( Logger log, ExecutionPlan executionPlan, Executor ioExecutor, + Executor recoveryExecutor, Configuration configuration, SlotPoolService slotPoolService, ScheduledExecutorService futureExecutor, @@ -134,6 +135,7 @@ public SchedulerNG createScheduler( log, executionPlan, ioExecutor, + recoveryExecutor, configuration, slotPoolService, futureExecutor, diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobManagerSharedServices.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobManagerSharedServices.java index 5dff09f5dfc96d..005433823b6735 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobManagerSharedServices.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobManagerSharedServices.java @@ -59,6 +59,8 @@ public class JobManagerSharedServices { private final ExecutorService ioExecutor; + private final ExecutorService recoveryExecutor; + private final LibraryCacheManager libraryCacheManager; private final ShuffleMaster shuffleMaster; @@ -71,9 +73,25 @@ public JobManagerSharedServices( LibraryCacheManager libraryCacheManager, ShuffleMaster shuffleMaster, @Nonnull BlobWriter blobWriter) { + this( + futureExecutor, + ioExecutor, + ioExecutor, + libraryCacheManager, + shuffleMaster, + blobWriter); + } + public JobManagerSharedServices( + ScheduledExecutorService futureExecutor, + ExecutorService ioExecutor, + ExecutorService recoveryExecutor, + LibraryCacheManager libraryCacheManager, + ShuffleMaster shuffleMaster, + @Nonnull BlobWriter blobWriter) { this.futureExecutor = checkNotNull(futureExecutor); this.ioExecutor = checkNotNull(ioExecutor); + this.recoveryExecutor = checkNotNull(recoveryExecutor); this.libraryCacheManager = checkNotNull(libraryCacheManager); this.shuffleMaster = checkNotNull(shuffleMaster); this.blobWriter = blobWriter; @@ -87,6 +105,15 @@ public Executor getIoExecutor() { return ioExecutor; } + /** + * Shared capacity for adaptive graph creation and state assignment, including synchronous + * storage calls within that work. Asynchronous checkpoint-trigger initialization and cleanup + * must continue to use the general IO executor to preserve their existing feedback. + */ + public Executor getRecoveryExecutor() { + return recoveryExecutor; + } + public LibraryCacheManager getLibraryCacheManager() { return libraryCacheManager; } @@ -113,8 +140,12 @@ public void shutdown() throws Exception { Throwable exception = null; try { + final ExecutorService[] executors = + recoveryExecutor == ioExecutor + ? new ExecutorService[] {futureExecutor, ioExecutor} + : new ExecutorService[] {futureExecutor, ioExecutor, recoveryExecutor}; ExecutorUtils.gracefulShutdown( - SHUTDOWN_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS, futureExecutor, ioExecutor); + SHUTDOWN_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS, executors); } catch (Throwable t) { exception = t; } @@ -183,7 +214,19 @@ public static JobManagerSharedServices fromConfiguration( .createShuffleMaster(shuffleMasterContext); shuffleMaster.start(); + // One fixed pool per JobManager, not per job or graph. The queue remains unbounded; + // BackgroundTask serializes suppliers within each adaptive scheduler. This separates + // worker availability, not storage clients or synchronous work inside a supplier. + final ExecutorService recoveryExecutor = + Executors.newFixedThreadPool( + jobManagerIoPoolSize, new ExecutorThreadFactory("jobmanager-recovery")); + return new JobManagerSharedServices( - futureExecutor, ioExecutor, libraryCacheManager, shuffleMaster, blobServer); + futureExecutor, + ioExecutor, + recoveryExecutor, + libraryCacheManager, + shuffleMaster, + blobServer); } } diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java index 39a29191b9761b..97b46365ffc5f1 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/JobMaster.java @@ -178,6 +178,8 @@ public class JobMaster extends FencedRpcEndpoint private final Executor ioExecutor; + private final Executor recoveryExecutor; + private final OnCompletionActions jobCompletionActions; private final FatalErrorHandler fatalErrorHandler; @@ -346,6 +348,8 @@ public void onUnknownDeploymentsOf( this.futureExecutor = MdcUtils.scopeToJob(jid, jobManagerSharedServices.getFutureExecutor()); this.ioExecutor = MdcUtils.scopeToJob(jid, jobManagerSharedServices.getIoExecutor()); + this.recoveryExecutor = + MdcUtils.scopeToJob(jid, jobManagerSharedServices.getRecoveryExecutor()); this.jobCompletionActions = checkNotNull(jobCompletionActions); this.fatalErrorHandler = checkNotNull(fatalErrorHandler); this.userCodeLoader = checkNotNull(userCodeLoader); @@ -431,6 +435,7 @@ private SchedulerNG createScheduler( log, executionPlan, ioExecutor, + recoveryExecutor, jobMasterConfiguration.getConfiguration(), slotPoolService, futureExecutor, diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/SlotPoolServiceSchedulerFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/SlotPoolServiceSchedulerFactory.java index d85f9b96552d38..ee5c51c8c3c385 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/SlotPoolServiceSchedulerFactory.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/SlotPoolServiceSchedulerFactory.java @@ -71,7 +71,9 @@ SlotPoolService createSlotPoolService( JobManagerOptions.SchedulerType getSchedulerType(); /** - * Creates a {@link SchedulerNG}. + * Creates a {@link SchedulerNG}. All executors are borrowed. The recovery executor is only for + * adaptive graph creation and state assignment, not asynchronous checkpoint-trigger + * initialization or cleanup. * * @return created SchedulerNG * @throws Exception if the scheduler creation fails @@ -80,6 +82,7 @@ SchedulerNG createScheduler( Logger log, ExecutionPlan executionPlan, Executor ioExecutor, + Executor recoveryExecutor, Configuration configuration, SlotPoolService slotPoolService, ScheduledExecutorService futureExecutor, diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultSchedulerFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultSchedulerFactory.java index ef02a90a875e87..7824baf89fe076 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultSchedulerFactory.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultSchedulerFactory.java @@ -65,6 +65,7 @@ public SchedulerNG createInstance( final Logger log, final ExecutionPlan executionPlan, final Executor ioExecutor, + final Executor recoveryExecutor, final Configuration jobMasterConfiguration, final SlotPoolService slotPoolService, final ScheduledExecutorService futureExecutor, diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SchedulerNGFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SchedulerNGFactory.java index 605870593ff241..e0ed44f411ec45 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SchedulerNGFactory.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SchedulerNGFactory.java @@ -45,10 +45,16 @@ /** Factory for {@link SchedulerNG}. */ public interface SchedulerNGFactory { + /** + * Creates a scheduler borrowing the supplied executors. The recovery executor is reserved for + * adaptive graph creation and state assignment; it must not replace the graph/coordinator IO + * executor or be shut down by the scheduler. Other scheduler types do not use it. + */ SchedulerNG createInstance( Logger log, ExecutionPlan executionPlan, Executor ioExecutor, + Executor recoveryExecutor, Configuration jobMasterConfiguration, SlotPoolService slotPoolService, ScheduledExecutorService futureExecutor, 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 b75f0d09f29e1e..cbc480273fd10b 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 @@ -421,6 +421,7 @@ public JobRescaleConfigInfo toJobRescaleConfigInfo() { private final long initializationTimestamp; private final Executor ioExecutor; + private final Executor recoveryExecutor; private final ClassLoader userCodeClassLoader; private final CheckpointsCleaner checkpointsCleaner; @@ -479,6 +480,7 @@ public AdaptiveScheduler( DeclarativeSlotPool declarativeSlotPool, SlotAllocator slotAllocator, Executor ioExecutor, + Executor recoveryExecutor, ClassLoader userCodeClassLoader, CheckpointsCleaner checkpointsCleaner, CheckpointRecoveryFactory checkpointRecoveryFactory, @@ -506,6 +508,7 @@ public AdaptiveScheduler( declarativeSlotPool, slotAllocator, ioExecutor, + recoveryExecutor, userCodeClassLoader, checkpointsCleaner, checkpointRecoveryFactory, @@ -531,6 +534,7 @@ public AdaptiveScheduler( DeclarativeSlotPool declarativeSlotPool, SlotAllocator slotAllocator, Executor ioExecutor, + Executor recoveryExecutor, ClassLoader userCodeClassLoader, CheckpointsCleaner checkpointsCleaner, CheckpointRecoveryFactory checkpointRecoveryFactory, @@ -573,6 +577,7 @@ public AdaptiveScheduler( this.declarativeSlotPool = declarativeSlotPool; this.initializationTimestamp = initializationTimestamp; this.ioExecutor = ioExecutor; + this.recoveryExecutor = Preconditions.checkNotNull(recoveryExecutor); this.userCodeClassLoader = userCodeClassLoader; this.restartBackoffTimeStrategy = restartBackoffTimeStrategy; this.fatalErrorHandler = fatalErrorHandler; @@ -1595,7 +1600,7 @@ private CompletableFuture createExecutionGraphAndRestoreStateAsy backgroundTask = backgroundTask.runAfter( () -> createExecutionGraphAndRestoreState(adjustedParallelismStore), - ioExecutor); + recoveryExecutor); return FutureUtils.switchExecutor( backgroundTask.getResultFuture(), getMainThreadExecutor()); diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerFactory.java index a6e06aa5ac2009..090f7bdb98331b 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerFactory.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerFactory.java @@ -69,6 +69,7 @@ public SchedulerNG createInstance( Logger log, ExecutionPlan executionPlan, Executor ioExecutor, + Executor recoveryExecutor, Configuration jobMasterConfiguration, SlotPoolService slotPoolService, ScheduledExecutorService futureExecutor, @@ -150,6 +151,7 @@ public SchedulerNG createInstance( declarativeSlotPool, slotAllocator, ioExecutor, + recoveryExecutor, userCodeLoader, new CheckpointsCleaner(), checkpointRecoveryFactory, diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptivebatch/AdaptiveBatchSchedulerFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptivebatch/AdaptiveBatchSchedulerFactory.java index 0970ffc0cbfd6d..e6220724e1217f 100644 --- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptivebatch/AdaptiveBatchSchedulerFactory.java +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptivebatch/AdaptiveBatchSchedulerFactory.java @@ -103,6 +103,7 @@ public SchedulerNG createInstance( Logger log, ExecutionPlan executionPlan, Executor ioExecutor, + Executor recoveryExecutor, Configuration jobMasterConfiguration, SlotPoolService slotPoolService, ScheduledExecutorService futureExecutor, diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobManagerSharedServicesTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobManagerSharedServicesTest.java index 1777edf0cd0767..5e2017a927504a 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobManagerSharedServicesTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobManagerSharedServicesTest.java @@ -23,21 +23,38 @@ import org.apache.flink.core.testutils.OneShotLatch; import org.apache.flink.runtime.blob.BlobServer; import org.apache.flink.runtime.blob.VoidBlobStore; +import org.apache.flink.runtime.blob.VoidBlobWriter; +import org.apache.flink.runtime.execution.librarycache.TestingLibraryCacheManager; +import org.apache.flink.runtime.shuffle.TestingShuffleMaster; import org.apache.flink.runtime.util.Hardware; import org.apache.flink.runtime.util.TestingFatalErrorHandler; import org.apache.flink.util.function.ThrowingRunnable; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import javax.annotation.Nonnull; import java.io.File; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import static org.apache.flink.runtime.testutils.CommonTestUtils.waitUntilCondition; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** * Tests auxiliary shared services created by {@link JobManagerSharedServices} and used by the @@ -49,6 +66,15 @@ class JobManagerSharedServicesTest { @TempDir private File TEMPORARY_FOLDER; + private final List blobServers = new ArrayList<>(); + + @AfterEach + void closeBlobServers() throws Exception { + for (BlobServer blobServer : blobServers) { + blobServer.close(); + } + } + @Test void testFutureExecutorNoConfiguration() throws Exception { final Configuration config = new Configuration(); @@ -74,9 +100,11 @@ void testFutureExecutorConfiguration() throws Exception { final JobManagerSharedServices jobManagerSharedServices = buildJobManagerSharedServices(config); - assertExecutorPoolSize(jobManagerSharedServices.getFutureExecutor(), futurePoolSize); - - jobManagerSharedServices.shutdown(); + try { + assertExecutorPoolSize(jobManagerSharedServices.getFutureExecutor(), futurePoolSize); + } finally { + jobManagerSharedServices.shutdown(); + } } @Test @@ -109,13 +137,306 @@ void testIoExecutorConfiguration() throws Exception { } } + @Test + void testRecoveryExecutorDefaultsToCpuCount() throws Exception { + final JobManagerSharedServices services = + buildJobManagerSharedServices(new Configuration()); + try { + assertExecutorPoolSize(services.getRecoveryExecutor(), CPU_CORES); + } finally { + services.shutdown(); + } + } + + @ParameterizedTest + @ValueSource(ints = {1, 2}) + void testRecoveryAndGeneralIoHaveIndependentFixedCapacity(int poolSize) throws Exception { + final Configuration configuration = new Configuration(); + configuration.set(JobManagerOptions.JOB_MANAGER_IO_POOL_SIZE, poolSize); + configuration.set(JobManagerOptions.JOB_MANAGER_FUTURE_POOL_SIZE, 1); + final JobManagerSharedServices services = buildJobManagerSharedServices(configuration); + try { + final ThreadPoolExecutor recovery = (ThreadPoolExecutor) services.getRecoveryExecutor(); + final ThreadPoolExecutor io = (ThreadPoolExecutor) services.getIoExecutor(); + assertThat(recovery).isNotSameAs(io).isNotSameAs(services.getFutureExecutor()); + assertThat(recovery.getCorePoolSize()).isEqualTo(poolSize); + assertThat(recovery.getMaximumPoolSize()).isEqualTo(poolSize); + assertThat(io.getMaximumPoolSize()).isEqualTo(poolSize); + assertIndependentCapacity(recovery, io, services.getFutureExecutor(), poolSize); + assertIndependentCapacity(io, recovery, services.getFutureExecutor(), poolSize); + assertThat( + CompletableFuture.supplyAsync( + () -> Thread.currentThread().getName(), recovery) + .get(10, TimeUnit.SECONDS)) + .startsWith("jobmanager-recovery"); + } finally { + services.shutdown(); + } + } + + @Test + void testShutdownDrainsAcceptedWorkAndRejectsNewRecoveryWork() throws Exception { + final Configuration configuration = new Configuration(); + configuration.set(JobManagerOptions.JOB_MANAGER_IO_POOL_SIZE, 1); + final JobManagerSharedServices services = buildJobManagerSharedServices(configuration); + final CountDownLatch started = new CountDownLatch(2); + final CountDownLatch release = new CountDownLatch(1); + final AtomicInteger completed = new AtomicInteger(); + final List> tasks = new ArrayList<>(); + CompletableFuture shutdown = null; + try { + for (Executor executor : + new Executor[] {services.getIoExecutor(), services.getRecoveryExecutor()}) { + tasks.add( + CompletableFuture.runAsync( + ThrowingRunnable.unchecked( + () -> { + started.countDown(); + assertThat(release.await(30, TimeUnit.SECONDS)) + .isTrue(); + completed.incrementAndGet(); + }), + executor)); + tasks.add(CompletableFuture.runAsync(completed::incrementAndGet, executor)); + } + assertThat(started.await(10, TimeUnit.SECONDS)).isTrue(); + shutdown = CompletableFuture.runAsync(ThrowingRunnable.unchecked(services::shutdown)); + waitUntilCondition( + () -> ((ExecutorService) services.getRecoveryExecutor()).isShutdown(), + 10L, + 1000); + assertThat(shutdown).isNotDone(); + final AtomicInteger rejectedWork = new AtomicInteger(); + assertThatThrownBy( + () -> + services.getRecoveryExecutor() + .execute(rejectedWork::incrementAndGet)) + .isInstanceOf(RejectedExecutionException.class); + assertThat(rejectedWork).hasValue(0); + release.countDown(); + shutdown.get(10, TimeUnit.SECONDS); + CompletableFuture.allOf(tasks.toArray(new CompletableFuture[0])) + .get(10, TimeUnit.SECONDS); + assertThat(completed).hasValue(4); + assertThat(((ExecutorService) services.getRecoveryExecutor()).isTerminated()).isTrue(); + assertThat(((ExecutorService) services.getIoExecutor()).isTerminated()).isTrue(); + assertThat(services.getFutureExecutor().isTerminated()).isTrue(); + } finally { + release.countDown(); + if (shutdown == null) { + services.shutdown(); + } else { + shutdown.get(10, TimeUnit.SECONDS); + } + } + } + + @Test + void testInterruptedShutdownInterruptsAllPoolsWithoutRunningQueuedRecovery() throws Exception { + final Configuration configuration = new Configuration(); + configuration.set(JobManagerOptions.JOB_MANAGER_IO_POOL_SIZE, 2); + configuration.set(JobManagerOptions.JOB_MANAGER_FUTURE_POOL_SIZE, 1); + final JobManagerSharedServices services = buildJobManagerSharedServices(configuration); + final CountDownLatch started = new CountDownLatch(5); + final CountDownLatch release = new CountDownLatch(1); + final AtomicInteger interrupted = new AtomicInteger(); + final AtomicInteger queuedWork = new AtomicInteger(); + final CompletableFuture shutdownResult = new CompletableFuture<>(); + final Runnable blocked = + () -> { + started.countDown(); + try { + assertThat(release.await(30, TimeUnit.SECONDS)).isTrue(); + } catch (InterruptedException expected) { + interrupted.incrementAndGet(); + Thread.currentThread().interrupt(); + } + }; + final Thread shutdownThread = + new Thread( + () -> { + Thread.currentThread().interrupt(); + try { + services.shutdown(); + shutdownResult.complete(Thread.currentThread().isInterrupted()); + } catch (Throwable error) { + shutdownResult.completeExceptionally(error); + } + }); + try { + services.getFutureExecutor().execute(blocked); + for (int i = 0; i < 2; i++) { + services.getIoExecutor().execute(blocked); + services.getRecoveryExecutor().execute(blocked); + } + assertThat(started.await(10, TimeUnit.SECONDS)).isTrue(); + services.getRecoveryExecutor().execute(queuedWork::incrementAndGet); + shutdownThread.start(); + assertThat(shutdownResult.get(10, TimeUnit.SECONDS)).isTrue(); + assertThat(services.getFutureExecutor().awaitTermination(10, TimeUnit.SECONDS)) + .isTrue(); + assertThat( + ((ExecutorService) services.getIoExecutor()) + .awaitTermination(10, TimeUnit.SECONDS)) + .isTrue(); + assertThat( + ((ExecutorService) services.getRecoveryExecutor()) + .awaitTermination(10, TimeUnit.SECONDS)) + .isTrue(); + assertThat(interrupted).hasValue(5); + assertThat(queuedWork).hasValue(0); + } finally { + release.countDown(); + shutdownThread.join(10000); + services.shutdown(); + } + } + + @Test + void testShutdownStartsAllPoolsBeforeAwaitingTermination() throws Exception { + final List events = new ArrayList<>(); + final RecordingExecutor future = new RecordingExecutor("future", events); + final RecordingExecutor io = new RecordingExecutor("io", events); + final RecordingExecutor recovery = new RecordingExecutor("recovery", events); + final JobManagerSharedServices services = + new JobManagerSharedServices( + future, + io, + recovery, + TestingLibraryCacheManager.newBuilder().build(), + new TestingShuffleMaster(), + VoidBlobWriter.getInstance()); + + try { + services.shutdown(); + + assertThat(events) + .containsExactly( + "future.shutdown", + "io.shutdown", + "recovery.shutdown", + "future.awaitTermination", + "io.awaitTermination", + "recovery.awaitTermination"); + assertThat(future.isTerminated()).isTrue(); + assertThat(io.isTerminated()).isTrue(); + assertThat(recovery.isTerminated()).isTrue(); + } finally { + future.shutdownNow(); + io.shutdownNow(); + recovery.shutdownNow(); + } + } + + @Test + void testLegacyConstructorOwnsIoAliasOnlyOnce() throws Exception { + final List events = new ArrayList<>(); + final RecordingExecutor future = new RecordingExecutor("future", events); + final RecordingExecutor io = new RecordingExecutor("io", events); + final AtomicInteger cacheCloses = new AtomicInteger(); + final AtomicInteger shuffleCloses = new AtomicInteger(); + final JobManagerSharedServices services = + new JobManagerSharedServices( + future, + io, + TestingLibraryCacheManager.newBuilder() + .setShutdownRunnable(cacheCloses::incrementAndGet) + .build(), + new TestingShuffleMaster() { + @Override + public void close() { + shuffleCloses.incrementAndGet(); + } + }, + VoidBlobWriter.getInstance()); + + try { + assertThat(services.getRecoveryExecutor()).isSameAs(io); + CompletableFuture.runAsync(() -> {}, services.getRecoveryExecutor()) + .get(10, TimeUnit.SECONDS); + CompletableFuture.runAsync(() -> {}, future).get(10, TimeUnit.SECONDS); + services.shutdown(); + + assertThat(events) + .containsExactly( + "future.shutdown", + "io.shutdown", + "future.awaitTermination", + "io.awaitTermination"); + assertThat(future.isTerminated()).isTrue(); + assertThat(io.isTerminated()).isTrue(); + assertThat(cacheCloses).hasValue(1); + assertThat(shuffleCloses).hasValue(1); + } finally { + future.shutdownNow(); + io.shutdownNow(); + } + } + + private static final class RecordingExecutor extends ScheduledThreadPoolExecutor { + private final String name; + private final List events; + + private RecordingExecutor(String name, List events) { + super(1); + this.name = name; + this.events = events; + } + + @Override + public void shutdown() { + events.add(name + ".shutdown"); + super.shutdown(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { + events.add(name + ".awaitTermination"); + return super.awaitTermination(timeout, unit); + } + } + + private static void assertIndependentCapacity( + Executor blockedExecutor, Executor otherExecutor, Executor futureExecutor, int poolSize) + throws Exception { + final CountDownLatch started = new CountDownLatch(poolSize); + final CountDownLatch release = new CountDownLatch(1); + final List> blockers = new ArrayList<>(); + try { + for (int i = 0; i < poolSize; i++) { + blockers.add( + CompletableFuture.runAsync( + ThrowingRunnable.unchecked( + () -> { + started.countDown(); + assertThat(release.await(30, TimeUnit.SECONDS)) + .isTrue(); + }), + blockedExecutor)); + } + assertThat(started.await(10, TimeUnit.SECONDS)).isTrue(); + final CompletableFuture queued = + CompletableFuture.runAsync(() -> {}, blockedExecutor); + CompletableFuture.runAsync(() -> {}, otherExecutor).get(10, TimeUnit.SECONDS); + CompletableFuture.runAsync(() -> {}, futureExecutor).get(10, TimeUnit.SECONDS); + assertThat(queued).isNotDone(); + release.countDown(); + queued.get(10, TimeUnit.SECONDS); + CompletableFuture.allOf(blockers.toArray(new CompletableFuture[0])) + .get(10, TimeUnit.SECONDS); + } finally { + release.countDown(); + } + } + @Nonnull private JobManagerSharedServices buildJobManagerSharedServices(Configuration configuration) throws Exception { + final BlobServer blobServer = + new BlobServer(configuration, TEMPORARY_FOLDER, new VoidBlobStore()); + blobServers.add(blobServer); return JobManagerSharedServices.fromConfiguration( - configuration, - new BlobServer(configuration, TEMPORARY_FOLDER, new VoidBlobStore()), - new TestingFatalErrorHandler()); + configuration, blobServer, new TestingFatalErrorHandler()); } private void assertExecutorPoolSize(Executor executor, int expectedPoolSize) @@ -131,18 +452,18 @@ private void assertExecutorPoolSize(Executor executor, int expectedPoolSize) expectedPoolSizePlusOneLatch.countDown(); expectedPoolSizeLatch.countDown(); // block the runnable to keep the thread occupied - releaseLatch.await(); + releaseLatch.await(30, TimeUnit.SECONDS); }; for (int i = 0; i < expectedPoolSizePlusOne; i++) { executor.execute(ThrowingRunnable.unchecked(countsDown)); } - // the expected pool size latch should complete since we expect to have enough threads - expectedPoolSizeLatch.await(); - assertThat(expectedPoolSizePlusOneLatch.getCount()).isOne(); - - // unblock the runnables - releaseLatch.trigger(); + try { + assertThat(expectedPoolSizeLatch.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(expectedPoolSizePlusOneLatch.getCount()).isOne(); + } finally { + releaseLatch.trigger(); + } } } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterSchedulerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterSchedulerTest.java index 61ad15d3cfd6dc..d0bf515fc528e2 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterSchedulerTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterSchedulerTest.java @@ -106,6 +106,7 @@ public SchedulerNG createInstance( Logger log, ExecutionPlan executionPlan, Executor ioExecutor, + Executor recoveryExecutor, Configuration jobMasterConfiguration, SlotPoolService slotPoolService, ScheduledExecutorService futureExecutor, diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterTest.java index 24a7be65256ce4..e0d423a369adec 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmaster/JobMasterTest.java @@ -33,6 +33,7 @@ import org.apache.flink.configuration.JobManagerOptions; import org.apache.flink.configuration.RestartStrategyOptions; import org.apache.flink.core.execution.SavepointFormatType; +import org.apache.flink.core.failure.FailureEnricher; import org.apache.flink.core.io.InputSplit; import org.apache.flink.core.io.InputSplitAssigner; import org.apache.flink.core.io.InputSplitSource; @@ -40,7 +41,11 @@ import org.apache.flink.events.Event; import org.apache.flink.events.EventBuilder; import org.apache.flink.events.Events; +import org.apache.flink.runtime.blob.BlobServer; +import org.apache.flink.runtime.blob.BlobWriter; +import org.apache.flink.runtime.blob.VoidBlobStore; import org.apache.flink.runtime.blocklist.BlockedNode; +import org.apache.flink.runtime.blocklist.BlocklistOperations; import org.apache.flink.runtime.blocklist.DefaultBlocklistHandler; import org.apache.flink.runtime.checkpoint.CheckpointProperties; import org.apache.flink.runtime.checkpoint.CheckpointRecoveryFactory; @@ -62,12 +67,14 @@ import org.apache.flink.runtime.executiongraph.AccessExecutionVertex; import org.apache.flink.runtime.executiongraph.ArchivedExecutionGraph; import org.apache.flink.runtime.executiongraph.ExecutionAttemptID; +import org.apache.flink.runtime.executiongraph.JobStatusListener; import org.apache.flink.runtime.executiongraph.ResultPartitionBytes; import org.apache.flink.runtime.executiongraph.failover.FailoverStrategyFactoryLoader; import org.apache.flink.runtime.heartbeat.HeartbeatServices; import org.apache.flink.runtime.heartbeat.HeartbeatServicesImpl; import org.apache.flink.runtime.heartbeat.TestingHeartbeatServices; import org.apache.flink.runtime.highavailability.TestingHighAvailabilityServices; +import org.apache.flink.runtime.io.network.partition.JobMasterPartitionTracker; import org.apache.flink.runtime.io.network.partition.ResultPartitionID; import org.apache.flink.runtime.io.network.partition.ResultPartitionType; import org.apache.flink.runtime.io.network.partition.TestingJobMasterPartitionTracker; @@ -102,11 +109,14 @@ import org.apache.flink.runtime.resourcemanager.ResourceManagerGateway; import org.apache.flink.runtime.resourcemanager.ResourceManagerId; import org.apache.flink.runtime.resourcemanager.utils.TestingResourceManagerGateway; +import org.apache.flink.runtime.rpc.FatalErrorHandler; import org.apache.flink.runtime.rpc.RpcUtils; import org.apache.flink.runtime.rpc.TestingRpcService; import org.apache.flink.runtime.rpc.exceptions.RecipientUnreachableException; import org.apache.flink.runtime.scheduler.DefaultSchedulerFactory; import org.apache.flink.runtime.scheduler.ExecutionGraphInfo; +import org.apache.flink.runtime.scheduler.SchedulerNG; +import org.apache.flink.runtime.scheduler.SchedulerNGFactory; import org.apache.flink.runtime.scheduler.SchedulerTestingUtils; import org.apache.flink.runtime.scheduler.TestingPhysicalSlot; import org.apache.flink.runtime.scheduler.TestingSchedulerNG; @@ -116,6 +126,7 @@ import org.apache.flink.runtime.shuffle.DefaultShuffleMetrics; import org.apache.flink.runtime.shuffle.NettyShuffleDescriptor; import org.apache.flink.runtime.shuffle.PartitionWithMetrics; +import org.apache.flink.runtime.shuffle.ShuffleMaster; import org.apache.flink.runtime.state.CompletedCheckpointStorageLocation; import org.apache.flink.runtime.state.StreamStateHandle; import org.apache.flink.runtime.taskexecutor.TaskExecutorGateway; @@ -136,8 +147,10 @@ import org.apache.flink.util.FlinkException; import org.apache.flink.util.FlinkRuntimeException; import org.apache.flink.util.InstantiationUtil; +import org.apache.flink.util.MdcUtils; import org.apache.flink.util.TestLoggerExtension; import org.apache.flink.util.concurrent.FutureUtils; +import org.apache.flink.util.function.ThrowingRunnable; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; @@ -147,6 +160,8 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.api.io.TempDir; +import org.slf4j.Logger; +import org.slf4j.MDC; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -175,6 +190,9 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; @@ -280,6 +298,164 @@ static void teardownAll() { } } + @Test + void testJobMastersBorrowSharedRecoveryPoolWithIsolatedJobContext() throws Exception { + configuration.set(JobManagerOptions.JOB_MANAGER_IO_POOL_SIZE, 1); + configuration.set(JobManagerOptions.JOB_MANAGER_FUTURE_POOL_SIZE, 1); + final JobGraph firstJob = JobGraphTestUtils.singleNoOpJobGraph(); + final JobGraph secondJob = JobGraphTestUtils.singleNoOpJobGraph(); + final TestingSchedulerFactory firstFactory = new TestingSchedulerFactory(); + final TestingSchedulerFactory secondFactory = new TestingSchedulerFactory(); + final CountDownLatch release = new CountDownLatch(1); + try (BlobServer blobServer = + new BlobServer(configuration, temporaryFolder.toFile(), new VoidBlobStore())) { + final JobManagerSharedServices services = + JobManagerSharedServices.fromConfiguration( + configuration, blobServer, testingFatalErrorHandler); + JobMaster firstMaster = null; + JobMaster secondMaster = null; + try { + final Thread recoveryWorker = + CompletableFuture.supplyAsync( + () -> { + MDC.setContextMap( + Collections.singletonMap("worker", "recovery")); + return Thread.currentThread(); + }, + services.getRecoveryExecutor()) + .get(10, TimeUnit.SECONDS); + firstMaster = + new JobMasterBuilder(firstJob, rpcService) + .withConfiguration(configuration) + .withJobManagerSharedServices(services) + .withSlotPoolServiceSchedulerFactory( + DefaultSlotPoolServiceSchedulerFactory.create( + TestingSlotPoolServiceBuilder.newBuilder(), + firstFactory)) + .createJobMaster(); + secondMaster = + new JobMasterBuilder(secondJob, rpcService) + .withConfiguration(configuration) + .withJobManagerSharedServices(services) + .withSlotPoolServiceSchedulerFactory( + DefaultSlotPoolServiceSchedulerFactory.create( + TestingSlotPoolServiceBuilder.newBuilder(), + secondFactory)) + .createJobMaster(); + final CompletableFuture> firstContext = + new CompletableFuture<>(); + final CompletableFuture firstRecovery = + CompletableFuture.runAsync( + ThrowingRunnable.unchecked( + () -> { + firstContext.complete( + Tuple2.of( + Thread.currentThread(), + MDC.get(MdcUtils.JOB_ID))); + assertThat(release.await(30, TimeUnit.SECONDS)) + .isTrue(); + }), + firstFactory.recoveryExecutor); + assertThat(firstContext.get(10, TimeUnit.SECONDS)) + .isEqualTo(Tuple2.of(recoveryWorker, firstJob.getJobID().toHexString())); + final CompletableFuture> secondContext = + executorContext(secondFactory.recoveryExecutor); + final Tuple2 firstIoContext = + executorContext(firstFactory.ioExecutor).get(10, TimeUnit.SECONDS); + final Tuple2 secondIoContext = + executorContext(secondFactory.ioExecutor).get(10, TimeUnit.SECONDS); + assertThat(firstIoContext.f0.getName()).startsWith("jobmanager-io"); + assertThat(firstIoContext.f0).isNotSameAs(recoveryWorker); + assertThat(firstIoContext.f1).isEqualTo(firstJob.getJobID().toHexString()); + assertThat(secondIoContext) + .isEqualTo( + Tuple2.of(firstIoContext.f0, secondJob.getJobID().toHexString())); + assertThat(secondContext).isNotDone(); + + firstMaster.closeAsync().get(10, TimeUnit.SECONDS); + + assertThat(((ExecutorService) services.getRecoveryExecutor()).isShutdown()) + .isFalse(); + assertThat(secondContext).isNotDone(); + release.countDown(); + firstRecovery.get(10, TimeUnit.SECONDS); + assertThat(secondContext.get(10, TimeUnit.SECONDS)) + .isEqualTo(Tuple2.of(recoveryWorker, secondJob.getJobID().toHexString())); + assertThat(recoveryWorker.getName()).startsWith("jobmanager-recovery"); + secondMaster.closeAsync().get(10, TimeUnit.SECONDS); + + assertThat( + CompletableFuture.supplyAsync( + MDC::getCopyOfContextMap, + services.getRecoveryExecutor()) + .get(10, TimeUnit.SECONDS)) + .containsExactlyEntriesOf(Collections.singletonMap("worker", "recovery")); + assertThat(executorContext(services.getIoExecutor()).get(10, TimeUnit.SECONDS).f1) + .isNull(); + CompletableFuture.runAsync(() -> {}, services.getFutureExecutor()) + .get(10, TimeUnit.SECONDS); + } finally { + release.countDown(); + try { + CompletableFuture.allOf( + firstMaster == null + ? CompletableFuture.completedFuture(null) + : firstMaster.closeAsync(), + secondMaster == null + ? CompletableFuture.completedFuture(null) + : secondMaster.closeAsync()) + .get(10, TimeUnit.SECONDS); + } finally { + services.shutdown(); + } + } + } + } + + private static CompletableFuture> executorContext(Executor executor) { + return CompletableFuture.supplyAsync( + () -> Tuple2.of(Thread.currentThread(), MDC.get(MdcUtils.JOB_ID)), executor); + } + + private static final class TestingSchedulerFactory implements SchedulerNGFactory { + private Executor ioExecutor; + private Executor recoveryExecutor; + + @Override + public SchedulerNG createInstance( + Logger log, + ExecutionPlan executionPlan, + Executor ioExecutor, + Executor recoveryExecutor, + Configuration jobMasterConfiguration, + SlotPoolService slotPoolService, + ScheduledExecutorService futureExecutor, + ClassLoader userCodeLoader, + CheckpointRecoveryFactory checkpointRecoveryFactory, + Duration rpcTimeout, + BlobWriter blobWriter, + JobManagerJobMetricGroup jobManagerJobMetricGroup, + Duration slotRequestTimeout, + ShuffleMaster shuffleMaster, + JobMasterPartitionTracker partitionTracker, + ExecutionDeploymentTracker executionDeploymentTracker, + long initializationTimestamp, + ComponentMainThreadExecutor mainThreadExecutor, + FatalErrorHandler fatalErrorHandler, + JobStatusListener jobStatusListener, + Collection failureEnrichers, + BlocklistOperations blocklistOperations) { + this.ioExecutor = ioExecutor; + this.recoveryExecutor = recoveryExecutor; + return TestingSchedulerNG.newBuilder().build(); + } + + @Override + public SchedulerType getSchedulerType() { + return SchedulerType.Default; + } + } + @Test void testTaskManagerRegistrationTriggersHeartbeating() throws Exception { final CompletableFuture heartbeatResourceIdFuture = new CompletableFuture<>(); diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingSchedulerNGFactory.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingSchedulerNGFactory.java index c3ff872dd5270b..717413fa5ce632 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingSchedulerNGFactory.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingSchedulerNGFactory.java @@ -57,6 +57,7 @@ public SchedulerNG createInstance( Logger log, ExecutionPlan executionPlan, Executor ioExecutor, + Executor recoveryExecutor, Configuration jobMasterConfiguration, SlotPoolService slotPoolService, ScheduledExecutorService futureExecutor, diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveRecoveryIoIsolationITCase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveRecoveryIoIsolationITCase.java new file mode 100644 index 00000000000000..cf52d625fb7c32 --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveRecoveryIoIsolationITCase.java @@ -0,0 +1,580 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.flink.runtime.scheduler.adaptive; + +import org.apache.flink.api.common.JobID; +import org.apache.flink.api.common.JobStatus; +import org.apache.flink.configuration.CheckpointingOptions; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.JobManagerOptions; +import org.apache.flink.core.execution.CheckpointType; +import org.apache.flink.core.fs.Path; +import org.apache.flink.runtime.checkpoint.CheckpointMetaData; +import org.apache.flink.runtime.checkpoint.CheckpointMetrics; +import org.apache.flink.runtime.checkpoint.CheckpointOptions; +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState; +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot; +import org.apache.flink.runtime.execution.Environment; +import org.apache.flink.runtime.execution.ExecutionState; +import org.apache.flink.runtime.executiongraph.ArchivedExecutionVertex; +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.jobgraph.OperatorID; +import org.apache.flink.runtime.jobgraph.tasks.AbstractInvokable; +import org.apache.flink.runtime.jobgraph.tasks.CheckpointCoordinatorConfiguration; +import org.apache.flink.runtime.jobgraph.tasks.JobCheckpointingSettings; +import org.apache.flink.runtime.minicluster.MiniCluster; +import org.apache.flink.runtime.minicluster.MiniClusterConfiguration; +import org.apache.flink.runtime.state.CheckpointStorage; +import org.apache.flink.runtime.state.CheckpointStorageAccess; +import org.apache.flink.runtime.state.CheckpointStorageLocation; +import org.apache.flink.runtime.state.CompletedCheckpointStorageLocation; +import org.apache.flink.runtime.state.IncrementalKeyedStateHandle.HandleAndLocalPath; +import org.apache.flink.runtime.state.IncrementalRemoteKeyedStateHandle; +import org.apache.flink.runtime.state.KeyGroupRange; +import org.apache.flink.runtime.state.KeyGroupRangeAssignment; +import org.apache.flink.runtime.state.KeyedStateHandle; +import org.apache.flink.runtime.state.StreamStateHandle; +import org.apache.flink.runtime.state.filesystem.FileStateHandle; +import org.apache.flink.runtime.state.memory.ByteStreamStateHandle; +import org.apache.flink.runtime.state.memory.MemoryBackendCheckpointStorageAccess; +import org.apache.flink.runtime.state.storage.JobManagerCheckpointStorage; +import org.apache.flink.streaming.util.RestartStrategyUtils; +import org.apache.flink.util.SerializedValue; +import org.apache.flink.util.TernaryBoolean; + +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.file.Files; +import java.time.Duration; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.apache.flink.runtime.testutils.CommonTestUtils.waitUntilCondition; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Exercises recovery and checkpoint backpressure through the MiniCluster's production factories. + */ +class AdaptiveRecoveryIoIsolationITCase { + + private static final int MAX_PARALLELISM = 4; + private static final long ADVANCE = -1; + private static final long CANCEL = -2; + private static final Duration WAIT = Duration.ofSeconds(20); + private static final Map CONTROLS = new ConcurrentHashMap<>(); + + @TempDir private java.nio.file.Path temporaryDirectory; + + @ParameterizedTest + @ValueSource(ints = {1, 2}) + @Timeout(120) + void restoresStateWhileAllGeneralIoWorkersAreDeleting(int ioWorkers) throws Exception { + final Configuration configuration = new Configuration(); + configuration.set(JobManagerOptions.SCHEDULER, JobManagerOptions.SchedulerType.Adaptive); + configuration.set(JobManagerOptions.JOB_MANAGER_IO_POOL_SIZE, ioWorkers); + configuration.set(CheckpointingOptions.MAX_RETAINED_CHECKPOINTS, 2); + configuration.set(CheckpointingOptions.CLEANER_PARALLEL_MODE, true); + configuration.set( + JobManagerOptions.SCHEDULER_SUBMISSION_RESOURCE_STABILIZATION_TIMEOUT, + Duration.ofMillis(1)); + configuration.set( + JobManagerOptions.SCHEDULER_EXECUTING_RESOURCE_STABILIZATION_TIMEOUT, + Duration.ofMillis(1)); + configuration.set( + JobManagerOptions.SCHEDULER_EXECUTING_COOLDOWN_AFTER_RESCALING, + Duration.ofMillis(1)); + + final JobVertex vertex = new JobVertex("Checkpointed counters"); + vertex.setInvokableClass(CounterInvokable.class); + vertex.setParallelism(2); + vertex.setMaxParallelism(MAX_PARALLELISM); + final JobGraph job = JobGraphTestUtils.streamingJobGraph(vertex); + job.setSnapshotSettings( + new JobCheckpointingSettings( + CheckpointCoordinatorConfiguration.builder() + .setCheckpointInterval(Long.MAX_VALUE) + .setCheckpointTimeout(60_000) + .build(), + null, + TernaryBoolean.FALSE, + new SerializedValue<>(new ObservedCheckpointStorage()), + null)); + RestartStrategyUtils.configureFixedDelayRestartStrategy(job, 2, 0); + + final Control control = new Control(temporaryDirectory, ioWorkers); + CONTROLS.put(job.getJobID(), control); + final MiniCluster cluster = + new MiniCluster( + new MiniClusterConfiguration.Builder() + .setConfiguration(configuration) + .setNumTaskManagers(1) + .setNumSlotsPerTaskManager(2) + .build()); + try { + cluster.start(); + cluster.submitJob(job).get(WAIT.toMillis(), TimeUnit.MILLISECONDS); + control.expectProgress(2, -1, 0); + control.advance(2, -1, 1); + final long first = checkpoint(cluster, job); + assertThat(first).isOne(); + control.advance(2, -1, 2); + final long second = checkpoint(cluster, job); + control.assertRetained(first, second); + control.advance(2, -1, 3); + final long third = checkpoint(cluster, job); + + assertThat(control.deletionsEntered.await(WAIT.toMillis(), TimeUnit.MILLISECONDS)) + .as("every configured general IO worker is blocked in real checkpoint disposal") + .isTrue(); + assertThat(control.blockedWorkers).hasSize(ioWorkers); + assertThat(control.blockedWorkers.keySet()) + .allSatisfy(name -> assertThat(name).startsWith("jobmanager-io")); + assertThat(control.blockedWorkers.values()).contains("shared"); + if (ioWorkers == 2) { + assertThat(control.blockedWorkers.values()).contains("private"); + } + control.assertRetained(second, third); + + rescale(cluster, job, vertex, 1); + control.expectProgress(1, third, 3); + control.advance(1, third, 4); + assertThat(control.graphCreationThreads) + .hasSize(2) + .allSatisfy(name -> assertThat(name).startsWith("jobmanager-recovery")); + + awaitRunning(cluster, job); + final CompletableFuture blockedCheckpoint = + cluster.triggerCheckpoint(job.getJobID(), CheckpointType.CONFIGURED); + assertThatThrownBy(() -> blockedCheckpoint.get(300, TimeUnit.MILLISECONDS)) + .isInstanceOf(TimeoutException.class); + assertThat(control.initializations).containsOnlyKeys(first, second, third); + control.advance(1, third, 5); + assertThat(blockedCheckpoint).isNotDone(); + assertThat(control.deletionGate.getCount()).isOne(); + control.assertRetained(second, third); + + control.deletionGate.countDown(); + final long fourth = blockedCheckpoint.get(WAIT.toMillis(), TimeUnit.MILLISECONDS); + assertThat(fourth).isGreaterThan(third); + control.awaitDiscarded(first, second); + control.assertRetained(third, fourth); + assertThat(control.initializations.get(fourth)).startsWith("jobmanager-io"); + + rescale(cluster, job, vertex, 2); + control.expectProgress(2, fourth, 5); + control.advance(2, fourth, 6); + final long fifth = checkpoint(cluster, job); + assertThat(fifth).isGreaterThan(fourth); + control.awaitDiscarded(third); + control.assertRetained(fourth, fifth); + + cluster.cancelJob(job.getJobID()).get(WAIT.toMillis(), TimeUnit.MILLISECONDS); + cluster.requestJobResult(job.getJobID()).get(WAIT.toMillis(), TimeUnit.MILLISECONDS); + assertThat( + cluster.getJobStatus(job.getJobID()) + .get(WAIT.toMillis(), TimeUnit.MILLISECONDS)) + .isEqualTo(JobStatus.CANCELED); + control.awaitDiscarded(first, second, third, fourth, fifth); + assertThat(control.mailboxes).isEmpty(); + } finally { + control.deletionGate.countDown(); + try { + cluster.closeAsync().get(30, TimeUnit.SECONDS); + } finally { + CONTROLS.remove(job.getJobID()); + } + } + } + + private static long checkpoint(MiniCluster cluster, JobGraph job) throws Exception { + awaitRunning(cluster, job); + return cluster.triggerCheckpoint(job.getJobID(), CheckpointType.CONFIGURED) + .get(WAIT.toMillis(), TimeUnit.MILLISECONDS); + } + + private static void awaitRunning(MiniCluster cluster, JobGraph job) throws Exception { + waitUntilCondition( + () -> { + boolean found = false; + for (ArchivedExecutionVertex vertex : + cluster.getArchivedExecutionGraph(job.getJobID()) + .get(WAIT.toMillis(), TimeUnit.MILLISECONDS) + .getAllExecutionVertices()) { + found = true; + if (vertex.getExecutionState() != ExecutionState.RUNNING) { + return false; + } + } + return found; + }, + 10L, + (int) WAIT.toMillis() / 10); + } + + private static void rescale( + MiniCluster cluster, JobGraph job, JobVertex vertex, int parallelism) throws Exception { + cluster.updateJobResourceRequirements( + job.getJobID(), + JobResourceRequirements.newBuilder() + .setParallelismForJobVertex( + vertex.getID(), parallelism, parallelism) + .build()) + .get(WAIT.toMillis(), TimeUnit.MILLISECONDS); + } + + private static final class Control { + private final java.nio.file.Path directory; + private final CountDownLatch deletionGate = new CountDownLatch(1); + private final CountDownLatch deletionsEntered; + private final Map blockedWorkers = new ConcurrentHashMap<>(); + private final Map initializations = new ConcurrentHashMap<>(); + private final BlockingQueue graphCreationThreads = new LinkedBlockingQueue<>(); + private final Map> mailboxes = new ConcurrentHashMap<>(); + private final BlockingQueue progress = new LinkedBlockingQueue<>(); + private final Map files = new ConcurrentHashMap<>(); + + private Control(java.nio.file.Path directory, int ioWorkers) { + this.directory = directory; + deletionsEntered = new CountDownLatch(ioWorkers); + } + + private void advance(int parallelism, long restoredCheckpoint, long expected) + throws Exception { + assertThat(mailboxes).hasSize(parallelism); + mailboxes.values().forEach(mailbox -> mailbox.add(ADVANCE)); + expectProgress(parallelism, restoredCheckpoint, expected); + } + + private void expectProgress(int parallelism, long restoredCheckpoint, long expected) + throws Exception { + final Map counters = new HashMap<>(); + final Set subtasks = new HashSet<>(); + for (int i = 0; i < parallelism; i++) { + final Progress next = progress.poll(WAIT.toMillis(), TimeUnit.MILLISECONDS); + assertThat(next) + .as( + "physical task progress after restoring checkpoint %s", + restoredCheckpoint) + .isNotNull(); + assertThat(next.parallelism).isEqualTo(parallelism); + assertThat(next.restoredCheckpoint).isEqualTo(restoredCheckpoint); + assertThat(subtasks.add(next.subtask)).isTrue(); + next.counters.forEach( + (key, value) -> { + assertThat(counters.put(key, value)).isNull(); + assertThat(value).isEqualTo(key * 1000L + expected); + }); + } + assertThat(counters).containsOnlyKeys(0, 1, 2, 3); + } + + private void assertRetained(long... checkpointIds) { + for (long checkpointId : checkpointIds) { + assertThat(files.values()).anyMatch(file -> file.checkpointId == checkpointId); + files.values().stream() + .filter(file -> file.checkpointId == checkpointId) + .forEach( + file -> { + assertThat(file.path).exists(); + assertThat(file.discardAttempts.get()).isZero(); + }); + } + } + + private void awaitDiscarded(long... checkpointIds) throws Exception { + for (long checkpointId : checkpointIds) { + assertThat(files.values()).anyMatch(file -> file.checkpointId == checkpointId); + waitUntilCondition( + () -> + files.values().stream() + .filter(file -> file.checkpointId == checkpointId) + .allMatch(file -> file.discardCompletions.get() > 0), + 10L, + (int) WAIT.toMillis() / 10); + files.values().stream() + .filter(file -> file.checkpointId == checkpointId) + .forEach( + file -> { + assertThat(file.path).doesNotExist(); + assertThat(file.discardAttempts.get()).isOne(); + assertThat(file.discardCompletions.get()).isOne(); + }); + } + } + } + + private static final class TrackedFile { + private final java.nio.file.Path path; + private final long checkpointId; + private final int subtask; + private final String kind; + private final AtomicInteger discardAttempts = new AtomicInteger(); + private final AtomicInteger discardCompletions = new AtomicInteger(); + + private TrackedFile(java.nio.file.Path path, long checkpointId, int subtask, String kind) { + this.path = path; + this.checkpointId = checkpointId; + this.subtask = subtask; + this.kind = kind; + } + } + + private static final class Progress { + private final int subtask; + private final int parallelism; + private final long restoredCheckpoint; + private final Map counters; + + private Progress( + int subtask, + int parallelism, + long restoredCheckpoint, + Map counters) { + this.subtask = subtask; + this.parallelism = parallelism; + this.restoredCheckpoint = restoredCheckpoint; + this.counters = new HashMap<>(counters); + } + } + + /** + * A task with file-backed keyed counters and a mailbox shared by input and checkpoint triggers. + */ + public static final class CounterInvokable extends AbstractInvokable { + private final BlockingQueue mailbox = new LinkedBlockingQueue<>(); + private final Map counters = new HashMap<>(); + + public CounterInvokable(Environment environment) { + super(environment); + } + + @Override + public void invoke() throws Exception { + final Environment environment = getEnvironment(); + final Control control = CONTROLS.get(environment.getJobID()); + final int subtask = getIndexInSubtaskGroup(); + final int parallelism = environment.getTaskInfo().getNumberOfParallelSubtasks(); + final KeyGroupRange range = + KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex( + MAX_PARALLELISM, parallelism, subtask); + final OperatorID operator = OperatorID.fromJobVertexID(environment.getJobVertexId()); + final long restored = + environment.getTaskStateManager().getRestoreCheckpointId().orElse(-1L); + if (restored < 0) { + for (int key : range) { + counters.put(key, key * 1000L); + } + } else { + final OperatorSubtaskState state = + environment + .getTaskStateManager() + .getSubtaskJobManagerRestoredState(operator) + .orElseThrow( + () -> + new IllegalStateException( + "Missing restored counter state")); + for (KeyedStateHandle keyed : state.getManagedKeyedState()) { + final IncrementalRemoteKeyedStateHandle incremental = + (IncrementalRemoteKeyedStateHandle) keyed; + assertThat(incremental.getCheckpointId()).isEqualTo(restored); + for (HandleAndLocalPath handle : incremental.getSharedState()) { + restore(handle.getHandle(), range, restored); + } + for (HandleAndLocalPath handle : incremental.getPrivateState()) { + try (DataInputStream input = + new DataInputStream(handle.getHandle().openInputStream())) { + assertThat(input.readLong()).isEqualTo(restored); + } + } + } + assertThat(counters).hasSize(range.getNumberOfKeyGroups()); + } + + control.mailboxes.put(subtask, mailbox); + try { + control.progress.add(new Progress(subtask, parallelism, restored, counters)); + long signal; + while ((signal = mailbox.take()) != CANCEL) { + if (signal == ADVANCE) { + counters.replaceAll((key, value) -> value + 1); + control.progress.add( + new Progress(subtask, parallelism, restored, counters)); + } else { + final FileStateHandle shared = + writeState(control, signal, subtask, "shared"); + final FileStateHandle privateState = + writeState(control, signal, subtask, "private"); + final IncrementalRemoteKeyedStateHandle snapshot = + new IncrementalRemoteKeyedStateHandle( + UUID.randomUUID(), + range, + signal, + Collections.singletonList( + HandleAndLocalPath.of(shared, "counters")), + Collections.singletonList( + HandleAndLocalPath.of(privateState, "manifest")), + new ByteStreamStateHandle( + "meta-" + signal + "-" + subtask, new byte[0])); + final TaskStateSnapshot taskState = new TaskStateSnapshot(); + taskState.putSubtaskStateByOperatorID( + operator, + OperatorSubtaskState.builder() + .setManagedKeyedState(snapshot) + .build()); + environment.acknowledgeCheckpoint( + signal, new CheckpointMetrics(), taskState); + } + } + } finally { + control.mailboxes.remove(subtask, mailbox); + } + } + + private void restore(StreamStateHandle handle, KeyGroupRange range, long checkpointId) + throws IOException { + try (DataInputStream input = new DataInputStream(handle.openInputStream())) { + assertThat(input.readLong()).isEqualTo(checkpointId); + final int size = input.readInt(); + for (int i = 0; i < size; i++) { + final int key = input.readInt(); + final long value = input.readLong(); + if (range.contains(key)) { + assertThat(counters.put(key, value)).isNull(); + } + } + } + } + + private FileStateHandle writeState( + Control control, long checkpointId, int subtask, String kind) throws IOException { + final java.nio.file.Path path = + Files.createTempFile( + control.directory, checkpointId + "-" + subtask + "-" + kind, ".state"); + try (DataOutputStream output = new DataOutputStream(Files.newOutputStream(path))) { + output.writeLong(checkpointId); + output.writeInt(counters.size()); + for (Map.Entry entry : counters.entrySet()) { + output.writeInt(entry.getKey()); + output.writeLong(entry.getValue()); + } + } + final String id = path.toUri().toString(); + control.files.put(id, new TrackedFile(path, checkpointId, subtask, kind)); + return new GatedFileStateHandle( + getEnvironment().getJobID(), new Path(path.toUri()), Files.size(path)); + } + + @Override + public void cancel() { + mailbox.add(CANCEL); + } + + @Override + public CompletableFuture triggerCheckpointAsync( + CheckpointMetaData metadata, CheckpointOptions options) { + mailbox.add(metadata.getCheckpointId()); + return CompletableFuture.completedFuture(true); + } + + @Override + public Future notifyCheckpointCompleteAsync(long checkpointId) { + return CompletableFuture.completedFuture(null); + } + + @Override + public Future notifyCheckpointSubsumedAsync(long checkpointId) { + return CompletableFuture.completedFuture(null); + } + } + + private static final class GatedFileStateHandle extends FileStateHandle { + private static final long serialVersionUID = 1L; + private final JobID jobId; + + private GatedFileStateHandle(JobID jobId, Path path, long size) { + super(path, size); + this.jobId = jobId; + } + + @Override + public void discardState() throws Exception { + final Control control = CONTROLS.get(jobId); + final TrackedFile file = control.files.get(getFilePath().toUri().toString()); + file.discardAttempts.incrementAndGet(); + if (file.checkpointId == 1 && file.subtask == 0) { + control.blockedWorkers.put(Thread.currentThread().getName(), file.kind); + control.deletionsEntered.countDown(); + if (!control.deletionGate.await(60, TimeUnit.SECONDS)) { + throw new IOException("Checkpoint deletion gate timed out"); + } + } + super.discardState(); + file.discardCompletions.incrementAndGet(); + } + } + + private static final class ObservedCheckpointStorage implements CheckpointStorage { + private static final long serialVersionUID = 1L; + + @Override + public CompletedCheckpointStorageLocation resolveCheckpoint(String pointer) + throws IOException { + return new JobManagerCheckpointStorage().resolveCheckpoint(pointer); + } + + @Override + public CheckpointStorageAccess createCheckpointStorage(JobID jobId) throws IOException { + final Control control = CONTROLS.get(jobId); + control.graphCreationThreads.add(Thread.currentThread().getName()); + return new MemoryBackendCheckpointStorageAccess( + jobId, null, null, true, JobManagerCheckpointStorage.DEFAULT_MAX_STATE_SIZE) { + @Override + public CheckpointStorageLocation initializeLocationForCheckpoint(long checkpointId) + throws IOException { + control.initializations.put(checkpointId, Thread.currentThread().getName()); + return super.initializeLocationForCheckpoint(checkpointId); + } + }; + } + } +} diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveRecoveryLifecycleTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveRecoveryLifecycleTest.java new file mode 100644 index 00000000000000..88f1b9a76f4a4a --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveRecoveryLifecycleTest.java @@ -0,0 +1,300 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.flink.runtime.scheduler.adaptive; + +import org.apache.flink.api.common.JobID; +import org.apache.flink.api.common.JobStatus; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.JobManagerOptions; +import org.apache.flink.runtime.blob.PermanentBlobKey; +import org.apache.flink.runtime.blob.TestingBlobWriter; +import org.apache.flink.runtime.checkpoint.CheckpointsCleaner; +import org.apache.flink.runtime.checkpoint.StandaloneCheckpointIDCounter; +import org.apache.flink.runtime.checkpoint.StandaloneCompletedCheckpointStore; +import org.apache.flink.runtime.checkpoint.TestingCheckpointRecoveryFactory; +import org.apache.flink.runtime.clusterframework.types.ResourceProfile; +import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutor; +import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutorServiceAdapter; +import org.apache.flink.runtime.deployment.TaskDeploymentDescriptor; +import org.apache.flink.runtime.executiongraph.ExecutionGraph; +import org.apache.flink.runtime.jobgraph.JobGraph; +import org.apache.flink.runtime.jobgraph.tasks.CheckpointCoordinatorConfiguration; +import org.apache.flink.runtime.jobgraph.tasks.JobCheckpointingSettings; +import org.apache.flink.runtime.jobmanager.slots.TaskManagerGateway; +import org.apache.flink.runtime.jobmaster.slotpool.DefaultAllocatedSlotPool; +import org.apache.flink.runtime.jobmaster.slotpool.DefaultDeclarativeSlotPool; +import org.apache.flink.runtime.messages.Acknowledge; +import org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder; +import org.apache.flink.runtime.util.ResourceCounter; +import org.apache.flink.testutils.executor.TestExecutorExtension; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import java.io.IOException; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeoutException; + +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.apache.flink.runtime.jobgraph.JobGraphTestUtils.singleNoOpJobGraph; +import static org.apache.flink.runtime.jobmaster.slotpool.SlotPoolTestUtils.createSlotOffersForResourceRequirements; +import static org.apache.flink.runtime.jobmaster.slotpool.SlotPoolTestUtils.createTaskManagerGateway; +import static org.apache.flink.runtime.jobmaster.slotpool.SlotPoolTestUtils.offerSlots; +import static org.assertj.core.api.Assertions.assertThat; + +/** Lifecycle regression for a real graph supplier running on the borrowed recovery executor. */ +class AdaptiveRecoveryLifecycleTest { + + private static final long TIMEOUT_SECONDS = 10; + + @RegisterExtension + private static final TestExecutorExtension mainExecutor = + new TestExecutorExtension<>(Executors::newSingleThreadScheduledExecutor); + + @RegisterExtension + private static final TestExecutorExtension generalExecutor = + new TestExecutorExtension<>(Executors::newSingleThreadScheduledExecutor); + + @RegisterExtension + private static final TestExecutorExtension recoveryExecutor = + new TestExecutorExtension<>(Executors::newSingleThreadScheduledExecutor); + + @RegisterExtension + private static final TestExecutorExtension futureExecutor = + new TestExecutorExtension<>(Executors::newSingleThreadScheduledExecutor); + + @Test + void closeWaitsForRunningRecoveryAndFencesSuccessfulStaleGraph() throws Exception { + final ComponentMainThreadExecutor mainThreadExecutor = + ComponentMainThreadExecutorServiceAdapter.forSingleThreadExecutor( + mainExecutor.getExecutor()); + final Thread recoveryThread = + recoveryExecutor + .getExecutor() + .submit(Thread::currentThread) + .get(TIMEOUT_SECONDS, SECONDS); + final CompletableFuture supplierEntered = new CompletableFuture<>(); + final CompletableFuture releaseSupplier = new CompletableFuture<>(); + final CompletableFuture supplierSucceeded = new CompletableFuture<>(); + final CompletableFuture storeShutdown = new CompletableFuture<>(); + final CompletableFuture counterShutdown = new CompletableFuture<>(); + final CompletableFuture deployment = new CompletableFuture<>(); + final CompletableFuture fatalError = new CompletableFuture<>(); + final List jobStatuses = new CopyOnWriteArrayList<>(); + + final StandaloneCompletedCheckpointStore store = + new StandaloneCompletedCheckpointStore(1) { + @Override + public void shutdown(JobStatus status, CheckpointsCleaner cleaner) + throws Exception { + storeShutdown.complete(status); + super.shutdown(status, cleaner); + } + }; + final StandaloneCheckpointIDCounter counter = + new StandaloneCheckpointIDCounter() { + @Override + public CompletableFuture shutdown(JobStatus status) { + counterShutdown.complete(status); + return super.shutdown(status); + } + }; + final TestingBlobWriter blobWriter = + new TestingBlobWriter(0) { + @Override + public PermanentBlobKey putPermanent(JobID jobId, byte[] value) + throws IOException { + supplierEntered.complete(Thread.currentThread()); + try { + releaseSupplier.get(TIMEOUT_SECONDS, SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } catch (ExecutionException | TimeoutException e) { + throw new IOException(e); + } + return super.putPermanent(jobId, value); + } + }; + + final JobGraph jobGraph = singleNoOpJobGraph(); + jobGraph.setSnapshotSettings( + new JobCheckpointingSettings( + CheckpointCoordinatorConfiguration.builder().build(), null)); + final Configuration configuration = new Configuration(); + configuration.set( + JobManagerOptions.SCHEDULER_SUBMISSION_RESOURCE_STABILIZATION_TIMEOUT, + Duration.ZERO); + final DefaultDeclarativeSlotPool slotPool = + new DefaultDeclarativeSlotPool( + jobGraph.getJobID(), + new DefaultAllocatedSlotPool(), + ignored -> {}, + Duration.ofMinutes(5), + Duration.ofSeconds(TIMEOUT_SECONDS), + Duration.ZERO, + mainThreadExecutor); + final TaskManagerGateway taskManagerGateway = + createTaskManagerGateway( + new TestingTaskExecutorGatewayBuilder() + .setSubmitTaskConsumer( + (descriptor, jobMasterId) -> { + deployment.complete(descriptor); + return CompletableFuture.completedFuture( + Acknowledge.get()); + }) + .createTestingTaskExecutorGateway()); + final AdaptiveScheduler scheduler = + new AdaptiveSchedulerBuilder( + jobGraph, mainThreadExecutor, generalExecutor.getExecutor()) + .setRecoveryExecutor(recoveryExecutor.getExecutor()) + .setFutureExecutor(futureExecutor.getExecutor()) + .setJobMasterConfiguration(configuration) + .setDeclarativeSlotPool(slotPool) + .setCheckpointRecoveryFactory( + new TestingCheckpointRecoveryFactory(store, counter)) + .setBlobWriter(blobWriter) + .setExecutionGraphFactoryDecorator( + delegate -> + (graph, + completedStore, + cleaner, + idCounter, + stats, + partitionConstraint, + timestamp, + attempts, + parallelism, + listener, + markFinished, + schedulingContext, + logger) -> { + final ExecutionGraph created = + delegate.createAndRestoreExecutionGraph( + graph, + completedStore, + cleaner, + idCounter, + stats, + partitionConstraint, + timestamp, + attempts, + parallelism, + listener, + markFinished, + schedulingContext, + logger); + supplierSucceeded.complete(created); + return created; + }) + .setFatalErrorHandler(fatalError::complete) + .setJobStatusListener((id, status, timestamp) -> jobStatuses.add(status)) + .build(); + CompletableFuture closeFuture = null; + try { + CompletableFuture.runAsync( + () -> { + scheduler.startScheduling(); + offerSlots( + slotPool, + createSlotOffersForResourceRequirements( + ResourceCounter.withResource( + ResourceProfile.UNKNOWN, 1)), + taskManagerGateway); + }, + mainThreadExecutor) + .get(TIMEOUT_SECONDS, SECONDS); + assertThat(supplierEntered.get(TIMEOUT_SECONDS, SECONDS)).isSameAs(recoveryThread); + + // Synchronous storage can still hold recovery even while the other lanes are idle. + final Thread generalThread = + generalExecutor + .getExecutor() + .submit(Thread::currentThread) + .get(TIMEOUT_SECONDS, SECONDS); + final Thread futureThread = + futureExecutor + .getExecutor() + .submit(Thread::currentThread) + .get(TIMEOUT_SECONDS, SECONDS); + assertThat(generalThread).isNotSameAs(recoveryThread).isNotSameAs(futureThread); + assertThat(futureThread).isNotSameAs(recoveryThread); + closeFuture = + CompletableFuture.supplyAsync( + () -> { + assertThat(scheduler.getState()) + .isInstanceOf(CreatingExecutionGraph.class); + return scheduler.closeAsync(); + }, + mainThreadExecutor) + .get(TIMEOUT_SECONDS, SECONDS); + drain(mainExecutor.getExecutor()); + assertThat(closeFuture).isNotDone(); + assertThat(storeShutdown).isNotDone(); + assertThat(counterShutdown).isNotDone(); + assertThat(deployment).isNotDone(); + + releaseSupplier.complete(null); + closeFuture.get(TIMEOUT_SECONDS, SECONDS); + // Also proves scheduler shutdown did not shut down its borrowed recovery executor. + assertThat( + recoveryExecutor + .getExecutor() + .submit(Thread::currentThread) + .get(TIMEOUT_SECONDS, SECONDS)) + .isSameAs(recoveryThread); + // The result handoff schedules another main-thread callback; drain both generations. + drain(mainExecutor.getExecutor()); + drain(mainExecutor.getExecutor()); + drain(futureExecutor.getExecutor()); + + // Closing also succeeds for a failed supplier: require a successful late graph here. + assertThat(supplierSucceeded.get(TIMEOUT_SECONDS, SECONDS)).isNotNull(); + assertThat(storeShutdown.get(TIMEOUT_SECONDS, SECONDS)).isEqualTo(JobStatus.SUSPENDED); + assertThat(counterShutdown.get(TIMEOUT_SECONDS, SECONDS)) + .isEqualTo(JobStatus.SUSPENDED); + assertThat( + CompletableFuture.supplyAsync(scheduler::getState, mainThreadExecutor) + .get(TIMEOUT_SECONDS, SECONDS)) + .isInstanceOf(Finished.class); + assertThat(jobStatuses).doesNotContain(JobStatus.RUNNING); + assertThat(deployment).isNotDone(); + assertThat(fatalError).isNotDone(); + } finally { + releaseSupplier.complete(null); + if (closeFuture == null) { + closeFuture = + CompletableFuture.supplyAsync(scheduler::closeAsync, mainThreadExecutor) + .get(TIMEOUT_SECONDS, SECONDS); + } + closeFuture.get(TIMEOUT_SECONDS, SECONDS); + } + } + + private static void drain(ScheduledExecutorService executor) throws Exception { + executor.submit(() -> {}).get(TIMEOUT_SECONDS, SECONDS); + } +} diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerBuilder.java index e8718be0ead69c..4f22e123b181ae 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerBuilder.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerBuilder.java @@ -59,9 +59,11 @@ import java.time.Duration; import java.util.Collection; import java.util.Collections; +import java.util.concurrent.Executor; import java.util.concurrent.ScheduledExecutorService; import java.util.function.BiFunction; import java.util.function.Function; +import java.util.function.UnaryOperator; /** Builder for {@link AdaptiveScheduler}. */ public class AdaptiveSchedulerBuilder { @@ -71,6 +73,8 @@ public class AdaptiveSchedulerBuilder { private final ComponentMainThreadExecutor mainThreadExecutor; private final ScheduledExecutorService executorService; + private Executor recoveryExecutor; + private ScheduledExecutorService futureExecutor; @Nullable private JobResourceRequirements jobResourceRequirements; @@ -82,6 +86,8 @@ public class AdaptiveSchedulerBuilder { private DeclarativeSlotPool declarativeSlotPool; private Duration rpcTimeout = DEFAULT_TIMEOUT; private BlobWriter blobWriter = VoidBlobWriter.getInstance(); + private UnaryOperator executionGraphFactoryDecorator = + UnaryOperator.identity(); private JobManagerJobMetricGroup jobManagerJobMetricGroup = UnregisteredMetricGroups.createUnregisteredJobManagerJobMetricGroup(); private ShuffleMaster shuffleMaster = ShuffleTestUtils.DEFAULT_SHUFFLE_MASTER; @@ -130,6 +136,18 @@ public AdaptiveSchedulerBuilder( Duration.ZERO, mainThreadExecutor); this.executorService = executorService; + this.recoveryExecutor = executorService; + this.futureExecutor = executorService; + } + + public AdaptiveSchedulerBuilder setRecoveryExecutor(Executor recoveryExecutor) { + this.recoveryExecutor = recoveryExecutor; + return this; + } + + public AdaptiveSchedulerBuilder setFutureExecutor(ScheduledExecutorService futureExecutor) { + this.futureExecutor = futureExecutor; + return this; } public AdaptiveSchedulerBuilder setJobResourceRequirements( @@ -250,13 +268,19 @@ public AdaptiveSchedulerBuilder setCheckpointStatsTrackerFactory( return this; } + public AdaptiveSchedulerBuilder setExecutionGraphFactoryDecorator( + UnaryOperator executionGraphFactoryDecorator) { + this.executionGraphFactoryDecorator = executionGraphFactoryDecorator; + return this; + } + public AdaptiveScheduler build() throws Exception { final ExecutionGraphFactory executionGraphFactory = new DefaultExecutionGraphFactory( jobMasterConfiguration, userCodeLoader, new DefaultExecutionDeploymentTracker(), - executorService, + futureExecutor, executorService, rpcTimeout, jobManagerJobMetricGroup, @@ -288,6 +312,7 @@ public AdaptiveScheduler build() throws Exception { TaskManagerOptions.TASK_MANAGER_LOAD_BALANCE_MODE)) : slotAllocator, executorService, + recoveryExecutor, userCodeLoader, checkpointsCleaner, checkpointRecoveryFactory, @@ -298,6 +323,6 @@ public AdaptiveScheduler build() throws Exception { fatalErrorHandler, jobStatusListener, failureEnrichers, - executionGraphFactory); + executionGraphFactoryDecorator.apply(executionGraphFactory)); } } diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/BackgroundTaskTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/BackgroundTaskTest.java index a2369a25a9ea1b..bb7ad3003b09b8 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/BackgroundTaskTest.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/BackgroundTaskTest.java @@ -18,19 +18,27 @@ package org.apache.flink.runtime.scheduler.adaptive; -import org.apache.flink.core.testutils.OneShotLatch; +import org.apache.flink.core.testutils.ManuallyTriggeredScheduledExecutorService; import org.apache.flink.testutils.executor.TestExecutorExtension; +import org.apache.flink.util.FlinkException; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; -import java.util.concurrent.ArrayBlockingQueue; -import java.util.concurrent.BlockingQueue; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; -import static org.apache.flink.core.testutils.FlinkAssertions.assertThatFuture; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -42,34 +50,34 @@ class BackgroundTaskTest { new TestExecutorExtension<>(() -> Executors.newFixedThreadPool(2)); @Test - void testFinishedBackgroundTaskIsTerminated() { + void testFinishedBackgroundTaskIsTerminated() throws Exception { final BackgroundTask finishedBackgroundTask = BackgroundTask.finishedBackgroundTask(); - assertThatFuture(finishedBackgroundTask.getTerminationFuture()).isDone(); - finishedBackgroundTask.getTerminationFuture().join(); + assertThat(finishedBackgroundTask.getTerminationFuture()).isDone(); + finishedBackgroundTask.getTerminationFuture().get(10, TimeUnit.SECONDS); } @Test void testFinishedBackgroundTaskDoesNotContainAResult() { final BackgroundTask finishedBackgroundTask = BackgroundTask.finishedBackgroundTask(); - assertThatFuture(finishedBackgroundTask.getResultFuture()).isCompletedExceptionally(); + assertThat(finishedBackgroundTask.getResultFuture()).isCompletedExceptionally(); } @Test - void testNormalCompletionOfBackgroundTask() { + void testNormalCompletionOfBackgroundTask() throws Exception { final String expectedValue = "foobar"; final BackgroundTask backgroundTask = BackgroundTask.finishedBackgroundTask() .runAfter(() -> expectedValue, TEST_EXECUTOR_EXTENSION.getExecutor()); - assertThat(backgroundTask.getResultFuture().join()).isEqualTo(expectedValue); - // check that the termination future has completed normally - backgroundTask.getTerminationFuture().join(); + assertThat(backgroundTask.getResultFuture().get(10, TimeUnit.SECONDS)) + .isEqualTo(expectedValue); + backgroundTask.getTerminationFuture().get(10, TimeUnit.SECONDS); } @Test - void testExceptionalCompletionOfBackgroundTask() throws InterruptedException { + void testExceptionalCompletionOfBackgroundTask() throws Exception { final Exception expectedException = new Exception("Test exception"); final BackgroundTask backgroundTask = BackgroundTask.finishedBackgroundTask() @@ -79,73 +87,195 @@ void testExceptionalCompletionOfBackgroundTask() throws InterruptedException { }, TEST_EXECUTOR_EXTENSION.getExecutor()); - assertThatThrownBy(() -> backgroundTask.getResultFuture().get()) - .withFailMessage("Expected an exceptionally completed result future.") + assertThatThrownBy(() -> backgroundTask.getResultFuture().get(10, TimeUnit.SECONDS)) .isInstanceOf(ExecutionException.class) .hasCause(expectedException); - // check that the termination future has completed normally - backgroundTask.getTerminationFuture().join(); + backgroundTask.getTerminationFuture().get(10, TimeUnit.SECONDS); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testSuccessorWaitsForPredecessorRegardlessOfOutcome(boolean failPredecessor) + throws Exception { + final ManuallyTriggeredScheduledExecutorService executor = + new ManuallyTriggeredScheduledExecutorService(); + final List completions = new ArrayList<>(); + final Exception failure = new Exception("predecessor failed"); + final BackgroundTask predecessor = + BackgroundTask.initialBackgroundTask( + () -> { + completions.add(1); + if (failPredecessor) { + throw failure; + } + return 1; + }, + executor); + final BackgroundTask successor = + predecessor.runAfter( + () -> { + completions.add(2); + return 2; + }, + executor); + + assertThat(executor.numQueuedRunnables()).isOne(); + assertThat(successor.getResultFuture()).isNotDone(); + executor.trigger(); + assertThat(completions).containsExactly(1); + assertThat(executor.numQueuedRunnables()).isOne(); + assertThat(successor.getResultFuture()).isNotDone(); + executor.trigger(); + + if (failPredecessor) { + assertThatThrownBy(() -> predecessor.getResultFuture().get(10, TimeUnit.SECONDS)) + .hasCause(failure); + } else { + assertThat(predecessor.getResultFuture().get(10, TimeUnit.SECONDS)).isEqualTo(1); + } + assertThat(successor.getResultFuture().get(10, TimeUnit.SECONDS)).isEqualTo(2); + successor.getTerminationFuture().get(10, TimeUnit.SECONDS); + assertThat(completions).containsExactly(1, 2); } @Test - void testRunAfterExecutesBackgroundTaskAfterPreviousHasCompleted() { - final OneShotLatch blockingLatch = new OneShotLatch(); - final BlockingQueue taskCompletions = new ArrayBlockingQueue<>(2); - final BackgroundTask backgroundTask = + void testAbortDoesNotBypassUnfinishedPredecessor() throws Exception { + final ManuallyTriggeredScheduledExecutorService executor = + new ManuallyTriggeredScheduledExecutorService(); + final List completions = new ArrayList<>(); + final BackgroundTask predecessor = BackgroundTask.initialBackgroundTask( - () -> { - blockingLatch.await(); - taskCompletions.offer(1); - return null; - }, - TEST_EXECUTOR_EXTENSION.getExecutor()) - .runAfter( - () -> { - taskCompletions.offer(2); - return null; - }, - TEST_EXECUTOR_EXTENSION.getExecutor()); + () -> { + completions.add(1); + return 1; + }, + executor); + final BackgroundTask aborted = + predecessor.runAfter( + () -> { + completions.add(2); + return 2; + }, + executor); + final BackgroundTask successor = + aborted.runAfter( + () -> { + completions.add(3); + return 3; + }, + executor); - blockingLatch.trigger(); + aborted.abort(); - backgroundTask.getTerminationFuture().join(); + assertThat(aborted.getTerminationFuture()).isNotDone(); + assertThat(successor.getResultFuture()).isNotDone(); + assertThat(executor.numQueuedRunnables()).isOne(); + executor.triggerAll(); - assertThat(taskCompletions).contains(1, 2); + assertAborted(aborted); + assertThat(successor.getResultFuture().get(10, TimeUnit.SECONDS)).isEqualTo(3); + successor.getTerminationFuture().get(10, TimeUnit.SECONDS); + assertThat(completions).containsExactly(1, 3); } @Test - void testAbortSkipsTasksWhichHaveNotBeenStarted() { - final OneShotLatch blockingLatch = new OneShotLatch(); - final BlockingQueue taskCompletions = new ArrayBlockingQueue<>(2); - final BackgroundTask backgroundTask = - BackgroundTask.initialBackgroundTask( - () -> { - blockingLatch.await(); - taskCompletions.offer(1); - return null; - }, - TEST_EXECUTOR_EXTENSION.getExecutor()) - .runAfter( - () -> { - taskCompletions.offer(2); - return null; - }, - TEST_EXECUTOR_EXTENSION.getExecutor()); + void testAbortSkipsTaskAlreadyQueuedOnExecutor() throws Exception { + final ManuallyTriggeredScheduledExecutorService executor = + new ManuallyTriggeredScheduledExecutorService(); + final AtomicBoolean invoked = new AtomicBoolean(); + final BackgroundTask task = + BackgroundTask.finishedBackgroundTask() + .runAfter(() -> invoked.getAndSet(true), executor); + + assertThat(executor.numQueuedRunnables()).isOne(); + task.abort(); + assertThat(task.getTerminationFuture()).isNotDone(); + executor.triggerAll(); - final BackgroundTask finalTask = - backgroundTask.runAfter( + assertAborted(task); + assertThat(invoked).isFalse(); + } + + @Test + void testAbortDoesNotInterruptRunningTaskOrReleaseSuccessor() throws Exception { + final CountDownLatch started = new CountDownLatch(1); + final CountDownLatch release = new CountDownLatch(1); + final AtomicBoolean predecessorFinished = new AtomicBoolean(); + final AtomicBoolean interrupted = new AtomicBoolean(); + final ExecutorService executor = TEST_EXECUTOR_EXTENSION.getExecutor(); + final BackgroundTask predecessor = + BackgroundTask.initialBackgroundTask( () -> { - taskCompletions.offer(3); - return null; + started.countDown(); + try { + assertThat(release.await(30, TimeUnit.SECONDS)).isTrue(); + } catch (InterruptedException e) { + interrupted.set(true); + throw e; + } + predecessorFinished.set(true); + return 1; }, - TEST_EXECUTOR_EXTENSION.getExecutor()); + executor); + try { + assertThat(started.await(10, TimeUnit.SECONDS)).isTrue(); + predecessor.abort(); + final BackgroundTask successor = + predecessor.runAfter( + () -> { + assertThat(predecessorFinished).isTrue(); + return 2; + }, + executor); + CompletableFuture.runAsync(() -> {}, executor).get(10, TimeUnit.SECONDS); + assertThat(predecessor.getTerminationFuture()).isNotDone(); + assertThat(successor.getResultFuture()).isNotDone(); + release.countDown(); - backgroundTask.abort(); + assertThat(predecessor.getResultFuture().get(10, TimeUnit.SECONDS)).isEqualTo(1); + assertThat(successor.getResultFuture().get(10, TimeUnit.SECONDS)).isEqualTo(2); + successor.getTerminationFuture().get(10, TimeUnit.SECONDS); + assertThat(interrupted).isFalse(); + } finally { + release.countDown(); + } + } - blockingLatch.trigger(); + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void testRejectionFailsResultButAllowsSuccessor(boolean predecessorPending) throws Exception { + final ManuallyTriggeredScheduledExecutorService executor = + new ManuallyTriggeredScheduledExecutorService(); + final BackgroundTask predecessor = + predecessorPending + ? BackgroundTask.initialBackgroundTask(() -> null, executor) + : BackgroundTask.finishedBackgroundTask(); + final RejectedExecutionException rejection = new RejectedExecutionException("shut down"); + final Executor rejectingExecutor = + command -> { + throw rejection; + }; + final AtomicBoolean invoked = new AtomicBoolean(); + final BackgroundTask rejected = + predecessor.runAfter(() -> invoked.getAndSet(true), rejectingExecutor); + final BackgroundTask successor = rejected.runAfter(() -> 42, executor); - finalTask.getTerminationFuture().join(); + executor.triggerAll(); - assertThat(taskCompletions).contains(1, 3); + assertThatThrownBy(() -> rejected.getResultFuture().get(10, TimeUnit.SECONDS)) + .isInstanceOf(ExecutionException.class) + .hasCause(rejection); + rejected.getTerminationFuture().get(10, TimeUnit.SECONDS); + assertThat(invoked).isFalse(); + assertThat(successor.getResultFuture().get(10, TimeUnit.SECONDS)).isEqualTo(42); + successor.getTerminationFuture().get(10, TimeUnit.SECONDS); + } + + private static void assertAborted(BackgroundTask task) throws Exception { + assertThatThrownBy(() -> task.getResultFuture().get(10, TimeUnit.SECONDS)) + .isInstanceOf(ExecutionException.class) + .hasCauseInstanceOf(FlinkException.class) + .hasRootCauseMessage("Background task has been aborted."); + task.getTerminationFuture().get(10, TimeUnit.SECONDS); } }