Shutdown hooks must run after joining non-daemon threads

Author: Christian Wimmer

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/JavaMainWrapper.java

@@ -29,6 +29,8 @@
 import static com.oracle.svm.core.option.SubstrateOptionsParser.BooleanOptionFormat.NAME_VALUE;
 import static com.oracle.svm.core.option.SubstrateOptionsParser.BooleanOptionFormat.PLUS_MINUS;
 
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
 //Checkstyle: allow reflection
 import java.lang.reflect.Method;
 import java.util.ArrayList;
@@ -85,21 +87,18 @@ public class JavaMainWrapper {
 
     public static class JavaMainSupport {
 
-        private final Method javaMainMethod;
+        final MethodHandle javaMainHandle;
+        final String javaMainClassName;
 
         @Platforms(Platform.HOSTED_ONLY.class)
-        public JavaMainSupport(Method javaMainMethod) {
-            this.javaMainMethod = javaMainMethod;
-        }
-
-        private Method getJavaMainMethod() {
-            assert javaMainMethod != null;
-            return javaMainMethod;
+        public JavaMainSupport(Method javaMainMethod) throws IllegalAccessException {
+            this.javaMainHandle = MethodHandles.lookup().unreflect(javaMainMethod);
+            this.javaMainClassName = javaMainMethod.getDeclaringClass().getName();
         }
 
         public String getJavaCommand() {
-            if (javaMainMethod != null && mainArgs != null) {
-                StringBuilder commandLine = new StringBuilder(javaMainMethod.getDeclaringClass().getName());
+            if (mainArgs != null) {
+                StringBuilder commandLine = new StringBuilder(javaMainClassName);
 
                 for (String arg : mainArgs) {
                     commandLine.append(' ');
@@ -160,28 +159,40 @@ public static int run(int paramArgc, CCharPointerPointer paramArgv) throws Excep
             args = RuntimePropertyParser.parse(args);
         }
         mainArgs = args;
-        final RuntimeSupport rs = RuntimeSupport.getRuntimeSupport();
+
         try {
-            try {
-                if (AllocationSite.Options.AllocationProfiling.getValue()) {
-                    Runtime.getRuntime().addShutdownHook(new AllocationSite.AllocationProfilingShutdownHook());
-                }
-                if (SubstrateOptions.PrintGCSummary.getValue()) {
-                    Runtime.getRuntime().addShutdownHook(new PrintGCSummaryShutdownHook());
-                }
-                rs.executeStartupHooks();
-                ImageSingletons.lookup(JavaMainSupport.class).getJavaMainMethod().invoke(null, (Object) mainArgs);
-            } catch (Throwable ex) {
-                JavaThreads.dispatchUncaughtException(Thread.currentThread(), ex);
+            if (AllocationSite.Options.AllocationProfiling.getValue()) {
+                Runtime.getRuntime().addShutdownHook(new AllocationSite.AllocationProfilingShutdownHook());
             }
-        } finally {
-            /* Shutdown before joining non-daemon threads. */
-            rs.shutdown();
-        }
+            if (SubstrateOptions.PrintGCSummary.getValue()) {
+                Runtime.getRuntime().addShutdownHook(new PrintGCSummaryShutdownHook());
+            }
+            RuntimeSupport.getRuntimeSupport().executeStartupHooks();
+
+            /*
+             * Invoke the application's main method. Invoking the main method via a method handle
+             * preserves exceptions, while invoking the main method via reflection would wrap
+             * exceptions in a InvocationTargetException.
+             */
+            ImageSingletons.lookup(JavaMainSupport.class).javaMainHandle.invokeExact(args);
 
-        JavaThreads.singleton().joinAllNonDaemons();
-        Counter.logValues();
+        } catch (Throwable ex) {
+            JavaThreads.dispatchUncaughtException(Thread.currentThread(), ex);
 
+        } finally {
+            /*
+             * Shutdown sequence: First wait for all non-daemon threads to exit.
+             */
+            JavaThreads.singleton().joinAllNonDaemons();
+            /*
+             * Run shutdown hooks (both our own hooks and application-registered hooks. Note that
+             * this can start new non-daemon threads. We are not responsible to wait until they have
+             * exited.
+             */
+            RuntimeSupport.getRuntimeSupport().shutdown();
+
+            Counter.logValues();
+        }
         return 0;
     }
 

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/JavaLangSubstitutions.java

@@ -48,6 +48,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Predicate;
 
 import org.graalvm.compiler.api.replacements.Fold;
 import org.graalvm.compiler.core.common.SuppressFBWarnings;
@@ -65,6 +66,7 @@
 import org.graalvm.word.WordFactory;
 
 import com.oracle.svm.core.MonitorSupport;
+import com.oracle.svm.core.SubstrateOptions;
 import com.oracle.svm.core.UnsafeAccess;
 import com.oracle.svm.core.annotate.Alias;
 import com.oracle.svm.core.annotate.AutomaticFeature;
@@ -545,6 +547,20 @@ static Object command(Object arg) {
     }
 }
 
+final class IsSingleThreaded implements Predicate<Class<?>> {
+    @Override
+    public boolean test(Class<?> t) {
+        return !SubstrateOptions.MultiThreaded.getValue();
+    }
+}
+
+final class IsMultiThreaded implements Predicate<Class<?>> {
+    @Override
+    public boolean test(Class<?> t) {
+        return SubstrateOptions.MultiThreaded.getValue();
+    }
+}
+
 @TargetClass(className = "java.lang.ApplicationShutdownHooks")
 final class Target_java_lang_ApplicationShutdownHooks {
 
@@ -559,9 +575,14 @@ final class Target_java_lang_ApplicationShutdownHooks {
     /**
      * Instead of starting all the threads in {@link #hooks}, just run the {@link Runnable}s one
      * after another.
+     *
+     * We need this substitution in single-threaded mode, where we cannot start new threads but
+     * still want to support shutdown hooks. In multi-threaded mode, this substitution is not
+     * present, i.e., the original JDK code runs the shutdown hooks in separate threads.
      */
     @Substitute
-    static void runHooks() {
+    @TargetElement(name = "runHooks", onlyWith = IsSingleThreaded.class)
+    static void runHooksSingleThreaded() {
         /* Claim all the hooks. */
         final Collection<Thread> threads;
         /* Checkstyle: allow synchronization. */
@@ -588,6 +609,10 @@ static void runHooks() {
         }
     }
 
+    @Alias
+    @TargetElement(name = "runHooks", onlyWith = IsMultiThreaded.class)
+    static native void runHooksMultiThreaded();
+
     /**
      * Interpose so that the first time someone adds an ApplicationShutdownHook, I set up a shutdown
      * hook to run all the ApplicationShutdownHooks. Then the rest of this method is copied from
@@ -636,7 +661,11 @@ public static void initializeOnce() {
                                         new Runnable() {
                                             @Override
                                             public void run() {
-                                                Target_java_lang_ApplicationShutdownHooks.runHooks();
+                                                if (SubstrateOptions.MultiThreaded.getValue()) {
+                                                    Target_java_lang_ApplicationShutdownHooks.runHooksMultiThreaded();
+                                                } else {
+                                                    Target_java_lang_ApplicationShutdownHooks.runHooksSingleThreaded();
+                                                }
                                             }
                                         });
                     } catch (InternalError ie) {