Dispose thread state on context finalization

Author: fangerer

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/PythonCextBuiltins.java

@@ -251,6 +251,7 @@
 import com.oracle.graal.python.nodes.util.CastToJavaStringNode;
 import com.oracle.graal.python.runtime.ExecutionContext.IndirectCallContext;
 import com.oracle.graal.python.runtime.PythonContext;
+import com.oracle.graal.python.runtime.PythonContext.PythonThreadState;
 import com.oracle.graal.python.runtime.PythonCore;
 import com.oracle.graal.python.runtime.PythonOptions;
 import com.oracle.graal.python.runtime.exception.ExceptionUtils;
@@ -1657,12 +1658,18 @@ Object wrap(VirtualFrame frame, Object bufferStructPointer, Object ownerObj, Obj
 
     @Builtin(name = "PyThreadState_Get")
     @GenerateNodeFactory
-    abstract static class PyThreadState_Get extends NativeBuiltin {
+    abstract static class PyThreadStateGet extends NativeBuiltin {
 
         @Specialization
         PThreadState get() {
+            PythonThreadState threadState = getContext().getThreadState();
+            PThreadState nativeWrapper = threadState.getNativeWrapper();
+            if (nativeWrapper == null) {
+                nativeWrapper = new PThreadState(threadState);
+                threadState.setNativeWrapper(nativeWrapper);
+            }
             // does not require a 'to_sulong' since it is already a native wrapper type
-            return new PThreadState(getContext().getThreadState());
+            return nativeWrapper;
         }
     }
 

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/PythonContext.java

@@ -61,6 +61,9 @@
 import com.oracle.graal.python.builtins.objects.PythonAbstractObject;
 import com.oracle.graal.python.builtins.objects.cext.PythonNativeClass;
 import com.oracle.graal.python.builtins.objects.cext.capi.CApiContext;
+import com.oracle.graal.python.builtins.objects.cext.capi.PThreadState;
+import com.oracle.graal.python.builtins.objects.cext.capi.PyTruffleObjectFree.ReleaseHandleNode;
+import com.oracle.graal.python.builtins.objects.cext.capi.PyTruffleObjectFreeFactory.ReleaseHandleNodeGen;
 import com.oracle.graal.python.builtins.objects.cext.capi.PythonNativeWrapper;
 import com.oracle.graal.python.builtins.objects.cext.hpy.GraalHPyContext;
 import com.oracle.graal.python.builtins.objects.common.HashingStorage;
@@ -139,6 +142,14 @@ public static final class PythonThreadState {
         /* corresponds to 'PyThreadState.dict' */
         PDict dict;
 
+        /*
+         * This is the native wrapper object if we need to expose the thread state as PyThreadState
+         * object. We need to store it here because the wrapper may receive 'toNative' in which case
+         * a handle is allocated. In order to avoid leaks, the handle needs to be free'd when the
+         * owning thread (or the whole context) is disposed.
+         */
+        PThreadState nativeWrapper;
+
         /*
          * The constructor needs to have this particular signature such that we can use it for
          * ContextThreadLocal.
@@ -191,6 +202,26 @@ public PDict getDict() {
         public void setDict(PDict dict) {
             this.dict = dict;
         }
+
+        public PThreadState getNativeWrapper() {
+            return nativeWrapper;
+        }
+
+        public void setNativeWrapper(PThreadState nativeWrapper) {
+            this.nativeWrapper = nativeWrapper;
+        }
+
+        public void dispose() {
+            ReleaseHandleNode releaseHandleNode = ReleaseHandleNodeGen.getUncached();
+            if (dict != null && dict.getNativeWrapper() != null) {
+                releaseHandleNode.execute(dict.getNativeWrapper());
+            }
+            dict = null;
+            if (nativeWrapper != null) {
+                releaseHandleNode.execute(nativeWrapper);
+                nativeWrapper = null;
+            }
+        }
     }
 
     private static final class AtExitHook {
@@ -245,8 +276,8 @@ private static final class AtExitHook {
      */
     private final PythonThreadState buildThreadState;
 
-    /* map of thread IDs to indices for array 'threadStates' */
-    private Map<Thread, PythonThreadState> threadStateMapping = Collections.synchronizedMap(new WeakHashMap<>());
+    /* map of thread IDs to the corresponding 'threadStates' */
+    private final Map<Thread, PythonThreadState> threadStateMapping = Collections.synchronizedMap(new WeakHashMap<>());
 
     private final ReentrantLock importLock = new ReentrantLock();
     @CompilationFinal private boolean isInitialized = false;
@@ -860,6 +891,16 @@ public void runShutdownHooks() {
         for (ShutdownHook h : shutdownHooks) {
             h.call(this);
         }
+        assert threadStateMapping != null;
+        for (PythonThreadState threadState : threadStateMapping.values()) {
+            threadState.dispose();
+        }
+        ReleaseHandleNode releaseHandleNode = ReleaseHandleNodeGen.getUncached();
+        for (PythonNativeWrapper singletonNativeWrapper : singletonNativePtrs) {
+            if (singletonNativeWrapper != null) {
+                releaseHandleNode.execute(singletonNativeWrapper);
+            }
+        }
     }
 
     @TruffleBoundary
@@ -1166,6 +1207,7 @@ public synchronized void disposeThread(Thread thread) {
         }
         ts.shutdown();
         threadStateMapping.remove(thread);
+        ts.dispose();
         releaseSentinelLock(ts.sentinelLock);
     }