Free native memory of native heap types

Author: fangerer

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/PThreadState.java

@@ -69,7 +69,7 @@ public final class PThreadState extends PythonStructNativeWrapper {
     private PThreadState(PythonThreadState threadState) {
         super(threadState, true);
         // 'registerReplacement' will set the native pointer if not running LLVM managed mode.
-        replacement = registerReplacement(allocateCLayout(threadState), InteropLibrary.getUncached());
+        replacement = registerReplacement(allocateCLayout(threadState), false, InteropLibrary.getUncached());
     }
 
     public static Object getThreadState(PythonLanguage language, PythonContext context) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/PyMemoryViewWrapper.java

@@ -165,7 +165,7 @@ private static Object allocate(PMemoryView object) {
     public Object getReplacement(InteropLibrary lib) {
         if (replacement == null) {
             Object pointerObject = allocate((PMemoryView) getDelegate());
-            replacement = registerReplacement(pointerObject, lib);
+            replacement = registerReplacement(pointerObject, false, lib);
         }
         return replacement;
     }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/PythonClassNativeWrapper.java

@@ -44,7 +44,7 @@
 import com.oracle.graal.python.builtins.objects.cext.common.CArrayWrappers.CStringWrapper;
 import com.oracle.graal.python.builtins.objects.cext.structs.CFields;
 import com.oracle.graal.python.builtins.objects.cext.structs.CStructAccess;
-import com.oracle.graal.python.builtins.objects.cext.structs.CStructAccessFactory;
+import com.oracle.graal.python.builtins.objects.cext.structs.CStructAccess.AllocateNode;
 import com.oracle.graal.python.builtins.objects.cext.structs.CStructs;
 import com.oracle.graal.python.builtins.objects.type.PythonClass;
 import com.oracle.graal.python.builtins.objects.type.PythonManagedClass;
@@ -55,7 +55,6 @@
 import com.oracle.graal.python.builtins.objects.type.TypeNodesFactory.SetItemSizeNodeGen;
 import com.oracle.graal.python.nodes.HiddenAttr;
 import com.oracle.graal.python.nodes.PGuards;
-import com.oracle.graal.python.nodes.attributes.WriteAttributeToObjectNode;
 import com.oracle.graal.python.util.PythonUtils;
 import com.oracle.truffle.api.CompilerAsserts;
 import com.oracle.truffle.api.CompilerDirectives;
@@ -111,7 +110,6 @@ public static void wrapNative(PythonManagedClass clazz, TruffleString name, Obje
 
         CStructAccess.ReadI64Node readI64 = CStructAccess.ReadI64Node.getUncached();
         CStructAccess.ReadPointerNode readPointer = CStructAccess.ReadPointerNode.getUncached();
-        WriteAttributeToObjectNode writeAttr = WriteAttributeToObjectNode.getUncached();
         InteropLibrary lib = InteropLibrary.getUncached();
 
         // some values are retained from the native representation
@@ -158,7 +156,7 @@ public static void wrapNative(PythonManagedClass clazz, TruffleString name, Obje
          * 'getReplacement' for more explanation).
          */
         wrapper.replacement = pointer;
-        wrapper.registerReplacement(pointer, lib);
+        wrapper.registerReplacement(pointer, false, lib);
     }
 
     public static void initNative(PythonManagedClass clazz, Object pointer) {
@@ -190,8 +188,8 @@ public Object getReplacement(InteropLibrary lib) {
              */
             PythonManagedClass clazz = (PythonManagedClass) getDelegate();
             boolean heaptype = (GetTypeFlagsNode.executeUncached(clazz) & TypeFlags.HEAPTYPE) != 0;
-            Object pointerObject = CStructAccessFactory.AllocateNodeGen.getUncached().alloc(heaptype ? CStructs.PyHeapTypeObject : CStructs.PyTypeObject);
-            replacement = registerReplacement(pointerObject, lib);
+            Object pointerObject = AllocateNode.allocUncached(heaptype ? CStructs.PyHeapTypeObject : CStructs.PyTypeObject);
+            replacement = registerReplacement(pointerObject, true, lib);
 
             ToNativeTypeNode.initializeType(this, pointerObject, heaptype);
         }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/PythonNativeWrapper.java

@@ -139,19 +139,19 @@ public Object getReplacement(@SuppressWarnings("unused") InteropLibrary lib) {
     }
 
     @TruffleBoundary
-    public final Object registerReplacement(Object pointer, InteropLibrary lib) {
+    public final Object registerReplacement(Object pointer, boolean freeAtCollection, InteropLibrary lib) {
         LOGGER.finest(() -> PythonUtils.formatJString("assigning %s with %s", getDelegate(), pointer));
         Object result;
         if (pointer instanceof Long lptr) {
             // need to convert to actual pointer
             result = PCallCapiFunction.callUncached(NativeCAPISymbol.FUN_CONVERT_POINTER, lptr);
-            CApiTransitions.createReference(this, lptr);
+            CApiTransitions.createReference(this, lptr, freeAtCollection);
         } else {
             result = pointer;
             if (lib.isPointer(pointer)) {
                 assert CApiContext.isPointerObject(pointer);
                 try {
-                    CApiTransitions.createReference(this, lib.asPointer(pointer));
+                    CApiTransitions.createReference(this, lib.asPointer(pointer), freeAtCollection);
                 } catch (UnsupportedMessageException e) {
                     throw CompilerDirectives.shouldNotReachHere(e);
                 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/transitions/CApiTransitions.java

@@ -207,29 +207,40 @@ public static final class PythonObjectReference extends IdReference<PythonNative
          * The index in the lookup table, where this reference is stored. This duplicates the native
          * field {@link CFields#GraalPyObject__handle_table_index} in order to save reading the
          * native field if we already have the reference object. The value may be {@code -1} in
-         * which case it means that the reference's referent is a managed object and
-         * {@link #pointer} points into the handle space.
+         * which case it means that {@link #pointer} is not a tagged pointer but a pointer to some
+         * other off-heap memory (e.g. {@code PyTypeObject} or other Python structs).
          */
         private final int handleTableIndex;
 
-        private PythonObjectReference(HandleContext handleContext, PythonNativeWrapper referent, boolean strong, long pointer, int handleTableIndex) {
+        /**
+         * Indicates if the native memory {@link #pointer} should be freed (using {@link FreeNode})
+         * if this reference was enqueued because the referent was collected. For example, this will
+         * be {@code true} if the referent is
+         * {@link com.oracle.graal.python.builtins.objects.cext.capi.PythonClassNativeWrapper} and
+         * the class native replacement was allocated on the heap (usually a heap type). It will be
+         * {@code false} if the class wraps a static type.
+         */
+        private final boolean freeAtCollection;
+
+        private PythonObjectReference(HandleContext handleContext, PythonNativeWrapper referent, boolean strong, long pointer, int handleTableIndex, boolean freeAtCollection) {
             super(handleContext, referent);
             this.pointer = pointer;
             this.strongReference = strong ? referent : null;
             referent.ref = this;
             this.handleTableIndex = handleTableIndex;
+            this.freeAtCollection = freeAtCollection;
             if (LOGGER.isLoggable(Level.FINE)) {
                 LOGGER.fine(PythonUtils.formatJString("new %s", toString()));
             }
         }
 
         static PythonObjectReference create(HandleContext handleContext, PythonAbstractObjectNativeWrapper referent, boolean strong, long pointer, int idx) {
             assert HandlePointerConverter.pointsToPyHandleSpace(pointer);
-            return new PythonObjectReference(handleContext, referent, strong, pointer, idx);
+            return new PythonObjectReference(handleContext, referent, strong, pointer, idx, true);
         }
 
-        static PythonObjectReference create(HandleContext handleContext, PythonNativeWrapper referent, long pointer) {
-            return new PythonObjectReference(handleContext, referent, true, pointer, -1);
+        static PythonObjectReference create(HandleContext handleContext, PythonNativeWrapper referent, long pointer, boolean freeAtCollection) {
+            return new PythonObjectReference(handleContext, referent, true, pointer, -1, freeAtCollection);
         }
 
         public boolean isStrongReference() {
@@ -244,6 +255,10 @@ public int getHandleTableIndex() {
             return handleTableIndex;
         }
 
+        public boolean isFreeAtCollection() {
+            return freeAtCollection;
+        }
+
         @Override
         @TruffleBoundary
         public String toString() {
@@ -390,6 +405,9 @@ public static void pollReferenceQueue() {
                         } else {
                             assert nativeLookupGet(context, reference.pointer) != null : Long.toHexString(reference.pointer);
                             nativeLookupRemove(context, reference.pointer);
+                            if (reference.freeAtCollection) {
+                                freeNativeStruct(reference);
+                            }
                         }
                     } else if (entry instanceof NativeObjectReference reference) {
                         nativeLookupRemove(context, reference.pointer);
@@ -454,21 +472,16 @@ public static void freeClassReplacements(HandleContext handleContext) {
         assert PythonContext.get(null).ownsGil();
         handleContext.nativeLookup.forEach((l, ref) -> {
             if (ref instanceof PythonObjectReference reference) {
-                /*
-                 * If a PythonObjectReference is in this table, it either refers to an object that
-                 * is backed by static native memory (e.g. static PyTypeObject), or it refers to
-                 * some non-object structure (e.g. PyThreadState). In any case, the objects are
-                 * immutable and so the references are expected to be strong.
-                 */
-                assert reference.strongReference != null;
+                // We don't expect references to wrappers that would have a native object stub.
                 assert reference.handleTableIndex == -1;
                 /*
-                 * TODO(fa): For classes, distinguish between static and heap memory.
-                 *
-                 * The ref may denote class wrappers where some of them are backed by static native
-                 * memory and some of them were allocated in heap. There is currently no reliable
-                 * way to distinguish them and so we leak the heap types.
+                 * The ref may denote: (a) class wrappers, where some of them are backed by static
+                 * native memory and some of them were allocated in heap, and (b) struct wrappers,
+                 * which may be freed manually in a separate step.
                  */
+                if (reference.freeAtCollection) {
+                    freeNativeStruct(reference);
+                }
             }
         });
         handleContext.nativeLookup.clear();
@@ -479,7 +492,7 @@ public static void disableReferenceQueuePolling(HandleContext handleContext) {
     }
 
     private static void freeNativeStub(long stubPointer) {
-        LOGGER.fine(() -> PythonUtils.formatJString("freeing native object stub 0x%x", stubPointer));
+        LOGGER.fine(() -> PythonUtils.formatJString("releasing native object stub 0x%x", stubPointer));
         FreeNode.executeUncached(stubPointer);
     }
 
@@ -488,6 +501,13 @@ private static void freeNativeStub(PythonObjectReference ref) {
         freeNativeStub(HandlePointerConverter.pointerToStub(ref.pointer));
     }
 
+    private static void freeNativeStruct(PythonObjectReference ref) {
+        assert ref.handleTableIndex == -1;
+        assert ref.freeAtCollection;
+        LOGGER.fine(() -> PythonUtils.formatJString("releasing %s", ref.toString()));
+        FreeNode.executeUncached(ref.pointer);
+    }
+
     private static void freeNativeStorage(NativeStorageReference ref) {
         LOGGER.fine(() -> PythonUtils.formatJString("releasing %s", ref.toString()));
         FreeNode.executeUncached(ref.ptr);
@@ -617,7 +637,12 @@ public static IdReference<?> nativeLookupGet(HandleContext context, long pointer
     }
 
     @TruffleBoundary
-    public static IdReference<?> nativeLookupPut(HandleContext context, long pointer, IdReference<?> value) {
+    public static IdReference<?> nativeLookupPut(HandleContext context, long pointer, NativeObjectReference value) {
+        return context.nativeLookup.put(pointer, value);
+    }
+
+    @TruffleBoundary
+    public static IdReference<?> nativeLookupPut(HandleContext context, long pointer, PythonObjectReference value) {
         return context.nativeLookup.put(pointer, value);
     }
 
@@ -814,7 +839,7 @@ static long doGeneric(Node inliningTarget, PythonAbstractObjectNativeWrapper wra
      */
     @TruffleBoundary
     @SuppressWarnings("try")
-    public static void createReference(PythonNativeWrapper obj, long ptr) {
+    public static void createReference(PythonNativeWrapper obj, long ptr, boolean freeAtCollection) {
         try (GilNode.UncachedAcquire ignored = GilNode.uncachedAcquire()) {
             /*
              * The first test if '!obj.isNative()' in the caller is done on a fast-path but not
@@ -825,7 +850,7 @@ public static void createReference(PythonNativeWrapper obj, long ptr) {
                 obj.setNativePointer(ptr);
                 pollReferenceQueue();
                 HandleContext context = getContext();
-                nativeLookupPut(context, ptr, PythonObjectReference.create(context, obj, ptr));
+                nativeLookupPut(context, ptr, PythonObjectReference.create(context, obj, ptr, freeAtCollection));
             }
         }
     }