Use asyncgen hooks to call aclose()

amaranth/sim/_async.py

@@ -362,15 +362,12 @@ async def until(self, condition: ValueLike):
                 raise TypeError(f"The shape of a condition may only be `signed` or `unsigned`, "
                                 f"not {shape!r}")
         tick = self.sample(condition).__aiter__()
-        try:
-            done = False
-            while not done:
-                clk, rst, *values, done = await tick.__anext__()
-                if rst:
-                    raise DomainReset
-            return tuple(values)
-        finally:
-            await tick.aclose()
+        done = False
+        while not done:
+            clk, rst, *values, done = await tick.__anext__()
+            if rst:
+                raise DomainReset
+        return tuple(values)
 
     async def repeat(self, count: int):
         """Repeat this trigger a specific number of times.
@@ -403,15 +400,12 @@ async def repeat(self, count: int):
         if count <= 0:
             raise ValueError(f"Repeat count must be a positive integer, not {count!r}")
         tick = self.__aiter__()
-        try:
-            for _ in range(count):
-                clk, rst, *values = await tick.__anext__()
-                if rst:
-                    raise DomainReset
-                assert clk
-            return tuple(values)
-        finally:
-            await tick.aclose()
+        for _ in range(count):
+            clk, rst, *values = await tick.__anext__()
+            if rst:
+                raise DomainReset
+            assert clk
+        return tuple(values)
 
     def _collect_trigger(self):
         clk_polarity = (1 if self._domain.clk_edge == "pos" else 0)

amaranth/sim/core.py

@@ -1,4 +1,6 @@
+from contextlib import contextmanager
 import inspect
+import sys
 import warnings
 
 from .._utils import deprecated
@@ -189,6 +191,15 @@ async def testbench(ctx):
         which they were added. If two testbenches share state, or must manipulate the design in
         a coordinated way, they may rely on this execution order for correctness.
 
+        .. warning::
+
+            On Python 3.12 and earlier, async generators (:py:`async` functions that also
+            :py:`yield`) are not cleaned up reliably when the simulator exits. To make sure context
+            managers or :py:`finally` blocks inside async generators run properly, use Python 3.13
+            or later. See `PEP-525`_ for more information on async generator finalization.
+
+            .. _PEP-525: https://peps.python.org/pep-0525/#finalization
+
         Raises
         ------
         :exc:`RuntimeError`
@@ -246,6 +257,15 @@ async def process(ctx):
             :py:`await ctx.tick().sample(...)`. Such state is visible in a waveform viewer,
             simplifying debugging.
 
+        .. warning::
+
+            On Python 3.12 and earlier, async generators (:py:`async` functions that also
+            :py:`yield`) are not cleaned up reliably when the simulator exits. To make sure context
+            managers or :py:`finally` blocks inside async generators run properly, use Python 3.13
+            or later. See `PEP-525`_ for more information on async generator finalization.
+
+            .. _PEP-525: https://peps.python.org/pep-0525/#finalization
+
         Raises
         ------
         :exc:`RuntimeError`
@@ -341,6 +361,33 @@ def run_until(self, deadline, *, run_passive=None):
         while self._engine.now < deadline.femtoseconds:
             self.advance()
 
+    @contextmanager
+    def _replace_asyncgen_hooks(self):
+        # Async generators require hooks for lifetime management. Replace existing hooks with ours
+        # for the duration of this context manager.
+
+        def firstiter(agen):
+            # Prevent any outer event loop from seeing this generator.
+            pass
+
+        def finalizer(agen):
+            # Generators can't be closed if they are currently running.
+            if not agen.ag_running:
+                # Try to run aclose() once, but skip it if it awaits anything.
+                try:
+                    coroutine = agen.aclose()
+                    coroutine.send(None)
+                except StopIteration:
+                    # Success
+                    return
+
+        old_hooks = sys.get_asyncgen_hooks()
+        sys.set_asyncgen_hooks(firstiter=firstiter, finalizer=finalizer)
+        try:
+            yield
+        finally:
+            sys.set_asyncgen_hooks(*old_hooks)
+
     def advance(self):
         """Advance the simulation.
 
@@ -356,7 +403,8 @@ def advance(self):
         :py:`False` otherwise.
         """
         self._running = True
-        return self._engine.advance()
+        with self._replace_asyncgen_hooks():
+            return self._engine.advance()
 
     def write_vcd(self, vcd_file, gtkw_file=None, *, traces=(), fs_per_delta=0):
         # `fs_per_delta`` is not currently documented; it is not clear if we want to expose

tests/test_sim.py

@@ -2077,6 +2077,16 @@ async def change(ctx):
         sim.add_clock(Period(s=4))
         sim.run()
 
+    def test_abandon_repeat(self):
+        m = Module()
+        toggle = Signal(1)
+        m.d.sync += toggle.eq(~toggle)
+        sim = Simulator(m)
+        async def process(ctx):
+            await ctx.tick().repeat(100)
+        sim.add_process(process)
+        sim.run()
+
     def test_trigger_wrong(self):
         a = Signal(4)
         m = Module()