JBR-9552 Move dynamic call target tables into generated classes to allow garbage collection

Author: YaaZ

src/java.base/share/classes/com/jetbrains/internal/jbrapi/AccessContext.java

@@ -46,21 +46,29 @@
  */
 class AccessContext {
 
+    static final int DYNAMIC_CALL_TARGET_NAME_OFFSET = 128;
+    @SuppressWarnings("unchecked")
+    static Supplier<MethodHandle>[] getDynamicCallTargets(Lookup target) {
+        try {
+            return (Supplier<MethodHandle>[]) target.findStaticVarHandle(
+                    target.lookupClass(), "dynamicCallTargets", Supplier[].class).get();
+        } catch (NoSuchFieldException | IllegalAccessException e) {
+            throw new RuntimeException(e);
+        }
+    }
     private static final DirectMethodHandleDesc BOOTSTRAP_DYNAMIC_DESC = MethodHandleDesc.ofMethod(
             DirectMethodHandleDesc.Kind.STATIC, desc(JBRApiSupport.class), "bootstrapDynamic",
             desc(CallSite.class, Lookup.class, String.class, MethodType.class));
 
     private final Map<Class<?>, Boolean> accessibleClasses = new HashMap<>();
     final Map<Proxy, Boolean> dependencies = new HashMap<>(); // true for required, false for optional
-    final List<DynamicCallTarget> dynamicCallTargets = new ArrayList<>();
+    final List<Supplier<MethodHandle>> dynamicCallTargets = new ArrayList<>();
     final Lookup caller;
 
     AccessContext(Lookup caller) {
         this.caller = caller;
     }
 
-    record DynamicCallTarget(String name, MethodTypeDesc descriptor, Supplier<MethodHandle> futureHandle) {}
-
     class Method {
         final CodeBuilder writer;
         private final boolean methodRequired;
@@ -84,10 +92,9 @@ void invokeDynamic(MethodHandle handle) {
         }
 
         void invokeDynamic(MethodType type, Supplier<MethodHandle> futureHandle) {
-            MethodTypeDesc desc = desc(erase(type));
-            DynamicCallTarget t = new DynamicCallTarget("dynamic" + dynamicCallTargets.size(), desc, futureHandle);
-            dynamicCallTargets.add(t);
-            writer.invokedynamic(DynamicCallSiteDesc.of(BOOTSTRAP_DYNAMIC_DESC, t.name, desc));
+            String name = String.valueOf((char) (dynamicCallTargets.size() + DYNAMIC_CALL_TARGET_NAME_OFFSET));
+            dynamicCallTargets.add(futureHandle);
+            writer.invokedynamic(DynamicCallSiteDesc.of(BOOTSTRAP_DYNAMIC_DESC, name, desc(erase(type))));
         }
 
         void invokeDirect(MethodHandleInfo handleInfo) {

src/java.base/share/classes/com/jetbrains/internal/jbrapi/BytecodeUtils.java

@@ -37,6 +37,7 @@
 import java.lang.invoke.MethodType;
 import java.lang.reflect.Method;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 import static java.lang.classfile.ClassFile.ACC_FINAL;
 import static java.lang.classfile.ClassFile.ACC_PUBLIC;
@@ -49,6 +50,8 @@ class BytecodeUtils {
     public static final ClassDesc VOID_DESC = desc(void.class);
     public static final ClassDesc OBJECT_DESC = desc(Object.class);
     public static final ClassDesc OBJECT_ARRAY_DESC = OBJECT_DESC.arrayType();
+    public static final ClassDesc SUPPLIER_DESC = desc(Supplier.class);
+    public static final ClassDesc SUPPLIER_ARRAY_DESC = SUPPLIER_DESC.arrayType();
     public static final ClassDesc EXTENSION_ARRAY_DESC = desc(long[].class);
     public static final ClassDesc PROXY_INTERFACE_DESC = desc(com.jetbrains.exported.JBRApiSupport.Proxy.class);
     public static final MethodTypeDesc GET_PROXY_TARGET_DESC = MethodTypeDesc.of(OBJECT_DESC);

src/java.base/share/classes/com/jetbrains/internal/jbrapi/JBRApi.java

@@ -33,10 +33,7 @@
 import java.lang.invoke.MethodType;
 import java.lang.reflect.Method;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
 import java.util.function.Function;
-import java.util.function.Supplier;
 
 import static java.lang.invoke.MethodHandles.Lookup;
 
@@ -80,17 +77,6 @@ public class JBRApi {
      */
     private static final boolean EXTEND_REGISTRY = Utils.property("jetbrains.runtime.api.extendRegistry", false);
 
-    record DynamicCallTargetKey(Class<?> proxy, String name, String descriptor) {}
-    static final ConcurrentMap<DynamicCallTargetKey, Supplier<MethodHandle>> dynamicCallTargets = new ConcurrentHashMap<>();
-
-    public static MethodHandle bindDynamic(Lookup caller, String name, MethodType type) {
-        if (VERBOSE) {
-            System.out.println("Binding call site " + caller.lookupClass().getName() + "#" + name + ": " + type);
-        }
-        if (!caller.hasFullPrivilegeAccess()) throw new Error("Caller lookup must have full privilege access"); // Authenticity check.
-        return dynamicCallTargets.get(new DynamicCallTargetKey(caller.lookupClass(), name, type.descriptorString())).get().asType(type);
-    }
-
     private final ProxyRepository proxyRepository;
     private final Boolean[] supportedExtensions;
     private final long[] emptyExtensionsBitfield;
@@ -250,4 +236,15 @@ public <T> T getService(Class<T> interFace, Enum<?>... extensions) {
         }
         return null;
     }
+
+    public static MethodHandle bindDynamic(Lookup caller, String name, MethodType type) {
+        int index = name.charAt(0) - AccessContext.DYNAMIC_CALL_TARGET_NAME_OFFSET;
+        if (VERBOSE) {
+            System.out.println("Binding call site " + caller.lookupClass().getName() + " #" + index);
+        }
+        if (!caller.hasFullPrivilegeAccess()) {
+            throw new Error("Caller lookup must have full privilege access"); // Authenticity check.
+        }
+        return AccessContext.getDynamicCallTargets(caller)[index].get().asType(type);
+    }
 }

src/java.base/share/classes/com/jetbrains/internal/jbrapi/ProxyGenerator.java

@@ -195,10 +195,10 @@ void init() {
         if (JBRApi.VERBOSE) {
             System.out.println("Initializing proxy " + interFace.getName());
         }
-        for (var t : accessContext.dynamicCallTargets) {
-            JBRApi.dynamicCallTargets.put(new JBRApi.DynamicCallTargetKey(
-                    generatedProxy.lookupClass(), t.name(), t.descriptor().descriptorString()
-            ), t.futureHandle());
+        if (!accessContext.dynamicCallTargets.isEmpty()) {
+            var table = AccessContext.getDynamicCallTargets(generatedProxy);
+            for (int i = 0; i < table.length; i++) table[i] = accessContext.dynamicCallTargets.get(i);
+
         }
     }
 
@@ -210,7 +210,7 @@ MethodHandle define(boolean service) {
         try {
             Object sericeTarget = service && info.targetLookup != null ? createServiceTarget() : null;
             generatedProxy = proxyGenLookup.defineHiddenClass(
-                    bytecode, false, Lookup.ClassOption.STRONG, Lookup.ClassOption.NESTMATE);
+                    bytecode, false, Lookup.ClassOption.NESTMATE);
             MethodHandle constructor = findConstructor();
             if (sericeTarget != null) constructor = MethodHandles.insertArguments(constructor, 0, sericeTarget);
             return constructor;
@@ -244,10 +244,10 @@ boolean generate() {
             cb.withFlags(ACC_SUPER | ACC_FINAL | ACC_SYNTHETIC)
               .withSuperclass(superclassDesc)
               .withInterfaceSymbols(superinterfaceDescs);
-            generateFields(cb);
             generateConstructor(cb);
             generateTargetGetter(cb);
             generateMethods(cb);
+            generateFields(cb);
         });
         if (JBRApi.VERIFY_BYTECODE) {
             List<VerifyError> errors = ClassFile.of().verify(bytecode);
@@ -267,6 +267,14 @@ private void generateFields(ClassBuilder cb) {
         if (EXTENSIONS_ENABLED) {
             cb.withField("extensions", EXTENSION_ARRAY_DESC, ACC_PRIVATE | ACC_FINAL);
         }
+        if (!accessContext.dynamicCallTargets.isEmpty()) {
+            cb.withField("dynamicCallTargets", SUPPLIER_ARRAY_DESC, ACC_PRIVATE | ACC_FINAL | ACC_STATIC);
+            cb.withMethodBody("<clinit>", MethodTypeDesc.of(VOID_DESC), ACC_PRIVATE | ACC_STATIC, m -> m
+                    .loadConstant(accessContext.dynamicCallTargets.size())
+                    .anewarray(SUPPLIER_DESC)
+                    .putstatic(proxyDesc, "dynamicCallTargets", SUPPLIER_ARRAY_DESC)
+                    .return_());
+        }
     }
 
     private void generateConstructor(ClassBuilder cb) {

test/jdk/jb/java/api/backend/RealTest.java

@@ -32,6 +32,7 @@
 import com.jetbrains.internal.jbrapi.JBRApi;
 
 import java.io.IOException;
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -43,9 +44,16 @@
 public class RealTest {
 
     public static void main(String[] args) throws Throwable {
+        // Plain run.
         run();
-        for (int i = 0; i < 2; i++) {
+
+        // Run in an isolated classloader at least 2 times until the proxy gets GC'ed.
+        WeakReference<?> weak = null;
+        for (int i = 0; i < 2 || weak.get() != null; i++) {
+            System.gc();
             new IsolatedLoader().loadClass(RealTest.class.getName()).getMethod("run").invoke(null);
+            if (weak == null) weak = new WeakReference<>(getProxy(Proxy.class));
+            if (i > 300) throw new Error("Proxy was not collected after 300 iterations");
         }
     }