Layered serialization support

Author: Zeavee

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

@@ -123,7 +123,7 @@
 import com.oracle.svm.core.reflect.RuntimeMetadataDecoder.FieldDescriptor;
 import com.oracle.svm.core.reflect.RuntimeMetadataDecoder.MethodDescriptor;
 import com.oracle.svm.core.reflect.fieldaccessor.UnsafeFieldAccessorFactory;
-import com.oracle.svm.core.reflect.serialize.SerializationRegistry;
+import com.oracle.svm.core.reflect.serialize.SerializationSupport;
 import com.oracle.svm.core.reflect.target.Target_jdk_internal_reflect_ConstantPool;
 import com.oracle.svm.core.util.LazyFinalReference;
 import com.oracle.svm.core.util.VMError;
@@ -1140,7 +1140,7 @@ public boolean isLinked() {
     }
 
     public boolean isRegisteredForSerialization() {
-        return ImageSingletons.lookup(SerializationRegistry.class).isRegisteredForSerialization(DynamicHub.toClass(this));
+        return SerializationSupport.isRegisteredForSerialization(this);
     }
 
     @KeepOriginal
@@ -2335,8 +2335,7 @@ public FieldAccessor newFieldAccessor(Field field0, boolean override) {
 
     @Substitute
     private Constructor<?> generateConstructor(Class<?> cl, Constructor<?> constructorToCall) {
-        SerializationRegistry serializationRegistry = ImageSingletons.lookup(SerializationRegistry.class);
-        ConstructorAccessor acc = (ConstructorAccessor) serializationRegistry.getSerializationConstructorAccessor(cl, constructorToCall.getDeclaringClass());
+        ConstructorAccessor acc = (ConstructorAccessor) SerializationSupport.getSerializationConstructorAccessor(cl, constructorToCall.getDeclaringClass());
         /*
          * Unlike other root constructors, this constructor is not copied for mutation but directly
          * mutated, as it is not cached. To cache this constructor, setAccessible call must be done

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

@@ -162,7 +162,7 @@ private static void registerLambdaForReflection(Class<?> lambdaClass) {
          * lambda-class information from the capturing class.
          */
         if (Serializable.class.isAssignableFrom(lambdaClass) &&
-                        SerializationSupport.singleton().isLambdaCapturingClassRegistered(LambdaUtils.capturingClass(lambdaClass.getName()))) {
+                        SerializationSupport.currentLayer().isLambdaCapturingClassRegistered(LambdaUtils.capturingClass(lambdaClass.getName()))) {
             try {
                 Method serializeLambdaMethod = lambdaClass.getDeclaredMethod("writeReplace");
                 RuntimeReflection.register(serializeLambdaMethod);

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/reflect/serialize/SerializationRegistry.java

@@ -25,10 +25,10 @@
  */
 package com.oracle.svm.core.reflect.serialize;
 
-public interface SerializationRegistry {
-
-    boolean isRegisteredForSerialization(Class<?> cl);
+import com.oracle.svm.core.hub.DynamicHub;
 
-    Object getSerializationConstructorAccessor(Class<?> serializationTargetClass, Class<?> targetConstructorClass);
+public interface SerializationRegistry {
+    boolean isRegisteredForSerialization0(DynamicHub dynamicHub);
 
+    Object getSerializationConstructorAccessor0(Class<?> declaringClass, Class<?> targetConstructorClass);
 }

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/reflect/serialize/SerializationSupport.java

@@ -31,26 +31,38 @@
 import java.lang.invoke.SerializedLambda;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Modifier;
+import java.util.EnumSet;
 import java.util.Objects;
 
 import org.graalvm.collections.EconomicMap;
 import org.graalvm.collections.MapCursor;
-import org.graalvm.nativeimage.ImageSingletons;
 import org.graalvm.nativeimage.Platform;
 import org.graalvm.nativeimage.Platforms;
 import org.graalvm.nativeimage.impl.ConfigurationCondition;
 
+import com.oracle.svm.core.BuildPhaseProvider;
+import com.oracle.svm.core.SubstrateUtil;
 import com.oracle.svm.core.configure.RuntimeConditionSet;
+import com.oracle.svm.core.hub.DynamicHub;
+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.reflect.SubstrateConstructorAccessor;
 import com.oracle.svm.core.util.ImageHeapMap;
 import com.oracle.svm.core.util.VMError;
 
 import jdk.graal.compiler.java.LambdaUtils;
 
-public class SerializationSupport implements SerializationRegistry {
+public class SerializationSupport implements MultiLayeredImageSingleton, SerializationRegistry, UnsavedSingleton {
 
-    public static SerializationSupport singleton() {
-        return (SerializationSupport) ImageSingletons.lookup(SerializationRegistry.class);
+    @Platforms(Platform.HOSTED_ONLY.class)
+    public static SerializationSupport currentLayer() {
+        return LayeredImageSingletonSupport.singleton().lookup(SerializationSupport.class, false, true);
+    }
+
+    public static SerializationSupport[] layeredSingletons() {
+        return MultiLayeredImageSingleton.getAllLayers(SerializationSupport.class);
     }
 
     /**
@@ -184,19 +196,48 @@ public String getClassLoaderSerializationLookupKey(ClassLoader classLoader) {
         throw VMError.shouldNotReachHere("No constructor accessor uses the class loader %s", classLoader);
     }
 
-    private final EconomicMap<Class<?>, RuntimeConditionSet> classes = EconomicMap.create();
+    /**
+     * This class is used as key in maps that use {@link Class} as key at runtime in layered images,
+     * because the hash code of {@link Class} objects cannot be injected in extension layers and is
+     * thus inconsistent across layers. The state of those maps is then incorrect at run time. The
+     * {@link DynamicHub} cannot be used directly either as its hash code at run time is the one of
+     * the {@link Class} object.
+     * <p>
+     * Temporary key for maps ideally indexed by their {@link Class} or {@link DynamicHub}. At
+     * runtime, these maps should be indexed by {@link DynamicHub#getTypeID}
+     */
+    public record DynamicHubKey(DynamicHub hub) {
+        public int getTypeID() {
+            return hub.getTypeID();
+        }
+    }
+
+    private final EconomicMap<Object /* DynamicHubKey or DynamicHub.typeID */, RuntimeConditionSet> classes = EconomicMap.create();
     private final EconomicMap<String, RuntimeConditionSet> lambdaCapturingClasses = EconomicMap.create();
 
     @Platforms(Platform.HOSTED_ONLY.class)
-    public void registerSerializationTargetClass(ConfigurationCondition cnd, Class<?> serializationTargetClass) {
+    public void registerSerializationTargetClass(ConfigurationCondition cnd, DynamicHub hub) {
         synchronized (classes) {
-            var previous = classes.putIfAbsent(serializationTargetClass, RuntimeConditionSet.createHosted(cnd));
+            var previous = classes.putIfAbsent(BuildPhaseProvider.isHostedUniverseBuilt() ? hub.getTypeID() : new DynamicHubKey(hub), RuntimeConditionSet.createHosted(cnd));
             if (previous != null) {
                 previous.addCondition(cnd);
             }
         }
     }
 
+    public void replaceHubKeyWithTypeID() {
+        EconomicMap<Integer, RuntimeConditionSet> newEntries = EconomicMap.create();
+        var cursor = classes.getEntries();
+        while (cursor.advance()) {
+            Object key = cursor.getKey();
+            if (key instanceof DynamicHubKey hubKey) {
+                newEntries.put(hubKey.getTypeID(), cursor.getValue());
+                cursor.remove();
+            }
+        }
+        classes.putAll(newEntries);
+    }
+
     @Platforms(Platform.HOSTED_ONLY.class)
     public void registerLambdaCapturingClass(ConfigurationCondition cnd, String lambdaCapturingClass) {
         synchronized (lambdaCapturingClasses) {
@@ -212,37 +253,63 @@ public boolean isLambdaCapturingClassRegistered(String lambdaCapturingClass) {
         return lambdaCapturingClasses.containsKey(lambdaCapturingClass);
     }
 
-    @Override
-    public Object getSerializationConstructorAccessor(Class<?> rawDeclaringClass, Class<?> rawTargetConstructorClass) {
-        Class<?> declaringClass = rawDeclaringClass;
+    public static Object getSerializationConstructorAccessor(Class<?> serializationTargetClass, Class<?> targetConstructorClass) {
+        Class<?> declaringClass = serializationTargetClass;
 
         if (LambdaUtils.isLambdaClass(declaringClass)) {
             declaringClass = SerializedLambda.class;
         }
 
+        if (SubstrateUtil.HOSTED) {
+            Object constructorAccessor = currentLayer().getSerializationConstructorAccessor0(declaringClass, targetConstructorClass);
+            if (constructorAccessor != null) {
+                return constructorAccessor;
+            }
+        } else {
+            for (var singleton : layeredSingletons()) {
+                Object constructorAccessor = singleton.getSerializationConstructorAccessor0(declaringClass, targetConstructorClass);
+                if (constructorAccessor != null) {
+                    return constructorAccessor;
+                }
+            }
+        }
+
+        String targetConstructorClassName = targetConstructorClass.getName();
+        if (ThrowMissingRegistrationErrors.hasBeenSet()) {
+            MissingSerializationRegistrationUtils.missingSerializationRegistration(declaringClass,
+                            "type " + declaringClass.getTypeName() + " with target constructor class: " + targetConstructorClassName);
+        } else {
+            throw VMError.unsupportedFeature("SerializationConstructorAccessor class not found for declaringClass: " + declaringClass.getName() +
+                            " (targetConstructorClass: " + targetConstructorClassName + "). Usually adding " + declaringClass.getName() +
+                            " to serialization-config.json fixes the problem.");
+        }
+        return null;
+    }
+
+    @Override
+    public Object getSerializationConstructorAccessor0(Class<?> declaringClass, Class<?> rawTargetConstructorClass) {
         VMError.guarantee(stubConstructor != null, "Called too early, no stub constructor yet.");
         Class<?> targetConstructorClass = Modifier.isAbstract(declaringClass.getModifiers()) ? stubConstructor.getDeclaringClass() : rawTargetConstructorClass;
-        Object constructorAccessor = constructorAccessors.get(new SerializationLookupKey(declaringClass, targetConstructorClass));
+        return constructorAccessors.get(new SerializationLookupKey(declaringClass, targetConstructorClass));
+    }
 
-        if (constructorAccessor != null) {
-            return constructorAccessor;
-        } else {
-            String targetConstructorClassName = targetConstructorClass.getName();
-            if (ThrowMissingRegistrationErrors.hasBeenSet()) {
-                MissingSerializationRegistrationUtils.missingSerializationRegistration(declaringClass,
-                                "type " + declaringClass.getTypeName() + " with target constructor class: " + targetConstructorClassName);
-            } else {
-                throw VMError.unsupportedFeature("SerializationConstructorAccessor class not found for declaringClass: " + declaringClass.getName() +
-                                " (targetConstructorClass: " + targetConstructorClassName + "). Usually adding " + declaringClass.getName() +
-                                " to serialization-config.json fixes the problem.");
+    public static boolean isRegisteredForSerialization(DynamicHub hub) {
+        for (SerializationRegistry singleton : SerializationSupport.layeredSingletons()) {
+            if (singleton.isRegisteredForSerialization0(hub)) {
+                return true;
             }
-            return null;
         }
+        return false;
     }
 
     @Override
-    public boolean isRegisteredForSerialization(Class<?> clazz) {
-        var conditionSet = classes.get(clazz);
+    public boolean isRegisteredForSerialization0(DynamicHub dynamicHub) {
+        var conditionSet = classes.get(dynamicHub.getTypeID());
         return conditionSet != null && conditionSet.satisfied();
     }
+
+    @Override
+    public EnumSet<LayeredImageSingletonBuilderFlags> getImageBuilderFlags() {
+        return LayeredImageSingletonBuilderFlags.ALL_ACCESS;
+    }
 }

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

@@ -39,6 +39,7 @@
 
 public abstract class ConditionalConfigurationRegistry {
     private Feature.BeforeAnalysisAccess beforeAnalysisAccess;
+    private SVMHost hostVM;
     private final Map<Class<?>, Collection<Runnable>> pendingReachabilityHandlers = new ConcurrentHashMap<>();
 
     protected void registerConditionalConfiguration(ConfigurationCondition condition, Consumer<ConfigurationCondition> consumer) {
@@ -82,4 +83,12 @@ public void setAnalysisAccess(Feature.BeforeAnalysisAccess beforeAnalysisAccess)
         pendingReachabilityHandlers.clear();
     }
 
+    public void setHostVM(SVMHost hostVM) {
+        this.hostVM = hostVM;
+    }
+
+    public SVMHost getHostVM() {
+        return hostVM;
+    }
+
 }

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

@@ -73,6 +73,8 @@
 import org.graalvm.nativeimage.impl.AnnotationExtractor;
 import org.graalvm.nativeimage.impl.CConstantValueSupport;
 import org.graalvm.nativeimage.impl.RuntimeClassInitializationSupport;
+import org.graalvm.nativeimage.impl.RuntimeReflectionSupport;
+import org.graalvm.nativeimage.impl.RuntimeSerializationSupport;
 import org.graalvm.nativeimage.impl.SizeOfSupport;
 import org.graalvm.word.PointerBase;
 
@@ -1029,11 +1031,13 @@ protected void setupNativeImage(String imageName, OptionValues options, Map<Meth
                 if (imageLayerLoader != null) {
                     imageLayerLoader.setHostedValuesProvider(hostedValuesProvider);
                 }
+                SVMHost hostVM = (SVMHost) aUniverse.hostVM();
+                ((ConditionalConfigurationRegistry) ImageSingletons.lookup(RuntimeSerializationSupport.class)).setHostVM(hostVM);
                 SubstratePlatformConfigurationProvider platformConfig = getPlatformConfig(aMetaAccess);
                 HostedProviders aProviders = createHostedProviders(target, aUniverse, originalProviders, platformConfig, aMetaAccess, classInitializationSupport);
                 aUniverse.hostVM().initializeProviders(aProviders);
 
-                ImageSingletons.add(SimulateClassInitializerSupport.class, ((SVMHost) aUniverse.hostVM()).createSimulateClassInitializerSupport(aMetaAccess));
+                ImageSingletons.add(SimulateClassInitializerSupport.class, (hostVM).createSimulateClassInitializerSupport(aMetaAccess));
 
                 bb = createBigBang(debug, options, aUniverse, aMetaAccess, aProviders, annotationSubstitutions);
                 aUniverse.setBigBang(bb);
@@ -1098,7 +1102,7 @@ protected void setupNativeImage(String imageName, OptionValues options, Map<Meth
                 }
 
                 initializeBigBang(bb, options, featureHandler, nativeLibraries, debug, aMetaAccess, aUniverse.getSubstitutions(), loader, true,
-                                new SubstrateClassInitializationPlugin((SVMHost) aUniverse.hostVM()), this.isStubBasedPluginsSupported(), aProviders);
+                                new SubstrateClassInitializationPlugin(hostVM), this.isStubBasedPluginsSupported(), aProviders);
 
                 loader.classLoaderSupport.getClassesToIncludeUnconditionally().forEach(cls -> bb.registerTypeForBaseImage(cls));
                 PreserveOptionsSupport.registerPreservedClasses(loader.classLoaderSupport);

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/code/RuntimeMetadataEncoderImpl.java

@@ -203,7 +203,7 @@ private void registerClass(HostedType type, ClassMetadata metadata) {
 
     private void registerField(HostedType declaringType, Object field, FieldMetadata metadata) {
         if (ImageLayerBuildingSupport.buildingImageLayer() &&
-                        !LayeredRuntimeMetadataSingleton.singleton().registerField(field, declaringType.getWrapped().getUniverse().getBigbang().getMetaAccess())) {
+                        !LayeredRuntimeMetadataSingleton.singleton().shouldRegisterField(field, declaringType.getWrapped().getUniverse().getBigbang().getMetaAccess(), metadata)) {
             return;
         }
         addType(declaringType);
@@ -219,7 +219,7 @@ private FieldMetadata[] getFields(HostedType declaringType) {
 
     private void registerMethod(HostedType declaringType, Object method, MethodMetadata metadata) {
         if (ImageLayerBuildingSupport.buildingImageLayer() &&
-                        !LayeredRuntimeMetadataSingleton.singleton().registerMethod(method, declaringType.getWrapped().getUniverse().getBigbang().getMetaAccess())) {
+                        !LayeredRuntimeMetadataSingleton.singleton().shouldRegisterMethod(method, declaringType.getWrapped().getUniverse().getBigbang().getMetaAccess(), metadata)) {
             return;
         }
         addType(declaringType);
@@ -235,7 +235,7 @@ private MethodMetadata[] getMethods(HostedType declaringType) {
 
     private void registerConstructor(HostedType declaringType, Object constructor, ConstructorMetadata metadata) {
         if (ImageLayerBuildingSupport.buildingImageLayer() &&
-                        !LayeredRuntimeMetadataSingleton.singleton().registerMethod(constructor, declaringType.getWrapped().getUniverse().getBigbang().getMetaAccess())) {
+                        !LayeredRuntimeMetadataSingleton.singleton().shouldRegisterMethod(constructor, declaringType.getWrapped().getUniverse().getBigbang().getMetaAccess(), metadata)) {
             return;
         }
         addType(declaringType);
@@ -1249,7 +1249,15 @@ private static LayeredRuntimeMetadataSingleton singleton() {
             return ImageSingletons.lookup(LayeredRuntimeMetadataSingleton.class);
         }
 
-        public boolean registerMethod(Object method, AnalysisMetaAccess metaAccess) {
+        public boolean shouldRegisterMethod(Object method, AnalysisMetaAccess metaAccess, AccessibleObjectMetadata metadata) {
+            if (!metadata.complete) {
+                /*
+                 * The method should be added to the list of registered methods only if the metadata
+                 * is complete. Incomplete metadata should always be registered and should not be
+                 * counted as the only metadata registered for a given method.
+                 */
+                return true;
+            }
             if (method instanceof AnalysisMethod analysisMethod) {
                 return registeredMethods.add(analysisMethod.getId());
             } else if (method instanceof HostedMethod hostedMethod) {
@@ -1262,7 +1270,10 @@ public boolean registerMethod(Object method, AnalysisMetaAccess metaAccess) {
             return true;
         }
 
-        public boolean registerField(Object field, AnalysisMetaAccess metaAccess) {
+        public boolean shouldRegisterField(Object field, AnalysisMetaAccess metaAccess, AccessibleObjectMetadata metadata) {
+            if (!metadata.complete) {
+                return true;
+            }
             if (field instanceof AnalysisField analysisField) {
                 return registeredFields.add(analysisField.getId());
             } else if (field instanceof HostedField hostedField) {