From cfdc2074286e7df64a65f15c5a99589e786a3588 Mon Sep 17 00:00:00 2001
From: Christian Haeubl <christian.haeubl@oracle.com>
Date: Fri, 15 Dec 2023 17:55:49 +0100
Subject: [PATCH] Execute LogManager shutdown hook after all other shutdown
 hooks.

---
 .../svm/core/jdk/JavaUtilSubstitutions.java   | 17 +++++--------
 .../core/jdk/Target_java_lang_Shutdown.java   | 24 +++++++++++++------
 vm/mx.vm/mx_vm_benchmark.py                   |  3 +--
 3 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/JavaUtilSubstitutions.java b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/JavaUtilSubstitutions.java
index 5feb7f6b421a..1e79daf96c57 100644
--- a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/JavaUtilSubstitutions.java
+++ b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/JavaUtilSubstitutions.java
@@ -213,7 +213,7 @@ final class Target_java_util_Currency {
 @TargetClass(className = "java.util.logging.LogManager", onlyWith = JavaLoggingModule.IsPresent.class)
 final class Target_java_util_logging_LogManager {
 
-    @Inject @RecomputeFieldValue(kind = Kind.NewInstance, declClass = AtomicBoolean.class) private AtomicBoolean addedShutdownHook = new AtomicBoolean();
+    @Inject @RecomputeFieldValue(kind = Kind.NewInstance, declClass = AtomicBoolean.class, isFinal = true) private AtomicBoolean addedShutdownHook = new AtomicBoolean();
 
     @Alias static Target_java_util_logging_LogManager manager;
 
@@ -222,22 +222,17 @@ final class Target_java_util_logging_LogManager {
 
     @Substitute
     public static Target_java_util_logging_LogManager getLogManager() {
-        /* First performing logic originally in getLogManager. */
+        /* Logic from original JDK method. */
         if (manager == null) {
-            return manager;
+            return null;
         }
         manager.ensureLogManagerInitialized();
 
-        /* Logic for adding shutdown hook. */
+        /* Add a shutdown hook to close the global handlers. */
         if (!manager.addedShutdownHook.getAndSet(true)) {
-            /* Add a shutdown hook to close the global handlers. */
-            try {
-                Runtime.getRuntime().addShutdownHook(SubstrateUtil.cast(new Target_java_util_logging_LogManager_Cleaner(manager), Thread.class));
-            } catch (IllegalStateException e) {
-                /* If the VM is already shutting down, we do not need to register shutdownHook. */
-            }
+            Runnable hook = SubstrateUtil.cast(new Target_java_util_logging_LogManager_Cleaner(manager), Runnable.class);
+            Util_java_lang_Shutdown.registerLogManagerShutdownHook(hook);
         }
-
         return manager;
     }
 }
diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/Target_java_lang_Shutdown.java b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/Target_java_lang_Shutdown.java
index e02f8eea04bb..743a0eced415 100644
--- a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/Target_java_lang_Shutdown.java
+++ b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/Target_java_lang_Shutdown.java
@@ -29,6 +29,7 @@
 import com.oracle.svm.core.annotate.Substitute;
 import com.oracle.svm.core.annotate.TargetClass;
 import com.oracle.svm.core.annotate.TargetElement;
+import com.oracle.svm.core.util.VMError;
 
 @TargetClass(className = "java.lang.Shutdown")
 public final class Target_java_lang_Shutdown {
@@ -41,12 +42,7 @@ public final class Target_java_lang_Shutdown {
 
     static {
         hooks = new Runnable[Util_java_lang_Shutdown.MAX_SYSTEM_HOOKS];
-        /*
-         * We use the last system hook slot (index 9), which is currently not used by the JDK, for
-         * our own shutdown hooks that are registered during image generation. The JDK currently
-         * uses slots 0, 1, and 2.
-         */
-        hooks[hooks.length - 1] = RuntimeSupport::executeShutdownHooks;
+        hooks[Util_java_lang_Shutdown.NATIVE_IMAGE_SHUTDOWN_HOOKS_SLOT] = RuntimeSupport::executeShutdownHooks;
     }
 
     @Substitute
@@ -73,7 +69,21 @@ final class Util_java_lang_Shutdown {
 
     /**
      * Value *copied* from {@code java.lang.Shutdown.MAX_SYSTEM_HOOKS} so that the value can be used
-     * during image generation (@Alias values are only visible at run time).
+     * during image generation (@Alias values are only visible at run time). The JDK currently uses
+     * slots 0, 1, and 2.
      */
     static final int MAX_SYSTEM_HOOKS = 10;
+
+    static final int LOG_MANAGER_SHUTDOWN_HOOK_SLOT = MAX_SYSTEM_HOOKS - 1;
+    static final int NATIVE_IMAGE_SHUTDOWN_HOOKS_SLOT = LOG_MANAGER_SHUTDOWN_HOOK_SLOT - 1;
+
+    public static void registerLogManagerShutdownHook(Runnable hook) {
+        /*
+         * Execute the LogManager shutdown hook after all other shutdown hooks. This is a workaround
+         * for GR-39429.
+         */
+        Runnable[] hooks = Target_java_lang_Shutdown.hooks;
+        VMError.guarantee(hooks[LOG_MANAGER_SHUTDOWN_HOOK_SLOT] == null, "slot must not be used");
+        hooks[LOG_MANAGER_SHUTDOWN_HOOK_SLOT] = hook;
+    }
 }
diff --git a/vm/mx.vm/mx_vm_benchmark.py b/vm/mx.vm/mx_vm_benchmark.py
index c3f4d6d4d21e..38ee533e5455 100644
--- a/vm/mx.vm/mx_vm_benchmark.py
+++ b/vm/mx.vm/mx_vm_benchmark.py
@@ -216,8 +216,7 @@ def __init__(self, vm, bm_suite, args):
                 base_image_build_args += ['-R:+FlightRecorder',
                                           '-R:StartFlightRecording=filename=default.jfr',
                                           '--enable-monitoring=jfr',
-                                          # We should enable this flag, but after we fix GR-39429.
-                                          # '-R:+JfrBasedExecutionSamplerStatistics'
+                                          '-R:+JfrBasedExecutionSamplerStatistics'
                                           ]
                 for stage in ('instrument-image', 'instrument-run'):
                     if stage in self.stages: