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..6b16282fb --- /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, object sensor): + """Add sensor.""" + # connect beacons + sensor.beacon = self.app.beacon.new(sensor) + self._sensors.add(sensor) + + def remove(self, object 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..43475b154 --- /dev/null +++ b/faust/transport/_cython/scheduler.pyx @@ -0,0 +1,257 @@ +# 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 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 + + if self.to_remove: + for tp in self.to_remove: + self.buffers.pop(tp, None) + self.to_remove.clear() + 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(): + self.tps.append(tp) + self.iters.append(it) + self.n = PyList_GET_SIZE(self.tps) + 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. + + 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 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 + + 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 self._snapshot_is_current(): + self.ti = 0 + return 0 + 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 + 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.""" + 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/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"} 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..1f92597c6 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,156 @@ 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 + + @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. 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