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

a_sync/executor.py

@@ -0,0 +1,619 @@
+"""
+This module provides several executor classes that facilitate running synchronous functions asynchronously using `asyncio`.
+
+With these executors, you can run sync functions in your executor with `await executor.run(fn, *args, **kwargs)`.
+The `executor.submit(fn, *args, **kwargs)` method works similarly to the `concurrent.futures` implementation but
+returns an `asyncio.Future` instead of a `concurrent.futures.Future`.
+
+Executor Classes:
+    - :class:`AsyncProcessPoolExecutor`: A process pool executor providing asynchronous run and submit methods, with support for synchronous mode
+    - :class:`AsyncThreadPoolExecutor`: A thread pool executor providing asynchronous run and submit methods, with support for synchronous mode
+    - :class:`PruningThreadPoolExecutor`: An :class:`AsyncThreadPoolExecutor` that prunes inactive threads after a timeout, ensuring at least one thread remains active to prevent locks.
+
+See Also:
+    - :mod:`concurrent.futures` for the original synchronous executor implementations.
+"""
+
+import asyncio
+import atexit
+import concurrent.futures
+import multiprocessing.context
+import queue
+import signal
+import threading
+import weakref
+from asyncio import sleep
+from asyncio.futures import _convert_future_exc
+from concurrent.futures import _base, thread
+
+from a_sync._typing import *
+from a_sync.primitives._debug import _DebugDaemonMixin
+
+# === Executor Shutdown Logic ===
+# All executors (module-level and user-created) are registered for shutdown on interpreter exit and signals.
+# Signal handlers are chainable: after our cleanup, the previous handler is called (unless SIG_DFL or SIG_IGN).
+# This ensures compatibility with other libraries and deduplicates shutdown logic.
+
+_EXECUTORS = set()
+
+
+def register_executor(executor) -> None:
+    """Register an executor for shutdown on exit/signals."""
+    _EXECUTORS.add(executor)
+
+
+def _shutdown_all_executors(*args) -> None:
+    """Shutdown all registered executors (non-blocking)."""
+    for executor in list(_EXECUTORS):
+        try:
+            executor.shutdown(wait=False)
+        except Exception:
+            pass
+
+
+def _register_executor_shutdown() -> None:
+    """Register atexit and chainable signal handlers for executor shutdown."""
+    atexit.register(_shutdown_all_executors)
+
+    def make_chainable_signal_handler(signalnum):
+        prev_handler = signal.getsignal(signalnum)
+
+        def handler(signum, frame):
+            _shutdown_all_executors()
+            if callable(prev_handler) and prev_handler not in (signal.SIG_DFL, signal.SIG_IGN):
+                prev_handler(signum, frame)
+
+        signal.signal(signalnum, handler)
+
+    make_chainable_signal_handler(signal.SIGINT)
+    make_chainable_signal_handler(signal.SIGTERM)
+
+
+_register_executor_shutdown()
+
+TEN_MINUTES = 60 * 10
+
+Initializer = Callable[..., object]
+
+
+class _AsyncExecutorMixin(concurrent.futures.Executor, _DebugDaemonMixin):
+    """
+    A mixin for Executors to provide asynchronous run and submit methods.
+
+    This mixin allows executors to operate in both asynchronous (normal) mode and synchronous mode.
+    In asynchronous (normal) mode, functions are submitted to the executor and awaited.
+    In synchronous mode, functions are executed directly in the current thread.
+
+    Examples:
+        >>> async def example():
+        >>>     result = await executor.run(some_function, arg1, arg2, kwarg1=value1)
+        >>>     print(result)
+
+    See Also:
+        - :meth:`submit` for submitting functions to the executor.
+    """
+
+    sync_mode: bool
+    """
+    Indicates if the executor is in synchronous mode (max_workers == 0).
+
+    Examples:
+        >>> if executor.sync_mode:
+        >>>     print("Executor is in synchronous mode.")
+    """
+
+    _max_workers: int
+
+    _workers: str
+    """The type of workers used."""
+
+    __slots__ = "_max_workers", "_initializer", "_initargs", "_broken", "_shutdown_lock"
+
+    async def run(self, fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs):
+        """
+        A shorthand way to call `await asyncio.get_event_loop().run_in_executor(this_executor, fn, *args)`.
+        Doesn't `await this_executor.run(fn, *args)` look so much better?
+
+        In synchronous mode, the function is executed directly in the current thread.
+        In asynchronous mode, the function is submitted to the executor and awaited.
+
+        Args:
+            fn: The function to run.
+            *args: Positional arguments for the function.
+            **kwargs: Keyword arguments for the function.
+
+        Examples:
+            >>> async def example():
+            >>>     result = await executor.run(some_function, arg1, arg2, kwarg1=value1)
+            >>>     print(result)
+
+        See Also:
+            - :meth:`submit` for submitting functions to the executor.
+        """
+        return fn(*args, **kwargs) if self.sync_mode else await self.submit(fn, *args, **kwargs)
+
+    @overload
+    def submit(self, fn: Callable[P, T], *args: P.args, fire_and_forget: Literal[True], **kwargs: P.kwargs) -> None:  # type: ignore [override]
+        """
+        Submits a job to the executor without expecting a result back. The executor will execute the task silently.
+
+        Args:
+            fn: The function to submit.
+            *args: Positional arguments for the function.
+            fire_and_forget: True
+            **kwargs: Keyword arguments for the function.
+
+        Examples:
+            >>> executor.submit(some_function, arg1, arg2, fire_and_forget=True, kwarg1=value1)
+
+        See Also:
+            - :meth:`run` for running functions with the executor.
+        """
+
+    @overload
+    def submit(self, fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> "asyncio.Future[T]":  # type: ignore [override]
+        """
+        Submits a job to the executor and returns an :class:`asyncio.Future` that can be awaited for the result without blocking.
+
+        Args:
+            fn: The function to submit.
+            *args: Positional arguments for the function.
+            **kwargs: Keyword arguments for the function.
+
+        Examples:
+            >>> future = executor.submit(some_function, arg1, arg2, kwarg1=value1)
+            >>> result = await future
+            >>> print(result)
+
+        See Also:
+            - :meth:`run` for running functions with the executor.
+        """
+
+    def submit(self, fn: Callable[P, T], *args: P.args, fire_and_forget: bool = False, **kwargs: P.kwargs) -> Optional["asyncio.Future[T]"]:  # type: ignore [override]
+        # sourcery skip: simplify-boolean-comparison
+        """
+        Submits a job to the executor and returns an :class:`asyncio.Future` that can be awaited for the result without blocking.
+
+        If `fire_and_forget` is True, the executor will not return any data and intead of a :class:`~Future` this function will return `None`.
+
+        Args:
+            fn: The function to submit.
+            *args: Positional arguments for the function.
+            fire_and_forget (optional): Set True to send the job to the executor without expecting a result. If `fire_and_forget` is True, this function will return None instead of a Future. Default False.
+            **kwargs: Keyword arguments for the function.
+
+        Examples:
+            >>> future = executor.submit(some_function, arg1, arg2, kwarg1=value1)
+            >>> result = await future
+            >>> print(result)
+
+        See Also:
+            - :meth:`run` for running functions with the executor.
+        """
+        if fire_and_forget is True:
+            # Send the job to the executor and return without creating a future or setting up callbacks
+            if self.sync_mode:
+                fn(*args, **kwargs)
+            else:
+                self.__super_submit(fn, *args, **kwargs)
+            return None
+
+        loop = self._get_loop()
+        fut = loop.create_future()
+        if self.sync_mode:
+            try:
+                _set_fut_result(fut, fn(*args, **kwargs))
+            except Exception as e:
+                _set_fut_exception(fut, e)
+        else:
+            self._ensure_debug_daemon(fut, fn, *args, **kwargs)
+
+            cf_fut = self.__super_submit(fn, *args, **kwargs)
+
+            # TODO: implement logic to actually cancel the job, not just the future which is useless for our use case
+            # def _call_check_cancel(destination: asyncio.Future):
+            #     if _fut_is_cancelled(destination):
+            #         cf_fut.cancel()
+            #
+            # fut.add_done_callback(_call_check_cancel)
+
+            def _call_copy_future_state(cf_fut: "concurrent.futures.Future"):
+                if _fut_is_cancelled(fut):
+                    return
+                loop.call_soon_threadsafe(
+                    _copy_future_state,
+                    cf_fut,
+                    fut,
+                )
+
+            _add_done_callback(cf_fut, _call_copy_future_state)
+
+        return fut
+
+    def __repr__(self) -> str:
+        worker_info = f"[{self.worker_count_current}/{self._max_workers} {self._workers}]"
+        identifier = getattr(self, "_thread_name_prefix", None) or hex(id(self))
+        return f"<{self.__class__.__name__} {identifier} {worker_info}>"
+
+    def __len__(self) -> int:
+        # NOTE: should this be queue length instead? probably
+        return self.worker_count_current
+
+    @property
+    def worker_count_current(self) -> int:
+        """
+        Returns the current number of workers.
+
+        Examples:
+            >>> print(f"Current worker count: {executor.worker_count_current}")
+        """
+        return len(getattr(self, f"_{self._workers}"))
+
+    def __init_mixin__(self):
+        self.sync_mode = self._max_workers == 0
+        self.__super_submit = super().submit
+        register_executor(self)
+
+    async def _debug_daemon(self, fut: asyncio.Future, fn, *args, **kwargs) -> None:
+        """
+        Runs until manually cancelled by the finished work item.
+
+        This code will only run if `self.logger.isEnabledFor(logging.DEBUG)` is True. You do not need to include any level checks in your custom implementations.
+
+        Args:
+            fut: The future being debugged.
+            fn: The function being executed.
+            *args: Positional arguments for the function.
+            **kwargs: Keyword arguments for the function.
+
+        See Also:
+            - :meth:`_start_debug_daemon` to start the debug daemon.
+        """
+        # TODO: make prettier strings for other types
+        if type(fn).__name__ == "function":
+            fnid = getattr(fn, "__qualname__", fn.__name__)
+            if fn.__module__:
+                fnid = f"{fn.__module__}.{fnid}"
+        else:
+            fnid = fn
+
+        msg = f"%s processing %s{args}"
+        if kwargs:
+            msg = f"{msg[:-1]} {', '.join(f'{k}={v}' for k, v in kwargs.items())})"
+        else:
+            msg = f"{msg[:-2]})"
+
+        done = fut.done
+        log_debug = self.logger.debug
+
+        while not done():
+            await sleep(15)
+            if not done():
+                log_debug(msg, self, fnid)
+
+
+# Process
+
+
+class AsyncProcessPoolExecutor(_AsyncExecutorMixin, concurrent.futures.ProcessPoolExecutor):
+    """
+    A :class:`concurrent.futures.ProcessPoolExecutor' subclass providing asynchronous
+    run and submit methods that support kwargs, with support for synchronous mode
+
+    Examples:
+        >>> executor = AsyncProcessPoolExecutor(max_workers=4)
+        >>> future = executor.submit(some_function, arg1, arg2, kwarg1='kwarg1')
+        >>> result = await future
+    """
+
+    _workers = "processes"
+    """The type of workers used, set to "processes"."""
+
+    __slots__ = (
+        "_mp_context",
+        "_processes",
+        "_pending_work_items",
+        "_call_queue",
+        "_result_queue",
+        "_queue_management_thread",
+        "_queue_count",
+        "_shutdown_thread",
+        "_work_ids",
+        "_queue_management_thread_wakeup",
+    )
+
+    def __init__(
+        self,
+        max_workers: Optional[int] = None,
+        mp_context: Optional[multiprocessing.context.BaseContext] = None,
+        initializer: Optional[Initializer] = None,
+        initargs: Tuple[Any, ...] = (),
+    ) -> None:
+        """
+        Initializes the AsyncProcessPoolExecutor.
+
+        Args:
+            max_workers: The maximum number of workers. Defaults to None.
+            mp_context: The multiprocessing context. Defaults to None.
+            initializer: An initializer callable. Defaults to None.
+            initargs: Arguments for the initializer. Defaults to ().
+
+        Examples:
+            >>> executor = AsyncProcessPoolExecutor(max_workers=4)
+            >>> future = executor.submit(some_function, arg1, arg2)
+            >>> result = await future
+        """
+        if max_workers == 0:
+            concurrent.futures.ProcessPoolExecutor.__init__(
+                self, 1, mp_context, initializer, initargs
+            )
+            self._max_workers = 0
+        else:
+            concurrent.futures.ProcessPoolExecutor.__init__(
+                self, max_workers, mp_context, initializer, initargs
+            )
+        self.__init_mixin__()
+
+
+# Thread
+
+
+class AsyncThreadPoolExecutor(_AsyncExecutorMixin, concurrent.futures.ThreadPoolExecutor):
+    """
+    A :class:`concurrent.futures.ThreadPoolExecutor' subclass providing asynchronous
+    run and submit methods that support kwargs, with support for synchronous mode
+
+    Examples:
+        >>> executor = AsyncThreadPoolExecutor(max_workers=10, thread_name_prefix="MyThread")
+        >>> future = executor.submit(some_function, arg1, arg2, kwarg1='kwarg1')
+        >>> result = await future
+    """
+
+    _workers = "threads"
+    """The type of workers used, set to "threads"."""
+
+    __slots__ = (
+        "_work_queue",
+        "_idle_semaphore",
+        "_threads",
+        "_shutdown",
+        "_thread_name_prefix",
+    )
+
+    def __init__(
+        self,
+        max_workers: Optional[int] = None,
+        thread_name_prefix: str = "",
+        initializer: Optional[Initializer] = None,
+        initargs: Tuple[Any, ...] = (),
+    ) -> None:
+        """
+        Initializes the AsyncThreadPoolExecutor.
+
+        Args:
+            max_workers: The maximum number of workers. Defaults to None.
+            thread_name_prefix: Prefix for thread names. Defaults to ''.
+            initializer: An initializer callable. Defaults to None.
+            initargs: Arguments for the initializer. Defaults to ().
+
+        Examples:
+            >>> executor = AsyncThreadPoolExecutor(max_workers=10, thread_name_prefix="MyThread")
+            >>> future = executor.submit(some_function, arg1, arg2)
+            >>> result = await future
+        """
+        if max_workers == 0:
+            concurrent.futures.ThreadPoolExecutor.__init__(
+                self, 1, thread_name_prefix, initializer, initargs
+            )
+            self._max_workers = 0
+        else:
+            concurrent.futures.ThreadPoolExecutor.__init__(
+                self, max_workers, thread_name_prefix, initializer, initargs
+            )
+        self.__init_mixin__()
+
+
+AsyncExecutor = Union[AsyncThreadPoolExecutor, AsyncProcessPoolExecutor]
+
+# For backward-compatibility
+ThreadPoolExecutor = AsyncThreadPoolExecutor
+ProcessPoolExecutor = AsyncProcessPoolExecutor
+
+# Pruning thread pool
+
+
+def _worker(executor_reference, work_queue, initializer, initargs, timeout):  # NOTE: NEW 'timeout'
+    """
+    Worker function for the PruningThreadPoolExecutor.
+
+    Args:
+        executor_reference: A weak reference to the executor.
+        work_queue: The work queue.
+        initializer: The initializer function.
+        initargs: Arguments for the initializer.
+        timeout: Timeout duration for pruning inactive threads.
+
+    See Also:
+        - :class:`PruningThreadPoolExecutor` for more details on thread pruning.
+    """
+    if initializer is not None:
+        try:
+            initializer(*initargs)
+        except BaseException:
+            _base.LOGGER.critical("Exception in initializer:", exc_info=True)
+            executor = executor_reference()
+            if executor is not None:
+                executor._initializer_failed()
+            return
+
+    try:
+        while True:
+            try:  # NOTE: NEW
+                work_item = work_queue.get(block=True, timeout=timeout)  # NOTE: NEW
+            except queue.Empty:  # NOTE: NEW
+                # Its been 'timeout' seconds and there are no new work items.  # NOTE: NEW
+                # Let's suicide the thread.  # NOTE: NEW
+                executor = executor_reference()  # NOTE: NEW
+
+                with executor._adjusting_lock:  # NOTE: NEW
+                    # NOTE: We keep a minimum of one thread active to prevent locks
+                    if len(executor) > 1:  # NOTE: NEW
+                        t = threading.current_thread()  # NOTE: NEW
+                        executor._threads.remove(t)  # NOTE: NEW
+                        thread._threads_queues.pop(t)  # NOTE: NEW
+                        # Let the executor know we have one less idle thread available
+                        executor._idle_semaphore.acquire(blocking=False)  # NOTE: NEW
+                        return  # NOTE: NEW
+                continue
+
+            if work_item is not None:
+                work_item.run()
+                # Delete references to object. See issue16284
+                del work_item
+
+                # attempt to increment idle count
+                executor = executor_reference()
+                if executor is not None:
+                    executor._idle_semaphore.release()
+                del executor
+                continue
+
+            executor = executor_reference()
+            # Exit if:
+            #   - The interpreter is shutting down OR
+            #   - The executor that owns the worker has been collected OR
+            #   - The executor that owns the worker has been shutdown OR
+            if thread._shutdown or executor is None or executor._shutdown:
+                # Flag the executor as shutting down as early as possible if it is not gc-ed yet.
+                if executor is not None:
+                    executor._shutdown = True
+                # Notice other workers
+                work_queue.put(None)
+                return
+            del executor
+    except BaseException:
+        _base.LOGGER.critical("Exception in worker", exc_info=True)
+
+
+class PruningThreadPoolExecutor(AsyncThreadPoolExecutor):
+    """
+    This :class:`~AsyncThreadPoolExecutor` implementation prunes inactive threads after 'timeout' seconds without a work item.
+    Pruned threads will be automatically recreated as needed for future workloads. Up to 'max_threads' can be active at any one time.
+    The executor ensures that at least one active thread remains to prevent locks.
+
+    Note:
+        The `_worker` function includes a check (`len(executor) > 1`) to ensure that at least one thread remains active.
+        This prevents the executor from having zero active threads, which could lead to deadlocks.
+
+    Examples:
+        >>> executor = PruningThreadPoolExecutor(max_workers=5, timeout=300)
+        >>> future = executor.submit(some_function, arg1, arg2, kwarg1='kwarg1')
+        >>> result = await future
+    """
+
+    __slots__ = "_timeout", "_adjusting_lock"
+
+    def __init__(
+        self,
+        max_workers=None,
+        thread_name_prefix="",
+        initializer=None,
+        initargs=(),
+        timeout=TEN_MINUTES,
+    ):
+        """
+        Initializes the PruningThreadPoolExecutor.
+
+        Args:
+            max_workers: The maximum number of workers. Defaults to None.
+            thread_name_prefix: Prefix for thread names. Defaults to ''.
+            initializer: An initializer callable. Defaults to None.
+            initargs: Arguments for the initializer. Defaults to ().
+            timeout: Timeout duration for pruning inactive threads. Defaults to TEN_MINUTES.
+
+        Examples:
+            >>> executor = PruningThreadPoolExecutor(max_workers=5, timeout=300)
+            >>> future = executor.submit(some_function, arg1, arg2)
+            >>> result = await future
+        """
+
+        self._timeout = timeout
+        """Timeout duration for pruning inactive threads."""
+
+        self._adjusting_lock = threading.Lock()
+        """Lock used to adjust the number of threads."""
+
+        AsyncThreadPoolExecutor.__init__(
+            self, max_workers, thread_name_prefix, initializer, initargs
+        )
+
+    def __len__(self) -> int:
+        return len(self._threads)
+
+    def _adjust_thread_count(self):
+        """
+        Adjusts the number of threads based on workload and idle threads.
+
+        See Also:
+            - :func:`_worker` for the worker function that handles thread pruning.
+        """
+        with self._adjusting_lock:
+            # if idle threads are available, don't spin new threads
+            if self._idle_semaphore.acquire(timeout=0):
+                return
+
+            # When the executor gets lost, the weakref callback will wake up
+            # the worker threads.
+            def weakref_cb(_, q=self._work_queue):
+                q.put(None)
+
+            num_threads = len(self._threads)
+            if num_threads < self._max_workers:
+                thread_name = "%s_%d" % (self._thread_name_prefix or self, num_threads)
+                t = threading.Thread(
+                    name=thread_name,
+                    target=_worker,
+                    args=(
+                        weakref.ref(self, weakref_cb),
+                        self._work_queue,
+                        self._initializer,
+                        self._initargs,
+                        self._timeout,
+                    ),
+                )
+                t.daemon = True
+                t.start()
+                self._threads.add(t)
+                thread._threads_queues[t] = self._work_queue
+
+
+def _copy_future_state(cf_fut: concurrent.futures.Future, fut: asyncio.Future):
+    """Internal helper to copy state from another Future.
+
+    The other Future may be a concurrent.futures.Future.
+    """
+    # check this again in case it was cancelled since the last check
+    if _fut_is_cancelled(fut):
+        return
+    exception = _get_cf_fut_exception(cf_fut)
+    if exception is None:
+        _set_fut_result(fut, _get_cf_fut_result(cf_fut))
+    else:
+        _set_fut_exception(fut, _convert_future_exc(exception))
+
+
+_fut_is_cancelled = asyncio.Future.cancelled
+_get_cf_fut_result = concurrent.futures.Future.result
+_get_cf_fut_exception = concurrent.futures.Future.exception
+_set_fut_result = asyncio.Future.set_result
+_set_fut_exception = asyncio.Future.set_exception
+_add_done_callback = concurrent.futures.Future.add_done_callback
+
+
+executor = PruningThreadPoolExecutor(128)
+
+__all__ = [
+    "AsyncThreadPoolExecutor",
+    "AsyncProcessPoolExecutor",
+    "PruningThreadPoolExecutor",
+]