From c812ec3f21bbaf6acfaaa66747477a1d515a190d Mon Sep 17 00:00:00 2001 From: Vizonex Date: Fri, 16 Oct 2026 16:20:18 -0500 Subject: [PATCH] use a cythonized WeakSet rather than an interpreted one --- winloop/includes/stdlib.pxi | 7 +- winloop/loop.pxd | 6 +- winloop/loop.pyx | 3 +- winloop/weakref_set.pxd | 28 ++++ winloop/weakref_set.pyx | 271 ++++++++++++++++++++++++++++++++++++ 5 files changed, 312 insertions(+), 3 deletions(-) create mode 100644 winloop/weakref_set.pxd create mode 100644 winloop/weakref_set.pyx diff --git a/winloop/includes/stdlib.pxi b/winloop/includes/stdlib.pxi index 3635067..0f38931 100644 --- a/winloop/includes/stdlib.pxi +++ b/winloop/includes/stdlib.pxi @@ -60,6 +60,8 @@ cdef col_deque = collections.deque cdef col_Iterable = collections.abc.Iterable cdef col_Counter = collections.Counter cdef col_OrderedDict = collections.OrderedDict +cdef col_abc_MutableSet = collections.abc.MutableSet + cdef cc_ThreadPoolExecutor = concurrent.futures.ThreadPoolExecutor cdef cc_Future = concurrent.futures.Future @@ -192,7 +194,10 @@ cdef warnings_warn = warnings.warn # TODO: weakref classes here could be rewritten to Cython to prevent bottlenecks when passing # between interpreter code and compiled code. cdef weakref_WeakValueDictionary = weakref.WeakValueDictionary -cdef weakref_WeakSet = weakref.WeakSet + +# Removed in favor of the new compiler WeakSet which +# uses compiled code. +# cdef weakref_WeakSet = weakref.WeakSet cdef py_inf = float('inf') diff --git a/winloop/loop.pxd b/winloop/loop.pxd index e70d923..509ae4b 100644 --- a/winloop/loop.pxd +++ b/winloop/loop.pxd @@ -79,7 +79,10 @@ cdef class Loop: cdef object __weakref__ - object _asyncgens + # altered to a cythonized WeakSet in 0.9.0 + # it's mostly a performance enhancement + # that remains unexposed. + WeakSet _asyncgens bint _asyncgens_shutdown_called bint _executor_shutdown_called @@ -228,6 +231,7 @@ include "sslproto.pxd" include "handles/udp.pxd" include "server.pxd" +include "weakref_set.pxd" cdef class _SyncSocketReaderFuture(Future): diff --git a/winloop/loop.pyx b/winloop/loop.pyx index 7ecc796..47c1791 100644 --- a/winloop/loop.pyx +++ b/winloop/loop.pyx @@ -223,7 +223,7 @@ cdef class Loop: # A weak set of all asynchronous generators that are # being iterated by the loop. - self._asyncgens = weakref_WeakSet() + self._asyncgens = WeakSet() # Set to True when `loop.shutdown_asyncgens` is called. self._asyncgens_shutdown_called = False @@ -3453,6 +3453,7 @@ include "sslproto.pyx" include "handles/udp.pyx" include "server.pyx" +include "weakref_set.pyx" # Used in UVProcess diff --git a/winloop/weakref_set.pxd b/winloop/weakref_set.pxd new file mode 100644 index 0000000..6357409 --- /dev/null +++ b/winloop/weakref_set.pxd @@ -0,0 +1,28 @@ +# mostly needed as a speedup for winloop's internal loop +# which contains unneeded & costly interpreter logic + +# The only downside is that this WeakSet doesn't fully support +# copying but it's a worth-it tradeoff. + +cdef class RemovalSystem: + cdef: + object this_ref + +cdef class WeakSet: + cdef: + object __weakref__ + public set data + set _iterating + list _pending_removals + RemovalSystem _remove + + cpdef _commit_removals(self) + cpdef add(self, item) + cpdef clear(self) + cpdef WeakSet copy(self) + cpdef pop(self) + cpdef remove(self, item) + cpdef discard(self, item) + cpdef update(self, other) + cpdef difference_update(self, other) + cpdef isdisjoint(self, other) diff --git a/winloop/weakref_set.pyx b/winloop/weakref_set.pyx new file mode 100644 index 0000000..230ed47 --- /dev/null +++ b/winloop/weakref_set.pyx @@ -0,0 +1,271 @@ +# Fork of Python's WeakSet class but cythonized for extra speed. + +from cpython.weakref cimport PyWeakref_NewRef, PyWeakref_GetRef +from cpython.object cimport PyObject +from cpython.set cimport PySet_Add, PySet_Clear, PySet_Discard, PySet_Pop + +cdef extern from "Python.h": + """ +static inline PyObject* Weakref_NewRef(PyObject* obj){ + return PyWeakref_NewRef(obj, NULL); +} +#define weakref_steal(obj) obj + """ + + # Same as PyWeakref_NewRef but NULL is passed for callback. + object Weakref_NewRef(object obj) + object weakref_steal(PyObject* obj) + WeakSet weakref_steal_as_weakset "weakref_steal"(PyObject* obj) + + + +cdef class _IterationGuard: + # This context manager registers itself in the current iterators of the + # weak container, such as to delay all removals until the context manager + # exits. + # This technique should be relatively thread-safe (since sets are). + + cdef: + object weakcontainer + + def __init__(self, object weakcontainer): + # Don't create cycles + self.weakcontainer = Weakref_NewRef(weakcontainer) + + def __enter__(self): + cdef PyObject* w + cdef WeakSet ws + + if PyWeakref_GetRef(self.weakcontainer, &w): + ws = weakref_steal_as_weakset(w) + PySet_Add(ws._iterating, self) + return self + + def __exit__(self, *args): + cdef set s + cdef PyObject* wr + cdef WeakSet ws + if PyWeakref_GetRef(self.weakcontainer, &wr): + ws = weakref_steal_as_weakset(wr) + s = ws._iterating + s.remove(self) + if not s: + ws._commit_removals() + + +cdef class RemovalSystem: + + def __cinit__(self, WeakSet this) -> None: + self.this_ref = Weakref_NewRef(this) + + def __call__(self, object item, object selfref=None): + cdef PyObject* w = NULL + cdef WeakSet ws + cdef object this_ref = selfref or self.this_ref + + if PyWeakref_GetRef(this_ref, &w): + ws = weakref_steal_as_weakset(w) + if ws._iterating: + ws._pending_removals.append(item) + else: + ws.data.discard(item) + + + + +cdef class WeakSet: + def __init__(self, object data = None): + self.data = set() + self._remove = RemovalSystem(self) + # A list of keys to be removed + self._pending_removals = [] + self._iterating = set() + if data is not None: + self.update(data) + + cpdef _commit_removals(self): + pop = self._pending_removals.pop + discard = self.data.discard + while True: + try: + item = pop() + except IndexError: + return + discard(item) + + def __iter__(self): + cdef PyObject* w + with _IterationGuard(self): + for itemref in self.data: + if PyWeakref_GetRef(itemref, &w): + # Caveat: the iterator will keep a strong reference to + # `item` until it is resumed or closed. + yield weakref_steal(w) + + def __len__(self): + return len(self.data) - len(self._pending_removals) + + def __contains__(self, item): + try: + wr = Weakref_NewRef(item) + except TypeError: + return False + return wr in self.data + + def __reduce__(self): + return self.__class__, (list(self),), self.__getstate__() + + cpdef add(self, item): + if self._pending_removals: + self._commit_removals() + PySet_Add(self.data, PyWeakref_NewRef(item, self._remove)) + + cpdef clear(self): + if self._pending_removals: + self._commit_removals() + PySet_Clear(self.data) + + cpdef WeakSet copy(self): + return self.__class__(self) + + cpdef pop(self): + cdef PyObject* ref + if self._pending_removals: + self._commit_removals() + while True: + try: + itemref = PySet_Pop(self.data) + except KeyError: + raise KeyError('pop from empty WeakSet') from None + if PyWeakref_GetRef(itemref, &ref): + return weakref_steal(ref) + + cpdef remove(self, item): + if self._pending_removals: + self._commit_removals() + self.data.remove(Weakref_NewRef(item)) + + cpdef discard(self, item): + if self._pending_removals: + self._commit_removals() + PySet_Discard(self.data, Weakref_NewRef(item)) + + cpdef update(self, other): + if self._pending_removals: + self._commit_removals() + for element in other: + self.add(element) + + def __ior__(self, other): + self.update(other) + return self + + def difference(self, other): + newset = self.copy() + newset.difference_update(other) + return newset + + def __sub__(self, other): + newset = self.copy() + newset.difference_update(other) + return newset + + cpdef difference_update(self, other): + self.__isub__(other) + + def __isub__(self, other): + if self._pending_removals: + self._commit_removals() + if self is other: + PySet_Clear(self.data) + else: + self.data.difference_update(Weakref_NewRef(item) for item in other) + return self + + def intersection(self, other): + return self.__class__(item for item in other if item in self) + + def __and__(self, other): + return self.__class__(item for item in other if item in self) + + def intersection_update(self, other): + self.__iand__(other) + + def __iand__(self, other): + if self._pending_removals: + self._commit_removals() + self.data.intersection_update(Weakref_NewRef(item) for item in other) + return self + + def issubset(self, other): + return self.data.issubset(Weakref_NewRef(item) for item in other) + + + def __richcmp__(self, other, op): + if op == 0: # < + return self.data < {Weakref_NewRef(item) for item in other} + + if op == 1: # <= + return self.data.issubset(Weakref_NewRef(item) for item in other) + + if op == 2: # == + if not isinstance(other, self.__class__): + return NotImplemented + return self.data == {Weakref_NewRef(item) for item in other} + + if op == 3: # != + if not isinstance(other, self.__class__): + return NotImplemented + return self.data != {Weakref_NewRef(item) for item in other} + + if op == 4: # > + return self.data > {Weakref_NewRef(item) for item in other} + + if op == 5: # >= + return self.issuperset(other) + + def issuperset(self, other): + return self.data.issuperset(Weakref_NewRef(item) for item in other) + + def symmetric_difference(self, other): + newset = self.copy() + newset.symmetric_difference_update(other) + return newset + + def __xor__(self, other): + newset = self.copy() + newset.symmetric_difference_update(other) + return newset + + def symmetric_difference_update(self, object other): + if self._pending_removals: + self._commit_removals() + if self is other: + PySet_Clear(self.data) + else: + self.data.symmetric_difference_update(PyWeakref_NewRef(item, self._remove) for item in other) + + def __ixor__(self, other): + if self._pending_removals: + self._commit_removals() + if self is other: + PySet_Clear(self.data) + else: + self.data.symmetric_difference_update(PyWeakref_NewRef(item, self._remove) for item in other) + return self + + def union(self, other): + return self.__class__(e for s in (self, other) for e in s) + + def __or__(self, other): + return self.__class__(e for s in (self, other) for e in s) + + cpdef isdisjoint(self, other): + return len(self.intersection(other)) == 0 + + def __repr__(self): + return repr(self.data) + + __class_getitem__ = classmethod(types_GenericAlias) + +col_abc_MutableSet.register(WeakSet)