maintain node hygiene in the heap by passing the user objects as arguments to the atExitHook

Author: timfel

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

@@ -80,15 +80,8 @@ private static class AtExitCallTarget extends RootNode {
 
             private final ContextReference<PythonContext> contextRef = lookupContextReference(PythonLanguage.class);
 
-            private Object callable;
-            private Object[] arguments;
-            private PKeyword[] keywords;
-
-            protected AtExitCallTarget(TruffleLanguage<?> language, Object callable, Object[] arguments, PKeyword[] keywords) {
+            protected AtExitCallTarget(TruffleLanguage<?> language) {
                 super(language);
-                this.callable = callable;
-                this.arguments = arguments;
-                this.keywords = keywords;
             }
 
             @Override
@@ -98,6 +91,10 @@ public Object execute(VirtualFrame frame) {
                 context.setTopFrameInfo(PFrame.Reference.EMPTY);
                 context.setCaughtException(PException.NO_EXCEPTION);
 
+                Object callable = frame.getArguments()[0];
+                Object[] arguments = (Object[]) frame.getArguments()[1];
+                PKeyword[] keywords = (PKeyword[]) frame.getArguments()[2];
+
                 // We deliberately pass 'null' frame here, the execution state will then be taken
                 // from the context.
                 try {
@@ -112,8 +109,8 @@ public Object execute(VirtualFrame frame) {
         @Specialization
         Object register(Object callable, Object[] arguments, PKeyword[] keywords) {
             CompilerDirectives.transferToInterpreter();
-            AtExitCallTarget atExitCallTarget = new AtExitCallTarget(getContext().getLanguage(), callable, arguments, keywords);
-            getContext().registerShutdownHook(callable, PythonUtils.getOrCreateCallTarget(atExitCallTarget));
+            AtExitCallTarget atExitCallTarget = new AtExitCallTarget(getContext().getLanguage());
+            getContext().registerShutdownHook(callable, arguments, keywords, PythonUtils.getOrCreateCallTarget(atExitCallTarget));
             return callable;
         }
     }

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

@@ -51,6 +51,7 @@
 import java.util.logging.Level;
 
 import org.graalvm.collections.EconomicMap;
+import org.graalvm.collections.MapCursor;
 import org.graalvm.nativeimage.ImageInfo;
 import org.graalvm.options.OptionKey;
 
@@ -69,6 +70,7 @@
 import com.oracle.graal.python.builtins.objects.dict.PDict;
 import com.oracle.graal.python.builtins.objects.frame.PFrame;
 import com.oracle.graal.python.builtins.objects.frame.PFrame.Reference;
+import com.oracle.graal.python.builtins.objects.function.PKeyword;
 import com.oracle.graal.python.builtins.objects.list.PList;
 import com.oracle.graal.python.builtins.objects.module.PythonModule;
 import com.oracle.graal.python.builtins.objects.object.PythonObjectLibrary;
@@ -164,6 +166,18 @@ List<WeakReference<Thread>> getOwners() {
         }
     }
 
+    private static final class AtExitHook {
+        final Object[] arguments;
+        final PKeyword[] keywords;
+        final CallTarget ct;
+
+        AtExitHook(Object[] arguments, PKeyword[] keywords, CallTarget ct) {
+            this.arguments = arguments;
+            this.keywords = keywords;
+            this.ct = ct;
+        }
+    }
+
     static final String PREFIX = "/";
     static final String LIB_PYTHON_3 = "/lib-python/3";
     static final String LIB_GRAALPYTHON = "/lib-graalpython";
@@ -178,7 +192,7 @@ List<WeakReference<Thread>> getOwners() {
     private PythonModule mainModule;
     private final PythonCore core;
     private final List<ShutdownHook> shutdownHooks = new ArrayList<>();
-    private final EconomicMap<Object, CallTarget> atExitHooks = EconomicMap.create();
+    private final EconomicMap<Object, AtExitHook> atExitHooks = EconomicMap.create(0);
     private final HashMap<PythonNativeClass, CyclicAssumption> nativeClassStableAssumptions = new HashMap<>();
     private final AtomicLong globalId = new AtomicLong(Integer.MAX_VALUE * 2L + 4L);
     private final ThreadGroup threadGroup = new ThreadGroup(GRAALPYTHON_THREADS);
@@ -653,8 +667,8 @@ public void registerShutdownHook(ShutdownHook shutdownHook) {
     }
 
     @TruffleBoundary
-    public void registerShutdownHook(Object callable, CallTarget ct) {
-        atExitHooks.put(callable, ct);
+    public void registerShutdownHook(Object callable, Object[] arguments, PKeyword[] keywords, CallTarget ct) {
+        atExitHooks.put(callable, new AtExitHook(arguments, keywords, ct));
     }
 
     @TruffleBoundary
@@ -666,13 +680,16 @@ public void deregisterShutdownHook(Object callable) {
     public void runShutdownHooks() {
         handler.shutdown();
         // run atExitHooks in reverse order they were registered
-        CallTarget[] hooks = new CallTarget[atExitHooks.size()];
-        Iterator<CallTarget> hooksIter = atExitHooks.getValues().iterator();
+        MapCursor<Object, AtExitHook> cursor = atExitHooks.getEntries();
+        AtExitHook[] hooks = new AtExitHook[atExitHooks.size()];
+        Object[] callables = new Object[atExitHooks.size()];
         for (int i = 0; i < hooks.length; i++) {
-            hooks[i] = hooksIter.next();
+            cursor.advance();
+            callables[i] = cursor.getKey();
+            hooks[i] = cursor.getValue();
         }
         for (int i = hooks.length - 1; i >= 0; i--) {
-            hooks[i].call();
+            hooks[i].ct.call(callables[i], hooks[i].arguments, hooks[i].keywords);
         }
         for (ShutdownHook h : shutdownHooks) {
             h.call(this);