Add a more serious GC check.

Author: patrick-kidger

tests/test_core.py

@@ -2,6 +2,7 @@
 import gc
 import threading
 import time
+import weakref
 from typing import Any
 
 import pytest
@@ -261,7 +262,7 @@ def g():
     assert loop.run(g()) == 3
 
 
-def test_gc():
+def test_gc_simple():
     def _block_add_one(x):
         return x + 1
 
@@ -281,6 +282,53 @@ def _gc(x: int) -> tinyio.Coro[tuple[int, int]]:
     assert set(loop._results.keys()) == {coro}
 
 
+@pytest.mark.parametrize("yield_from", (False, True))
+@pytest.mark.parametrize("timeout", (None, 10))
+def test_gc_after_event(yield_from, timeout):
+    """The interesting case here is `yield_from=True`, `timeout=10`.
+    (The others are just for completeness.)
+
+    In this case we have that:
+    - `f1` has `timeout=10` but triggers immediately.
+    - `f2` has `timeout=2` but triggers at the end of our main coroutine.
+    And so we have that `f2` is before of `f1` in the internal heap of timeouts, but that `f1` will trigger first.
+    Thus when `f1` triggers it will remain in that heap even after it has triggered (until `f2` has triggered as well
+    and they can both be popped).
+    In this scenario, we don't want the generator object to remain in memory just because it's still sitting in that
+    heap!
+    This test checks that the generator can be cleaned up even whilst we wait for the `_Wait` object to get collected
+    later.
+    """
+
+    def wait(event, wait_time):
+        if yield_from:
+            yield from event.wait(wait_time)
+        else:
+            yield event.wait(wait_time)
+
+    def set_event(event):
+        for _ in range(20):
+            yield
+        event.set()
+
+    def baz():
+        event1 = tinyio.Event()
+        event2 = tinyio.Event()
+        f1 = wait(event1, timeout)
+        f2 = wait(event2, 2)
+        ref = weakref.ref(f1)
+        yield {f2}
+        yield [f1, set_event(event1)]
+        del f1
+        gc.collect()
+        assert ref() is None
+        event2.set()
+        return 3
+
+    loop = tinyio.Loop()
+    assert loop.run(baz()) == 3
+
+
 def test_event_fairness():
     """This checks that once one event unblocks, that we don't just keep chasing all the stuff downstream of that event,
     i.e. that we do schedule work from any other event that has finished.

tinyio/_core.py

@@ -350,7 +350,7 @@ def cleanup(self):
             assert self._waiting_for is not None
             assert self._event is not None
             self.state = _WaitState.DONE
-            self._waiting_for = None
+            self._waiting_for = None  # For GC purposes.
             del self._event._waits[self]
             self._event = None  # For GC purposes.