[FLINK-37607] Fix RpcEndpoint#MainThreadExecutor lost scheduling tasks when not running

Author: zhouhengtong

flink-rpc/flink-rpc-core/src/main/java/org/apache/flink/runtime/rpc/RpcEndpoint.java

@@ -50,6 +50,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 import static org.apache.flink.util.Preconditions.checkNotNull;
 
@@ -131,7 +132,7 @@ public abstract class RpcEndpoint implements RpcGateway, AutoCloseableAsync {
      * <p>IMPORTANT: the running state is not thread safe and can be used only in the main thread of
      * the rpc endpoint.
      */
-    private boolean isRunning;
+    private CompletableFuture<Void> isRunningFuture;
 
     /**
      * Initializes the RPC endpoint.
@@ -147,8 +148,14 @@ protected RpcEndpoint(
         this.rpcServer = rpcService.startServer(this, loggingContext);
         this.resourceRegistry = new CloseableRegistry();
 
+        this.isRunningFuture = new CompletableFuture<>();
         this.mainThreadExecutor =
-                new MainThreadExecutor(rpcServer, this::validateRunsInMainThread, endpointId);
+                new MainThreadExecutor(
+                        rpcServer,
+                        this::validateRunsInMainThread,
+                        Executors.newSingleThreadScheduledExecutor(
+                                new ExecutorThreadFactory(endpointId + "-main-scheduler")),
+                        () -> isRunningFuture);
         registerResource(this.mainThreadExecutor);
     }
 
@@ -187,7 +194,7 @@ public String getEndpointId() {
      */
     protected boolean isRunning() {
         validateRunsInMainThread();
-        return isRunning;
+        return isRunningFuture.isDone();
     }
 
     // ------------------------------------------------------------------------
@@ -210,7 +217,7 @@ public final void start() {
      */
     public final void internalCallOnStart() throws Exception {
         validateRunsInMainThread();
-        isRunning = true;
+        isRunningFuture.complete(null);
         onStart();
     }
 
@@ -253,7 +260,7 @@ public final CompletableFuture<Void> internalCallOnStop() {
                     new RuntimeException("Close resource registry fail", e));
         }
         stopFuture = CompletableFuture.allOf(stopFuture, onStop());
-        isRunning = false;
+        isRunningFuture = new CompletableFuture<>();
         return stopFuture;
     }
 
@@ -489,6 +496,13 @@ protected static class MainThreadExecutor implements ComponentMainThreadExecutor
          */
         private final ScheduledExecutorService mainScheduledExecutor;
 
+        /**
+         * The future indicate the gateway whether is running, NOTICE: can't change the state of
+         * future.
+         */
+        private final Supplier<CompletableFuture<Void>> getRunningFuture;
+
+        @VisibleForTesting
         MainThreadExecutor(
                 MainThreadExecutable gateway, Runnable mainThreadCheck, String endpointId) {
             this(
@@ -503,9 +517,22 @@ protected static class MainThreadExecutor implements ComponentMainThreadExecutor
                 MainThreadExecutable gateway,
                 Runnable mainThreadCheck,
                 ScheduledExecutorService mainScheduledExecutor) {
+            this(
+                    gateway,
+                    mainThreadCheck,
+                    mainScheduledExecutor,
+                    () -> CompletableFuture.completedFuture(null));
+        }
+
+        MainThreadExecutor(
+                MainThreadExecutable gateway,
+                Runnable mainThreadCheck,
+                ScheduledExecutorService mainScheduledExecutor,
+                Supplier<CompletableFuture<Void>> getRunningFuture) {
             this.gateway = Preconditions.checkNotNull(gateway);
             this.mainThreadCheck = Preconditions.checkNotNull(mainThreadCheck);
             this.mainScheduledExecutor = mainScheduledExecutor;
+            this.getRunningFuture = getRunningFuture;
         }
 
         @Override
@@ -526,14 +553,21 @@ public void execute(@Nonnull Runnable command) {
         public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
             final long delayMillis = TimeUnit.MILLISECONDS.convert(delay, unit);
             FutureTask<Void> ft = new FutureTask<>(command, null);
-            if (mainScheduledExecutor.isShutdown()) {
-                log.warn(
-                        "The scheduled executor service is shutdown and ignores the command {}",
-                        command);
-            } else {
-                mainScheduledExecutor.schedule(
-                        () -> gateway.runAsync(ft), delayMillis, TimeUnit.MILLISECONDS);
-            }
+            getRunningFuture
+                    .get()
+                    .thenAccept(
+                            ignore -> {
+                                if (mainScheduledExecutor.isShutdown()) {
+                                    log.warn(
+                                            "The scheduled executor service is shutdown and ignores the command {}",
+                                            command);
+                                } else {
+                                    mainScheduledExecutor.schedule(
+                                            () -> gateway.runAsync(ft),
+                                            delayMillis,
+                                            TimeUnit.MILLISECONDS);
+                                }
+                            });
             return new ScheduledFutureAdapter<>(ft, delayMillis, TimeUnit.MILLISECONDS);
         }
 
@@ -551,14 +585,21 @@ public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit)
         public <V> ScheduledFuture<V> schedule(Callable<V> callable, long delay, TimeUnit unit) {
             final long delayMillis = TimeUnit.MILLISECONDS.convert(delay, unit);
             FutureTask<V> ft = new FutureTask<>(callable);
-            if (mainScheduledExecutor.isShutdown()) {
-                log.warn(
-                        "The scheduled executor service is shutdown and ignores the callable {}",
-                        callable);
-            } else {
-                mainScheduledExecutor.schedule(
-                        () -> gateway.runAsync(ft), delayMillis, TimeUnit.MILLISECONDS);
-            }
+            getRunningFuture
+                    .get()
+                    .thenAccept(
+                            ignore -> {
+                                if (mainScheduledExecutor.isShutdown()) {
+                                    log.warn(
+                                            "The scheduled executor service is shutdown and ignores the callable {}",
+                                            callable);
+                                } else {
+                                    mainScheduledExecutor.schedule(
+                                            () -> gateway.runAsync(ft),
+                                            delayMillis,
+                                            TimeUnit.MILLISECONDS);
+                                }
+                            });
             return new ScheduledFutureAdapter<>(ft, delayMillis, TimeUnit.MILLISECONDS);
         }
 

flink-runtime/src/test/java/org/apache/flink/runtime/rpc/RpcEndpointTest.java

@@ -40,6 +40,8 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 /** Tests for the RpcEndpoint, its self gateways and MainThreadExecutor scheduling command. */
 class RpcEndpointTest {
@@ -472,6 +474,36 @@ void testCallAsyncTimeout() throws InterruptedException, ExecutionException, Tim
         }
     }
 
+    /** Test schedule task when the RPC is running. */
+    @Test
+    public void testScheduleTaskAfterStart() throws Exception {
+        final RpcEndpoint endpoint = new BaseEndpoint(rpcService);
+        final CompletableFuture<Void> taskCompletedFuture = new CompletableFuture<>();
+        try {
+            final Duration expectedDelay = Duration.ofSeconds(0);
+            ScheduledFuture<?> future =
+                    endpoint.getMainThreadExecutor()
+                            .schedule(
+                                    () -> taskCompletedFuture.complete(null),
+                                    expectedDelay.toMillis(),
+                                    TimeUnit.MILLISECONDS);
+            assertThrows(
+                    TimeoutException.class,
+                    () ->
+                            taskCompletedFuture.get(
+                                    expectedDelay.toMillis() + 3000L, TimeUnit.MILLISECONDS));
+            assertFalse(taskCompletedFuture.isDone());
+            assertFalse(future.isDone());
+
+            endpoint.start();
+
+            taskCompletedFuture.get(1000L, TimeUnit.MILLISECONDS);
+        } finally {
+            RpcUtils.terminateRpcEndpoint(endpoint);
+            endpoint.validateResourceClosed();
+        }
+    }
+
     private static class TestMainThreadExecutable implements MainThreadExecutable {
 
         private final Consumer<Runnable> scheduleRunAsyncConsumer;