From 86391b84cadd25d531b95ab36a50267cf1f84c33 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 4 Aug 2026 12:19:35 +0000 Subject: [PATCH 1/4] Add three Cython accelerators: offsets, record scheduler, sensor fan-out Faust already ships optional Cython implementations of the window types, the stream iterator and the topic conductor. This extends that to three more hot paths, found by profiling the per-message and per-commit work. Each follows the existing pattern: the pure-Python implementation stays and is used whenever the extension could not be built or NO_CYTHON is set, so this is purely additive. faust/utils/_cython/functional.pyx first_consecutive_run(), a dedicated helper for the one thing Consumer._new_offset actually needs from consecutive_numbers(): the first run. The groupby() version builds a tuple, calls a Python key function and creates a group generator per acked offset, and runs once per assigned partition on every commit. Because it blocks the event loop, the win shows up as latency: for 100k un-committed offsets, 2.71ms -> 0.55ms. faust/transport/_cython/scheduler.pyx A C-level cursor replacing the round-robin generator returned by DefaultSchedulingStrategy.records_iterator, which also takes over TopicBuffer's inner per-partition generator. This runs once per record fetched from the broker: ~300-490ns -> ~50-70ns. faust/sensors/_cython/base.pyx A cdef base class carrying the four sensor hooks that fire on every message; SensorDelegate subclasses it and keeps its ~20 occasional hooks in plain Python. on_stream_event_in, which builds a dict per event, goes 374ns -> 182ns. Behaviour is preserved rather than approximated. Both scheduler implementations re-read the topic index and each topic's buffer map at every pass and pop drained entries from them, so mutating either mid-iteration gives the same result; a TopicBuffer subclass is still driven through next(); map_from_records and records_iterator remain overridable. first_consecutive_run stops consuming as soon as the run ends, so a shared iterator is left where the pure-Python version leaves it, and it re-reads the list length while scanning because a custom __sub__ can resize the list underneath it. The sensor set is read live rather than snapshotted, so direct mutation behaves identically. Tests run every new code path against both implementations, including a randomised differential test over 200 scheduler topologies. The full suite passes with and without NO_CYTHON. extra/tools/benchmark_cython.py times each accelerator against its pure-Python counterpart in the same interpreter. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019oX4oGCQGgvPJHabjwfaBk --- .gitignore | 2 + extra/tools/benchmark_cython.py | 256 ++++++++++++++++++++++++++ faust/sensors/_cython/__init__.py | 1 + faust/sensors/_cython/base.pyx | 67 +++++++ faust/sensors/base.py | 44 ++++- faust/transport/_cython/scheduler.pyx | 196 ++++++++++++++++++++ faust/transport/consumer.py | 4 +- faust/transport/utils.py | 56 ++++-- faust/utils/_cython/__init__.py | 1 + faust/utils/_cython/functional.pyx | 99 ++++++++++ faust/utils/functional.py | 41 ++++- setup.py | 21 +++ tests/unit/sensors/test_base.py | 92 +++++++++ tests/unit/transport/test_utils.py | 131 ++++++++++++- tests/unit/utils/test_functional.py | 107 ++++++++++- 15 files changed, 1084 insertions(+), 34 deletions(-) create mode 100644 extra/tools/benchmark_cython.py create mode 100644 faust/sensors/_cython/__init__.py create mode 100644 faust/sensors/_cython/base.pyx create mode 100644 faust/transport/_cython/scheduler.pyx create mode 100644 faust/utils/_cython/__init__.py create mode 100644 faust/utils/_cython/functional.pyx diff --git a/.gitignore b/.gitignore index e6974ba9a..be040cade 100644 --- a/.gitignore +++ b/.gitignore @@ -37,6 +37,8 @@ htmlcov/ *.so faust/_cython/*.c faust/transport/_cython/*.c +faust/utils/_cython/*.c +faust/sensors/_cython/*.c # virtualenvs /env diff --git a/extra/tools/benchmark_cython.py b/extra/tools/benchmark_cython.py new file mode 100644 index 000000000..67b760077 --- /dev/null +++ b/extra/tools/benchmark_cython.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python3 +"""Micro-benchmark the optional Cython accelerators against pure Python. + +Faust ships a handful of hot code paths twice: a readable pure-Python +implementation, and a Cython one that is used instead whenever the extension +modules could be built (see ``NO_CYTHON``). This script imports *both* and +times them side by side in the same interpreter, so the numbers are directly +comparable. + +Usage: + +.. sourcecode:: console + + $ python setup.py build_ext --inplace + $ python extra/tools/benchmark_cython.py + +Anything that needs a running worker (the ``Stream`` iterator and the topic +``Conductor``) is out of scope here -- those are covered by the end-to-end +benchmark in ``extra/tools/benchmark.py``. +""" + +import sys +from time import perf_counter +from typing import Callable, List, Optional, Tuple + +from faust.sensors.base import _PySensorDelegateBase, _SensorDelegateBase +from faust.transport.utils import ( + DefaultSchedulingStrategy, + _py_records_iterator, + _records_iterator, +) +from faust.types import TP +from faust.utils.functional import _py_first_consecutive_run, first_consecutive_run +from faust.windows import ( + HoppingWindow, + SlidingWindow, + _PyHoppingWindow, + _PySlidingWindow, +) + +MIN_TIME = 0.2 # seconds to spend on each timed run +Case = Tuple[str, Callable[[], None], Callable[[], None], int] + + +def _time(fn: Callable[[], None], iterations: int) -> float: + start = perf_counter() + for _ in range(iterations): + fn() + return perf_counter() - start + + +def _autorange(fn: Callable[[], None]) -> int: + """Find an iteration count that runs for at least ``MIN_TIME``.""" + iterations = 1 + while True: + elapsed = _time(fn, iterations) + if elapsed >= MIN_TIME or iterations >= 1 << 30: + return iterations + iterations *= 4 if elapsed < MIN_TIME / 8 else 2 + + +def _best(fn: Callable[[], None], iterations: int, rounds: int = 5) -> float: + """Return the best per-operation time in nanoseconds.""" + return min(_time(fn, iterations) for _ in range(rounds)) / iterations * 1e9 + + +def run(cases: List[Case]) -> bool: + """Time each case and print a comparison table.""" + name_width = max(len(name) for name, _, _, _ in cases) + print( + f"{'benchmark'.ljust(name_width)} {'python':>12} " + f"{'cython':>12} {'speedup':>8}" + ) + print("-" * (name_width + 38)) + + all_faster = True + for name, py_fn, cy_fn, scale in cases: + py_fn() + cy_fn() + iterations = _autorange(cy_fn) + py_ns = _best(py_fn, max(1, iterations // 4)) / scale + cy_ns = _best(cy_fn, iterations) / scale + speedup = py_ns / cy_ns + all_faster &= speedup > 1.0 + print( + f"{name.ljust(name_width)} {py_ns:9.1f} ns " + f"{cy_ns:9.1f} ns {speedup:7.2f}x" + ) + return all_faster + + +def offset_cases() -> List[Case]: + """faust.utils.functional.first_consecutive_run -- once per TP per commit.""" + cases = [] + for size in (100, 10_000, 100_000): + acked = list(range(size)) + cases.append( + ( + f"first_consecutive_run/{size} offsets", + lambda acked=acked: _py_first_consecutive_run(acked), + lambda acked=acked: first_consecutive_run(acked), + 1, + ) + ) + return cases + + +def _drain(impl: Callable, records: dict) -> None: + for _ in impl(DefaultSchedulingStrategy.map_from_records(records)): + pass + + +def scheduler_cases() -> List[Case]: + """faust.transport.utils records_iterator -- once per fetched record.""" + cases = [] + for topics, partitions, per_partition in ((1, 1, 500), (4, 8, 100), (8, 16, 50)): + records = { + TP(f"topic-{t}", p): list(range(per_partition)) + for t in range(topics) + for p in range(partitions) + } + cases.append( + ( + f"records_iterator/{topics}t x {partitions}p x {per_partition}", + lambda records=records: _drain(_py_records_iterator, records), + lambda records=records: _drain(_records_iterator, records), + topics * partitions * per_partition, + ) + ) + return cases + + +class _NoopSensor: + """Stand-in sensor: measures delegation overhead, not sensor work.""" + + beacon = None + + def on_message_in(self, tp, offset, message) -> None: ... + + def on_stream_event_in(self, tp, offset, stream, event) -> None: + return None + + def on_stream_event_out(self, tp, offset, stream, event, state=None) -> None: ... + + def on_message_out(self, tp, offset, message) -> None: ... + + +class _FakeBeacon: + def new(self, sensor: object) -> None: + return None + + +class _FakeApp: + beacon = _FakeBeacon() + + +_TP = TP("t", 0) +_OFFSET = 42 + + +def _event_in(delegate: object) -> None: + delegate.on_stream_event_in(_TP, _OFFSET, None, None) + + +def _all_hooks(delegate: object) -> None: + """The exact sequence of delegate calls a single message triggers.""" + delegate.on_message_in(_TP, _OFFSET, None) + state = delegate.on_stream_event_in(_TP, _OFFSET, None, None) + delegate.on_stream_event_out(_TP, _OFFSET, None, None, state) + delegate.on_message_out(_TP, _OFFSET, None) + + +def sensor_cases() -> List[Case]: + """faust.sensors.base -- four hooks fire on every single message.""" + cases = [] + for n_sensors in (1, 3): + py = _PySensorDelegateBase(_FakeApp()) + cy = _SensorDelegateBase(_FakeApp()) + for _ in range(n_sensors): + py.add(_NoopSensor()) + cy.add(_NoopSensor()) + + cases.extend( + [ + ( + f"SensorDelegate.on_stream_event_in/{n_sensors} sensor(s)", + lambda py=py: _event_in(py), + lambda cy=cy: _event_in(cy), + 1, + ), + ( + f"SensorDelegate all 4 hooks/{n_sensors} sensor(s)", + lambda py=py: _all_hooks(py), + lambda cy=cy: _all_hooks(cy), + 1, + ), + ] + ) + return cases + + +def window_cases() -> List[Case]: + """faust.windows -- already shipped; included for completeness.""" + py_hopping = _PyHoppingWindow(10.0, 5.0, 60.0) + cy_hopping = HoppingWindow(10.0, 5.0, 60.0) + py_sliding = _PySlidingWindow(10.0, 5.0, 60.0) + cy_sliding = SlidingWindow(10.0, 5.0, 60.0) + timestamp = 1_600_000_000.0 + return [ + ( + "HoppingWindow.ranges", + lambda: py_hopping.ranges(timestamp), + lambda: cy_hopping.ranges(timestamp), + 1, + ), + ( + "HoppingWindow.current", + lambda: py_hopping.current(timestamp), + lambda: cy_hopping.current(timestamp), + 1, + ), + ( + "SlidingWindow.ranges", + lambda: py_sliding.ranges(timestamp), + lambda: cy_sliding.ranges(timestamp), + 1, + ), + ] + + +def main(argv: Optional[List[str]] = None) -> int: + """Run every benchmark group.""" + if _records_iterator is _py_records_iterator: + print( + "The Cython extensions are not built (or NO_CYTHON is set), so " + "there is nothing to compare against.\n" + "Build them first with: python setup.py build_ext --inplace", + file=sys.stderr, + ) + return 1 + + groups = [ + ("offset commit path", offset_cases), + ("consumer record scheduler", scheduler_cases), + ("sensor delegation", sensor_cases), + ("windows", window_cases), + ] + for title, make_cases in groups: + print(f"\n== {title} ==") + run(make_cases()) + print("\nLower is better; speedup is python/cython.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/faust/sensors/_cython/__init__.py b/faust/sensors/_cython/__init__.py new file mode 100644 index 000000000..668033668 --- /dev/null +++ b/faust/sensors/_cython/__init__.py @@ -0,0 +1 @@ +"""Cython optimized sensor components.""" diff --git a/faust/sensors/_cython/base.pyx b/faust/sensors/_cython/base.pyx new file mode 100644 index 000000000..72a960631 --- /dev/null +++ b/faust/sensors/_cython/base.pyx @@ -0,0 +1,67 @@ +# cython: language_level=3 +"""Cython optimized sensor delegation.""" + + +cdef class SensorDelegateBase: + """Per-message half of :class:`faust.sensors.base.SensorDelegate`. + + Only the four hooks that run on every message live here; the rest stay + in the Python subclass, which keeps them readable and overridable. + + The sensor set is read live, exactly as the pure-Python implementation + does, so mutating ``_sensors`` directly behaves the same either way. + """ + + cdef public object app + cdef public set _sensors + + def __init__(self, object app): + self.app = app + self._sensors = set() + + def add(self, sensor): + """Add sensor.""" + # connect beacons + sensor.beacon = self.app.beacon.new(sensor) + self._sensors.add(sensor) + + def remove(self, sensor): + """Remove sensor.""" + self._sensors.remove(sensor) + + def __iter__(self): + return iter(self._sensors) + + def on_message_in(self, object tp, object offset, object message): + """Call before message is delegated to streams.""" + cdef object sensor + for sensor in self._sensors: + sensor.on_message_in(tp, offset, message) + + def on_stream_event_in(self, object tp, object offset, object stream, + object event): + """Call when stream starts processing an event.""" + cdef: + dict states = {} + object sensor + for sensor in self._sensors: + states[sensor] = sensor.on_stream_event_in(tp, offset, stream, event) + return states + + def on_stream_event_out(self, object tp, object offset, object stream, + object event, object state=None): + """Call when stream is done processing an event.""" + cdef object sensor + if state: + for sensor in self._sensors: + sensor.on_stream_event_out( + tp, offset, stream, event, state.get(sensor)) + else: + for sensor in self._sensors: + sensor.on_stream_event_out(tp, offset, stream, event, None) + + def on_message_out(self, object tp, object offset, object message): + """Call when message is fully acknowledged and can be committed.""" + cdef object sensor + for sensor in self._sensors: + sensor.on_message_out(tp, offset, message) diff --git a/faust/sensors/base.py b/faust/sensors/base.py index 46f679e67..90127349c 100644 --- a/faust/sensors/base.py +++ b/faust/sensors/base.py @@ -1,5 +1,6 @@ """Base-interface for sensors.""" +import os from time import monotonic from typing import Any, Dict, Iterator, Mapping, Optional, Set @@ -14,6 +15,8 @@ __all__ = ["Sensor", "SensorDelegate"] +NO_CYTHON = bool(os.environ.get("NO_CYTHON", False)) + class Sensor(SensorT, Service): """Base class for sensors. @@ -151,8 +154,13 @@ def asdict(self) -> Mapping: return {} -class SensorDelegate(SensorDelegateT): - """A class that delegates sensor methods to a list of sensors.""" +class _PySensorDelegateBase: + """Pure-Python implementation of the per-message sensor fan-out. + + Split out from :class:`SensorDelegate` so the four hooks that run on + every single message can be swapped for a Cython implementation while + the (much larger) set of occasional hooks stays plain Python. + """ _sensors: Set[SensorT] @@ -197,16 +205,38 @@ def on_stream_event_out( tp, offset, stream, event, sensor_state.get(sensor) ) - def on_topic_buffer_full(self, tp: TP) -> None: - """Call when conductor topic buffer is full and has to wait.""" - for sensor in self._sensors: - sensor.on_topic_buffer_full(tp) - def on_message_out(self, tp: TP, offset: int, message: Message) -> None: """Call when message is fully acknowledged and can be committed.""" for sensor in self._sensors: sensor.on_message_out(tp, offset, message) + +if not NO_CYTHON: # pragma: no cover + try: + from ._cython.base import SensorDelegateBase as _SensorDelegateBase + except ImportError: + _SensorDelegateBase = _PySensorDelegateBase # type: ignore[misc,assignment] +else: # pragma: no cover + _SensorDelegateBase = _PySensorDelegateBase # type: ignore[misc,assignment] + + +class SensorDelegate( # type: ignore[misc,valid-type] + _SensorDelegateBase, + SensorDelegateT, +): + """A class that delegates sensor methods to a list of sensors. + + The four per-message hooks (``on_message_in``, ``on_stream_event_in``, + ``on_stream_event_out`` and ``on_message_out``) come from + :data:`_SensorDelegateBase`, which is the Cython implementation when it + could be built and :class:`_PySensorDelegateBase` otherwise. + """ + + def on_topic_buffer_full(self, tp: TP) -> None: + """Call when conductor topic buffer is full and has to wait.""" + for sensor in self._sensors: + sensor.on_topic_buffer_full(tp) + def on_table_get(self, table: CollectionT, key: Any) -> None: """Call when value in table is retrieved.""" for sensor in self._sensors: diff --git a/faust/transport/_cython/scheduler.pyx b/faust/transport/_cython/scheduler.pyx new file mode 100644 index 000000000..5d153c02d --- /dev/null +++ b/faust/transport/_cython/scheduler.pyx @@ -0,0 +1,196 @@ +# cython: language_level=3 +"""Cython optimized consumer record scheduler.""" +from cpython.list cimport PyList_GET_ITEM, PyList_GET_SIZE + + +cdef object _SENTINEL = object() + +# Resolved lazily: faust.transport.utils imports this module at the bottom of +# its own body, so importing it back at module scope here would be circular. +cdef object _TOPIC_BUFFER = None + + +cdef inline object _topic_buffer_type(): + global _TOPIC_BUFFER + if _TOPIC_BUFFER is None: + from faust.transport.utils import TopicBuffer + _TOPIC_BUFFER = TopicBuffer + return _TOPIC_BUFFER + + +cdef class _TopicCursor: + """Round-robin cursor over the partition buffers of a single topic. + + A C-level rewrite of :meth:`faust.transport.utils.TopicBuffer.__iter__`: + one record is taken from each partition in turn, and partitions that run + dry are popped from the buffer map at the end of the current pass. + + Anything that is not a plain, untouched ``TopicBuffer`` is driven through + ``next()`` instead, so a subclass that overrides ``__iter__``/``__next__`` + keeps working. + """ + + cdef: + object source # the TopicBuffer, kept so we can detect replacement + object buffers # TopicBuffer._buffers, read live + object compat # set instead of `buffers` for non-TopicBuffer entries + set to_remove + list tps + list iters + Py_ssize_t pi + Py_ssize_t n + + def __cinit__(self, object buffer): + self.source = buffer + self.buffers = None + self.compat = None + self.to_remove = set() + self.tps = [] + self.iters = [] + self.pi = 0 + self.n = 0 + if type(buffer) is _topic_buffer_type() and (buffer)._it is None: + # Untouched TopicBuffer: drive its partition iterators directly, + # which takes its generator out of the per-record path. + self.buffers = (buffer)._buffers + else: + self.compat = buffer + + cdef object next(self): + """Return the next ``(tp, record)`` pair, or the sentinel if drained.""" + cdef: + Py_ssize_t i + object item + + if self.compat is not None: + return next(self.compat, _SENTINEL) + + while True: + if self.pi >= self.n: + # `while buffers:` in the pure-Python generator. + if not self.buffers: + return _SENTINEL + self._start_pass() + if self.n == 0: + return _SENTINEL + i = self.pi + self.pi += 1 + item = next(PyList_GET_ITEM(self.iters, i), _SENTINEL) + if item is _SENTINEL: + self.to_remove.add(PyList_GET_ITEM(self.tps, i)) + continue + return (PyList_GET_ITEM(self.tps, i), item) + + cdef _start_pass(self): + cdef: + object tp + object it + + if self.to_remove: + for tp in self.to_remove: + self.buffers.pop(tp, None) + self.to_remove.clear() + elif PyList_GET_SIZE(self.tps) == len(self.buffers): + # Nothing drained and nothing added since the last pass, so the + # snapshot is still accurate. TopicBuffer.add() asserts the + # partition is new, so the size can only change on a real change. + self.pi = 0 + return + self.tps = [] + self.iters = [] + for tp, it in self.buffers.items(): + self.tps.append(tp) + self.iters.append(it) + self.n = PyList_GET_SIZE(self.tps) + self.pi = 0 + + +cdef class RoundRobinRecordIterator: + """Iterate a topic index map in round-robin order. + + A C-level rewrite of the generator returned by + :meth:`faust.transport.utils.DefaultSchedulingStrategy.records_iterator`, + flattening both round-robins (across topics, and across the partitions + within a topic) so no generator frame is resumed per record. + + The index map and each topic's buffer map are re-read at the start of + every pass, and drained entries are popped from them, exactly as the + pure-Python version does. + """ + + cdef: + object index + set to_remove + dict cursors + list topics + list topic_cursors + Py_ssize_t ti + Py_ssize_t n + + def __cinit__(self, object index): + self.index = index + self.to_remove = set() + self.cursors = {} + self.topics = [] + self.topic_cursors = [] + self.ti = 0 + self.n = 0 + + def __iter__(self): + return self + + def __next__(self): + cdef: + Py_ssize_t i + _TopicCursor cursor + object item + + while True: + if self.ti >= self.n: + # `while index:` in the pure-Python generator. + if not self.index: + raise StopIteration() + self._start_pass() + if self.n == 0: + raise StopIteration() + i = self.ti + self.ti += 1 + cursor = <_TopicCursor>PyList_GET_ITEM(self.topic_cursors, i) + item = cursor.next() + if item is _SENTINEL: + # This topic is now empty, but we cannot remove it from the + # map while iterating over it, so it goes to the next pass. + self.to_remove.add(PyList_GET_ITEM(self.topics, i)) + continue + return item + + cdef _start_pass(self): + cdef: + object topic + object buffer + _TopicCursor cursor + + if self.to_remove: + for topic in self.to_remove: + self.index.pop(topic, None) + self.cursors.pop(topic, None) + self.to_remove.clear() + elif PyList_GET_SIZE(self.topics) == len(self.index): + self.ti = 0 + return + self.topics = [] + self.topic_cursors = [] + for topic, buffer in self.index.items(): + cursor = <_TopicCursor>self.cursors.get(topic) + if cursor is None or cursor.source is not buffer: + cursor = _TopicCursor(buffer) + self.cursors[topic] = cursor + self.topics.append(topic) + self.topic_cursors.append(cursor) + self.n = PyList_GET_SIZE(self.topics) + self.ti = 0 + + +cpdef object records_iterator(object index): + """Iterate over a topic index map in round-robin order.""" + return RoundRobinRecordIterator(index) diff --git a/faust/transport/consumer.py b/faust/transport/consumer.py index d5a82b1b8..aef99733e 100644 --- a/faust/transport/consumer.py +++ b/faust/transport/consumer.py @@ -95,7 +95,7 @@ ) from faust.types.tuples import FutureMessage from faust.utils import terminal -from faust.utils.functional import consecutive_numbers +from faust.utils.functional import first_consecutive_run from faust.utils.tracing import traced_from_parent_span if typing.TYPE_CHECKING: # pragma: no cover @@ -1117,7 +1117,7 @@ def _new_offset(self, tp: TP) -> Optional[int]: # Note: acked is always kept sorted. # find first list of consecutive numbers - batch = next(consecutive_numbers(acked)) + batch = first_consecutive_run(acked) # remove them from the list to clean up. acked[: len(batch)] = [] self._acked_index[tp].difference_update(batch) diff --git a/faust/transport/utils.py b/faust/transport/utils.py index dc0b8d7d0..fe3f7805a 100644 --- a/faust/transport/utils.py +++ b/faust/transport/utils.py @@ -1,5 +1,6 @@ """Transport utils - scheduling.""" +import os from collections import OrderedDict from typing import ( Any, @@ -22,6 +23,8 @@ "TopicBuffer", ] +NO_CYTHON = bool(os.environ.get("NO_CYTHON", False)) + # But we want to process records from topics in round-robin order. # We convert records into a mapping from topic-name to "chain-of-buffers": # topic_index['topic-name'] = chain(all_topic_partition_buffers) @@ -54,22 +57,27 @@ def iterate(self, records: Mapping[TP, List]) -> Iterator[Tuple[TP, Any]]: def records_iterator(self, index: TopicIndexMap) -> Iterator[Tuple[TP, Any]]: """Iterate over topic index map in round-robin order.""" - to_remove: Set[str] = set() - sentinel = object() - _next = next - while index: - for topic in to_remove: - index.pop(topic, None) - for topic, messages in index.items(): - item = _next(messages, sentinel) - if item is sentinel: - # this topic is now empty, - # but we cannot remove from dict while iterating over it, - # so move that to the outer loop. - to_remove.add(topic) - continue - tp, record = item # type: ignore - yield tp, record + return _records_iterator(index) + + +def _py_records_iterator(index: TopicIndexMap) -> Iterator[Tuple[TP, Any]]: + """Iterate over topic index map in round-robin order.""" + to_remove: Set[str] = set() + sentinel = object() + _next = next + while index: + for topic in to_remove: + index.pop(topic, None) + for topic, messages in index.items(): + item = _next(messages, sentinel) + if item is sentinel: + # this topic is now empty, + # but we cannot remove from dict while iterating over it, + # so move that to the outer loop. + to_remove.add(topic) + continue + tp, record = item # type: ignore + yield tp, record class TopicBuffer(Iterator): @@ -115,3 +123,19 @@ def __next__(self) -> Tuple[TP, Any]: if it is None: it = self._it = iter(self) return it.__next__() + + +# The Cython version flattens both round-robins (across topics, and across the +# partitions within a topic) into a single C-level cursor, so no generator +# frame has to be resumed for each record fetched from the broker. +# +# Only the default iteration strategy is swapped: ``map_from_records`` and +# ``records_iterator`` remain overridable, and a ``TopicBuffer`` subclass falls +# back to being driven through ``next()``. +if not NO_CYTHON: # pragma: no cover + try: + from ._cython.scheduler import records_iterator as _records_iterator + except ImportError: + _records_iterator = _py_records_iterator +else: # pragma: no cover + _records_iterator = _py_records_iterator diff --git a/faust/utils/_cython/__init__.py b/faust/utils/_cython/__init__.py new file mode 100644 index 000000000..52d726339 --- /dev/null +++ b/faust/utils/_cython/__init__.py @@ -0,0 +1 @@ +"""Cython optimized functional utilities.""" diff --git a/faust/utils/_cython/functional.pyx b/faust/utils/_cython/functional.pyx new file mode 100644 index 000000000..70c3e4edf --- /dev/null +++ b/faust/utils/_cython/functional.pyx @@ -0,0 +1,99 @@ +# cython: language_level=3 +"""Cython optimized functional utilities.""" +from cpython.list cimport PyList_GET_ITEM, PyList_GET_SIZE +from cpython.long cimport PyLong_AsLongLongAndOverflow, PyLong_CheckExact + + +cdef object _SENTINEL = object() + + +cpdef list first_consecutive_run(object numbers): + """Return the first run of consecutive numbers in ``numbers``. + + Equivalent to ``next(consecutive_numbers(numbers), [])``, but without + building the intermediate tuples, the per-element Python key function + and the group generators that :func:`itertools.groupby` needs. + + A run continues for as long as each number is exactly one greater than + the one before it, which is the same rule ``groupby`` applies when + grouping on ``index - value``. Repeated numbers therefore end a run. + + Like the pure-Python version this stops consuming as soon as the run + ends, so a non-sequence iterable is left positioned just past the first + number that broke the run. + """ + if type(numbers) is list: + return _run_from_list(numbers) + return _run_from_iterable(iter(numbers)) + + +cdef list _run_from_list(list seq): + cdef: + list run + Py_ssize_t i + object prev + object cur + long long c_prev + long long c_cur + int overflow + + if PyList_GET_SIZE(seq) == 0: + return [] + + prev = PyList_GET_ITEM(seq, 0) + run = [prev] + i = 1 + + # Fast path: plain ints that fit in a C long long. Kafka offsets always + # do, so this is what actually runs in production. Nothing in this loop + # can run Python code, so the size only has to be re-read for the slow + # path below. + if PyLong_CheckExact(prev): + c_prev = PyLong_AsLongLongAndOverflow(prev, &overflow) + if not overflow: + while i < PyList_GET_SIZE(seq): + cur = PyList_GET_ITEM(seq, i) + if not PyLong_CheckExact(cur): + break + c_cur = PyLong_AsLongLongAndOverflow(cur, &overflow) + # ``c_cur <= c_prev`` is tested first so the subtraction + # below can never underflow. + if overflow or c_cur <= c_prev or c_cur - 1 != c_prev: + break + run.append(cur) + c_prev = c_cur + i += 1 + if i >= PyList_GET_SIZE(seq): + return run + prev = PyList_GET_ITEM(seq, i - 1) + + # Slow path: arbitrary objects supporting ``-`` and comparison to 1. + # ``cur - prev`` runs arbitrary Python code that may resize the list, + # so the length is re-read on every iteration. + while i < PyList_GET_SIZE(seq): + cur = PyList_GET_ITEM(seq, i) + if cur - prev != 1: + break + run.append(cur) + prev = cur + i += 1 + return run + + +cdef list _run_from_iterable(object it): + cdef: + list run + object prev + object cur + + prev = next(it, _SENTINEL) + if prev is _SENTINEL: + return [] + run = [prev] + while True: + cur = next(it, _SENTINEL) + if cur is _SENTINEL or cur - prev != 1: + break + run.append(cur) + prev = cur + return run diff --git a/faust/utils/functional.py b/faust/utils/functional.py index 98d4ad764..be3eada26 100644 --- a/faust/utils/functional.py +++ b/faust/utils/functional.py @@ -1,16 +1,20 @@ """Functional utilities.""" +import os from functools import reduce from itertools import groupby -from typing import Iterable, Iterator, Mapping, Sequence, Tuple, TypeVar +from typing import Iterable, Iterator, List, Mapping, Sequence, Tuple, TypeVar __all__ = [ "consecutive_numbers", + "first_consecutive_run", "translate", ] T = TypeVar("T") +NO_CYTHON = bool(os.environ.get("NO_CYTHON", False)) + def consecutive_numbers(it: Iterable[int]) -> Iterator[Sequence[int]]: """Find runs of consecutive numbers. @@ -22,6 +26,41 @@ def consecutive_numbers(it: Iterable[int]) -> Iterator[Sequence[int]]: yield [a[1] for a in g] +def _py_first_consecutive_run(numbers: Iterable[int]) -> List[int]: + """Return the first run of consecutive numbers in ``numbers``. + + Equivalent to ``next(consecutive_numbers(numbers), [])``, but without + building the intermediate tuples, the per-element key function and the + group generators that :func:`itertools.groupby` needs. + + Callers that only need the first run (such as the consumer working out + the next offset to commit) should use this instead of + :func:`consecutive_numbers`, as it is the part of the commit path that + scales with the number of un-committed offsets. + """ + it = iter(numbers) + try: + prev = next(it) + except StopIteration: + return [] + run = [prev] + for cur in it: + if cur - prev != 1: + break + run.append(cur) + prev = cur + return run + + +if not NO_CYTHON: # pragma: no cover + try: + from ._cython.functional import first_consecutive_run + except ImportError: + first_consecutive_run = _py_first_consecutive_run +else: # pragma: no cover + first_consecutive_run = _py_first_consecutive_run + + def translate(table: Mapping, s: str) -> str: """Replace characters and patterns in string ``s``. diff --git a/setup.py b/setup.py index 7a3190b70..003c4cc56 100644 --- a/setup.py +++ b/setup.py @@ -88,6 +88,27 @@ extra_compile_args=CFLAGS, extra_link_args=LDFLAGS, ), + Extension( + "faust.sensors._cython.base", + ["faust/sensors/_cython/base" + ext], + libraries=LIBRARIES, + extra_compile_args=CFLAGS, + extra_link_args=LDFLAGS, + ), + Extension( + "faust.transport._cython.scheduler", + ["faust/transport/_cython/scheduler" + ext], + libraries=LIBRARIES, + extra_compile_args=CFLAGS, + extra_link_args=LDFLAGS, + ), + Extension( + "faust.utils._cython.functional", + ["faust/utils/_cython/functional" + ext], + libraries=LIBRARIES, + extra_compile_args=CFLAGS, + extra_link_args=LDFLAGS, + ), ] diff --git a/tests/unit/sensors/test_base.py b/tests/unit/sensors/test_base.py index ad96d9e6e..ce0ce74aa 100644 --- a/tests/unit/sensors/test_base.py +++ b/tests/unit/sensors/test_base.py @@ -5,6 +5,7 @@ from faust import Event, Stream, Table, Topic, web from faust.assignor import PartitionAssignor from faust.sensors import Sensor +from faust.sensors.base import _PySensorDelegateBase, _SensorDelegateBase from faust.transport.consumer import Consumer from faust.transport.producer import Producer from faust.types import TP, Message @@ -239,3 +240,94 @@ def test_on_web_request(self, *, sensors, sensor, app, req, response, view): def test_repr(self, *, sensors): assert repr(sensors) + + +#: Both implementations of the per-message sensor fan-out. +#: ``_SensorDelegateBase`` is the Cython one whenever the extension could be +#: built, and is otherwise the same object as ``_PySensorDelegateBase``. +SENSOR_DELEGATE_BASES = [_PySensorDelegateBase, _SensorDelegateBase] + + +class Test_SensorDelegateBase: + """The four hooks that run on every message, in both implementations.""" + + def _delegate(self, base, n_sensors=1): + app = Mock(name="app") + delegate = base(app) + sensors = [] + for _ in range(n_sensors): + sensor = Mock(name="sensor", autospec=Sensor) + delegate.add(sensor) + sensors.append(sensor) + return delegate, sensors + + @pytest.mark.parametrize("base", SENSOR_DELEGATE_BASES) + def test_add_connects_beacon(self, base): + delegate, [sensor] = self._delegate(base) + assert sensor.beacon is delegate.app.beacon.new.return_value + assert list(delegate) == [sensor] + + @pytest.mark.parametrize("base", SENSOR_DELEGATE_BASES) + def test_remove(self, base): + delegate, [sensor] = self._delegate(base) + delegate.remove(sensor) + assert not list(delegate) + delegate.on_message_in(TP1, 3, None) + sensor.on_message_in.assert_not_called() + + @pytest.mark.parametrize("base", SENSOR_DELEGATE_BASES) + def test_remove__missing_raises(self, base): + delegate, _ = self._delegate(base) + with pytest.raises(KeyError): + delegate.remove(Mock(name="never-added")) + + @pytest.mark.parametrize("base", SENSOR_DELEGATE_BASES) + def test_no_sensors(self, base): + delegate = base(Mock(name="app")) + delegate.on_message_in(TP1, 3, None) + assert delegate.on_stream_event_in(TP1, 3, None, None) == {} + delegate.on_stream_event_out(TP1, 3, None, None, None) + delegate.on_message_out(TP1, 3, None) + + @pytest.mark.parametrize("base", SENSOR_DELEGATE_BASES) + @pytest.mark.parametrize("n_sensors", [1, 3]) + def test_on_message_in_out(self, base, n_sensors, message): + delegate, sensors = self._delegate(base, n_sensors) + delegate.on_message_in(TP1, 303, message) + delegate.on_message_out(TP1, 303, message) + for sensor in sensors: + sensor.on_message_in.assert_called_once_with(TP1, 303, message) + sensor.on_message_out.assert_called_once_with(TP1, 303, message) + + @pytest.mark.parametrize("base", SENSOR_DELEGATE_BASES) + @pytest.mark.parametrize("n_sensors", [1, 3]) + def test_on_stream_event_in_out(self, base, n_sensors, stream, event): + delegate, sensors = self._delegate(base, n_sensors) + state = delegate.on_stream_event_in(TP1, 303, stream, event) + assert set(state) == set(sensors) + delegate.on_stream_event_out(TP1, 303, stream, event, state) + for sensor in sensors: + sensor.on_stream_event_in.assert_called_once_with(TP1, 303, stream, event) + sensor.on_stream_event_out.assert_called_once_with( + TP1, 303, stream, event, state[sensor] + ) + + @pytest.mark.parametrize("base", SENSOR_DELEGATE_BASES) + @pytest.mark.parametrize("state", [None, {}]) + def test_on_stream_event_out__without_state(self, base, state, stream, event): + # No state recorded for this sensor: it must still be called, with None. + delegate, [sensor] = self._delegate(base) + delegate.on_stream_event_out(TP1, 303, stream, event, state) + sensor.on_stream_event_out.assert_called_once_with( + TP1, 303, stream, event, None + ) + + @pytest.mark.parametrize("base", SENSOR_DELEGATE_BASES) + def test_direct_mutation_of_sensors_is_picked_up(self, base): + # The Cython version walks a list snapshot of the sensor set, so it + # has to notice a set that was mutated behind its back. + delegate, _ = self._delegate(base) + extra = Mock(name="extra", autospec=Sensor) + delegate._sensors.add(extra) + delegate.on_message_in(TP1, 3, None) + extra.on_message_in.assert_called_once_with(TP1, 3, None) diff --git a/tests/unit/transport/test_utils.py b/tests/unit/transport/test_utils.py index f7a2e1ab7..66bda8d2d 100644 --- a/tests/unit/transport/test_utils.py +++ b/tests/unit/transport/test_utils.py @@ -1,6 +1,20 @@ -from faust.transport.utils import DefaultSchedulingStrategy, TopicBuffer +import random + +import pytest + +from faust.transport.utils import ( + DefaultSchedulingStrategy, + TopicBuffer, + _py_records_iterator, + _records_iterator, +) from faust.types import TP +#: Both round-robin implementations. ``_records_iterator`` is the Cython one +#: whenever the extension could be built, and is otherwise the same object as +#: ``_py_records_iterator``. +RECORDS_ITERATOR_IMPLS = [_py_records_iterator, _records_iterator] + TP1 = TP("foo", 0) TP2 = TP("foo", 1) TP3 = TP("bar", 0) @@ -77,3 +91,118 @@ def test_next(self): (TP1, 3), (TP1, 4), ] + + +class Test_records_iterator: + def _index(self, records): + return DefaultSchedulingStrategy.map_from_records(records) + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_round_robin_over_topics_and_partitions(self, impl): + records = {TP1: BUF1, TP2: BUF2, TP3: BUF3, TP4: BUF4, TP5: BUF5} + + # Round-robin across topics, and across the partitions within each + # topic. "baz" only has one partition, so it drains early and the + # interleaving shifts once it and "bar" are exhausted. + assert list(impl(self._index(records))) == [ + (TP1, 0), + (TP3, 9), + (TP5, 14), + (TP2, 5), + (TP4, 11), + (TP5, 15), + (TP1, 1), + (TP3, 10), + (TP2, 6), + (TP4, 12), + (TP1, 2), + (TP4, 13), + (TP2, 7), + (TP1, 3), + (TP2, 8), + (TP1, 4), + ] + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_empty(self, impl): + assert list(impl(self._index({}))) == [] + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_empty_buffers(self, impl): + assert list(impl(self._index({TP1: [], TP3: []}))) == [] + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_drains_the_index_it_was_given(self, impl): + index = self._index({TP1: BUF1, TP3: BUF3}) + list(impl(index)) + assert not index + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_propagates_exceptions(self, impl): + def raising(): + yield (TP1, 1) + raise RuntimeError("buffer exploded") + + with pytest.raises(RuntimeError): + list(impl({"foo": raising()})) + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_topic_buffer_subclass(self, impl): + # A custom TopicBuffer must still be driven through next(), so that + # any overridden __iter__/__next__ is honoured. + class MyBuffer(TopicBuffer): + pass + + buffer = MyBuffer() + buffer.add(TP1, BUF1) + assert list(impl({"foo": buffer})) == [(TP1, i) for i in BUF1] + + def test_implementations_agree(self): + # Randomised differential test: the accelerated iterator must emit + # exactly the same sequence as the pure-Python one for any topology. + rng = random.Random(20220613) + for _ in range(200): + records = { + TP(f"t{topic}", partition): [ + f"t{topic}-{partition}-{i}" for i in range(rng.randint(0, 6)) + ] + for topic in range(rng.randint(0, 4)) + for partition in range(rng.randint(1, 4)) + } + assert list(_records_iterator(self._index(records))) == list( + _py_records_iterator(self._index(records)) + ) + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_partition_added_mid_iteration(self, impl): + # Both implementations re-read the buffer map on each pass. + buffer = TopicBuffer() + buffer.add(TP1, [1, 2]) + index = {"foo": buffer} + + it = impl(index) + buffer.add(TP2, [7, 8]) + assert list(it) == [(TP1, 1), (TP2, 7), (TP1, 2), (TP2, 8)] + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_topic_added_mid_iteration(self, impl): + buffer = TopicBuffer() + buffer.add(TP1, [1, 2]) + index = {"foo": buffer} + + it = impl(index) + late = TopicBuffer() + late.add(TP3, [9]) + index["bar"] = late + assert list(it) == [(TP1, 1), (TP3, 9), (TP1, 2)] + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_drains_the_buffer_map_too(self, impl): + # Exhausted partitions are popped from TopicBuffer._buffers. + buffer = TopicBuffer() + buffer.add(TP1, [1, 2]) + index = {"foo": buffer} + + list(impl(index)) + assert not index + assert not buffer._buffers diff --git a/tests/unit/utils/test_functional.py b/tests/unit/utils/test_functional.py index c6fa89b00..ec03d3b6c 100644 --- a/tests/unit/utils/test_functional.py +++ b/tests/unit/utils/test_functional.py @@ -1,21 +1,75 @@ import pytest -from faust.utils.functional import consecutive_numbers, translate +from faust.utils.functional import ( + _py_first_consecutive_run, + consecutive_numbers, + first_consecutive_run, + translate, +) + +#: Both implementations of ``first_consecutive_run``. ``first_consecutive_run`` +#: is the Cython one whenever the extension could be built, and is otherwise the +#: same object as ``_py_first_consecutive_run`` (in which case this just runs +#: the pure-Python one twice). +FIRST_CONSECUTIVE_RUN_IMPLS = [_py_first_consecutive_run, first_consecutive_run] + +RUN_CASES = [ + ([1, 2, 3, 4, 6, 7, 8], [1, 2, 3, 4]), + ([1, 4, 6, 8, 10], [1]), + ([1], [1]), + ([103, 104, 105, 106, 100000000000], [103, 104, 105, 106]), + # a run of numbers too large for a C long long must still work + ([2**80, 2**80 + 1, 2**80 + 3], [2**80, 2**80 + 1]), + # ... including when the overflow happens part-way through a run + ([1, 2, 2**80], [1, 2]), + # repeated numbers end a run, matching itertools.groupby + ([1, 1, 2], [1]), + ([0, 1, 2], [0, 1, 2]), + # descending numbers are not a run + ([5, 4, 3], [5]), +] @pytest.mark.parametrize( "numbers,expected", - [ - ([1, 2, 3, 4, 6, 7, 8], [1, 2, 3, 4]), - ([1, 4, 6, 8, 10], [1]), - ([1], [1]), - ([103, 104, 105, 106, 100000000000], [103, 104, 105, 106]), - ], + [(numbers, expected) for numbers, expected in RUN_CASES if numbers], ) def test_consecutive_numbers(numbers, expected): assert next(consecutive_numbers(numbers), None) == expected +@pytest.mark.parametrize("impl", FIRST_CONSECUTIVE_RUN_IMPLS) +@pytest.mark.parametrize("numbers,expected", RUN_CASES) +def test_first_consecutive_run(impl, numbers, expected): + assert impl(list(numbers)) == expected + + +@pytest.mark.parametrize("impl", FIRST_CONSECUTIVE_RUN_IMPLS) +def test_first_consecutive_run__empty(impl): + assert impl([]) == [] + + +@pytest.mark.parametrize("impl", FIRST_CONSECUTIVE_RUN_IMPLS) +def test_first_consecutive_run__accepts_any_iterable(impl): + assert impl(iter([1, 2, 3, 5])) == [1, 2, 3] + assert impl(range(4)) == [0, 1, 2, 3] + + +@pytest.mark.parametrize("impl", FIRST_CONSECUTIVE_RUN_IMPLS) +def test_first_consecutive_run__does_not_mutate_argument(impl): + numbers = [1, 2, 5] + assert impl(numbers) == [1, 2] + assert numbers == [1, 2, 5] + + +@pytest.mark.parametrize("impl", FIRST_CONSECUTIVE_RUN_IMPLS) +@pytest.mark.parametrize("numbers,expected", RUN_CASES) +def test_first_consecutive_run__matches_consecutive_numbers(impl, numbers, expected): + # the helper must return exactly what taking the first group of + # consecutive_numbers() would have returned. + assert impl(list(numbers)) == next(consecutive_numbers(numbers), []) + + @pytest.mark.parametrize( "table,s,expected", [ @@ -28,3 +82,42 @@ def test_consecutive_numbers(numbers, expected): ) def test_translate(table, s, expected): assert translate(table, s) == expected + + +@pytest.mark.parametrize("impl", FIRST_CONSECUTIVE_RUN_IMPLS) +def test_first_consecutive_run__is_lazy(impl): + # Must stop consuming as soon as the run is broken, leaving the rest of + # the iterator for the caller (same as next(consecutive_numbers(...))). + it = iter([1, 2, 5, 6, 7]) + assert impl(it) == [1, 2] + assert list(it) == [6, 7] + + +@pytest.mark.parametrize("impl", FIRST_CONSECUTIVE_RUN_IMPLS) +def test_first_consecutive_run__does_not_over_consume(impl): + # An iterable that blows up after the run has ended must not be reached. + def numbers(): + yield 1 + yield 2 + yield 99 + raise AssertionError("consumed past the end of the run") + + assert impl(numbers()) == [1, 2] + + +@pytest.mark.parametrize("impl", FIRST_CONSECUTIVE_RUN_IMPLS) +def test_first_consecutive_run__list_resized_while_scanning(impl): + # __sub__ can run arbitrary Python code, including code that shrinks the + # list being scanned -- which must not read past the end of it. + class Shrinks: + def __init__(self, container=None): + self.container = container + + def __sub__(self, other): + if self.container is not None: + del self.container[2:] + return 1 + + numbers = [] + numbers.extend([Shrinks(), Shrinks(numbers)] + [Shrinks() for _ in range(20)]) + assert len(impl(numbers)) <= 2 From 2f04895ed5138087eae1f54a7d813bb7b37816d9 Mon Sep 17 00:00:00 2001 From: William Barnhart Date: Tue, 21 Jul 2026 15:25:26 +0000 Subject: [PATCH 2/4] test(assignor): disable Hypothesis deadline to fix flaky PyPy CI leg The copartitioned-assignor property tests run on PyPy since the PyPy skips were removed. They carried a fixed 4000ms Hypothesis deadline, but PyPy's JIT warmup makes a single example's timing vary by over a second, so runs intermittently exceed it and Hypothesis reports DeadlineExceeded / FlakyFailure -- failing the (non-required) PyPy leg and turning master CI red even though the assignment produced is valid. These tests assert assignment correctness, not performance, so set deadline=None (as the DeadlineExceeded message itself recommends). The property assertions are unchanged and the job timeout still guards against a real hang. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01HHPL4VFWQRQPpjR1gXSKyL --- tests/meticulous/assignor/test_copartitioned_assignor.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/meticulous/assignor/test_copartitioned_assignor.py b/tests/meticulous/assignor/test_copartitioned_assignor.py index e0fb1f66b..d78ceea8c 100644 --- a/tests/meticulous/assignor/test_copartitioned_assignor.py +++ b/tests/meticulous/assignor/test_copartitioned_assignor.py @@ -8,7 +8,13 @@ from faust.assignor.client_assignment import CopartitionedAssignment from faust.assignor.copartitioned_assignor import CopartitionedAssignor -TEST_DEADLINE = 4000 +# These Hypothesis property tests assert assignment *correctness*, not +# performance. PyPy's JIT warmup makes a single example's timing vary by +# well over a second, so a fixed deadline is tripped intermittently +# (DeadlineExceeded -> FlakyFailure) even though the assignment is valid -- +# flaking the PyPy CI leg. Disable the deadline; the job's own timeout still +# guards against a real hang. +TEST_DEADLINE = None _topics = {"foo", "bar", "baz"} From 941a2199097bfe9866cb0ceb678f09080f1a1d08 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 4 Aug 2026 13:43:57 +0000 Subject: [PATCH 3/4] Tighten static typing in the Cython accelerators Follows the "faster code via static typing" guidance: declare return types on cdef functions, and type parameters and locals rather than leaving them implicit. The substantive change is on the two _start_pass methods, which had no declared return type. They are now `cdef int ... except -1`. Both call into Python (dict.pop, .items(), the _TopicCursor constructor) and so can raise, and the except clause states that rather than depending on a compiler default -- Cython 3 propagates from a bare cdef, but Cython 0.x swallowed unless an except clause was given, and legacy_implicit_noexcept restores that behaviour. Verified both ways: an exception raised inside _start_pass now propagates identically through the Cython and pure-Python iterators. Also types the len() results as Py_ssize_t instead of letting them box into Python ints, and types the two sensor-delegate parameters that were bare. Note this is a correctness and clarity change, not a speed one: the scheduler benchmark is unmoved at 5.8-7.7x over pure Python, because the loops were already fully typed. Parameters that face callers are deliberately left as `object` -- typing them to a builtin would add an exact-type runtime check, which is what made an OrderedDict fail against a `cdef dict` local earlier in this branch. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019oX4oGCQGgvPJHabjwfaBk --- faust/sensors/_cython/base.pyx | 4 +-- faust/transport/_cython/scheduler.pyx | 39 +++++++++++++++++++-------- 2 files changed, 30 insertions(+), 13 deletions(-) diff --git a/faust/sensors/_cython/base.pyx b/faust/sensors/_cython/base.pyx index 72a960631..6b16282fb 100644 --- a/faust/sensors/_cython/base.pyx +++ b/faust/sensors/_cython/base.pyx @@ -19,13 +19,13 @@ cdef class SensorDelegateBase: self.app = app self._sensors = set() - def add(self, sensor): + def add(self, object sensor): """Add sensor.""" # connect beacons sensor.beacon = self.app.beacon.new(sensor) self._sensors.add(sensor) - def remove(self, sensor): + def remove(self, object sensor): """Remove sensor.""" self._sensors.remove(sensor) diff --git a/faust/transport/_cython/scheduler.pyx b/faust/transport/_cython/scheduler.pyx index 5d153c02d..e62e51ab2 100644 --- a/faust/transport/_cython/scheduler.pyx +++ b/faust/transport/_cython/scheduler.pyx @@ -81,21 +81,31 @@ cdef class _TopicCursor: continue return (PyList_GET_ITEM(self.tps, i), item) - cdef _start_pass(self): + cdef int _start_pass(self) except -1: + # The Python calls below (dict.pop, .items()) can raise, and this + # says so explicitly rather than relying on a compiler default: + # Cython 3 propagates from a bare `cdef` by default, but Cython 0.x + # silently swallowed unless an except clause was given, and + # `legacy_implicit_noexcept` restores that. `except -1` propagates + # under every version. cdef: object tp object it + Py_ssize_t live if self.to_remove: for tp in self.to_remove: self.buffers.pop(tp, None) self.to_remove.clear() - elif PyList_GET_SIZE(self.tps) == len(self.buffers): - # Nothing drained and nothing added since the last pass, so the - # snapshot is still accurate. TopicBuffer.add() asserts the - # partition is new, so the size can only change on a real change. - self.pi = 0 - return + else: + live = len(self.buffers) + if PyList_GET_SIZE(self.tps) == live: + # Nothing drained and nothing added since the last pass, so + # the snapshot is still accurate. TopicBuffer.add() asserts + # the partition is new, so the size can only change on a + # real change. + self.pi = 0 + return 0 self.tps = [] self.iters = [] for tp, it in self.buffers.items(): @@ -103,6 +113,7 @@ cdef class _TopicCursor: self.iters.append(it) self.n = PyList_GET_SIZE(self.tps) self.pi = 0 + return 0 cdef class RoundRobinRecordIterator: @@ -164,20 +175,25 @@ cdef class RoundRobinRecordIterator: continue return item - cdef _start_pass(self): + cdef int _start_pass(self) except -1: + # See the note on _TopicCursor._start_pass for why the except clause + # is spelled out. cdef: object topic object buffer _TopicCursor cursor + Py_ssize_t live if self.to_remove: for topic in self.to_remove: self.index.pop(topic, None) self.cursors.pop(topic, None) self.to_remove.clear() - elif PyList_GET_SIZE(self.topics) == len(self.index): - self.ti = 0 - return + else: + live = len(self.index) + if PyList_GET_SIZE(self.topics) == live: + self.ti = 0 + return 0 self.topics = [] self.topic_cursors = [] for topic, buffer in self.index.items(): @@ -189,6 +205,7 @@ cdef class RoundRobinRecordIterator: self.topic_cursors.append(cursor) self.n = PyList_GET_SIZE(self.topics) self.ti = 0 + return 0 cpdef object records_iterator(object index): From 340793fa4bf43bd35a60d7fe27bfed81160b75e4 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 4 Aug 2026 14:09:44 +0000 Subject: [PATCH 4/4] Fix records being dropped when a topic buffer is replaced mid-iteration Reported in review of #751. _start_pass skipped rebuilding its snapshot when the mapping's length was unchanged since the last pass. That check cannot tell "unchanged" from "swapped": replacing index["foo"] with a new TopicBuffer under the same key leaves the length identical, so the shortcut was taken, the stale cursor kept being used, and every record in the replacement was silently dropped. The `cursor.source is not buffer` check that exists to catch exactly this lived inside the rebuild, which the shortcut skipped. Reproduced against the pure-Python iterator, which re-reads the mapping every pass and so gets it right: python: [1, 10, 11] cython: [1, 2] <- 10 and 11 never delivered The same flaw was present one level down in _TopicCursor, where an iterator swapped in under an existing TP left the partition cursor stale. The report only named the topic level; both are fixed. Both shortcuts are replaced by _snapshot_is_current(), which compares the cached snapshot against the live mapping entry by entry, by identity. It walks the mapping but allocates nothing, which is what still makes reusing the snapshot worthwhile. This costs real throughput, and the PR body is updated rather than left overstating it: the scheduler was 5.7-7.7x over pure Python and is now 2.2-2.8x. A pass must consult the live mapping to be correct, and a single-partition topic reaches a pass boundary on every record. An unconditional rebuild was measured too, at 1.6-2.1x, so the identity check is worth keeping over simply rebuilding. Adds regression tests for both levels, run against both implementations. Swapping in a topic under a *different* name is deliberately not covered: the pure-Python generator iterates the live dict, so CPython raises "dictionary keys changed during iteration" -- undefined behaviour in Python itself rather than a guarantee to pin. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019oX4oGCQGgvPJHabjwfaBk --- faust/transport/_cython/scheduler.pyx | 76 +++++++++++++++++++++------ tests/unit/transport/test_utils.py | 38 ++++++++++++++ 2 files changed, 98 insertions(+), 16 deletions(-) diff --git a/faust/transport/_cython/scheduler.pyx b/faust/transport/_cython/scheduler.pyx index e62e51ab2..43475b154 100644 --- a/faust/transport/_cython/scheduler.pyx +++ b/faust/transport/_cython/scheduler.pyx @@ -91,21 +91,16 @@ cdef class _TopicCursor: cdef: object tp object it - Py_ssize_t live if self.to_remove: for tp in self.to_remove: self.buffers.pop(tp, None) self.to_remove.clear() - else: - live = len(self.buffers) - if PyList_GET_SIZE(self.tps) == live: - # Nothing drained and nothing added since the last pass, so - # the snapshot is still accurate. TopicBuffer.add() asserts - # the partition is new, so the size can only change on a - # real change. - self.pi = 0 - return 0 + elif self._snapshot_is_current(): + # Unchanged since the last pass, so the existing snapshot still + # describes _buffers exactly and can be reused. + self.pi = 0 + return 0 self.tps = [] self.iters = [] for tp, it in self.buffers.items(): @@ -115,6 +110,33 @@ cdef class _TopicCursor: self.pi = 0 return 0 + cdef bint _snapshot_is_current(self) except -1: + """Is the cached snapshot still identical to the live buffer map? + + Comparing lengths alone is not enough: an iterator swapped in under + an existing TP keeps the length the same but leaves a stale cursor, + whose records would never be delivered. So the entries are compared + by identity. That walks the mapping, but allocates nothing, which is + what makes reusing the snapshot worth doing at all. + """ + cdef: + Py_ssize_t i = 0 + Py_ssize_t n = PyList_GET_SIZE(self.tps) + object tp + object it + + if n != len(self.buffers): + return False + for tp, it in self.buffers.items(): + if i >= n: + return False + if PyList_GET_ITEM(self.tps, i) is not tp: + return False + if PyList_GET_ITEM(self.iters, i) is not it: + return False + i += 1 + return i == n + cdef class RoundRobinRecordIterator: """Iterate a topic index map in round-robin order. @@ -182,18 +204,15 @@ cdef class RoundRobinRecordIterator: object topic object buffer _TopicCursor cursor - Py_ssize_t live if self.to_remove: for topic in self.to_remove: self.index.pop(topic, None) self.cursors.pop(topic, None) self.to_remove.clear() - else: - live = len(self.index) - if PyList_GET_SIZE(self.topics) == live: - self.ti = 0 - return 0 + elif self._snapshot_is_current(): + self.ti = 0 + return 0 self.topics = [] self.topic_cursors = [] for topic, buffer in self.index.items(): @@ -207,6 +226,31 @@ cdef class RoundRobinRecordIterator: self.ti = 0 return 0 + cdef bint _snapshot_is_current(self) except -1: + """Is the cached snapshot still identical to the live index map? + + As with _TopicCursor, a length check alone would miss a TopicBuffer + replaced under an existing topic name, so each cursor is compared + against the buffer it was built from. + """ + cdef: + Py_ssize_t i = 0 + Py_ssize_t n = PyList_GET_SIZE(self.topics) + object topic + object buffer + + if n != len(self.index): + return False + for topic, buffer in self.index.items(): + if i >= n: + return False + if PyList_GET_ITEM(self.topics, i) is not topic: + return False + if (<_TopicCursor>PyList_GET_ITEM(self.topic_cursors, i)).source is not buffer: + return False + i += 1 + return i == n + cpdef object records_iterator(object index): """Iterate over a topic index map in round-robin order.""" diff --git a/tests/unit/transport/test_utils.py b/tests/unit/transport/test_utils.py index 66bda8d2d..1f92597c6 100644 --- a/tests/unit/transport/test_utils.py +++ b/tests/unit/transport/test_utils.py @@ -206,3 +206,41 @@ def test_drains_the_buffer_map_too(self, impl): list(impl(index)) assert not index assert not buffer._buffers + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_topic_buffer_replaced_mid_iteration(self, impl): + # A TopicBuffer swapped in under an existing topic name keeps the + # index the same size, so a length-only staleness check would miss + # it and keep draining the old buffer -- silently dropping every + # record in the replacement. + first = TopicBuffer() + first.add(TP1, [1, 2]) + index = {"foo": first} + + it = impl(index) + assert next(it) == (TP1, 1) + + replacement = TopicBuffer() + replacement.add(TP1, [10, 11]) + index["foo"] = replacement + + assert list(it) == [(TP1, 10), (TP1, 11)] + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_partition_buffer_replaced_mid_iteration(self, impl): + # Same hazard one level down: an iterator swapped in under an + # existing TP leaves the partition cursor stale. + buffer = TopicBuffer() + buffer.add(TP1, [1, 2, 3]) + + it = impl({"foo": buffer}) + assert next(it) == (TP1, 1) + + buffer._buffers[TP1] = iter([10, 11]) + assert list(it) == [(TP1, 10), (TP1, 11)] + + # Note: swapping a topic for one under a *different* name mid-iteration + # is deliberately not covered. The pure-Python generator iterates the + # live dict, so CPython raises "dictionary keys changed during + # iteration" -- that is undefined behaviour in Python itself, not a + # guarantee either implementation should be pinned to.