[GR-66292] Deduplicate ClassForNameSupport data in extension layers

PullRequest: graal/21186

Author: Zeavee

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/hub/ClassForNameSupport.java

@@ -28,11 +28,17 @@
 import static jdk.graal.compiler.options.OptionStability.EXPERIMENTAL;
 
 import java.lang.reflect.Modifier;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.BooleanSupplier;
+import java.util.function.Consumer;
 
 import org.graalvm.collections.EconomicMap;
 import org.graalvm.nativeimage.Platform;
@@ -46,22 +52,30 @@
 import com.oracle.svm.core.configure.RuntimeConditionSet;
 import com.oracle.svm.core.feature.AutomaticallyRegisteredImageSingleton;
 import com.oracle.svm.core.hub.registry.ClassRegistries;
+import com.oracle.svm.core.layeredimagesingleton.ImageSingletonLoader;
+import com.oracle.svm.core.layeredimagesingleton.ImageSingletonWriter;
 import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingletonBuilderFlags;
 import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingletonSupport;
 import com.oracle.svm.core.layeredimagesingleton.MultiLayeredImageSingleton;
-import com.oracle.svm.core.layeredimagesingleton.UnsavedSingleton;
 import com.oracle.svm.core.metadata.MetadataTracer;
 import com.oracle.svm.core.option.HostedOptionKey;
+import com.oracle.svm.core.option.LayerVerifiedOption;
 import com.oracle.svm.core.reflect.MissingReflectionRegistrationUtils;
 import com.oracle.svm.core.util.ImageHeapMap;
 import com.oracle.svm.core.util.VMError;
 
 import jdk.graal.compiler.options.Option;
 
 @AutomaticallyRegisteredImageSingleton
-public final class ClassForNameSupport implements MultiLayeredImageSingleton, UnsavedSingleton {
+public final class ClassForNameSupport implements MultiLayeredImageSingleton {
+
+    public static final String CLASSES_REGISTERED = "classes registered";
+    public static final String CLASSES_REGISTERED_STATES = "classes registered states";
+    public static final String UNSAFE_REGISTERED = "unsafe registered";
+    public static final String RESPECTS_CLASS_LOADER = "respects class loader";
 
     public static final class Options {
+        @LayerVerifiedOption(kind = LayerVerifiedOption.Kind.Changed, severity = LayerVerifiedOption.Severity.Error)//
         @Option(help = "Class.forName and similar respect their class loader argument.", stability = EXPERIMENTAL)//
         public static final HostedOptionKey<Boolean> ClassForNameRespectsClassLoader = new HostedOptionKey<>(false);
     }
@@ -98,7 +112,8 @@ private static ClassForNameSupport[] layeredSingletons() {
     }
 
     /**
-     * The map used to collect registered classes. Not used when respecting class loaders.
+     * The map used to collect registered classes. Not used when respecting class loaders. This map
+     * only collects data for the current layer.
      */
     private final EconomicMap<String, ConditionalRuntimeValue<Object>> knownClasses;
     /**
@@ -112,14 +127,36 @@ private static ClassForNameSupport[] layeredSingletons() {
      */
     private final EconomicMap<String, Throwable> knownExceptions;
     /**
-     * The map used to collect unsafe allocated classes.
+     * The map used to collect unsafe allocated classes. This map only collects data for the current
+     * layer.
      */
     private final EconomicMap<Class<?>, RuntimeConditionSet> unsafeInstantiatedClasses;
 
+    /**
+     * The map used to collect classes registered in previous layers. The boolean associated to each
+     * class is true if the registered value is complete and false in the case of a negative query.
+     * A complete data registered in the current layer will overwrite a negative query in previous
+     * layers. In this case, the data will be stored in {@link ClassForNameSupport#knownClasses} of
+     * the current layer and the boolean value will be changed in the map of the next extension
+     * layer.
+     */
+    @Platforms(HOSTED_ONLY.class) //
+    private final Map<String, Boolean> previousLayerClasses;
+
+    /**
+     * The set used to collect unsafe allocated classes in previous layers.
+     */
+    @Platforms(HOSTED_ONLY.class) //
+    private final Set<String> previousLayerUnsafe;
+
     private static final Object NEGATIVE_QUERY = new Object();
 
     public ClassForNameSupport() {
-        if (respectClassLoader()) {
+        this(Map.of(), Set.of(), respectClassLoader());
+    }
+
+    public ClassForNameSupport(Map<String, Boolean> previousLayerClasses, Set<String> previousLayerUnsafe, boolean respectsClassLoader) {
+        if (respectsClassLoader) {
             knownClasses = null;
             knownClassNames = ImageHeapMap.createNonLayeredMap();
             knownExceptions = ImageHeapMap.createNonLayeredMap();
@@ -129,6 +166,8 @@ public ClassForNameSupport() {
             knownExceptions = null;
         }
         unsafeInstantiatedClasses = ImageHeapMap.createNonLayeredMap();
+        this.previousLayerClasses = previousLayerClasses;
+        this.previousLayerUnsafe = previousLayerUnsafe;
     }
 
     public static boolean respectClassLoader() {
@@ -177,15 +216,15 @@ public void registerClass(ConfigurationCondition condition, Class<?> clazz, Clas
                                 currentValue == clazz) {
                     currentValue = clazz;
                     var cond = updateConditionalValue(existingEntry, currentValue, condition);
-                    knownClasses.put(name, cond);
+                    addKnownClass(name, cond);
                 } else if (currentValue instanceof Throwable) { // failed at linking time
                     var cond = updateConditionalValue(existingEntry, currentValue, condition);
                     /*
                      * If the class has already been seen as throwing an error, we don't overwrite
                      * this error. Nevertheless, we have to update the set of conditionals to be
                      * correct.
                      */
-                    knownClasses.put(name, cond);
+                    addKnownClass(name, cond);
                 } else {
                     throw VMError.shouldNotReachHere("""
                                     Invalid Class.forName value for %s: %s
@@ -199,6 +238,17 @@ accessible through the builder class loader, and it was already registered by na
         }
     }
 
+    private void addKnownClass(String name, ConditionalRuntimeValue<Object> cond) {
+        addKnownClass(name, (map) -> map.put(name, cond), cond);
+    }
+
+    private void addKnownClass(String name, Consumer<EconomicMap<String, ConditionalRuntimeValue<Object>>> callback, ConditionalRuntimeValue<Object> cond) {
+        Boolean previousLayerData = previousLayerClasses.get(name);
+        if (previousLayerData == null || (!previousLayerData && cond.getValueUnconditionally() != NEGATIVE_QUERY)) {
+            callback.accept(knownClasses);
+        }
+    }
+
     @Platforms(HOSTED_ONLY.class)
     private boolean isLibGraalClass(Class<?> clazz) {
         return libGraalLoader != null && clazz.getClassLoader() == libGraalLoader;
@@ -260,19 +310,24 @@ private void registerKnownClassName(ConfigurationCondition condition, String cla
     public void registerUnsafeAllocated(ConfigurationCondition condition, Class<?> clazz) {
         if (!clazz.isArray() && !clazz.isInterface() && !Modifier.isAbstract(clazz.getModifiers())) {
             /* Otherwise, UNSAFE.allocateInstance results in InstantiationException */
-            var conditionSet = unsafeInstantiatedClasses.putIfAbsent(clazz, RuntimeConditionSet.createHosted(condition));
-            if (conditionSet != null) {
-                conditionSet.addCondition(condition);
+            if (!previousLayerUnsafe.contains(clazz.getName())) {
+                var conditionSet = unsafeInstantiatedClasses.putIfAbsent(clazz, RuntimeConditionSet.createHosted(condition));
+                if (conditionSet != null) {
+                    conditionSet.addCondition(condition);
+                }
             }
         }
     }
 
     private void updateCondition(ConfigurationCondition condition, String className, Object value) {
         synchronized (knownClasses) {
-            var runtimeConditions = knownClasses.putIfAbsent(className, new ConditionalRuntimeValue<>(RuntimeConditionSet.createHosted(condition), value));
-            if (runtimeConditions != null) {
-                runtimeConditions.getConditions().addCondition(condition);
-            }
+            var cond = new ConditionalRuntimeValue<>(RuntimeConditionSet.createHosted(condition), value);
+            addKnownClass(className, (map) -> {
+                var runtimeConditions = map.putIfAbsent(className, cond);
+                if (runtimeConditions != null) {
+                    runtimeConditions.getConditions().addCondition(condition);
+                }
+            }, cond);
         }
     }
 
@@ -468,4 +523,59 @@ public static boolean canUnsafeInstantiateAsInstance(DynamicHub hub) {
     public EnumSet<LayeredImageSingletonBuilderFlags> getImageBuilderFlags() {
         return LayeredImageSingletonBuilderFlags.ALL_ACCESS;
     }
+
+    @Override
+    public PersistFlags preparePersist(ImageSingletonWriter writer) {
+        List<String> classNames = new ArrayList<>();
+        List<Boolean> classStates = new ArrayList<>();
+        Set<String> unsafeNames = new HashSet<>(previousLayerUnsafe);
+
+        var cursor = knownClasses.getEntries();
+        while (cursor.advance()) {
+            classNames.add(cursor.getKey());
+            boolean isNegativeQuery = cursor.getValue().getValueUnconditionally() == NEGATIVE_QUERY;
+            classStates.add(!isNegativeQuery);
+        }
+
+        for (var entry : previousLayerClasses.entrySet()) {
+            /*
+             * If a complete entry overwrites a negative query from a previous layer, the
+             * previousLayerClasses map entry needs to be skipped to register the new entry for
+             * extension layers.
+             */
+            if (!classNames.contains(entry.getKey())) {
+                classNames.add(entry.getKey());
+                classStates.add(entry.getValue());
+            }
+        }
+
+        unsafeInstantiatedClasses.getKeys().iterator().forEachRemaining(c -> unsafeNames.add(c.getName()));
+
+        writer.writeStringList(CLASSES_REGISTERED, classNames);
+        writer.writeBoolList(CLASSES_REGISTERED_STATES, classStates);
+        writer.writeStringList(UNSAFE_REGISTERED, unsafeNames.stream().toList());
+        /*
+         * The option is not accessible when the singleton is loaded, so the boolean needs to be
+         * persisted.
+         */
+        writer.writeInt(RESPECTS_CLASS_LOADER, respectClassLoader() ? 1 : 0);
+
+        return PersistFlags.CREATE;
+    }
+
+    @SuppressWarnings("unused")
+    public static Object createFromLoader(ImageSingletonLoader loader) {
+        List<String> previousLayerClassKeys = loader.readStringList(CLASSES_REGISTERED);
+        List<Boolean> previousLayerClassStates = loader.readBoolList(CLASSES_REGISTERED_STATES);
+
+        Map<String, Boolean> previousLayerClasses = new HashMap<>();
+        for (int i = 0; i < previousLayerClassKeys.size(); ++i) {
+            previousLayerClasses.put(previousLayerClassKeys.get(i), previousLayerClassStates.get(i));
+        }
+
+        Set<String> previousLayerUnsafe = Set.copyOf(loader.readStringList(UNSAFE_REGISTERED));
+        boolean respectsClassLoader = loader.readInt(RESPECTS_CLASS_LOADER) == 1;
+
+        return new ClassForNameSupport(Collections.unmodifiableMap(previousLayerClasses), previousLayerUnsafe, respectsClassLoader);
+    }
 }