[GR-54727] Fix leak of StructSequence builtin nodes and call targets.

PullRequest: graalpython/3376

Author: steve-s

graalpython/com.oracle.graal.python.test/src/com/oracle/graal/python/test/advanced/LeakTest.java

@@ -333,11 +333,13 @@ protected void launch(Builder contextBuilder) {
                     }
                 }
             }
-            // the check at the end will make a dump anyway, so no createHeapDump flag here
             long currentSize = getJavaHeapSize(false);
             System.out.printf("Heap size after all repetitions: %,d\n", currentSize);
             if (currentSize > initialSize * 1.1) {
                 System.err.printf("Heap size grew too much after repeated context creations and invocations. From %,d bytes to %,d bytes.\n", initialSize, currentSize);
+                if (keepDump) {
+                    dumpHeap(ManagementFactory.getPlatformMBeanServer(), true);
+                }
                 System.exit(255);
             }
         }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/PythonLanguage.java

@@ -37,6 +37,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.lang.invoke.VarHandle;
+import java.lang.ref.WeakReference;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.InvalidPathException;
 import java.util.ArrayList;
@@ -70,6 +71,10 @@
 import com.oracle.graal.python.builtins.objects.function.PArguments;
 import com.oracle.graal.python.builtins.objects.module.PythonModule;
 import com.oracle.graal.python.builtins.objects.object.PythonObject;
+import com.oracle.graal.python.builtins.objects.tuple.StructSequence;
+import com.oracle.graal.python.builtins.objects.tuple.StructSequence.BuiltinTypeDescriptor;
+import com.oracle.graal.python.builtins.objects.tuple.StructSequence.Descriptor;
+import com.oracle.graal.python.builtins.objects.tuple.StructSequence.DescriptorCallTargets;
 import com.oracle.graal.python.builtins.objects.type.MroShape;
 import com.oracle.graal.python.builtins.objects.type.PythonAbstractClass;
 import com.oracle.graal.python.builtins.objects.type.PythonManagedClass;
@@ -333,6 +338,46 @@ public boolean isSingleContext() {
 
     @CompilationFinal(dimensions = 1) private final RootCallTarget[] builtinSlotsCallTargets;
 
+    /**
+     * Weak hash map of call targets for builtin functions associated with named tuples generated at
+     * runtime from C extensions. We hold the cached call targets also weakly, because otherwise we
+     * would have a cycle from the value (call targets reference builtin nodes which wrap the
+     * descriptor) to the key. The key should be GC'ed when the corresponding generated named tuple
+     * class is GC'ed.
+     */
+    private final WeakHashMap<StructSequence.Descriptor, WeakReference<StructSequence.DescriptorCallTargets>> structSequenceTargets = new WeakHashMap<>();
+
+    /**
+     * The same as {@link #structSequenceTargets}, but for builtin named tuples. There is a bounded
+     * statically known number of builtin named tuples.
+     */
+    private final ConcurrentHashMap<StructSequence.Descriptor, StructSequence.DescriptorCallTargets> structSequenceBuiltinTargets = new ConcurrentHashMap<>();
+
+    public StructSequence.DescriptorCallTargets getOrCreateStructSequenceCallTargets(StructSequence.Descriptor descriptor,
+                    Function<StructSequence.Descriptor, StructSequence.DescriptorCallTargets> factory) {
+        if (singleContext) {
+            return factory.apply(descriptor);
+        }
+        if (descriptor instanceof BuiltinTypeDescriptor builtinDescriptor) {
+            // There must be finite set of objects initialized at build time, no need for a weak map
+            assert !ImageInfo.inImageCode() || builtinDescriptor.wasInitializedAtBuildTime();
+            return structSequenceBuiltinTargets.computeIfAbsent(builtinDescriptor, factory);
+        }
+        return getOrCreateStructSeqNonBuiltinTargets(descriptor, factory);
+    }
+
+    private DescriptorCallTargets getOrCreateStructSeqNonBuiltinTargets(Descriptor descriptor, Function<Descriptor, DescriptorCallTargets> factory) {
+        synchronized (structSequenceTargets) {
+            WeakReference<DescriptorCallTargets> weakResult = structSequenceTargets.computeIfAbsent(descriptor, d -> new WeakReference<>(factory.apply(d)));
+            DescriptorCallTargets result = weakResult.get();
+            if (result == null) {
+                result = factory.apply(descriptor);
+                structSequenceTargets.put(descriptor, new WeakReference<>(result));
+            }
+            return result;
+        }
+    }
+
     /**
      * We cannot initialize call targets in language ctor and the next suitable hook is context
      * initialization, but that is called multiple times. We use this flag to run the language
@@ -410,6 +455,9 @@ public MroShape getMroShapeRoot() {
     protected void finalizeContext(PythonContext context) {
         context.finalizeContext();
         super.finalizeContext(context);
+        // trigger cleanup of stale entries in weak hash maps
+        structSequenceTargets.size();
+        indirectCallDataMap.size();
     }
 
     @Override
@@ -1000,9 +1048,87 @@ public void setBuiltinSlotCallTarget(int index, RootCallTarget callTarget) {
     }
 
     /**
-     * Cache call targets that are created for every new context, based on a single key.
+     * Caches call target that wraps a node that is not parametrized, i.e., has only a parameterless
+     * ctor and all its instances implement the same logic. Parametrized nodes must include the
+     * parameters that alter their behavior as part of the cache key.
+     */
+    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Class<? extends Node> key) {
+        // It's complicated with RootNodes, but regular nodes should have only parameterless ctor to
+        // be appropriate keys for the cache
+        assert RootNode.class.isAssignableFrom(key) || key.getConstructors().length <= 1;
+        assert RootNode.class.isAssignableFrom(key) || key.getConstructors().length == 0 || key.getConstructors()[0].getParameterCount() == 0;
+        return createCachedCallTargetUnsafe(rootNodeFunction, key);
+    }
+
+    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Enum<?> key) {
+        return createCachedCallTargetUnsafe(rootNodeFunction, key);
+    }
+
+    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Class<? extends Node> nodeClass, String key) {
+        // for builtins: name is needed to distinguish builtins that share the same underlying node
+        // in general: a String may be parameter of the node wrapped in the root node or the root
+        // node itself, there must be finite number of strings that can appear here (i.e., must not
+        // be dynamically generated unless their number is bounded).
+        return createCachedCallTargetUnsafe(rootNodeFunction, nodeClass, key);
+    }
+
+    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Class<? extends Node> nodeClass, TruffleString key) {
+        // See the String overload
+        return createCachedCallTargetUnsafe(rootNodeFunction, nodeClass, key);
+    }
+
+    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Class<? extends Node> nodeClass, int key) {
+        return createCachedCallTargetUnsafe(rootNodeFunction, nodeClass, key);
+    }
+
+    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Class<? extends Node> nodeClass1, Class<?> nodeClass2, String name) {
+        // for slot wrappers: the root node may be wrapping a helper wrapper node implementing the
+        // wrapper logic and the bare slot node itself
+        return createCachedCallTargetUnsafe(rootNodeFunction, nodeClass1, nodeClass2, name);
+    }
+
+    /**
+     * Caches call targets for external C functions created by extensions at runtime.
+     * <p>
+     * For the time being, we assume finite/limited number of extensions and their external
+     * functions. This may hold onto call targets created by one extension used in a context that
+     * was closed in the meanwhile and no other context ever loads the extension.
      */
-    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Object key) {
+    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Class<? extends RootNode> klass, Enum<?> signature, TruffleString name,
+                    boolean doArgumentAndResultConversion) {
+        return createCachedCallTargetUnsafe(rootNodeFunction, klass, signature, name, doArgumentAndResultConversion);
+    }
+
+    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Enum<?> signature, TruffleString name,
+                    boolean doArgumentAndResultConversion) {
+        return createCachedCallTargetUnsafe(rootNodeFunction, signature, name, doArgumentAndResultConversion);
+    }
+
+    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Enum<?> signature, TruffleString name) {
+        return createCachedCallTargetUnsafe(rootNodeFunction, signature, name);
+    }
+
+    public RootCallTarget createStructSeqIndexedMemberAccessCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, int memberIndex) {
+        return createCachedCallTargetUnsafe(rootNodeFunction, StructSequence.class, memberIndex);
+    }
+
+    public RootCallTarget createCachedPropAccessCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Class<?> nodeClass, String name, int type, int offset) {
+        // For the time being, we assume finite/limited number of cext/hpy types members, their
+        // types and offsets
+        return createCachedCallTargetUnsafe(rootNodeFunction, nodeClass, name, type, offset);
+    }
+
+    /**
+     * Keys in any caches held by {@link PythonLanguage} must be context independent objects and
+     * there must be either finite number of their instances, or if the key is a context independent
+     * mirror of some runtime data structure, it must be cached weakly. This call targets cache is
+     * strong.
+     * <p>
+     * To avoid memory leaks, all key types must be known to have finite number of possible
+     * instances. Public methods for adding to the cache must take concrete key type(s) so that all
+     * possible cache keys are explicit and documented.
+     */
+    private RootCallTarget createCachedCallTargetUnsafe(Function<PythonLanguage, RootNode> rootNodeFunction, Object key) {
         CompilerAsserts.neverPartOfCompilation();
         if (!singleContext) {
             return cachedCallTargets.computeIfAbsent(key, k -> PythonUtils.getOrCreateCallTarget(rootNodeFunction.apply(this)));
@@ -1011,11 +1137,8 @@ public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode>
         }
     }
 
-    /**
-     * Cache call targets that are created for every new context, based on a list of keys.
-     */
-    public RootCallTarget createCachedCallTarget(Function<PythonLanguage, RootNode> rootNodeFunction, Object... cacheKeys) {
-        return createCachedCallTarget(rootNodeFunction, Arrays.asList(cacheKeys));
+    private RootCallTarget createCachedCallTargetUnsafe(Function<PythonLanguage, RootNode> rootNodeFunction, Object... cacheKeys) {
+        return createCachedCallTargetUnsafe(rootNodeFunction, Arrays.asList(cacheKeys));
     }
 
     public void registerBuiltinDescriptorCallTarget(BuiltinMethodDescriptor descriptor, RootCallTarget callTarget) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/Builtin.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017, 2023, Oracle and/or its affiliates.
+ * Copyright (c) 2017, 2024, Oracle and/or its affiliates.
  * Copyright (c) 2013, Regents of the University of California
  *
  * All rights reserved.
@@ -109,4 +109,14 @@
     String raiseErrorName() default StringLiterals.J_EMPTY_STRING;
 
     boolean forceSplitDirectCalls() default false;
+
+    /**
+     * If set to {@code true}, then this builtin will be initialized and registered in
+     * {@link PythonBuiltins#initialize(Python3Core)}. Otherwise, it will be ignored even when it
+     * has a generated node factory. This is useful when we want to initialize the builtin manually
+     * in different way (e.g., wrap it in method, descriptor, ...). By convention set this to
+     * {@code false} also for builtins declared outside of {@link PythonBuiltins} subclass to
+     * document the intent to not initialize them automatically.
+     */
+    boolean autoRegister() default true;
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/PythonBuiltins.java

@@ -244,7 +244,12 @@ private void initializeEachFactoryWith(BiConsumer<NodeFactory<? extends PythonBu
         PythonOS currentOs = PythonOS.getPythonOS();
         for (NodeFactory<? extends PythonBuiltinBaseNode> factory : factories) {
             Boolean needsFrame = null;
+            boolean initialized = false;
             for (Builtin builtin : factory.getNodeClass().getAnnotationsByType(Builtin.class)) {
+                if (!builtin.autoRegister()) {
+                    assert !initialized : "Builtin annotations on " + factory.getNodeClass().getName() + " do not agree on 'autoInitialize' property.";
+                    break;
+                }
                 if (builtin.os() == PythonOS.PLATFORM_ANY || builtin.os() == currentOs) {
                     if (needsFrame == null) {
                         needsFrame = builtin.needsFrame();