Break the initialization cycle in NIO/cgroups code in a simpler and more thread-safe way.

christianhaeubl · zakkak · commit df95ad5257ff · 2023-10-10T13:26:32.000+03:00
(cherry picked from commit af74d9f)
diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/heap/PhysicalMemory.java b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/heap/PhysicalMemory.java
@@ -24,8 +24,6 @@
  */
 package com.oracle.svm.core.heap;
 
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
 
 import org.graalvm.nativeimage.ImageSingletons;
@@ -34,7 +32,6 @@
 
 import com.oracle.svm.core.Containers;
 import com.oracle.svm.core.SubstrateOptions;
-import com.oracle.svm.core.jdk.UninterruptibleUtils.AtomicInteger;
 import com.oracle.svm.core.stack.StackOverflowCheck;
 import com.oracle.svm.core.thread.PlatformThreads;
 import com.oracle.svm.core.thread.VMOperation;
@@ -58,22 +55,16 @@ default boolean hasSize() {
         UnsignedWord size();
     }
 
-    private static final CountDownLatch CACHED_SIZE_AVAIL_LATCH = new CountDownLatch(1);
-    private static final AtomicInteger INITIALIZING = new AtomicInteger(0);
     private static final ReentrantLock LOCK = new ReentrantLock();
     private static final UnsignedWord UNSET_SENTINEL = UnsignedUtils.MAX_VALUE;
     private static UnsignedWord cachedSize = UNSET_SENTINEL;
 
     public static boolean isInitialized() {
-        return INITIALIZING.get() > 1;
+        return cachedSize != UNSET_SENTINEL;
     }
 
-    /**
-     *
-     * @return {@code true} when PhycialMemory.size() is still initializing
-     */
-    private static boolean isInitializing() {
-        return INITIALIZING.get() == 1;
+    public static boolean isInitializationInProgress() {
+        return LOCK.isHeldByCurrentThread();
     }
 
     /**
@@ -92,42 +83,23 @@ public static UnsignedWord size() {
             throw VMError.shouldNotReachHere("Accessing the physical memory size may require allocation and synchronization");
         }
 
-        LOCK.lock();
-        try {
-            if (!isInitialized()) {
-                if (isInitializing()) {
-                    /*
-                     * Recursive initializations need to wait for the one initializing thread to
-                     * finish so as to get correct reads of the cachedSize value.
-                     */
-                    try {
-                        boolean expired = !CACHED_SIZE_AVAIL_LATCH.await(1L, TimeUnit.SECONDS);
-                        if (expired) {
-                            throw new InternalError("Expired latch!");
-                        }
-                        VMError.guarantee(cachedSize != UNSET_SENTINEL, "Expected cached size to be set");
-                        return cachedSize;
-                    } catch (InterruptedException e) {
-                        throw VMError.shouldNotReachHere("Interrupt on countdown latch!");
+        if (!isInitialized()) {
+            long memoryLimit = SubstrateOptions.MaxRAM.getValue();
+            if (memoryLimit > 0) {
+                cachedSize = WordFactory.unsigned(memoryLimit);
+            } else {
+                LOCK.lock();
+                try {
+                    if (!isInitialized()) {
+                        memoryLimit = Containers.memoryLimitInBytes();
+                        cachedSize = memoryLimit > 0
+                                        ? WordFactory.unsigned(memoryLimit)
+                                        : ImageSingletons.lookup(PhysicalMemorySupport.class).size();
                     }
+                } finally {
+                    LOCK.unlock();
                 }
-                INITIALIZING.incrementAndGet();
-                long memoryLimit = SubstrateOptions.MaxRAM.getValue();
-                if (memoryLimit > 0) {
-                    cachedSize = WordFactory.unsigned(memoryLimit);
-                } else {
-                    memoryLimit = Containers.memoryLimitInBytes();
-                    cachedSize = memoryLimit > 0
-                                    ? WordFactory.unsigned(memoryLimit)
-                                    : ImageSingletons.lookup(PhysicalMemorySupport.class).size();
-                }
-                // Now that we have set the cachedSize let other threads know it's
-                // available to use.
-                INITIALIZING.incrementAndGet();
-                CACHED_SIZE_AVAIL_LATCH.countDown();
             }
-        } finally {
-            LOCK.unlock();
         }
 
         return cachedSize;
diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/Target_jdk_internal_misc_VM.java b/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/Target_jdk_internal_misc_VM.java
@@ -25,7 +25,6 @@
 package com.oracle.svm.core.jdk;
 
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import com.oracle.svm.core.NeverInline;
 import com.oracle.svm.core.SubstrateOptions;
@@ -38,7 +37,6 @@
 import com.oracle.svm.core.annotate.TargetClass;
 import com.oracle.svm.core.heap.PhysicalMemory;
 import com.oracle.svm.core.snippets.KnownIntrinsics;
-import com.oracle.svm.core.util.VMError;
 
 import jdk.internal.misc.Unsafe;
 
@@ -72,44 +70,26 @@ public static ClassLoader latestUserDefinedLoader0() {
     private static long directMemory;
     @Alias @InjectAccessors(PageAlignDirectMemoryAccessors.class) //
     private static boolean pageAlignDirectMemory;
-
-    @Alias //
-    public static native void initLevel(int newVal);
-
-    @Alias //
-    public static native int initLevel();
 }
 
 final class DirectMemoryAccessors {
+    private static final long DIRECT_MEMORY_DURING_INITIALIZATION = 25 * 1024 * 1024;
 
     /*
-     * Full initialization is two-staged. First, we init directMemory to a static value (25MB) so
-     * that initialization of PhysicalMemory has a chance to finish. At that point isInintialized
-     * will be false, since we need to (potentially) set the value to the actual configured heap
-     * size. That can only be done once PhysicalMemory init completed. We'd introduce a cycle
-     * otherwise.
+     * Not volatile to avoid a memory barrier when reading the values. Instead, an explicit barrier
+     * is inserted when writing the values.
      */
-    private static boolean isInitialized;
-    private static final int INITIALIZING = 1;
-    private static final int INITIALIZED = 2;
-    private static final AtomicInteger INIT_COUNT = new AtomicInteger();
-    private static final long STATIC_DIRECT_MEMORY_AMOUNT = 25 * 1024 * 1024;
+    private static boolean initialized;
     private static long directMemory;
 
     static long getDirectMemory() {
-        if (!isInitialized) {
-            initialize();
+        if (!initialized) {
+            return tryInitialize();
         }
         return directMemory;
     }
 
-    private static void initialize() {
-        if (INIT_COUNT.get() == INITIALIZED) {
-            /*
-             * Safeguard for recursive init
-             */
-            return;
-        }
+    private static long tryInitialize() {
         /*
          * The JDK method VM.saveAndRemoveProperties looks at the system property
          * "sun.nio.MaxDirectMemorySize". However, that property is always set by the Java HotSpot
@@ -120,72 +100,41 @@ private static void initialize() {
         if (newDirectMemory == 0) {
             /*
              * No value explicitly specified. The default in the JDK in this case is the maximum
-             * heap size. However, we cannot rely on Runtime.maxMemory() until PhysicalMemory has
-             * fully initialized. Runtime.maxMemory() has a dependency on PhysicalMemory.size()
-             * which in turn depends on container support which might use NIO. To avoid this cycle,
-             * we first initialize the 'directMemory' field to an arbitrary value (25MB), and only
-             * use the Runtime.maxMemory() API once PhysicalMemory has fully initialized.
+             * heap size.
              */
-            if (!PhysicalMemory.isInitialized()) {
+            if (PhysicalMemory.isInitializationInProgress()) {
                 /*
-                 * While initializing physical memory we might end up back here with an INIT_COUNT
-                 * of 1, since we read the directMemory field during container support code
-                 * execution which runs when PhysicalMemory is still initializing.
+                 * When initializing PhysicalMemory, we use NIO/cgroups code that calls
+                 * VM.getDirectMemory(). When this initialization is in progress, we need to prevent
+                 * that Runtime.maxMemory() is called below because it would trigger a recursive
+                 * initialization of PhysicalMemory. So, we return a temporary value.
                  */
-                VMError.guarantee(INIT_COUNT.get() <= INITIALIZING, "Initial run needs to have init count 0 or 1");
-                newDirectMemory = STATIC_DIRECT_MEMORY_AMOUNT; // Static value during initialization
-                INIT_COUNT.setRelease(INITIALIZING);
-            } else {
-                VMError.guarantee(INIT_COUNT.get() <= INITIALIZING, "Runtime.maxMemory() invariant");
-                /*
-                 * Once we know PhysicalMemory has been properly initialized we can use
-                 * Runtime.maxMemory(). Note that we might end up in this branch for code explicitly
-                 * using the JDK cgroups code. At that point PhysicalMemory has likely been
-                 * initialized.
-                 */
-                INIT_COUNT.setRelease(INITIALIZED);
-                newDirectMemory = Runtime.getRuntime().maxMemory();
-            }
-        } else {
-            /*
-             * For explicitly set direct memory we are done
-             */
-            Unsafe.getUnsafe().storeFence();
-            directMemory = newDirectMemory;
-            isInitialized = true;
-            if (Target_jdk_internal_misc_VM.initLevel() < 1) {
-                // only the first accessor needs to set this
-                Target_jdk_internal_misc_VM.initLevel(1);
+                return DIRECT_MEMORY_DURING_INITIALIZATION;
             }
-            return;
+            newDirectMemory = Runtime.getRuntime().maxMemory();
         }
-        VMError.guarantee(newDirectMemory > 0, "New direct memory should be initialized");
 
-        Unsafe.getUnsafe().storeFence();
+        /*
+         * The initialization is not synchronized, so multiple threads can race. Usually this will
+         * lead to the same value, unless the runtime options are modified concurrently - which is
+         * possible but not a case we care about.
+         */
         directMemory = newDirectMemory;
-        if (PhysicalMemory.isInitialized() && INITIALIZED == INIT_COUNT.get()) {
-            /*
-             * Complete initialization hand-shake once PhysicalMemory is properly initialized. Also
-             * set the VM init level to 1 so as to provoke the NIO code to re-set the internal
-             * MAX_MEMORY field.
-             */
-            isInitialized = true;
-            if (Target_jdk_internal_misc_VM.initLevel() < 1) {
-                // only the first accessor needs to set this
-                Target_jdk_internal_misc_VM.initLevel(1);
-            }
-        }
+
+        /* Ensure values are published to other threads before marking fields as initialized. */
+        Unsafe.getUnsafe().storeFence();
+        initialized = true;
+
+        return newDirectMemory;
     }
 }
 
 final class PageAlignDirectMemoryAccessors {
-
     /*
      * Not volatile to avoid a memory barrier when reading the values. Instead, an explicit barrier
      * is inserted when writing the values.
      */
     private static boolean initialized;
-
     private static boolean pageAlignDirectMemory;
 
     static boolean getPageAlignDirectMemory() {
