fix: adds a handler to gracefully shutdown (#5895)

* fix: adds a handler to gracefully shutdown service

Couple changes related to application lifecycle management. Firstly, adds a shutdown handler
that gracefully shuts down the service when the jvm determines its time to shut down (e.g.
when it receives a termination signal). Secondly, this patch reorganizes some of the startup,
steady-state, and shutdown code to make shutdown easier to reason about. Specifically, all
these methods are now called from the same thread, so both the thread-safety and order of
execution are guaranteed. All the shutdown hook does is notify the main thread and then wait
for it to exit.

ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/ConnectExecutable.java

@@ -21,6 +21,7 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
 import org.apache.kafka.common.utils.Utils;
 import org.apache.kafka.connect.cli.ConnectDistributed;
 import org.apache.kafka.connect.errors.ConnectException;
@@ -35,6 +36,7 @@ public final class ConnectExecutable implements Executable {
   private final ConnectDistributed connectDistributed;
   private final Map<String, String> workerProps;
   private Connect connect;
+  private final CountDownLatch terminateLatch = new CountDownLatch(1);
 
   public static ConnectExecutable of(final String configFile) throws IOException {
     final Map<String, String> workerProps = !configFile.isEmpty()
@@ -64,16 +66,19 @@ public void startAsync() {
   }
 
   @Override
-  public void triggerShutdown() {
+  public void shutdown() {
     if (connect != null) {
       connect.stop();
     }
   }
 
   @Override
-  public void awaitTerminated() {
-    if (connect != null) {
-      connect.awaitStop();
-    }
+  public void notifyTerminated() {
+    terminateLatch.countDown();
+  }
+
+  @Override
+  public void awaitTerminated() throws InterruptedException {
+    terminateLatch.await();
   }
 }

ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/Executable.java

@@ -21,19 +21,24 @@
 public interface Executable {
 
   /**
-   * Starts the executable asynchronously.
+   * Starts the executable asynchronously. Guaranteed to be called before shutdown.
    */
   default void startAsync() throws Exception {}
 
   /**
-   * Triggers a shutdown asynchronously, in order to ensure that the shutdown
-   * has finished use {@link #awaitTerminated()}
+   * Called to notify threads awaiting termination (see #awaitTerminated)
+   * that it's time to shutdown.
    */
-  default void triggerShutdown() throws Exception {}
+  default void notifyTerminated() {}
 
   /**
-   * Awaits the {@link #triggerShutdown()} to finish. This is a blocking
-   * operation.
+   * Shutdown the service.
+   */
+  default void shutdown() throws Exception {}
+
+  /**
+   * Awaits the {@link #notifyTerminated()} notification. This is a blocking
+   * operation. Guaranteed to be called before shutdown.
    */
   default void awaitTerminated() throws InterruptedException {}
 }

ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlRestApplication.java

@@ -128,7 +128,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -169,7 +169,6 @@ public final class KsqlRestApplication implements Executable {
   private final Optional<HeartbeatAgent> heartbeatAgent;
   private final Optional<LagReportingAgent> lagReportingAgent;
   private final PullQueryExecutor pullQueryExecutor;
-  private final AtomicBoolean shuttingDown = new AtomicBoolean(false);
   private final ServerInfoResource serverInfoResource;
   private final Optional<HeartbeatResource> heartbeatResource;
   private final Optional<ClusterStatusResource> clusterStatusResource;
@@ -185,7 +184,7 @@ public final class KsqlRestApplication implements Executable {
 
   // The startup thread that can be interrupted if necessary during shutdown.  This should only
   // happen if startup hangs.
-  private volatile Thread startAsyncThread;
+  private AtomicReference<Thread> startAsyncThreadRef = new AtomicReference<>(null);
 
   public static SourceName getCommandsStreamName() {
     return COMMANDS_STREAM_NAME;
@@ -305,7 +304,7 @@ public void startAsync() {
         pullQueryExecutor
     );
 
-    startAsyncThread = Thread.currentThread();
+    startAsyncThreadRef.set(Thread.currentThread());
     try {
       final Endpoints endpoints = new KsqlServerEndpoints(
           ksqlEngine,
@@ -340,7 +339,7 @@ public void startAsync() {
     } catch (AbortApplicationStartException e) {
       log.error("Aborting application start", e);
     } finally {
-      startAsyncThread = null;
+      startAsyncThreadRef.set(null);
     }
   }
 
@@ -392,14 +391,14 @@ private void waitForPreconditions() {
           1000,
           30000,
           this::checkPreconditions,
-          shuttingDown::get,
+          terminatedFuture::isDone,
           predicates
       );
     } catch (KsqlFailedPrecondition e) {
       log.error("Failed to meet preconditions. Exiting...", e);
     }
 
-    if (shuttingDown.get()) {
+    if (terminatedFuture.isDone()) {
       throw new AbortApplicationStartException(
           "Shutting down application during waitForPreconditions");
     }
@@ -441,17 +440,19 @@ private void initialize(final KsqlConfig configWithApplicationServer) {
     serverState.setReady();
   }
 
-  @SuppressWarnings("checkstyle:NPathComplexity")
   @Override
-  public void triggerShutdown() {
-    log.debug("ksqlDB triggerShutdown called");
-    // First, make sure the server wasn't stuck in startup.  Set the shutdown flag and interrupt the
-    // startup thread if it's been hanging.
-    shuttingDown.set(true);
+  public void notifyTerminated() {
+    terminatedFuture.complete(null);
+    final Thread startAsyncThread = startAsyncThreadRef.get();
     if (startAsyncThread != null) {
       startAsyncThread.interrupt();
     }
+  }
 
+  @SuppressWarnings("checkstyle:NPathComplexity")
+  @Override
+  public void shutdown() {
+    log.info("ksqlDB shutdown called");
     try {
       streamedQueryResource.closeMetrics();
     } catch (final Exception e) {
@@ -496,9 +497,7 @@ public void triggerShutdown() {
 
     shutdownAdditionalAgents();
 
-    log.debug("ksqlDB triggerShutdown complete");
-
-    terminatedFuture.complete(null);
+    log.info("ksqlDB shutdown complete");
   }
 
   @Override

ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlServerMain.java

@@ -24,6 +24,8 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executor;
 import org.apache.kafka.streams.StreamsConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -32,6 +34,7 @@ public class KsqlServerMain {
 
   private static final Logger log = LoggerFactory.getLogger(KsqlServerMain.class);
 
+  private final Executor shutdownHandler;
   private final Executable executable;
 
   public static void main(final String[] args) {
@@ -55,29 +58,46 @@ public static void main(final String[] args) {
       final Optional<String> queriesFile = serverOptions.getQueriesFile(properties);
       final Executable executable = createExecutable(
           properties, queriesFile, installDir, ksqlConfig);
-      new KsqlServerMain(executable).tryStartApp();
+      new KsqlServerMain(
+          executable,
+          r -> Runtime.getRuntime().addShutdownHook(new Thread(r))
+      ).tryStartApp();
     } catch (final Exception e) {
       log.error("Failed to start KSQL", e);
       System.exit(-1);
     }
   }
 
-  KsqlServerMain(final Executable executable) {
+  KsqlServerMain(final Executable executable, final Executor shutdownHandler) {
     this.executable = Objects.requireNonNull(executable, "executable");
+    this.shutdownHandler = Objects.requireNonNull(shutdownHandler, "shutdownHandler");
   }
 
   void tryStartApp() throws Exception {
+    final CountDownLatch latch = new CountDownLatch(1);
+    shutdownHandler.execute(() -> {
+      executable.notifyTerminated();
+      try {
+        latch.await();
+      } catch (final InterruptedException e) {
+        throw new RuntimeException(e);
+      }
+    });
     try {
-      log.info("Starting server");
-      executable.startAsync();
-      log.info("Server up and running");
-      executable.awaitTerminated();
-    } catch (Throwable t) {
-      log.error("Unhandled exception in server startup", t);
-      throw t;
+      try {
+        log.info("Starting server");
+        executable.startAsync();
+        log.info("Server up and running");
+        executable.awaitTerminated();
+      } catch (Throwable t) {
+        log.error("Unhandled exception in server startup", t);
+        throw t;
+      } finally {
+        log.info("Server shutting down");
+        executable.shutdown();
+      }
     } finally {
-      log.info("Server shutting down");
-      executable.triggerShutdown();
+      latch.countDown();
     }
   }
 

ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/MultiExecutable.java

@@ -43,8 +43,13 @@ public void startAsync() throws Exception {
   }
 
   @Override
-  public void triggerShutdown() throws Exception {
-    doAction(Executable::triggerShutdown);
+  public void shutdown() throws Exception {
+    doAction(Executable::shutdown);
+  }
+
+  @Override
+  public void notifyTerminated() {
+    doAction(Executable::notifyTerminated);
   }
 
   @Override

ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/StandaloneExecutor.java

@@ -129,12 +129,17 @@ public void startAsync() {
       versionChecker.start(KsqlModuleType.SERVER, properties);
     } catch (final Exception e) {
       log.error("Failed to start KSQL Server with query file: " + queriesFile, e);
-      triggerShutdown();
       throw e;
     }
   }
 
-  public void triggerShutdown() {
+  @Override
+  public void notifyTerminated() {
+    shutdownLatch.countDown();
+  }
+
+  @Override
+  public void shutdown() {
     try {
       ksqlEngine.close();
     } catch (final Exception e) {
@@ -145,7 +150,6 @@ public void triggerShutdown() {
     } catch (final Exception e) {
       log.warn("Failed to cleanly shutdown services", e);
     }
-    shutdownLatch.countDown();
   }
 
   @Override

ksqldb-rest-app/src/test/java/io/confluent/ksql/rest/server/ConnectIntegrationTest.java

@@ -95,7 +95,7 @@ public static void setUpClass() {
 
   @AfterClass
   public static void tearDownClass() {
-    CONNECT.triggerShutdown();
+    CONNECT.shutdown();
   }
 
   private KsqlRestClient ksqlRestClient;

ksqldb-rest-app/src/test/java/io/confluent/ksql/rest/server/KsqlRestApplicationTest.java

@@ -202,7 +202,7 @@ public void tearDown() {
   @Test
   public void shouldCloseServiceContextOnClose() {
     // When:
-    app.triggerShutdown();
+    app.shutdown();
 
     // Then:
     verify(serviceContext).close();
@@ -211,7 +211,7 @@ public void shouldCloseServiceContextOnClose() {
   @Test
   public void shouldCloseSecurityExtensionOnClose() {
     // When:
-    app.triggerShutdown();
+    app.shutdown();
 
     // Then:
     verify(securityExtension).close();

ksqldb-rest-app/src/test/java/io/confluent/ksql/rest/server/KsqlServerMainTest.java

@@ -15,7 +15,9 @@
 
 package io.confluent.ksql.rest.server;
 
+import static org.easymock.EasyMock.capture;
 import static org.easymock.EasyMock.expectLastCall;
+import static org.easymock.EasyMock.newCapture;
 import static org.easymock.EasyMock.replay;
 import static org.easymock.EasyMock.verify;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -27,6 +29,8 @@
 
 import io.confluent.ksql.util.KsqlServerException;
 import java.io.File;
+import java.util.concurrent.Executor;
+import org.easymock.Capture;
 import org.easymock.EasyMockRunner;
 import org.easymock.Mock;
 import org.easymock.MockType;
@@ -40,12 +44,14 @@ public class KsqlServerMainTest {
 
   @Mock(MockType.NICE)
   private Executable executable;
+  @Mock(MockType.NICE)
+  private Executor shutdownHandler;
 
   private final File mockStreamsStateDir = mock(File.class);
 
   @Before
   public void setUp() {
-    main = new KsqlServerMain(executable);
+    main = new KsqlServerMain(executable, shutdownHandler);
     when(mockStreamsStateDir.exists()).thenReturn(true);
     when(mockStreamsStateDir.mkdirs()).thenReturn(true);
     when(mockStreamsStateDir.isDirectory()).thenReturn(true);
@@ -57,7 +63,7 @@ public void setUp() {
   @Test
   public void shouldStopAppOnJoin() throws Exception {
     // Given:
-    executable.triggerShutdown();
+    executable.shutdown();
     expectLastCall();
     replay(executable);
 
@@ -74,7 +80,7 @@ public void shouldStopAppOnErrorStarting() throws Exception {
     executable.startAsync();
     expectLastCall().andThrow(new RuntimeException("Boom"));
 
-    executable.triggerShutdown();
+    executable.shutdown();
     expectLastCall();
     replay(executable);
 
@@ -90,6 +96,24 @@ public void shouldStopAppOnErrorStarting() throws Exception {
     verify(executable);
   }
 
+  @Test
+  public void shouldNotifyAppOnTerminate() throws Exception {
+    // Given:
+    final Capture<Runnable> captureShutdownHandler = newCapture();
+    shutdownHandler.execute(capture(captureShutdownHandler));
+    executable.notifyTerminated();
+    expectLastCall();
+    replay(shutdownHandler, executable);
+    main.tryStartApp();
+    final Runnable handler = captureShutdownHandler.getValue();
+
+    // When:
+    handler.run();
+
+    // Then:
+    verify(executable);
+  }
+
   @Test
   public void shouldFailIfStreamsStateDirectoryCannotBeCreated() {
     // Given: