Fix asyncgen close

Lib/test/test_coroutines.py

@@ -1668,7 +1668,6 @@ class CoroutineTest(unittest.TestCase):
 
         self.assertEqual(sys.getrefcount(aiter), refs_before)
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; refcount leak in async for/with
     def test_for_6(self):
         I = 0
 

crates/vm/src/builtins/asyncgenerator.rs

@@ -180,6 +180,7 @@ impl PyRef<PyAsyncGen> {
         exc_tb: OptionalArg,
         vm: &VirtualMachine,
     ) -> PyResult<PyAsyncGenAThrow> {
+        warn_deprecated_throw_signature(&exc_val, &exc_tb, vm)?;
         PyAsyncGen::init_hooks(&self, vm)?;
         Ok(PyAsyncGenAThrow {
             ag: self,
@@ -328,7 +329,7 @@ impl PyAsyncGenASend {
         let res = self.ag.inner.send(self.ag.as_object(), val, vm);
         let res = PyAsyncGenWrappedValue::unbox(&self.ag, res, vm);
         if res.is_err() {
-            self.close();
+            self.set_closed();
         }
         res
     }
@@ -341,8 +342,23 @@ impl PyAsyncGenASend {
         exc_tb: OptionalArg,
         vm: &VirtualMachine,
     ) -> PyResult {
-        if let AwaitableState::Closed = self.state.load() {
-            return Err(vm.new_runtime_error("cannot reuse already awaited __anext__()/asend()"));
+        match self.state.load() {
+            AwaitableState::Closed => {
+                return Err(
+                    vm.new_runtime_error("cannot reuse already awaited __anext__()/asend()")
+                );
+            }
+            AwaitableState::Init => {
+                if self.ag.running_async.load() {
+                    self.state.store(AwaitableState::Closed);
+                    return Err(
+                        vm.new_runtime_error("anext(): asynchronous generator is already running")
+                    );
+                }
+                self.ag.running_async.store(true);
+                self.state.store(AwaitableState::Iter);
+            }
+            AwaitableState::Iter => {}
         }
 
         warn_deprecated_throw_signature(&exc_val, &exc_tb, vm)?;
@@ -355,13 +371,36 @@ impl PyAsyncGenASend {
         );
         let res = PyAsyncGenWrappedValue::unbox(&self.ag, res, vm);
         if res.is_err() {
-            self.close();
+            self.set_closed();
         }
         res
     }
 
     #[pymethod]
-    fn close(&self) {
+    fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
+        if matches!(self.state.load(), AwaitableState::Closed) {
+            return Ok(());
+        }
+        let result = self.throw(
+            vm.ctx.exceptions.generator_exit.to_owned().into(),
+            OptionalArg::Missing,
+            OptionalArg::Missing,
+            vm,
+        );
+        match result {
+            Ok(_) => Err(vm.new_runtime_error("coroutine ignored GeneratorExit")),
+            Err(e)
+                if e.fast_isinstance(vm.ctx.exceptions.stop_iteration)
+                    || e.fast_isinstance(vm.ctx.exceptions.stop_async_iteration)
+                    || e.fast_isinstance(vm.ctx.exceptions.generator_exit) =>
+            {
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
+    }
+
+    fn set_closed(&self) {
         self.state.store(AwaitableState::Closed);
     }
 }
@@ -472,6 +511,32 @@ impl PyAsyncGenAThrow {
         exc_tb: OptionalArg,
         vm: &VirtualMachine,
     ) -> PyResult {
+        match self.state.load() {
+            AwaitableState::Closed => {
+                return Err(
+                    vm.new_runtime_error("cannot reuse already awaited aclose()/athrow()")
+                );
+            }
+            AwaitableState::Init => {
+                if self.ag.running_async.load() {
+                    self.state.store(AwaitableState::Closed);
+                    let msg = if self.aclose {
+                        "aclose(): asynchronous generator is already running"
+                    } else {
+                        "athrow(): asynchronous generator is already running"
+                    };
+                    return Err(vm.new_runtime_error(msg.to_owned()));
+                }
+                if self.ag.inner.closed() {
+                    self.state.store(AwaitableState::Closed);
+                    return Err(vm.new_stop_iteration(None));
+                }
+                self.ag.running_async.store(true);
+                self.state.store(AwaitableState::Iter);
+            }
+            AwaitableState::Iter => {}
+        }
+
         warn_deprecated_throw_signature(&exc_val, &exc_tb, vm)?;
         let ret = self.ag.inner.throw(
             self.ag.as_object(),
@@ -493,8 +558,27 @@ impl PyAsyncGenAThrow {
     }
 
     #[pymethod]
-    fn close(&self) {
-        self.state.store(AwaitableState::Closed);
+    fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
+        if matches!(self.state.load(), AwaitableState::Closed) {
+            return Ok(());
+        }
+        let result = self.throw(
+            vm.ctx.exceptions.generator_exit.to_owned().into(),
+            OptionalArg::Missing,
+            OptionalArg::Missing,
+            vm,
+        );
+        match result {
+            Ok(_) => Err(vm.new_runtime_error("coroutine ignored GeneratorExit")),
+            Err(e)
+                if e.fast_isinstance(vm.ctx.exceptions.stop_iteration)
+                    || e.fast_isinstance(vm.ctx.exceptions.stop_async_iteration)
+                    || e.fast_isinstance(vm.ctx.exceptions.generator_exit) =>
+            {
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
     }
 
     fn ignored_close(&self, res: &PyResult<PyIterReturn>) -> bool {