lionherd-core 1.0.0a3__py3-none-any.whl

lionherd_core/libs/concurrency/_patterns.py

@@ -0,0 +1,238 @@
+# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Iterable, Sequence
+from typing import TypeVar
+
+import anyio
+
+# ExceptionGroup is built-in in Python 3.11+
+from ._cancel import effective_deadline, move_on_at
+from ._errors import non_cancel_subgroup
+from ._primitives import CapacityLimiter
+from ._task import create_task_group
+from ._utils import current_time
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+__all__ = (
+    "CompletionStream",
+    "bounded_map",
+    "gather",
+    "race",
+    "retry",
+)
+
+
+async def gather(*aws: Awaitable[T], return_exceptions: bool = False) -> list[T | BaseException]:
+    """Run awaitables concurrently."""
+    if not aws:
+        return []
+
+    results: list[T | BaseException | None] = [None] * len(aws)
+
+    async def _runner(idx: int, aw: Awaitable[T]) -> None:
+        try:
+            results[idx] = await aw
+        except BaseException as exc:
+            results[idx] = exc
+            if not return_exceptions:
+                raise  # Propagate to the TaskGroup
+
+    try:
+        async with create_task_group() as tg:
+            for i, aw in enumerate(aws):
+                tg.start_soon(_runner, i, aw)
+    except ExceptionGroup as eg:
+        if not return_exceptions:
+            # Separate cancellations from real failures while preserving structure/tracebacks
+            rest = non_cancel_subgroup(eg)
+            if rest is not None:
+                raise rest
+            # All were cancellations -> propagate cancellation group
+            raise
+
+    return results  # type: ignore
+
+
+async def race(*aws: Awaitable[T]) -> T:
+    """Return first completion."""
+    if not aws:
+        raise ValueError("race() requires at least one awaitable")
+    send, recv = anyio.create_memory_object_stream(0)
+
+    async def _runner(aw: Awaitable[T]) -> None:
+        try:
+            res = await aw
+            await send.send((True, res))
+        except BaseException as exc:
+            await send.send((False, exc))
+
+    async with send, recv, create_task_group() as tg:
+        for aw in aws:
+            tg.start_soon(_runner, aw)
+        ok, payload = await recv.receive()
+        tg.cancel_scope.cancel()
+
+    # Raise outside the TaskGroup context to avoid ExceptionGroup wrapping
+    if ok:
+        return payload  # type: ignore[return-value]
+    raise payload  # type: ignore[misc]
+
+
+async def bounded_map(
+    func: Callable[[T], Awaitable[R]],
+    items: Iterable[T],
+    *,
+    limit: int,
+    return_exceptions: bool = False,
+) -> list[R | BaseException]:
+    """Apply async function to items with concurrency limit."""
+    if limit <= 0:
+        raise ValueError("limit must be >= 1")
+
+    seq = list(items)
+    if not seq:
+        return []
+
+    out: list[R | BaseException | None] = [None] * len(seq)
+    limiter = CapacityLimiter(limit)
+
+    async def _runner(i: int, x: T) -> None:
+        async with limiter:
+            try:
+                out[i] = await func(x)
+            except BaseException as exc:
+                out[i] = exc
+                if not return_exceptions:
+                    raise  # Propagate to the TaskGroup
+
+    try:
+        async with create_task_group() as tg:
+            for i, x in enumerate(seq):
+                tg.start_soon(_runner, i, x)
+    except ExceptionGroup as eg:
+        if not return_exceptions:
+            # Separate cancellations from real failures while preserving structure/tracebacks
+            rest = non_cancel_subgroup(eg)
+            if rest is not None:
+                raise rest
+            # All were cancellations -> propagate cancellation group
+            raise
+
+    return out  # type: ignore
+
+
+class CompletionStream:
+    """Async completion stream with structured concurrency and explicit lifecycle."""
+
+    def __init__(self, aws: Sequence[Awaitable[T]], *, limit: int | None = None):
+        self.aws = aws
+        self.limit = limit
+        self._task_group: anyio.abc.TaskGroup | None = None
+        self._send: anyio.abc.ObjectSendStream[tuple[int, T]] | None = None
+        self._recv: anyio.abc.ObjectReceiveStream[tuple[int, T]] | None = None
+        self._completed_count = 0
+        self._total_count = len(aws)
+
+    async def __aenter__(self):
+        n = len(self.aws)
+        self._send, self._recv = anyio.create_memory_object_stream(n)
+        self._task_group = anyio.create_task_group()
+        await self._task_group.__aenter__()
+
+        limiter = CapacityLimiter(self.limit) if self.limit else None
+
+        async def _runner(i: int, aw: Awaitable[T]) -> None:
+            if limiter:
+                await limiter.acquire()
+            try:
+                res = await aw
+                try:
+                    assert self._send is not None
+                    await self._send.send((i, res))  # type: ignore[arg-type]
+                except anyio.ClosedResourceError:
+                    # Stream was closed (e.g., early break from iteration)
+                    # Swallow the error gracefully
+                    pass
+            finally:
+                if limiter:
+                    limiter.release()
+
+        # Start all tasks
+        for i, aw in enumerate(self.aws):
+            self._task_group.start_soon(_runner, i, aw)
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        # Cancel remaining tasks and clean up
+        if self._task_group:
+            await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
+        if self._send:
+            await self._send.aclose()
+        if self._recv:
+            await self._recv.aclose()
+        return False
+
+    def __aiter__(self):
+        if not self._recv:
+            raise RuntimeError("CompletionStream must be used as async context manager")
+        return self
+
+    async def __anext__(self):
+        if self._completed_count >= self._total_count:
+            raise StopAsyncIteration
+
+        try:
+            result = await self._recv.receive()
+            self._completed_count += 1
+            return result
+        except anyio.EndOfStream:
+            raise StopAsyncIteration
+
+
+async def retry(
+    fn: Callable[[], Awaitable[T]],
+    *,
+    attempts: int = 3,
+    base_delay: float = 0.1,
+    max_delay: float = 2.0,
+    retry_on: tuple[type[BaseException], ...] = (Exception,),
+    jitter: float = 0.1,
+) -> T:
+    """Deadline-aware exponential backoff retry."""
+    attempt = 0
+    deadline = effective_deadline()
+    while True:
+        try:
+            return await fn()
+        except retry_on:
+            attempt += 1
+            if attempt >= attempts:
+                raise
+
+            delay = min(max_delay, base_delay * (2 ** (attempt - 1)))
+            if jitter:
+                import random
+
+                delay *= 1 + random.random() * jitter
+
+            # Cap by ambient deadline if one exists
+            if deadline is not None:
+                remaining = deadline - current_time()
+                if remaining <= 0:
+                    # Out of time; surface the last error
+                    raise
+                # Use move_on_at to avoid TOCTOU race between deadline check and sleep
+                with move_on_at(deadline):
+                    await anyio.sleep(delay)
+                # If we were cancelled by deadline, re-raise the original exception
+                if current_time() >= deadline:
+                    raise
+            else:
+                await anyio.sleep(delay)

lionherd_core/libs/concurrency/_primitives.py

@@ -0,0 +1,253 @@
+# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Generic, Self, TypeVar
+
+import anyio
+import anyio.abc
+
+T = TypeVar("T")
+
+
+__all__ = (
+    "CapacityLimiter",
+    "Condition",
+    "Event",
+    "Lock",
+    "Queue",
+    "Semaphore",
+)
+
+
+class Lock:
+    """Async mutex lock."""
+
+    __slots__ = ("_lock",)
+
+    def __init__(self) -> None:
+        self._lock = anyio.Lock()
+
+    async def acquire(self) -> None:
+        """Acquire lock."""
+        await self._lock.acquire()
+
+    def release(self) -> None:
+        """Release lock."""
+        self._lock.release()
+
+    async def __aenter__(self) -> Self:
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        self.release()
+
+
+class Semaphore:
+    """Async semaphore."""
+
+    __slots__ = ("_sem",)
+
+    def __init__(self, initial_value: int) -> None:
+        if initial_value < 0:
+            raise ValueError("initial_value must be >= 0")
+        self._sem = anyio.Semaphore(initial_value)
+
+    async def acquire(self) -> None:
+        """Acquire slot."""
+        await self._sem.acquire()
+
+    def release(self) -> None:
+        """Release slot."""
+        self._sem.release()
+
+    async def __aenter__(self) -> Self:
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        self.release()
+
+
+class CapacityLimiter:
+    """Async capacity limiter."""
+
+    __slots__ = ("_lim",)
+
+    def __init__(self, total_tokens: float) -> None:
+        """Initialize capacity limiter."""
+        if total_tokens <= 0:
+            raise ValueError("total_tokens must be > 0")
+        self._lim = anyio.CapacityLimiter(total_tokens)
+
+    async def acquire(self) -> None:
+        """Acquire capacity."""
+        await self._lim.acquire()
+
+    def release(self) -> None:
+        """Release capacity."""
+        self._lim.release()
+
+    @property
+    def remaining_tokens(self) -> float:
+        """Available capacity (deprecated)."""
+        return self._lim.available_tokens
+
+    @property
+    def total_tokens(self) -> float:
+        """Get capacity limit."""
+        return self._lim.total_tokens
+
+    @total_tokens.setter
+    def total_tokens(self, value: float) -> None:
+        """Set capacity limit."""
+        if value <= 0:
+            raise ValueError("total_tokens must be > 0")
+        self._lim.total_tokens = value
+
+    @property
+    def borrowed_tokens(self) -> float:
+        """Get borrowed tokens."""
+        return self._lim.borrowed_tokens
+
+    @property
+    def available_tokens(self) -> float:
+        """Get available tokens."""
+        return self._lim.available_tokens
+
+    async def acquire_on_behalf_of(self, borrower: object) -> None:
+        """Acquire for borrower."""
+        await self._lim.acquire_on_behalf_of(borrower)
+
+    def release_on_behalf_of(self, borrower: object) -> None:
+        """Release for borrower."""
+        self._lim.release_on_behalf_of(borrower)
+
+    # Support idiomatic AnyIO usage: `async with limiter: ...`
+    async def __aenter__(self) -> Self:
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        self.release()
+
+
+@dataclass(slots=True)
+class Queue(Generic[T]):
+    """Async FIFO queue."""
+
+    _send: anyio.abc.ObjectSendStream[T]
+    _recv: anyio.abc.ObjectReceiveStream[T]
+
+    @classmethod
+    def with_maxsize(cls, maxsize: int) -> Queue[T]:
+        """Create queue with maxsize."""
+        send, recv = anyio.create_memory_object_stream(maxsize)
+        return cls(send, recv)
+
+    async def put(self, item: T) -> None:
+        """Put item."""
+        await self._send.send(item)
+
+    def put_nowait(self, item: T) -> None:
+        """Put without blocking."""
+        self._send.send_nowait(item)  # type: ignore[attr-defined]
+
+    async def get(self) -> T:
+        """Get item."""
+        return await self._recv.receive()
+
+    def get_nowait(self) -> T:
+        """Get without blocking."""
+        return self._recv.receive_nowait()  # type: ignore[attr-defined]
+
+    async def close(self) -> None:
+        """Close streams."""
+        await self._send.aclose()
+        await self._recv.aclose()
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        await self.close()
+
+    @property
+    def sender(self) -> anyio.abc.ObjectSendStream[T]:
+        """Get send stream."""
+        return self._send
+
+    @property
+    def receiver(self) -> anyio.abc.ObjectReceiveStream[T]:
+        """Get receive stream."""
+        return self._recv
+
+
+class Event:
+    """Async event for task signaling."""
+
+    __slots__ = ("_event",)
+
+    def __init__(self) -> None:
+        self._event = anyio.Event()
+
+    def set(self) -> None:
+        """Set flag."""
+        self._event.set()
+
+    def is_set(self) -> bool:
+        """Check if set."""
+        return self._event.is_set()
+
+    async def wait(self) -> None:
+        """Wait for flag."""
+        await self._event.wait()
+
+    def statistics(self) -> anyio.EventStatistics:
+        """Get statistics."""
+        return self._event.statistics()
+
+
+class Condition:
+    """Async condition variable."""
+
+    __slots__ = ("_condition",)
+
+    def __init__(self, lock: Lock | None = None) -> None:
+        """Initialize condition variable."""
+        _lock = lock._lock if lock else None
+        self._condition = anyio.Condition(_lock)
+
+    async def acquire(self) -> None:
+        """Acquire lock."""
+        await self._condition.acquire()
+
+    def release(self) -> None:
+        """Release lock."""
+        self._condition.release()
+
+    async def __aenter__(self) -> Self:
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        self.release()
+
+    async def wait(self) -> None:
+        """Wait until notified."""
+        await self._condition.wait()
+
+    def notify(self, n: int = 1) -> None:
+        """Wake n tasks."""
+        self._condition.notify(n)
+
+    def notify_all(self) -> None:
+        """Wake all tasks."""
+        self._condition.notify_all()
+
+    def statistics(self) -> anyio.ConditionStatistics:
+        """Get statistics."""
+        return self._condition.statistics()

lionherd_core/libs/concurrency/_priority_queue.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import heapq
+from typing import Generic, TypeVar
+
+from ._primitives import Condition
+
+T = TypeVar("T")
+
+__all__ = ("PriorityQueue", "QueueEmpty", "QueueFull")
+
+
+class QueueEmpty(Exception):  # noqa: N818
+    """Exception raised when queue.get_nowait() is called on empty queue."""
+
+
+class QueueFull(Exception):  # noqa: N818
+    """Exception raised when queue.put_nowait() is called on full queue."""
+
+
+class PriorityQueue(Generic[T]):
+    """Async priority queue (heapq + anyio.Condition).
+
+    API: Similar to asyncio.PriorityQueue, but nowait methods are async.
+
+    Attributes:
+        maxsize: Maximum queue size (0 = unlimited)
+    """
+
+    def __init__(self, maxsize: int = 0):
+        """Initialize priority queue.
+
+        Args:
+            maxsize: Max size (0 = unlimited)
+        """
+        self.maxsize = maxsize
+        self._queue: list[T] = []
+        self._condition = Condition()
+
+    async def put(self, item: T) -> None:
+        """Put item into queue (blocks if full).
+
+        Args:
+            item: Item (tuple with priority as first element)
+        """
+        async with self._condition:
+            # Wait if queue is full
+            while self.maxsize > 0 and len(self._queue) >= self.maxsize:
+                await self._condition.wait()
+
+            heapq.heappush(self._queue, item)
+            self._condition.notify()
+
+    async def put_nowait(self, item: T) -> None:
+        """Put item without blocking (async, unlike asyncio).
+
+        Args:
+            item: Item (tuple with priority as first element)
+
+        Raises:
+            QueueFull: If queue is at maxsize
+        """
+        async with self._condition:
+            if self.maxsize > 0 and len(self._queue) >= self.maxsize:
+                raise QueueFull("Queue is full")
+
+            heapq.heappush(self._queue, item)
+            # Notify waiting getters that item is available
+            self._condition.notify()
+
+    async def get(self) -> T:
+        """Get highest priority item (blocks if empty).
+
+        Returns:
+            Highest priority item (lowest value first)
+        """
+        async with self._condition:
+            # Wait if queue is empty
+            while not self._queue:
+                await self._condition.wait()
+
+            item = heapq.heappop(self._queue)
+            self._condition.notify()
+            return item
+
+    async def get_nowait(self) -> T:
+        """Get item without blocking (async, unlike asyncio).
+
+        Returns:
+            Highest priority item
+
+        Raises:
+            QueueEmpty: If queue is empty
+        """
+        async with self._condition:
+            if not self._queue:
+                raise QueueEmpty("Queue is empty")
+
+            item = heapq.heappop(self._queue)
+            # Notify waiting putters that space is available
+            self._condition.notify()
+            return item
+
+    def qsize(self) -> int:
+        """Approximate queue size (unlocked, racy).
+
+        Note: Value may be stale immediately. Use for monitoring only.
+
+        Returns:
+            Number of items in queue
+        """
+        return len(self._queue)
+
+    def empty(self) -> bool:
+        """Check if queue is empty (unlocked, racy).
+
+        Note: Value may be stale immediately. Use for monitoring only.
+
+        Returns:
+            True if queue is empty
+        """
+        return len(self._queue) == 0
+
+    def full(self) -> bool:
+        """Check if queue is full (unlocked, racy).
+
+        Note: Value may be stale immediately. Use for monitoring only.
+
+        Returns:
+            True if queue is at maxsize
+        """
+        return self.maxsize > 0 and len(self._queue) >= self.maxsize

lionherd_core/libs/concurrency/_resource_tracker.py

@@ -0,0 +1,66 @@
+# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import threading
+import time
+import weakref
+from dataclasses import dataclass
+
+__all__ = (
+    "LeakInfo",
+    "LeakTracker",
+    "track_resource",
+    "untrack_resource",
+)
+
+
+@dataclass(frozen=True, slots=True)
+class LeakInfo:
+    name: str
+    kind: str | None
+    created_at: float
+
+
+class LeakTracker:
+    """Track live objects for leak detection."""
+
+    def __init__(self) -> None:
+        self._live: dict[int, LeakInfo] = {}
+        self._lock = threading.Lock()
+
+    def track(self, obj: object, *, name: str | None, kind: str | None) -> None:
+        info = LeakInfo(name=name or f"obj-{id(obj)}", kind=kind, created_at=time.time())
+        key = id(obj)
+
+        def _finalizer(_key: int = key) -> None:
+            with self._lock:
+                self._live.pop(_key, None)
+
+        with self._lock:
+            self._live[key] = info
+        weakref.finalize(obj, _finalizer)
+
+    def untrack(self, obj: object) -> None:
+        with self._lock:
+            self._live.pop(id(obj), None)
+
+    def live(self) -> list[LeakInfo]:
+        with self._lock:
+            return list(self._live.values())
+
+    def clear(self) -> None:
+        with self._lock:
+            self._live.clear()
+
+
+_TRACKER = LeakTracker()
+
+
+def track_resource(obj: object, name: str | None = None, kind: str | None = None) -> None:
+    _TRACKER.track(obj, name=name, kind=kind)
+
+
+def untrack_resource(obj: object) -> None:
+    _TRACKER.untrack(obj)