reset callbacks

quarkusio · Jul 31, 2024 · e6710e8 · e6710e8
1 parent ba20b72
commit e6710e8
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 39 deletions.
diff --git a/core/deployment/src/main/java/io/quarkus/runner/bootstrap/StartupActionImpl.java b/core/deployment/src/main/java/io/quarkus/runner/bootstrap/StartupActionImpl.java
@@ -212,34 +212,40 @@ public int runMainClassBlocking(String... args) throws Exception {
             AtomicInteger result = new AtomicInteger();
             Class<?> lifecycleManager = Class.forName(ApplicationLifecycleManager.class.getName(), true, runtimeClassLoader);
             AtomicBoolean alreadyStarted = new AtomicBoolean();
-            lifecycleManager.getDeclaredMethod("setDefaultExitCodeHandler", Consumer.class).invoke(null,
-                    (Consumer<Integer>) result::set);
-            lifecycleManager.getDeclaredMethod("setAlreadyStartedCallback", Consumer.class).invoke(null,
-                    (Consumer<Boolean>) alreadyStarted::set);
-            // force init here
-            Class<?> appClass = Class.forName(className, true, runtimeClassLoader);
-            Method start = appClass.getMethod("main", String[].class);
-            start.invoke(null, (Object) (args == null ? new String[0] : args));
+            Method setDefaultExitCodeHandler = lifecycleManager.getDeclaredMethod("setDefaultExitCodeHandler", Consumer.class);
+            Method setAlreadyStartedCallback = lifecycleManager.getDeclaredMethod("setAlreadyStartedCallback", Consumer.class);
 
-            CountDownLatch latch = new CountDownLatch(1);
-            new Thread(() -> {
-                try {
-                    Class<?> q = Class.forName(Quarkus.class.getName(), true, runtimeClassLoader);
-                    q.getMethod("blockingExit").invoke(null);
-                } catch (Exception e) {
-                    throw new RuntimeException(e);
-                } finally {
-                    latch.countDown();
-                }
-            }).start();
-            latch.await();
+            try {
+                setDefaultExitCodeHandler.invoke(null, (Consumer<Integer>) result::set);
+                setAlreadyStartedCallback.invoke(null, (Consumer<Boolean>) alreadyStarted::set);
+                // force init here
+                Class<?> appClass = Class.forName(className, true, runtimeClassLoader);
+                Method start = appClass.getMethod("main", String[].class);
+                start.invoke(null, (Object) (args == null ? new String[0] : args));
+
+                CountDownLatch latch = new CountDownLatch(1);
+                new Thread(() -> {
+                    try {
+                        Class<?> q = Class.forName(Quarkus.class.getName(), true, runtimeClassLoader);
+                        q.getMethod("blockingExit").invoke(null);
+                    } catch (Exception e) {
+                        throw new RuntimeException(e);
+                    } finally {
+                        latch.countDown();
+                    }
+                }).start();
+                latch.await();
 
-            if (alreadyStarted.get()) {
-                //quarkus was not actually started by the main method
-                //just return
-                return 0;
+                if (alreadyStarted.get()) {
+                    //quarkus was not actually started by the main method
+                    //just return
+                    return 0;
+                }
+                return result.get();
+            } finally {
+                setDefaultExitCodeHandler.invoke(null, (Consumer<?>) null);
+                setAlreadyStartedCallback.invoke(null, (Consumer<?>) null);
             }
-            return result.get();
         } finally {
             for (var closeTask : runtimeCloseTasks) {
                 try {

diff --git a/core/runtime/src/main/java/io/quarkus/runtime/ApplicationLifecycleManager.java b/core/runtime/src/main/java/io/quarkus/runtime/ApplicationLifecycleManager.java
@@ -2,7 +2,6 @@
 
 import java.util.List;
 import java.util.Locale;
-import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
@@ -51,17 +50,15 @@ public class ApplicationLifecycleManager {
 
     // used by ShutdownEvent to propagate the information about shutdown reason
     public static volatile ShutdownEvent.ShutdownReason shutdownReason = ShutdownEvent.ShutdownReason.STANDARD;
-    private static volatile BiConsumer<Integer, Throwable> defaultExitCodeHandler = new BiConsumer<Integer, Throwable>() {
-        @Override
-        public void accept(Integer integer, Throwable cause) {
-            Logger logger = Logger.getLogger(Application.class);
-            logger.debugf("Shutting down with exit code %s", integer);
-            if (logger.isTraceEnabled()) {
-                logger.tracef(new RuntimeException("Shutdown Stack Trace"), "Shutdown triggered");
-            }
-            System.exit(integer);
+    private static final BiConsumer<Integer, Throwable> MAIN_EXIT_CODE_HANDLER = (integer, cause) -> {
+        Logger logger = Logger.getLogger(Application.class);
+        logger.debugf("Shutting down with exit code %s", integer);
+        if (logger.isTraceEnabled()) {
+            logger.tracef(new RuntimeException("Shutdown Stack Trace"), "Shutdown triggered");
         }
+        System.exit(integer);
     };
+    private static volatile BiConsumer<Integer, Throwable> defaultExitCodeHandler = MAIN_EXIT_CODE_HANDLER;
 
     private ApplicationLifecycleManager() {
 
@@ -357,7 +354,9 @@ public static boolean isVmShuttingDown() {
      * @param defaultExitCodeHandler the new default exit handler
      */
     public static void setDefaultExitCodeHandler(BiConsumer<Integer, Throwable> defaultExitCodeHandler) {
-        Objects.requireNonNull(defaultExitCodeHandler);
+        if (defaultExitCodeHandler == null) {
+            defaultExitCodeHandler = MAIN_EXIT_CODE_HANDLER;
+        }
         ApplicationLifecycleManager.defaultExitCodeHandler = defaultExitCodeHandler;
     }
 
@@ -370,12 +369,18 @@ public static void setDefaultExitCodeHandler(BiConsumer<Integer, Throwable> defa
      *
      * @param defaultExitCodeHandler the new default exit handler
      */
+    // Used by StartupActionImpl via reflection
     public static void setDefaultExitCodeHandler(Consumer<Integer> defaultExitCodeHandler) {
-        setDefaultExitCodeHandler((exitCode, cause) -> defaultExitCodeHandler.accept(exitCode));
+        BiConsumer<Integer, Throwable> biConsumer = defaultExitCodeHandler == null ? null
+                : (exitCode, cause) -> defaultExitCodeHandler.accept(exitCode);
+        setDefaultExitCodeHandler(biConsumer);
     }
 
-    public static void setAlreadyStartedCallback(Consumer<Boolean> callback) {
-        alreadyStartedCallback = callback;
+    @SuppressWarnings("unused")
+    // Used by StartupActionImpl via reflection
+    public static void setAlreadyStartedCallback(Consumer<Boolean> alreadyStartedCallback) {
+        ApplicationLifecycleManager.alreadyStartedCallback = alreadyStartedCallback != null ? alreadyStartedCallback : b -> {
+        };
     }
 
     /**