Recover scheduler throughput lost to the snapshot-staleness check

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 <object>PyList_GET_ITEM(self.tps, i) is not <object>pkey:
+                    return False
+                if <object>PyList_GET_ITEM(self.iters, i) is not <object>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 <object>PyList_GET_ITEM(self.topics, i) is not <object>pkey:
+                    return False
+                if (<_TopicCursor>PyList_GET_ITEM(
+                        self.topic_cursors, i)).source is not <object>pvalue:
+                    return False
+                i += 1
+            return i == n
         for topic, buffer in self.index.items():
             if i >= n:
                 return False

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

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