fix: full PyPy support — deterministic stream cleanup + assignor livelock fix, all PyPy skips removed

.github/workflows/python-package.yml

@@ -70,7 +70,10 @@ jobs:
   test-pypy:
     name: 'Python pypy3.11/Cython: false'
     runs-on: ubuntu-latest
-    timeout-minutes: 10
+    # PyPy runs the pure-Python paths and is markedly slower than CPython, and
+    # this leg now runs the formerly-skipped tests too; give it headroom over
+    # the old 10-minute cap, which already clipped the run at ~96%.
+    timeout-minutes: 15
     continue-on-error: true
     env:
       USE_CYTHON: 'false'

faust/assignor/copartitioned_assignor.py

@@ -259,8 +259,39 @@ def _assign_round_robin(self, unassigned: Iterable[int], active: bool) -> None:
                         and assigment.can_assign(partition, active)
                     )
                 )  # By above assertion, should never throw error
-                unassigned_partition = assign_to.pop_partition(active)
+                unassigned_partition = self._free_up_partition(assign_to, active)
                 unassigned.append(unassigned_partition)
 
             # Assign partition
             assign_to.assign_partition(partition, active)
+
+    def _free_up_partition(
+        self, assignment: CopartitionedAssignment, active: bool
+    ) -> int:
+        # Evict a partition from an exhausted client so the partition we are
+        # placing can take its slot.  Prefer a victim that some *other*
+        # not-yet-exhausted client can immediately take, so the re-queued
+        # victim makes forward progress instead of bouncing straight back
+        # here.  ``CopartitionedAssignment.pop_partition`` pops an arbitrary
+        # set element, and when the arbitrary choice is the victim's only
+        # possible home two partitions ping-pong between the queue and this
+        # client forever.  Whether that happens depends on set iteration
+        # order, so the livelock strikes on PyPy while CPython (with a
+        # different order) escapes it for the same input.  Choosing a
+        # re-homable victim in deterministic (sorted) order makes the loop
+        # terminate on every interpreter.
+        held = assignment.get_assigned_partitions(active)
+        for candidate in sorted(held):
+            for other in self._client_assignments.values():
+                if other is assignment:
+                    continue
+                if not self._client_exhausted(other, active) and other.can_assign(
+                    candidate, active
+                ):
+                    assignment.unassign_partition(candidate, active)
+                    return candidate
+        # No immediately re-homable victim: fall back to a deterministic
+        # arbitrary choice (still correct, just not progress-guaranteed).
+        candidate = min(held)
+        assignment.unassign_partition(candidate, active)
+        return candidate

faust/streams.py

@@ -7,8 +7,10 @@
 import weakref
 from asyncio import CancelledError
 from contextvars import ContextVar
+from functools import wraps
 from typing import (
     Any,
+    AsyncGenerator,
     AsyncIterable,
     AsyncIterator,
     Callable,
@@ -88,6 +90,30 @@ async def maybe_forward(value: Any, channel: ChannelT) -> Any:
     return value
 
 
+def _tracks_buffer_agen(
+    fun: Callable[..., AsyncIterable],
+) -> Callable[..., AsyncIterable]:
+    """Register buffering generators (``take()`` and friends) on the stream.
+
+    The cleanup for these generators -- acking consumed events, restoring
+    ``enable_acks`` and detaching the buffering processor -- lives in
+    ``finally`` blocks that only run once the generator is finalized.  When a
+    caller abandons the generator (``break`` inside ``async for``), CPython's
+    reference counting finalizes it right away, but PyPy defers finalization
+    to the next major GC cycle, which may never come.  Tracking every
+    generator in a WeakSet lets ``Stream.on_stop`` close leftovers
+    explicitly, making cleanup deterministic on any interpreter.
+    """
+
+    @wraps(fun)
+    def _create_agen(self: StreamT, *args: Any, **kwargs: Any) -> AsyncIterable:
+        agen = fun(self, *args, **kwargs)
+        cast("Stream", self)._active_agens.add(agen)
+        return agen
+
+    return _create_agen
+
+
 class _LinkedListDirection(NamedTuple):
     attr: str
     getter: Callable[[StreamT], Optional[StreamT]]
@@ -149,6 +175,13 @@ def __init__(
         self._processors = list(processors) if processors else []
         self._on_start = on_start
 
+        # Live async generators handed out by this stream --
+        # __aiter__ iterators and the take() family -- closed explicitly
+        # in on_stop().  A WeakSet so that on CPython an abandoned
+        # generator is still finalized (and thus cleaned up) promptly by
+        # reference counting.
+        self._active_agens: "weakref.WeakSet[AsyncGenerator]" = weakref.WeakSet()
+
         # attach beacon to channel, or if iterable attach to current task.
         task = current_task(loop=self.loop)
         if task is not None:
@@ -300,6 +333,7 @@ async def events(self) -> AsyncIterable[EventT]:
             if self.current_event is not None:
                 yield self.current_event
 
+    @_tracks_buffer_agen
     async def take(self, max_: int, within: Seconds) -> AsyncIterable[Sequence[T_co]]:
         """Buffer n values at a time and yield a list of buffered values.
 
@@ -392,6 +426,7 @@ async def add_to_buffer(value: T) -> T:
             self.enable_acks = stream_enable_acks
             self._processors.remove(add_to_buffer)
 
+    @_tracks_buffer_agen
     async def take_events(
         self, max_: int, within: Seconds
     ) -> AsyncIterable[Sequence[EventT]]:
@@ -485,6 +520,7 @@ async def add_to_buffer(value: T) -> T:
             self.enable_acks = stream_enable_acks
             self._processors.remove(add_to_buffer)
 
+    @_tracks_buffer_agen
     async def take_with_timestamp(
         self, max_: int, within: Seconds, timestamp_field_name: str
     ) -> AsyncIterable[Sequence[T_co]]:
@@ -592,6 +628,7 @@ def enumerate(self, start: int = 0) -> AsyncIterable[Tuple[int, T_co]]:
         """
         return aenumerate(self, start)
 
+    @_tracks_buffer_agen
     async def noack_take(
         self, max_: int, within: Seconds
     ) -> AsyncIterable[Sequence[T_co]]:
@@ -1033,6 +1070,23 @@ async def stop(self) -> None:
 
     async def on_stop(self) -> None:
         """Signal that the stream is stopping."""
+        # Close any buffering generator (take() and friends) the caller
+        # abandoned, so their ``finally`` blocks -- acking consumed events,
+        # restoring ``enable_acks`` and detaching the buffering processor --
+        # run deterministically now instead of whenever the interpreter
+        # finalizes the generator (which on PyPy waits for a GC cycle that
+        # may never come).  ``aclose()`` on an exhausted generator is a
+        # harmless no-op.
+        for agen in list(self._active_agens):
+            try:
+                await agen.aclose()
+            except (RuntimeError, ValueError):
+                # Generator is still running in another task (e.g. an agent
+                # actor iterating this stream during a rebalance); its owner
+                # will clean it up.  CPython raises RuntimeError for this,
+                # PyPy raises ValueError ("async generator already
+                # executing") -- catch both so on_stop never aborts here.
+                pass
         self._passive = False
         self._passive_started.clear()
         for table_or_stream in self.combined:
@@ -1046,9 +1100,14 @@ def __next__(self) -> T:
 
     def __aiter__(self) -> AsyncIterator[T_co]:  # pragma: no cover
         if _CStreamIterator is not None:
-            return self._c_aiter()
+            it = self._c_aiter()
         else:
-            return self._py_aiter()
+            it = self._py_aiter()
+        # Track the iterator so on_stop() can close it if the caller
+        # abandons it (see _tracks_buffer_agen for why this matters on
+        # interpreters without reference counting, e.g. PyPy).
+        self._active_agens.add(cast(AsyncGenerator, it))
+        return it
 
     async def _c_aiter(self) -> AsyncIterator[T_co]:  # pragma: no cover
         self.log.dev("Using Cython optimized __aiter__")

tests/conftest.py

@@ -1,4 +1,7 @@
+import asyncio
+import gc
 import os
+import platform
 import threading
 import time
 from http import HTTPStatus
@@ -44,6 +47,27 @@ def loop(event_loop):
     return event_loop
 
 
+@pytest.fixture()
+def event_loop(request):
+    """Per-test event loop (overrides pytest-asyncio's default).
+
+    Identical to the pytest-asyncio 0.21 fixture, except that on PyPy we
+    finalize leftover async generators while the loop is still alive.
+    PyPy has no reference counting, so an abandoned async generator is
+    only finalized at a later GC cycle -- by which time this loop is
+    closed and its cleanup lands on a dead loop (leaking tasks into later
+    tests and misattributing warnings).  A gc pass plus
+    ``shutdown_asyncgens()`` inside the owning test keeps the cleanup
+    here, matching CPython's prompt-finalization behavior.
+    """
+    loop = asyncio.get_event_loop_policy().new_event_loop()
+    yield loop
+    if platform.python_implementation() == "PyPy":
+        gc.collect()
+        loop.run_until_complete(loop.shutdown_asyncgens())
+    loop.close()
+
+
 class _patching(object):
     def __init__(self, monkeypatch, request):
         self.monkeypatch = monkeypatch