fix(java): Correct the async meta-shared layer bootstrap in ObjectStreamSerializer

java/fory-core/src/main/java/org/apache/fory/builder/MetaSharedLayerCodecBuilder.java

@@ -31,14 +31,11 @@
 import org.apache.fory.memory.MemoryBuffer;
 import org.apache.fory.meta.TypeDef;
 import org.apache.fory.reflect.TypeRef;
-import org.apache.fory.serializer.CodegenSerializer;
 import org.apache.fory.serializer.MetaSharedLayerSerializer;
 import org.apache.fory.serializer.MetaSharedLayerSerializerBase;
-import org.apache.fory.serializer.Serializers;
 import org.apache.fory.type.Descriptor;
 import org.apache.fory.type.DescriptorGrouper;
 import org.apache.fory.util.ExceptionUtils;
-import org.apache.fory.util.GraalvmSupport;
 import org.apache.fory.util.Preconditions;
 import org.apache.fory.util.StringUtils;
 
@@ -56,18 +53,16 @@
  */
 public class MetaSharedLayerCodecBuilder extends ObjectCodecBuilder {
   private final TypeDef layerTypeDef;
-  private final Class<?> layerMarkerClass;
 
   public MetaSharedLayerCodecBuilder(
       TypeRef<?> beanType, Fory fory, TypeDef layerTypeDef, Class<?> layerMarkerClass) {
     super(beanType, fory, GeneratedMetaSharedLayerSerializer.class);
     Preconditions.checkArgument(
         !fory.getConfig().checkClassVersion(),
         "Class version check should be disabled when compatible mode is enabled.");
-    this.layerTypeDef = layerTypeDef;
-    this.layerMarkerClass = layerMarkerClass;
+    this.layerTypeDef = fory.getTypeResolver().cacheTypeDef(layerTypeDef);
     DescriptorGrouper grouper =
-        typeResolver(r -> r.createDescriptorGrouper(layerTypeDef, beanClass));
+        typeResolver(r -> r.createDescriptorGrouper(this.layerTypeDef, beanClass));
     objectCodecOptimizer = new ObjectCodecOptimizer(beanClass, grouper, false, ctx);
   }
 
@@ -102,13 +97,15 @@ public String genCode() {
             ""
                 + "super(${fory}, ${cls});\n"
                 + "this.${fory} = ${fory};\n"
-                + "${serializer} = ${builderClass}.setCodegenSerializer(${fory}, ${cls}, this);\n",
+                + "${serializer} = ${builderClass}.setCodegenSerializer(${fory}, ${cls}, ${layerTypeDefId});\n",
             "fory",
             FORY_NAME,
             "cls",
             POJO_CLASS_TYPE_NAME,
             "builderClass",
             MetaSharedLayerCodecBuilder.class.getName(),
+            "layerTypeDefId",
+            layerTypeDef.getId() + "L",
             "serializer",
             SERIALIZER_FIELD_NAME);
     ctx.clearExprState();
@@ -130,23 +127,13 @@ protected void addCommonImports() {
   // Invoked by JIT.
   @SuppressWarnings({"unchecked", "rawtypes"})
   public static MetaSharedLayerSerializerBase setCodegenSerializer(
-      Fory fory, Class<?> cls, GeneratedMetaSharedLayerSerializer s) {
-    if (GraalvmSupport.isGraalRuntime()) {
-      return (MetaSharedLayerSerializerBase) typeResolver(fory, r -> r.getSerializer(s.getType()));
-    }
-    // This method hold jit lock, so create jit serializer async to avoid block serialization.
-    // Use MetaSharedLayerSerializer as fallback since it's compatible with
-    // MetaSharedLayerSerializerBase
-    Class serializerClass =
-        fory.getJITContext()
-            .registerSerializerJITCallback(
-                () -> MetaSharedLayerSerializer.class,
-                () -> CodegenSerializer.loadCodegenSerializer(fory, s.getType()),
-                c ->
-                    s.serializer =
-                        (MetaSharedLayerSerializerBase)
-                            Serializers.newSerializer(fory, s.getType(), c));
-    return (MetaSharedLayerSerializerBase) Serializers.newSerializer(fory, cls, serializerClass);
+      Fory fory, Class<?> cls, long layerTypeDefId) {
+    TypeDef layerTypeDef =
+        Preconditions.checkNotNull(
+            fory.getTypeResolver().getTypeDefById(layerTypeDefId),
+            "Missing cached layer TypeDef for id " + layerTypeDefId + " and class " + cls);
+    return new MetaSharedLayerSerializer(
+        fory, cls, layerTypeDef, LayerMarkerClassGenerator.getOrCreate(fory, cls, 0));
   }
 
   @Override

java/fory-core/src/main/java/org/apache/fory/resolver/TypeResolver.java

@@ -1119,6 +1119,10 @@ public final TypeDef cacheTypeDef(TypeDef typeDef) {
     return sharedRegistry.getOrCreateTypeDef(typeDef);
   }
 
+  public final TypeDef getTypeDefById(long typeDefId) {
+    return sharedRegistry.typeDefById.get(typeDefId);
+  }
+
   public final boolean isSerializable(Class<?> cls) {
     // Enums are always serializable, even if abstract (enums with abstract methods)
     if (cls.isEnum()) {

java/fory-core/src/test/java/org/apache/fory/serializer/ObjectStreamSerializerTest.java

@@ -41,16 +41,20 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.TreeMap;
+import java.util.TreeSet;
 import java.util.Vector;
 import java.util.concurrent.ConcurrentHashMap;
 import lombok.EqualsAndHashCode;
 import org.apache.fory.Fory;
 import org.apache.fory.ForyTestBase;
+import org.apache.fory.builder.Generated;
 import org.apache.fory.collection.LongMap;
 import org.apache.fory.config.CompatibleMode;
 import org.apache.fory.config.ForyBuilder;
 import org.apache.fory.config.Language;
 import org.apache.fory.memory.MemoryBuffer;
+import org.apache.fory.reflect.ReflectionUtils;
 import org.apache.fory.resolver.MetaContext;
 import org.apache.fory.resolver.SharedRegistry;
 import org.apache.fory.resolver.TypeInfo;
@@ -1228,6 +1232,107 @@ public void testNestedObjectSerialization(CompatibleMode compatible) {
     assertEquals(result.nestedList.get(1).nestedValue, "list2");
   }
 
+  public static class AsyncTreeSetSubclass extends TreeSet<String> {
+    public AsyncTreeSetSubclass() {}
+  }
+
+  public static class AsyncTreeMapSubclass extends TreeMap<String, String> {
+    public AsyncTreeMapSubclass() {}
+  }
+
+  @EqualsAndHashCode
+  public static class AsyncLayerJitContainer implements Serializable {
+    private String name;
+    private AsyncTreeSetSubclass values;
+    private AsyncTreeMapSubclass attributes;
+
+    public AsyncLayerJitContainer() {}
+
+    public AsyncLayerJitContainer(
+        String name, AsyncTreeSetSubclass values, AsyncTreeMapSubclass attributes) {
+      this.name = name;
+      this.values = values;
+      this.attributes = attributes;
+    }
+
+    private void writeObject(ObjectOutputStream s) throws IOException {
+      s.defaultWriteObject();
+    }
+
+    private void readObject(ObjectInputStream s) throws IOException, ClassNotFoundException {
+      s.defaultReadObject();
+    }
+  }
+
+  @Test(timeOut = 60000)
+  public void testAsyncCompilationNestedTreeCollectionsCompatibleMode()
+      throws InterruptedException {
+    Fory fory = newCompatibleAsyncObjectStreamFory(true);
+    fory.registerSerializer(
+        AsyncLayerJitContainer.class,
+        new ObjectStreamSerializer(fory, AsyncLayerJitContainer.class));
+    fory.registerSerializer(
+        AsyncTreeSetSubclass.class, new ObjectStreamSerializer(fory, AsyncTreeSetSubclass.class));
+    fory.registerSerializer(
+        AsyncTreeMapSubclass.class, new ObjectStreamSerializer(fory, AsyncTreeMapSubclass.class));
+
+    AsyncTreeSetSubclass values = new AsyncTreeSetSubclass();
+    values.add("one");
+    values.add("two");
+    AsyncTreeMapSubclass attributes = new AsyncTreeMapSubclass();
+    attributes.put("alpha", "A");
+    attributes.put("beta", "B");
+    AsyncLayerJitContainer obj = new AsyncLayerJitContainer("container", values, attributes);
+
+    serDeCheckSerializer(fory, obj, "ObjectStreamSerializer");
+
+    waitForGeneratedLayerSerializer(fory, AsyncLayerJitContainer.class);
+    waitForGeneratedLayerSerializer(fory, AsyncTreeSetSubclass.class);
+    waitForGeneratedLayerSerializer(fory, AsyncTreeMapSubclass.class);
+
+    serDeCheckSerializer(fory, obj, "ObjectStreamSerializer");
+  }
+
+  @Test(timeOut = 60000)
+  public void testAsyncCompilationTreeSetSubclassObjectStreamSerializer()
+      throws InterruptedException {
+    Fory fory = newCompatibleAsyncObjectStreamFory(true);
+    fory.registerSerializer(
+        AsyncTreeSetSubclass.class, new ObjectStreamSerializer(fory, AsyncTreeSetSubclass.class));
+
+    AsyncTreeSetSubclass values = new AsyncTreeSetSubclass();
+    values.add("one");
+    values.add("two");
+
+    serDeCheckSerializer(fory, values, "ObjectStreamSerializer");
+    waitForGeneratedLayerSerializer(fory, AsyncTreeSetSubclass.class);
+    serDeCheckSerializer(fory, values, "ObjectStreamSerializer");
+  }
+
+  @Test
+  public void testTreeCollectionsStillWorkWithoutAsyncCompilation() {
+    Fory fory = newCompatibleAsyncObjectStreamFory(false);
+    fory.registerSerializer(
+        AsyncLayerJitContainer.class,
+        new ObjectStreamSerializer(fory, AsyncLayerJitContainer.class));
+    fory.registerSerializer(
+        AsyncTreeSetSubclass.class, new ObjectStreamSerializer(fory, AsyncTreeSetSubclass.class));
+    fory.registerSerializer(
+        AsyncTreeMapSubclass.class, new ObjectStreamSerializer(fory, AsyncTreeMapSubclass.class));
+
+    AsyncTreeSetSubclass values = new AsyncTreeSetSubclass();
+    values.add("one");
+    values.add("two");
+    AsyncTreeMapSubclass attributes = new AsyncTreeMapSubclass();
+    attributes.put("alpha", "A");
+    attributes.put("beta", "B");
+
+    serDeCheckSerializer(
+        fory,
+        new AsyncLayerJitContainer("container", values, attributes),
+        "ObjectStreamSerializer");
+  }
+
   // ==================== Circular Reference in Custom Serialization ====================
 
   /** Class with potential circular reference. */
@@ -1371,4 +1476,45 @@ public void testAllPrimitiveTypes(CompatibleMode compatible) {
     assertEquals(result.charVal, 'A');
     assertEquals(result.boolVal, true);
   }
+
+  private Fory newCompatibleAsyncObjectStreamFory(boolean asyncCompilation) {
+    return Fory.builder()
+        .withLanguage(Language.JAVA)
+        .requireClassRegistration(false)
+        .withRefTracking(true)
+        .withCodegen(true)
+        .withCompatibleMode(CompatibleMode.COMPATIBLE)
+        .withAsyncCompilation(asyncCompilation)
+        .build();
+  }
+
+  private void waitForGeneratedLayerSerializer(Fory fory, Class<?> type)
+      throws InterruptedException {
+    long deadline = System.currentTimeMillis() + 30_000;
+    while (System.currentTimeMillis() < deadline) {
+      if (hasGeneratedLayerSerializer(fory, type)) {
+        return;
+      }
+      Thread.sleep(10);
+    }
+    Assert.fail("Timed out waiting for generated layer serializer for " + type.getName());
+  }
+
+  private boolean hasGeneratedLayerSerializer(Fory fory, Class<?> type) {
+    Serializer<?> serializer = fory.getTypeResolver().getSerializer(type);
+    if (!(serializer instanceof ObjectStreamSerializer)) {
+      return false;
+    }
+    Object[] slotsInfos = (Object[]) ReflectionUtils.getObjectFieldValue(serializer, "slotsInfos");
+    if (slotsInfos.length == 0) {
+      return false;
+    }
+    for (Object slotsInfo : slotsInfos) {
+      Object slotsSerializer = ReflectionUtils.getObjectFieldValue(slotsInfo, "slotsSerializer");
+      if (!(slotsSerializer instanceof Generated)) {
+        return false;
+      }
+    }
+    return true;
+  }
 }