ez-a-sync 0.32.29__cp310-cp310-win_amd64.whl

a_sync/primitives/locks/prio_semaphore.pyx

@@ -0,0 +1,597 @@
+# cython: boundscheck=False
+"""
+This module provides priority-based semaphore implementations. These semaphores allow 
+waiters to be assigned priorities, ensuring that higher priority waiters are 
+processed before lower priority ones.
+"""
+
+import asyncio
+import collections
+import heapq
+from logging import getLogger
+from typing import List, Literal, Optional, Protocol, Set, Type, TypeVar
+
+from cpython.unicode cimport PyUnicode_CompareWithASCIIString
+
+from a_sync.primitives.locks.semaphore cimport Semaphore
+
+
+# cdef collections
+cdef object deque = collections.deque
+del collections
+
+# cdef heapq
+cdef object heappush = heapq.heappush
+cdef object heappop = heapq.heappop
+del heapq
+
+# cdef logging
+cdef public object logger = getLogger(__name__)
+cdef object c_logger = logger
+cdef object DEBUG = 10
+cdef object _logger_log = logger._log
+cdef object _logger_is_enabled = logger.isEnabledFor
+del getLogger
+
+class Priority(Protocol):
+    def __lt__(self, other) -> bool: ...
+
+
+PT = TypeVar("PT", bound=Priority)
+
+CM = TypeVar("CM", bound="_AbstractPrioritySemaphoreContextManager[Priority]")
+
+
+cdef class _AbstractPrioritySemaphore(Semaphore):
+    """
+    A semaphore that allows prioritization of waiters.
+
+    This semaphore manages waiters with associated priorities, ensuring that waiters with higher
+    priorities are processed before those with lower priorities. Subclasses must define the
+    `_top_priority` attribute to specify the default top priority behavior.
+
+    The `_context_manager_class` attribute should specify the class used for managing semaphore contexts.
+
+    See Also:
+        :class:`PrioritySemaphore` for an implementation using numeric priorities.
+    """
+
+    def __cinit__(self) -> None:
+        self._context_managers = {}
+        """A dictionary mapping priorities to their context managers."""
+    
+        self._Semaphore__waiters = []
+        """A heap queue of context managers, sorted by priority."""
+
+        # NOTE: This should (hopefully) be temporary
+        self._potential_lost_waiters = set()
+        """A list of futures representing waiters that might have been lost."""
+    
+    def __init__(
+        self, 
+        context_manager_class: Type[_AbstractPrioritySemaphoreContextManager],
+        top_priority: object,
+        value: int = 1, 
+        *, 
+        name: Optional[str] = None, 
+    ) -> None:
+        """Initializes the priority semaphore.
+
+        Args:
+            value: The initial capacity of the semaphore.
+            name: An optional name for the semaphore, used for debugging.
+
+        Examples:
+            >>> semaphore = _AbstractPrioritySemaphore(5, name="test_semaphore")
+        """
+        # context manager class is some temporary hacky shit, just ignore this
+        Semaphore.__init__(self, value, name=name)
+
+        self._capacity = value
+        """The initial capacity of the semaphore."""
+
+        self._top_priority = top_priority
+        self._top_priority_manager = context_manager_class(
+                self, top_priority, name=self.name
+            )
+        self._context_manager_class = context_manager_class
+
+    def __repr__(self) -> str:
+        """Returns a string representation of the semaphore."""
+        return f"<{self.__class__.__name__} name={self.name} capacity={self._capacity} value={self._Semaphore__value} waiters={self._count_waiters()}>"
+
+    async def __aenter__(self) -> None:
+        """Enters the semaphore context, acquiring it with the top priority.
+
+        This method is part of the asynchronous context management protocol.
+
+        Examples:
+            >>> semaphore = _AbstractPrioritySemaphore(5)
+            >>> async with semaphore:
+            ...     await do_stuff()
+        """
+        await self._top_priority_manager.acquire()
+
+    async def __aexit__(self, *_) -> None:
+        """Exits the semaphore context, releasing it with the top priority.
+
+        This method is part of the asynchronous context management protocol.
+
+        Examples:
+            >>> semaphore = _AbstractPrioritySemaphore(5)
+            >>> async with semaphore:
+            ...     await do_stuff()
+        """
+        self._top_priority_manager.release()
+
+    cpdef object acquire(self):
+        """Acquires the semaphore with the top priority.
+
+        This method overrides :meth:`Semaphore.acquire` to handle priority-based logic.
+
+        Examples:
+            >>> semaphore = _AbstractPrioritySemaphore(5)
+            >>> await semaphore.acquire()
+        """
+        return self._top_priority_manager.acquire()
+
+    def __getitem__(
+        self, priority: Optional[PT]
+    ) -> "_AbstractPrioritySemaphoreContextManager[PT]":
+        """Gets the context manager for a given priority.
+
+        Args:
+            priority: The priority for which to get the context manager. If None, uses the top priority.
+
+        Returns:
+            The context manager associated with the given priority.
+
+        Examples:
+            >>> semaphore = _AbstractPrioritySemaphore(5)
+            >>> context_manager = semaphore[priority]
+        """
+        return self.c_getitem(priority)
+    
+    cdef _AbstractPrioritySemaphoreContextManager c_getitem(self, object priority):
+        cdef _AbstractPrioritySemaphoreContextManager context_manager
+        cdef dict[object, _AbstractPrioritySemaphoreContextManager] context_managers
+
+        if priority is None or priority == self._top_priority:
+            return self._top_priority_manager
+        
+        context_managers = self._context_managers
+        context_manager = context_managers.get(priority)
+        if context_manager is None:
+            context_manager = self._context_manager_class(
+                self, priority, name=self.name
+            )
+            heappush(self._Semaphore__waiters, context_manager)
+            context_managers[priority] = context_manager
+        return context_manager
+
+    cpdef bint locked(self):
+        """Checks if the semaphore is locked.
+
+        Returns:
+            True if the semaphore cannot be acquired immediately, False otherwise.
+
+        Examples:
+            >>> semaphore = _AbstractPrioritySemaphore(5)
+            >>> semaphore.locked()
+        """
+        cdef list waiters
+        if self._Semaphore__value == 0:
+            return True
+        for cm in self._context_managers.values():
+            waiters = (<Semaphore>cm).__waiters
+            for waiter in waiters:
+                if _is_not_cancelled(waiter):
+                    return True
+        return False
+
+    cdef dict[object, Py_ssize_t] _count_waiters(self):
+        """Counts the number of waiters for each priority.
+
+        Returns:
+            A dictionary mapping each priority to the number of waiters.
+
+        Examples:
+            >>> semaphore = _AbstractPrioritySemaphore(5)
+            >>> semaphore._count_waiters()
+        """
+        cdef _AbstractPrioritySemaphoreContextManager manager
+        cdef list[_AbstractPrioritySemaphoreContextManager] waiters = self._Semaphore__waiters
+        return {manager._priority: len(manager._Semaphore__waiters) for manager in sorted(waiters)}
+        
+    cpdef void _wake_up_next(self):
+        """Wakes up the next waiter in line.
+
+        This method handles the waking of waiters based on priority. It includes an emergency
+        procedure to handle potential lost waiters, ensuring that no waiter is left indefinitely
+        waiting.
+
+        The emergency procedure is a temporary measure to address potential issues with lost waiters.
+
+        Examples:
+            >>> semaphore = _AbstractPrioritySemaphore(5)
+            >>> semaphore._wake_up_next()
+        """
+        cdef _AbstractPrioritySemaphoreContextManager manager
+        cdef object manager_waiters, get_next
+        cdef Py_ssize_t start_len, end_len
+        cdef set potential_lost_waiters
+        cdef bint woke_up, debug_logs
+
+        manager = self._top_priority_manager
+        manager_waiters = manager._Semaphore__waiters
+        potential_lost_waiters = self._potential_lost_waiters
+        debug_logs = _logger_is_enabled(DEBUG)
+        if manager_waiters:
+            get_next = manager_waiters.popleft
+            while manager_waiters:
+                waiter = get_next()
+                potential_lost_waiters.discard(waiter)
+                if _is_not_done(waiter):
+                    waiter.set_result(None)
+                    if debug_logs:
+                        log_debug("woke up %s", (waiter, ))
+                    return
+        
+        while manager_waiters:
+            waiter = get_next()
+            potential_lost_waiters.discard(waiter)
+            if _is_not_done(waiter):
+                waiter.set_result(None)
+                if debug_logs:
+                    log_debug("woke up %s", (waiter, ))
+                return
+
+        cdef list self_waiters = self._Semaphore__waiters
+        while self_waiters:
+            manager = heappop(self_waiters)
+            if len(manager) == 0:
+                # There are no more waiters, get rid of the empty manager
+                if debug_logs:
+                    log_debug(
+                        "manager %s has no more waiters, popping from %s",
+                        (manager._repr_no_parent_(), self),
+                    )
+                self._context_managers.pop(manager._priority)
+                continue
+
+            woke_up = False
+            start_len = len(manager)
+
+            manager_waiters = manager._Semaphore__waiters
+            if debug_logs:
+                log_debug("waking up next for %s", (manager._repr_no_parent_(), ))
+                if not manager_waiters:
+                    log_debug("not manager._Semaphore__waiters", ())
+
+            if manager_waiters:
+                get_next = manager_waiters.popleft
+                waiter = get_next()
+                potential_lost_waiters.discard(waiter)
+                if _is_not_done(waiter):
+                    waiter.set_result(None)
+                    woke_up = True
+                    if debug_logs:
+                        log_debug("woke up %s", (waiter, ))
+                    break
+
+                if not woke_up:
+                    while manager_waiters:
+                        waiter = get_next()
+                        potential_lost_waiters.discard(waiter)
+                        if _is_not_done(waiter):
+                            waiter.set_result(None)
+                            woke_up = True
+                            if debug_logs:
+                                log_debug("woke up %s", (waiter, ))
+                            break
+
+            if not woke_up:
+                self._context_managers.pop(manager._priority)
+                continue
+
+            end_len = len(manager)
+
+            assert start_len > end_len, f"start {start_len} end {end_len}"
+
+            if end_len:
+                # There are still waiters, put the manager back
+                heappush(self_waiters, manager)
+            else:
+                # There are no more waiters, get rid of the empty manager
+                self._context_managers.pop(manager._priority)
+            return
+
+        # emergency procedure (hopefully temporary):
+        if not debug_logs:
+            while potential_lost_waiters:
+                waiter = potential_lost_waiters.pop()
+                if _is_not_done(waiter):
+                    waiter.set_result(None)
+                    return
+            return
+            
+        while potential_lost_waiters:
+            waiter = potential_lost_waiters.pop()
+            log_debug("we found a lost waiter %s", (waiter, ))
+            if _is_not_done(waiter):
+                waiter.set_result(None)
+                log_debug("woke up lost waiter %s", (waiter, ))
+                return
+
+        log_debug("%s has no waiters to wake", (self, ))
+
+
+cdef class _AbstractPrioritySemaphoreContextManager(Semaphore):
+    """
+    A context manager for priority semaphore waiters.
+
+    This context manager is associated with a specific priority and handles
+    the acquisition and release of the semaphore for waiters with that priority.
+    """
+
+    def __init__(
+        self,
+        parent: _AbstractPrioritySemaphore,
+        priority: PT,
+        name: Optional[str] = None,
+    ) -> None:
+        """Initializes the context manager for a specific priority.
+
+        Args:
+            parent: The parent semaphore.
+            priority: The priority associated with this context manager.
+            name: An optional name for the context manager, used for debugging.
+
+        Examples:
+            >>> parent_semaphore = _AbstractPrioritySemaphore(5)
+            >>> context_manager = _AbstractPrioritySemaphoreContextManager(parent_semaphore, priority=1)
+        """
+
+        Semaphore.__init__(self, 0, name=name)
+
+        self._parent = parent
+        """The parent semaphore."""
+
+        self._priority = priority
+        """The priority associated with this context manager."""
+
+    def __repr__(self) -> str:
+        """Returns a string representation of the context manager."""
+        return f"<{self.__class__.__name__} parent={self._parent} {self._priority_name}={self._priority} waiters={len(self)}>"
+
+    cpdef str _repr_no_parent_(self):
+        """Returns a string representation of the context manager without the parent."""
+        return f"<{self.__class__.__name__} parent_name={self._parent.name} {self._priority_name}={self._priority} waiters={len(self)}>"
+
+    def __richcmp__(self, _AbstractPrioritySemaphoreContextManager other, int op) -> bint:
+        """Rich comparison special method. Compares this context manager with another based on priority.
+
+        Args:
+            other: The other context manager to compare with.
+            op: The operation being performed. One of:
+                0 -> Py_LT (a < b)
+                1 -> Py_LE (a <= b)
+                2 -> Py_EQ (a == b)
+                3 -> Py_NE (a != b)
+                4 -> Py_GT (a > b)
+                5 -> Py_GE (a >= b)
+
+        Returns:
+            The boolean result of the comparison.
+
+        Raises:
+            TypeError: If the other object is not an instance of :class:`~_AbstractPrioritySemaphoreContextManager`.
+
+        Examples:
+            >>> cm1 = _AbstractPrioritySemaphoreContextManager(parent, priority=1)
+            >>> cm2 = _AbstractPrioritySemaphoreContextManager(parent, priority=2)
+            >>> cm1 < cm2
+        """
+        if op == 0:  # Py_LT
+            return self._priority < other._priority
+        elif op == 1:  # Py_LE
+            return self._priority <= other._priority
+        elif op == 2:  # Py_EQ
+            return self is other
+        elif op == 3:  # Py_NE
+            return self is not other
+        elif op == 4:  # Py_GT
+            return self._priority > other._priority
+        elif op == 5:  # Py_GE
+            return self._priority >= other._priority
+        return NotImplemented
+
+    cpdef object acquire(self):
+        """Acquires the semaphore for this context manager.
+
+        If the internal counter is larger than zero on entry,
+        decrement it by one and return True immediately. If it is
+        zero on entry, block, waiting until some other coroutine has
+        called release() to make it larger than 0, and then return
+        True.
+
+        This method overrides :meth:`Semaphore.acquire` to handle priority-based logic.
+
+        Examples:
+            >>> context_manager = _AbstractPrioritySemaphoreContextManager(parent, priority=1)
+            >>> await context_manager.acquire()
+        """
+        if self._parent._Semaphore__value <= 0:
+            self._c_ensure_debug_daemon((),{})
+        return self.__acquire()
+    
+    async def __acquire(self) -> Literal[True]:
+        cdef object fut
+        cdef _AbstractPrioritySemaphore parent = self._parent
+        while parent._Semaphore__value <= 0:
+            if self._Semaphore__waiters is None:
+                self._Semaphore__waiters = deque()
+            fut = self._c_get_loop().create_future()
+            self._Semaphore__waiters.append(fut)
+            parent._potential_lost_waiters.add(fut)
+            try:
+                await fut
+            except:
+                # See the similar code in Queue.get.
+                fut.cancel()
+                if parent._Semaphore__value > 0 and _is_not_cancelled(fut):
+                    parent._wake_up_next()
+                raise
+        parent._Semaphore__value -= 1
+        return True
+
+    cpdef void release(self):
+        """Releases the semaphore for this context manager.
+
+        This method overrides :meth:`Semaphore.release` to handle priority-based logic.
+
+        Examples:
+            >>> context_manager = _AbstractPrioritySemaphoreContextManager(parent, priority=1)
+            >>> context_manager.release()
+        """
+        self._parent.release()
+
+
+cdef inline bint _is_not_done(fut: asyncio.Future):
+    return PyUnicode_CompareWithASCIIString(fut._state, b"PENDING") == 0
+
+cdef inline bint _is_not_cancelled(fut: asyncio.Future):
+    return PyUnicode_CompareWithASCIIString(fut._state, b"CANCELLED") != 0
+
+
+del asyncio
+
+
+cdef class _PrioritySemaphoreContextManager(_AbstractPrioritySemaphoreContextManager):
+    """Context manager for numeric priority semaphores."""
+
+    def __cinit__(self):
+        self._priority_name = "priority"
+        # Semaphore.__cinit__(self)
+        self._Semaphore__waiters = deque()
+        self._decorated: Set[str] = set()
+
+    def __richcmp__(self, _PrioritySemaphoreContextManager other, int op) -> bint:
+        """Rich comparison special method. Compares this context manager with another based on priority.
+
+        Args:
+            other: The other context manager to compare with.
+            op: The operation being performed. One of:
+                0 -> Py_LT (a < b)
+                1 -> Py_LE (a <= b)
+                2 -> Py_EQ (a == b)
+                3 -> Py_NE (a != b)
+                4 -> Py_GT (a > b)
+                5 -> Py_GE (a >= b)
+                
+        Returns:
+            The boolean result of the comparison.
+
+        Raises:
+            TypeError: If the other object is not an instance of :class:`~_PrioritySemaphoreContextManager`.
+
+        Examples:
+            >>> cm1 = _AbstractPrioritySemaphoreContextManager(parent, priority=1)
+            >>> cm2 = _AbstractPrioritySemaphoreContextManager(parent, priority=2)
+            >>> cm1 < cm2
+        """
+        if op == 0:  # Py_LT
+            return <int>self._priority < <int>other._priority
+        elif op == 1:  # Py_LE
+            return <int>self._priority <= <int>other._priority
+        elif op == 2:  # Py_EQ
+            return self is other
+        elif op == 3:  # Py_NE
+            return self is not other
+        elif op == 4:  # Py_GT
+            return <int>self._priority > <int>other._priority
+        elif op == 5:  # Py_GE
+            return <int>self._priority >= <int>other._priority
+        return NotImplemented
+
+cdef class PrioritySemaphore(_AbstractPrioritySemaphore):
+    """Semaphore that uses numeric priorities for waiters.
+
+    This class extends :class:`_AbstractPrioritySemaphore` and provides a concrete implementation
+    using numeric priorities. The `_context_manager_class` is set to :class:`_PrioritySemaphoreContextManager`,
+    and the `_top_priority` is set to -1, which is the highest priority.
+
+    Examples:
+        The primary way to use this semaphore is by specifying a priority.
+
+        >>> priority_semaphore = PrioritySemaphore(10)
+        >>> async with priority_semaphore[priority]:
+        ...     await do_stuff()
+
+        You can also enter and exit this semaphore without specifying a priority, and it will use the top priority by default:
+
+        >>> priority_semaphore = PrioritySemaphore(10)
+        >>> async with priority_semaphore:
+        ...     await do_stuff()
+
+    See Also:
+        :class:`_AbstractPrioritySemaphore` for the base class implementation.
+    """
+
+    def __cinit__(self):        
+        # _AbstractPrioritySemaphore.__cinit__(self)
+
+        self._context_managers = {}
+        """A dictionary mapping priorities to their context managers."""
+    
+        self._Semaphore__waiters = []
+        """A heap queue of context managers, sorted by priority."""
+
+        # NOTE: This should (hopefully) be temporary
+        self._potential_lost_waiters = set()
+        """A list of futures representing waiters that might have been lost."""
+    
+    def __init__(
+        self, 
+        value: int = 1, 
+        *, 
+        name: Optional[str] = None, 
+    ) -> None:
+        # context manager class is some temporary hacky shit, just ignore this
+        _AbstractPrioritySemaphore.__init__(self, _PrioritySemaphoreContextManager, -1, value, name=name)
+
+    def __getitem__(
+        self, priority: Optional[PT]
+    ) -> "_PrioritySemaphoreContextManager[PT]":
+        """Gets the context manager for a given priority.
+
+        Args:
+            priority: The priority for which to get the context manager. If None, uses the top priority.
+
+        Returns:
+            The context manager associated with the given priority.
+
+        Examples:
+            >>> semaphore = _AbstractPrioritySemaphore(5)
+            >>> context_manager = semaphore[priority]
+        """
+        return self.c_getitem(priority or -1)
+    
+    cdef _PrioritySemaphoreContextManager c_getitem(self, object priority):
+        if self._Semaphore__value is None:
+            raise ValueError(self._Semaphore__value)
+        
+        cdef dict[int, _PrioritySemaphoreContextManager] context_managers
+        cdef _PrioritySemaphoreContextManager context_manager
+        
+        context_managers = self._context_managers
+        context_manager = context_managers.get(priority)
+        if context_manager is None:
+            context_manager = _PrioritySemaphoreContextManager(self, <int>priority, name=self.name)
+            heappush(self._Semaphore__waiters, context_manager)
+            context_managers[<int>priority] = context_manager
+        return context_manager
+    
+
+cdef void log_debug(str message, tuple args):
+    _logger_log(DEBUG, message, args)