Add firstiter async hook

This can be used by various async frameworks to ensure asyncgens shut
down - e.g. asyncio.shutdown_asyncgens(). This is even more important
in GraalPy, since we do not run the finalizers.

Author: Adam Hrbac

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

@@ -1136,6 +1136,12 @@ Object getProfile() {
     abstract static class SetAsyncgenHooks extends PythonBuiltinNode {
         @Specialization
         Object setAsyncgenHooks(Object firstIter, Object finalizer) {
+            if (firstIter != PNone.NO_VALUE && firstIter != PNone.NONE) {
+                getContext().getThreadState(getLanguage()).setAsyncgenFirstIter(firstIter);
+            } else if (firstIter == PNone.NONE) {
+                getContext().getThreadState(getLanguage()).setAsyncgenFirstIter(null);
+            }
+            // Ignore finalizer, since we don't have a useful place to call it
             return PNone.NONE;
         }
     }
@@ -1146,7 +1152,9 @@ abstract static class GetAsyncgenHooks extends PythonBuiltinNode {
         @Specialization
         Object setAsyncgenHooks() {
             // TODO: use asyncgen_hooks object
-            return factory().createTuple(new Object[]{PNone.NONE, PNone.NONE});
+            PythonContext.PythonThreadState threadState = getContext().getThreadState(getLanguage());
+            Object firstiter = threadState.getAsyncgenFirstIter();
+            return factory().createTuple(new Object[]{firstiter == null ? PNone.NONE : firstiter, PNone.NONE});
         }
     }
 

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/asyncio/AsyncGeneratorBuiltins.java

@@ -40,17 +40,24 @@
  */
 package com.oracle.graal.python.builtins.objects.asyncio;
 
+import static com.oracle.graal.python.nodes.SpecialMethodNames.J___AITER__;
+import static com.oracle.graal.python.nodes.SpecialMethodNames.J___ANEXT__;
+
+import java.util.List;
+
 import com.oracle.graal.python.builtins.Builtin;
 import com.oracle.graal.python.builtins.CoreFunctions;
 import com.oracle.graal.python.builtins.PythonBuiltinClassType;
 import com.oracle.graal.python.builtins.PythonBuiltins;
 import com.oracle.graal.python.builtins.objects.PNone;
 import com.oracle.graal.python.builtins.objects.generator.GeneratorBuiltins;
+import com.oracle.graal.python.nodes.call.special.CallUnaryMethodNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonBinaryBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
 import com.oracle.graal.python.runtime.PAsyncGen;
+import com.oracle.graal.python.runtime.PythonContext;
 import com.oracle.truffle.api.dsl.Bind;
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.GenerateNodeFactory;
@@ -60,13 +67,20 @@
 import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.profiles.InlinedConditionProfile;
 
-import java.util.List;
-
-import static com.oracle.graal.python.nodes.SpecialMethodNames.J___AITER__;
-import static com.oracle.graal.python.nodes.SpecialMethodNames.J___ANEXT__;
-
 @CoreFunctions(extendClasses = PythonBuiltinClassType.PAsyncGenerator)
 public class AsyncGeneratorBuiltins extends PythonBuiltins {
+    private static void callHooks(VirtualFrame frame, PAsyncGen self, PythonContext.PythonThreadState state, CallUnaryMethodNode invokeFirstIter) {
+        Object firstIter = state.getAsyncgenFirstIter();
+        if (firstIter == null) {
+            return;
+        }
+        if (self.isHookCalled()) {
+            return;
+        }
+        self.setHookCalled(true);
+        invokeFirstIter.executeObject(frame, firstIter, self);
+    }
+
     @Override
     protected List<? extends NodeFactory<? extends PythonBuiltinBaseNode>> getNodeFactories() {
         return AsyncGeneratorBuiltinsFactory.getFactories();
@@ -116,7 +130,9 @@ public boolean isRunning(PAsyncGen self) {
     @GenerateNodeFactory
     public abstract static class ASend extends PythonBinaryBuiltinNode {
         @Specialization
-        public Object aSend(PAsyncGen self, Object sent) {
+        public Object aSend(VirtualFrame frame, PAsyncGen self, Object sent,
+                        @Cached CallUnaryMethodNode callFirstIter) {
+            callHooks(frame, self, getContext().getThreadState(getLanguage()), callFirstIter);
             return factory().createAsyncGeneratorASend(self, sent);
         }
     }
@@ -127,7 +143,9 @@ public abstract static class AThrow extends PythonBuiltinNode {
         public abstract Object execute(VirtualFrame frame, PAsyncGen self, Object arg1, Object arg2, Object arg3);
 
         @Specialization
-        public Object athrow(PAsyncGen self, Object arg1, Object arg2, Object arg3) {
+        public Object athrow(VirtualFrame frame, PAsyncGen self, Object arg1, Object arg2, Object arg3,
+                        @Cached CallUnaryMethodNode callFirstIter) {
+            callHooks(frame, self, getContext().getThreadState(getLanguage()), callFirstIter);
             return factory().createAsyncGeneratorAThrow(self, arg1, arg2, arg3);
         }
     }
@@ -145,7 +163,9 @@ public Object aIter(PAsyncGen self) {
     @GenerateNodeFactory
     public abstract static class ANext extends PythonUnaryBuiltinNode {
         @Specialization
-        public Object aNext(PAsyncGen self) {
+        public Object aNext(VirtualFrame frame, PAsyncGen self,
+                        @Cached CallUnaryMethodNode callFirstIter) {
+            callHooks(frame, self, getContext().getThreadState(getLanguage()), callFirstIter);
             return factory().createAsyncGeneratorASend(self, PNone.NONE);
         }
     }
@@ -154,7 +174,9 @@ public Object aNext(PAsyncGen self) {
     @GenerateNodeFactory
     public abstract static class AClose extends PythonUnaryBuiltinNode {
         @Specialization
-        public Object aClose(PAsyncGen self) {
+        public Object aClose(VirtualFrame frame, PAsyncGen self,
+                        @Cached CallUnaryMethodNode callFirstIter) {
+            callHooks(frame, self, getContext().getThreadState(getLanguage()), callFirstIter);
             return factory().createAsyncGeneratorAThrow(self, null, PNone.NO_VALUE, PNone.NO_VALUE);
         }
     }

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

@@ -49,7 +49,7 @@
 
 public final class PAsyncGen extends PGenerator {
     private boolean closed = false;
-    private boolean hooksCalled = false;
+    private boolean hookCalled = false;
     private boolean runningAsync = false;
 
     public static PAsyncGen create(PythonLanguage lang, TruffleString name, TruffleString qualname, PBytecodeRootNode rootNode, RootCallTarget[] callTargets, Object[] arguments) {
@@ -69,12 +69,12 @@ public void markClosed() {
         this.closed = true;
     }
 
-    public boolean isHooksCalled() {
-        return hooksCalled;
+    public boolean isHookCalled() {
+        return hookCalled;
     }
 
-    public void setHooksCalled(boolean hooksCalled) {
-        this.hooksCalled = hooksCalled;
+    public void setHookCalled(boolean hookCalled) {
+        this.hookCalled = hookCalled;
     }
 
     public boolean isRunningAsync() {

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

@@ -300,6 +300,11 @@ public static final class PythonThreadState {
          */
         Object runningEventLoop;
 
+        /*
+         * A callable that should be called for the first iteration of an async generators.
+         */
+        Object asyncgenFirstIter;
+
         /*
          * Counter for C-level recursion depth used for Py_(Enter/Leave)RecursiveCall.
          */
@@ -477,6 +482,14 @@ public void profilingStart() {
         public void profilingStop() {
             this.profiling = false;
         }
+
+        public Object getAsyncgenFirstIter() {
+            return asyncgenFirstIter;
+        }
+
+        public void setAsyncgenFirstIter(Object asyncgenFirstIter) {
+            this.asyncgenFirstIter = asyncgenFirstIter;
+        }
     }
 
     private static final class AtExitHook {