diff --git a/amaranth/sim/_async.py b/amaranth/sim/_async.py index d21aa9bdd..3bf7021ed 100644 --- a/amaranth/sim/_async.py +++ b/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) diff --git a/amaranth/sim/core.py b/amaranth/sim/core.py index c96335904..e8ebf313e 100644 --- a/amaranth/sim/core.py +++ b/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 diff --git a/tests/test_sim.py b/tests/test_sim.py index 3745e058f..9a1f95d14 100644 --- a/tests/test_sim.py +++ b/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()