[GR-65971] Introduce onSingletonRegistration layered callback.

PullRequest: graal/22049

Author: teshull

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/heap/GCCause.java

@@ -25,7 +25,6 @@
 package com.oracle.svm.core.heap;
 
 import java.util.Collections;
-import java.util.EnumSet;
 import java.util.List;
 import java.util.function.Function;
 
@@ -37,15 +36,18 @@
 import com.oracle.svm.core.feature.AutomaticallyRegisteredFeature;
 import com.oracle.svm.core.feature.AutomaticallyRegisteredImageSingleton;
 import com.oracle.svm.core.feature.InternalFeature;
-import com.oracle.svm.core.imagelayer.BuildingImageLayerPredicate;
 import com.oracle.svm.core.imagelayer.ImageLayerBuildingSupport;
 import com.oracle.svm.core.layeredimagesingleton.ImageSingletonLoader;
 import com.oracle.svm.core.layeredimagesingleton.ImageSingletonWriter;
 import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingleton;
-import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingletonBuilderFlags;
 import com.oracle.svm.core.traits.BuiltinTraits.AllAccess;
 import com.oracle.svm.core.traits.BuiltinTraits.SingleLayer;
+import com.oracle.svm.core.traits.SingletonLayeredCallbacks;
+import com.oracle.svm.core.traits.SingletonLayeredCallbacksSupplier;
+import com.oracle.svm.core.traits.SingletonLayeredInstallationKind.Independent;
 import com.oracle.svm.core.traits.SingletonLayeredInstallationKind.InitialLayerOnly;
+import com.oracle.svm.core.traits.SingletonTrait;
+import com.oracle.svm.core.traits.SingletonTraitKind;
 import com.oracle.svm.core.traits.SingletonTraits;
 import com.oracle.svm.core.util.DuplicatedInNativeCode;
 import com.oracle.svm.core.util.ImageHeapList;
@@ -128,7 +130,13 @@ void installGCCause(GCCause gcCause) {
 }
 
 @AutomaticallyRegisteredFeature
+@SingletonTraits(access = AllAccess.class, layeredCallbacks = GCCauseFeature.LayeredCallbacks.class, layeredInstallationKind = Independent.class)
 class GCCauseFeature implements InternalFeature {
+    /**
+     * In layered builds all {@link GCCause}s are registered and installed in the initial layer.
+     * Here we track which {@link GCCause}s were installed in the initial layer to detect issues.
+     */
+    List<String> registeredGCCauses;
 
     @Override
     public void duringSetup(DuringSetupAccess access) {
@@ -158,8 +166,7 @@ public void duringSetup(DuringSetupAccess access) {
                 var gcCauseList = GCCause.getGCCauses();
                 idToGCCauseName = (idx) -> gcCauseList.get(idx).getName();
             } else {
-                var gcCauseNames = LayeredGCCauseTracker.getRegisteredGCCauses();
-                idToGCCauseName = gcCauseNames::get;
+                idToGCCauseName = registeredGCCauses::get;
             }
             access.registerObjectReplacer(obj -> {
                 if (obj instanceof GCCause gcCause) {
@@ -172,49 +179,37 @@ public void duringSetup(DuringSetupAccess access) {
             });
         }
     }
-}
-
-/**
- * In layered builds all {@link GCCause}s are registered and installed in the initial layer. Here we
- * track which {@link GCCause}s were installed in the initial layer to detect issues.
- */
-@AutomaticallyRegisteredImageSingleton(onlyWith = BuildingImageLayerPredicate.class)
-class LayeredGCCauseTracker implements LayeredImageSingleton {
-    List<String> registeredGCCauses;
 
-    public static List<String> getRegisteredGCCauses() {
-        assert ImageLayerBuildingSupport.buildingExtensionLayer();
-        return ImageSingletons.lookup(LayeredGCCauseTracker.class).registeredGCCauses;
-    }
+    static class LayeredCallbacks extends SingletonLayeredCallbacksSupplier {
+
+        @Override
+        public SingletonTrait getLayeredCallbacksTrait() {
+            return new SingletonTrait(SingletonTraitKind.LAYERED_CALLBACKS, new SingletonLayeredCallbacks() {
+                @Override
+                public LayeredImageSingleton.PersistFlags doPersist(ImageSingletonWriter writer, Object singleton) {
+                    List<String> gcCauses;
+                    if (ImageLayerBuildingSupport.buildingInitialLayer()) {
+                        gcCauses = GCCause.getGCCauses().stream().map(gcCause -> {
+                            if (gcCause == null) {
+                                return "";
+                            } else {
+                                assert !gcCause.getName().isEmpty() : Assertions.errorMessage("Empty string is reserved for non-existent GCCauses", gcCause);
+                                return gcCause.getName();
+                            }
+                        }).toList();
+                    } else {
+                        gcCauses = ((GCCauseFeature) singleton).registeredGCCauses;
+                    }
+                    writer.writeStringList("registeredGCCauses", gcCauses);
 
-    @Override
-    public EnumSet<LayeredImageSingletonBuilderFlags> getImageBuilderFlags() {
-        return LayeredImageSingletonBuilderFlags.BUILDTIME_ACCESS_ONLY;
-    }
+                    return LayeredImageSingleton.PersistFlags.CALLBACK_ON_REGISTRATION;
+                }
 
-    @Override
-    public PersistFlags preparePersist(ImageSingletonWriter writer) {
-        List<String> gcCauses;
-        if (ImageLayerBuildingSupport.buildingInitialLayer()) {
-            gcCauses = GCCause.getGCCauses().stream().map(gcCause -> {
-                if (gcCause == null) {
-                    return "";
-                } else {
-                    assert !gcCause.getName().isEmpty() : Assertions.errorMessage("Empty string is reserved for non-existent GCCauses", gcCause);
-                    return gcCause.getName();
+                @Override
+                public void onSingletonRegistration(ImageSingletonLoader loader, Object singleton) {
+                    ((GCCauseFeature) singleton).registeredGCCauses = Collections.unmodifiableList(loader.readStringList("registeredGCCauses"));
                 }
-            }).toList();
-        } else {
-            gcCauses = registeredGCCauses;
+            });
         }
-        writer.writeStringList("registeredGCCauses", gcCauses);
-        return PersistFlags.CREATE;
-    }
-
-    @SuppressWarnings("unused")
-    public static Object createFromLoader(ImageSingletonLoader loader) {
-        var causeTracker = new LayeredGCCauseTracker();
-        causeTracker.registeredGCCauses = Collections.unmodifiableList(loader.readStringList("registeredGCCauses"));
-        return causeTracker;
     }
 }

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/layeredimagesingleton/LayeredImageSingleton.java

@@ -31,6 +31,7 @@
 import org.graalvm.nativeimage.Platforms;
 
 import com.oracle.svm.core.imagelayer.ImageLayerBuildingSupport;
+import com.oracle.svm.core.traits.SingletonLayeredCallbacks;
 
 /**
  * In additional to the traditional singleton model, i.e. a key-value map whose lookups are constant
@@ -79,7 +80,18 @@ enum PersistFlags {
          * Indicates in a subsequent image a new singleton should be created and linked via calling
          * {@code Object createFromLoader(ImageSingletonLoader)}.
          */
-        CREATE
+        CREATE,
+        /**
+         * Indicates that when and/or if this singleton is registered in the next image via
+         * {@link ImageSingletons#add}, then
+         * {@link SingletonLayeredCallbacks#onSingletonRegistration} should be called on this
+         * singleton.
+         * <p>
+         * If a singleton is registered under multiple keys, then
+         * {@link SingletonLayeredCallbacks#onSingletonRegistration} should only be called once and
+         * will have finished executing before the singleton is installed into the registry.
+         */
+        CALLBACK_ON_REGISTRATION
     }
 
     /*

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/traits/SingletonLayeredCallbacks.java

@@ -70,7 +70,16 @@ public interface LayeredSingletonInstantiator {
      * determine how to instantiate the singleton in the next layer.
      */
     public Class<? extends LayeredSingletonInstantiator> getSingletonInstantiator() {
-        throw VMError.shouldNotReachHere("getSingletonInstantiator is not implemented. This method must only be implemented if doPersist returns PersistFlag.CREATE");
+        throw VMError.shouldNotReachHere("getSingletonInstantiator is not implemented. This method must be implemented if doPersist returns PersistFlag.CREATE");
+    }
+
+    /**
+     * See description in {@link PersistFlags#CALLBACK_ON_REGISTRATION} for more details. Note this
+     * method will be called at most once for each registered singleton object.
+     */
+    @SuppressWarnings("unused")
+    public void onSingletonRegistration(ImageSingletonLoader loader, Object singleton) {
+        throw VMError.shouldNotReachHere("onSingletonRegistration is not implemented. This method must be implemented if doPersist returns PersistFlag.CALLBACK_ON_REGISTRATION");
     }
 
 }

substratevm/src/com.oracle.svm.hosted/resources/SharedLayerSnapshotCapnProtoSchema.capnp

@@ -12,6 +12,7 @@ using MethodId = Int32;
 using FieldId = Int32;
 using ConstantId = Int32;
 using SingletonObjId = Int32;
+using KeyStoreId = Int32;
 using HostedMethodIndex = Int32;
 
 struct PersistedAnalysisType {
@@ -243,17 +244,23 @@ struct ImageSingletonKey {
   objectId @2 :SingletonObjId;
   constantId @3 :ConstantId;
   isInitialLayerOnly @4 :Bool;
+  keyStoreId @5 :KeyStoreId;
 }
 
 struct ImageSingletonObject {
   id @0 :SingletonObjId;
   className @1 :Text;
-  store @2 :List(KeyStoreEntry);
+  keyStoreId @2 :KeyStoreId;
   recreateClass @3 :Text;
   # GR-66792 remove once no custom persist actions exist
   recreateMethod @4 :Text;
 }
 
+struct KeyStoreInstance {
+  id @0 :KeyStoreId;
+  keyStore @1 :List(KeyStoreEntry);
+}
+
 struct Annotation {
   typeName @0 :Text;
   values @1 :List(AnnotationValue);
@@ -303,6 +310,7 @@ struct SharedLayerSnapshot {
   sharedLayerBootLayerModules @21 :List(Text);
   layeredModule @22 :LayeredModule;
   cGlobals @23 :List(CGlobalDataInfo);
+  keyStoreInstances @24 :List(KeyStoreInstance);
 }
 
 struct StaticFinalFieldFoldingSingleton {

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/ImageSingletonsSupportImpl.java

@@ -25,14 +25,18 @@
 package com.oracle.svm.hosted;
 
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Comparator;
 import java.util.EnumMap;
 import java.util.EnumSet;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -355,7 +359,7 @@ public static void clear() {
             singletonDuringImageBuild = null;
         }
 
-        public static void persist() {
+        public static void persistSingletonInfo() {
             var list = singletonDuringImageBuild.configObjects.entrySet().stream().filter(e -> e.getValue().traitMap.getTrait(SingletonTraitKind.LAYERED_CALLBACKS).isPresent())
                             .sorted(Comparator.comparing(e -> e.getKey().getName()))
                             .toList();
@@ -364,12 +368,21 @@ public static void persist() {
 
         private final Map<Class<?>, SingletonInfo> configObjects;
         private final Map<Object, SingletonTraitMap> singletonToTraitMap;
+        /**
+         * Tracks the status of singletons for which a registration callback needs to be executed
+         * upon installation. The key will always be the singleton object, and the value will be
+         * either a {@link Boolean} or {@link Lock} based on whether the callback's execution is
+         * still in progress or has completed.
+         */
+        private final Map<Object, Object> singletonRegistrationCallbackStatus;
 
         private final EnumSet<SingletonLayeredInstallationKind.InstallationKind> forbiddenInstallationKinds;
         private Set<Class<?>> multiLayeredImageSingletonKeys;
         private final boolean layeredBuild;
+        private final boolean extensionLayerBuild;
         private final AnnotationExtractor extractor;
         private final Function<Class<?>, SingletonTrait[]> singletonTraitInjector;
+        private final SVMImageLayerSingletonLoader singletonLoader;
 
         public HostedManagement() {
             this(null, null);
@@ -382,7 +395,10 @@ public HostedManagement(HostedImageLayerBuildingSupport support, AnnotationExtra
             forbiddenInstallationKinds = EnumSet.of(SingletonLayeredInstallationKind.InstallationKind.DISALLOWED);
             if (support != null) {
                 this.layeredBuild = support.buildingImageLayer;
+                this.extensionLayerBuild = support.buildingImageLayer && !support.buildingInitialLayer;
                 this.singletonTraitInjector = support.getSingletonTraitInjector();
+                this.singletonLoader = support.getSingletonLoader();
+                this.singletonRegistrationCallbackStatus = extensionLayerBuild ? new ConcurrentIdentityHashMap<>() : null;
                 if (support.buildingImageLayer) {
                     if (!support.buildingApplicationLayer) {
                         forbiddenInstallationKinds.add(SingletonLayeredInstallationKind.InstallationKind.APP_LAYER_ONLY);
@@ -393,7 +409,10 @@ public HostedManagement(HostedImageLayerBuildingSupport support, AnnotationExtra
                 }
             } else {
                 this.layeredBuild = false;
+                this.extensionLayerBuild = false;
                 this.singletonTraitInjector = null;
+                this.singletonLoader = null;
+                this.singletonRegistrationCallbackStatus = null;
             }
             this.extractor = extractor;
         }
@@ -484,12 +503,71 @@ private void addSingletonToMap(Class<?> key, Object value, SingletonTraitMap tra
                 }
             });
 
+            /* Run onSingletonRegistration hook if needed. */
+            if (extensionLayerBuild) {
+                if (singletonLoader.hasRegistrationCallback(key)) {
+                    synchronizeRegistrationCallbackExecution(value, () -> {
+                        var trait = traitMap.getTrait(SingletonTraitKind.LAYERED_CALLBACKS).get();
+                        var callbacks = ((SingletonLayeredCallbacks) trait.metadata());
+                        callbacks.onSingletonRegistration(singletonLoader.getImageSingletonLoader(key), value);
+                    });
+                }
+            }
+
             Object prevValue = configObjects.putIfAbsent(key, new SingletonInfo(value, traitMap));
             if (prevValue != null) {
                 throw UserError.abort("ImageSingletons.add must not overwrite existing key %s%nExisting value: %s%nNew value: %s", key.getTypeName(), prevValue, value);
             }
         }
 
+        /**
+         * Ensures the provided registrationCallback will execute only once per a singleton.
+         * Regardless of which thread executes the registrationCallback, this method will not return
+         * until the registrationCallback has been executed.
+         */
+        private void synchronizeRegistrationCallbackExecution(Object singleton, Runnable registrationCallback) {
+            while (true) {
+                var status = singletonRegistrationCallbackStatus.get(singleton);
+                if (status == null) {
+                    // create a lock for other threads to wait on
+                    ReentrantLock lock = new ReentrantLock();
+                    lock.lock();
+                    try {
+                        status = singletonRegistrationCallbackStatus.computeIfAbsent(singleton, _ -> lock);
+                        if (status != lock) {
+                            // failed to install lock. Repeat loop.
+                            continue;
+                        }
+
+                        // Run registrationCallback
+                        registrationCallback.run();
+
+                        // the registrationCallback has finished - update its status
+                        var prev = singletonRegistrationCallbackStatus.put(singleton, Boolean.TRUE);
+                        VMError.guarantee(prev == lock);
+                    } finally {
+                        lock.unlock();
+                    }
+                } else if (status instanceof Lock lock) {
+                    lock.lock();
+                    try {
+                        /*
+                         * Once the lock can be acquired we know the registrationCallback has been
+                         * completed and we can proceed.
+                         */
+                        assert singletonRegistrationCallbackStatus.get(singleton) == Boolean.TRUE;
+                    } finally {
+                        lock.unlock();
+                    }
+                } else {
+                    // the registrationCallback has already completed
+                    assert status == Boolean.TRUE;
+                }
+                /* At this point the registrationCallback has executed so it is safe to proceed. */
+                break;
+            }
+        }
+
         Collection<Class<?>> getMultiLayeredImageSingletonKeys() {
             return multiLayeredImageSingletonKeys;
         }
@@ -499,7 +577,7 @@ void freezeLayeredImageSingletonMetadata() {
         }
 
         Set<Object> getSingletonsWithTrait(SingletonLayeredInstallationKind.InstallationKind kind) {
-            return configObjects.values().stream().filter(singletonInfo -> {
+            return Collections.unmodifiableSet(configObjects.values().stream().filter(singletonInfo -> {
                 /*
                  * We must filter out forbidden objects, as they are not actually installed in this
                  * image.
@@ -511,7 +589,7 @@ Set<Object> getSingletonsWithTrait(SingletonLayeredInstallationKind.Installation
                     }
                 }
                 return false;
-            }).map(SingletonInfo::singleton).collect(Collectors.toUnmodifiableSet());
+            }).map(SingletonInfo::singleton).collect(Collectors.toCollection(() -> Collections.newSetFromMap(new IdentityHashMap<>()))));
         }
 
         void forbidNewTraitInstallations(SingletonLayeredInstallationKind.InstallationKind kind) {

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/NativeImageGenerator.java

@@ -757,7 +757,7 @@ protected void doRun(Map<Method, CEntryPointData> entryPoints, JavaMainSupport j
             }
             try (StopTimer t = TimerCollection.createTimerAndStart(TimerCollection.Registry.ARCHIVE_LAYER)) {
                 if (ImageLayerBuildingSupport.buildingSharedLayer()) {
-                    ImageSingletonsSupportImpl.HostedManagement.persist();
+                    ImageSingletonsSupportImpl.HostedManagement.persistSingletonInfo();
                     HostedImageLayerBuildingSupport.singleton().archiveLayer();
                 }
             }