diff --git a/faust/transport/_cython/scheduler.pyx b/faust/transport/_cython/scheduler.pyx index 43475b154..071f81646 100644 --- a/faust/transport/_cython/scheduler.pyx +++ b/faust/transport/_cython/scheduler.pyx @@ -1,6 +1,8 @@ # cython: language_level=3 """Cython optimized consumer record scheduler.""" +from cpython.dict cimport PyDict_CheckExact, PyDict_Next from cpython.list cimport PyList_GET_ITEM, PyList_GET_SIZE +from cpython.ref cimport PyObject cdef object _SENTINEL = object() @@ -122,11 +124,28 @@ cdef class _TopicCursor: cdef: Py_ssize_t i = 0 Py_ssize_t n = PyList_GET_SIZE(self.tps) + Py_ssize_t pos = 0 + PyObject *pkey + PyObject *pvalue object tp object it if n != len(self.buffers): return False + if PyDict_CheckExact(self.buffers): + # Walks the dict without allocating an items view, an iterator + # or a tuple per entry -- which is what makes this cheaper than + # simply rebuilding the snapshot. The borrowed references are + # safe because nothing here mutates the dict. + while PyDict_Next(self.buffers, &pos, &pkey, &pvalue): + if i >= n: + return False + if PyList_GET_ITEM(self.tps, i) is not pkey: + return False + if PyList_GET_ITEM(self.iters, i) is not pvalue: + return False + i += 1 + return i == n for tp, it in self.buffers.items(): if i >= n: return False @@ -236,11 +255,25 @@ cdef class RoundRobinRecordIterator: cdef: Py_ssize_t i = 0 Py_ssize_t n = PyList_GET_SIZE(self.topics) + Py_ssize_t pos = 0 + PyObject *pkey + PyObject *pvalue object topic object buffer if n != len(self.index): return False + if PyDict_CheckExact(self.index): + while PyDict_Next(self.index, &pos, &pkey, &pvalue): + if i >= n: + return False + if PyList_GET_ITEM(self.topics, i) is not pkey: + return False + if (<_TopicCursor>PyList_GET_ITEM( + self.topic_cursors, i)).source is not pvalue: + return False + i += 1 + return i == n for topic, buffer in self.index.items(): if i >= n: return False diff --git a/faust/transport/utils.py b/faust/transport/utils.py index fe3f7805a..0267c11df 100644 --- a/faust/transport/utils.py +++ b/faust/transport/utils.py @@ -1,7 +1,6 @@ """Transport utils - scheduling.""" import os -from collections import OrderedDict from typing import ( Any, Dict, @@ -87,9 +86,11 @@ class TopicBuffer(Iterator): _it: Optional[Iterator] def __init__(self) -> None: - # note: this is a regular dict, but ordered on Python 3.6 - # we use this alias to signify it must be ordered. - self._buffers = OrderedDict() + # Insertion-ordered by language guarantee since 3.7, and Faust + # requires 3.10, so a plain dict is enough. Using one rather than + # OrderedDict also lets the Cython scheduler walk it with + # PyDict_Next, which allocates nothing per entry. + self._buffers = {} # getmany calls next(_TopicBuffer), and does not call iter(), # so the first call to next caches an iterator. self._it = None diff --git a/tests/unit/transport/test_utils.py b/tests/unit/transport/test_utils.py index 1f92597c6..48b009625 100644 --- a/tests/unit/transport/test_utils.py +++ b/tests/unit/transport/test_utils.py @@ -244,3 +244,22 @@ def test_partition_buffer_replaced_mid_iteration(self, impl): # 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. + + @pytest.mark.parametrize("impl", RECORDS_ITERATOR_IMPLS) + def test_non_dict_index_mapping(self, impl): + # The Cython snapshot check has a PyDict_Next fast path for exact + # dicts and falls back to .items() for anything else, so a mapping + # that is not a plain dict has to give the same answer. + from collections import OrderedDict + + buffer = TopicBuffer() + buffer.add(TP1, BUF1) + index = OrderedDict([("foo", buffer)]) + assert list(impl(index)) == [(TP1, i) for i in BUF1] + + def test_buffers_is_a_plain_dict(self): + # PyDict_Next requires an exact dict; TopicBuffer._buffers used to be + # an OrderedDict, which is a subclass and would take the slow path. + buffer = TopicBuffer() + buffer.add(TP1, BUF1) + assert type(buffer._buffers) is dict