tigrcorn-protocols 0.3.16.dev5__py3-none-any.whl

tigrcorn_protocols/scheduler/dispatch.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable
+
+from .policy import SchedulerPolicy
+
+
+class TaskDispatcher:
+    def __init__(self, policy: SchedulerPolicy | None = None) -> None:
+        self.policy = policy or SchedulerPolicy()
+        self.tasks: set[asyncio.Task] = set()
+
+    def spawn(self, coro: Awaitable):
+        if len(self.tasks) >= self.policy.max_tasks:
+            close = getattr(coro, 'close', None)
+            if callable(close):
+                close()
+            raise RuntimeError('task quota exceeded')
+        task = asyncio.create_task(coro)
+        self.tasks.add(task)
+        task.add_done_callback(self.tasks.discard)
+        return task
+
+
+async def spawn(coro: Awaitable):
+    return asyncio.create_task(coro)

tigrcorn_protocols/scheduler/fairness.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass, field
+from typing import Generic, TypeVar
+
+T = TypeVar('T')
+
+
+@dataclass(slots=True)
+class FairnessPolicy(Generic[T]):
+    round_robin: bool = True
+    _queue: deque[T] = field(default_factory=deque)
+
+    def push(self, item: T) -> None:
+        self._queue.append(item)
+
+    def pop(self) -> T | None:
+        if not self._queue:
+            return None
+        return self._queue.popleft()

tigrcorn_protocols/scheduler/policy.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class SchedulerPolicy:
+    max_connections: int = 10_000
+    max_tasks: int = 50_000
+    max_streams_per_session: int = 128
+    limit_concurrency: int | None = None
+    drain_on_shutdown: bool = True

tigrcorn_protocols/scheduler/priorities.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True, order=True)
+class Priority:
+    value: int = 0

tigrcorn_protocols/scheduler/quotas.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class Quotas:
+    max_connections: int = 10_000
+    max_streams_per_connection: int = 128
+    current_connections: int = 0
+
+    def acquire_connection(self) -> bool:
+        if self.current_connections >= self.max_connections:
+            return False
+        self.current_connections += 1
+        return True
+
+    def release_connection(self) -> None:
+        self.current_connections = max(0, self.current_connections - 1)

tigrcorn_protocols/scheduler/runtime.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable
+from dataclasses import dataclass
+from typing import Any
+
+from .dispatch import TaskDispatcher
+from .policy import SchedulerPolicy
+from .quotas import Quotas
+from .tasks import TaskSet
+
+
+@dataclass(slots=True)
+class ConnectionLease:
+    scheduler: "ProductionScheduler"
+    released: bool = False
+
+    def release(self) -> None:
+        if self.released:
+            return
+        self.scheduler.release_connection()
+        self.released = True
+
+    def __enter__(self) -> "ConnectionLease":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.release()
+
+
+@dataclass(slots=True)
+class WorkLease:
+    scheduler: "ProductionScheduler"
+    released: bool = False
+
+    def release(self) -> None:
+        if self.released:
+            return
+        self.scheduler.release_work()
+        self.released = True
+
+    def __enter__(self) -> "WorkLease":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.release()
+
+
+class ProductionScheduler:
+    """Package-owned runtime scheduler for connection admission and task draining.
+
+    The scheduler keeps protocol code out of ad-hoc concurrency decisions. It owns:
+    - connection quotas
+    - task quotas
+    - global in-flight work admission (`limit_concurrency`)
+    - graceful shutdown / drain behavior
+    - task tracking for server-internal relay work
+    """
+
+    def __init__(self, policy: SchedulerPolicy | None = None) -> None:
+        self.policy = policy or SchedulerPolicy()
+        self.dispatcher = TaskDispatcher(self.policy)
+        self.quotas = Quotas(
+            max_connections=self.policy.max_connections,
+            max_streams_per_connection=self.policy.max_streams_per_session,
+        )
+        self.tasks = TaskSet()
+        self._closed = False
+        self._draining = False
+        self._owners: dict[asyncio.Task[Any], str | None] = {}
+        self._inflight = 0
+
+    @property
+    def open_connections(self) -> int:
+        return self.quotas.current_connections
+
+    @property
+    def active_tasks(self) -> int:
+        return len(self.dispatcher.tasks)
+
+    @property
+    def current_inflight(self) -> int:
+        return self._inflight
+
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    def _can_admit_work(self) -> bool:
+        if self._closed or self._draining:
+            return False
+        limit = self.policy.limit_concurrency
+        return limit is None or self._inflight < limit
+
+    def acquire_connection(self) -> ConnectionLease | None:
+        if self._closed or self._draining:
+            return None
+        if not self.quotas.acquire_connection():
+            return None
+        return ConnectionLease(self)
+
+    def release_connection(self) -> None:
+        self.quotas.release_connection()
+
+    def acquire_work(self) -> WorkLease | None:
+        if not self._can_admit_work():
+            return None
+        self._inflight += 1
+        return WorkLease(self)
+
+    def release_work(self) -> None:
+        self._inflight = max(0, self._inflight - 1)
+
+    def spawn(self, coro: Awaitable[Any], *, owner: str | None = None) -> asyncio.Task[Any]:
+        if self._closed or self._draining:
+            close = getattr(coro, 'close', None)
+            if callable(close):
+                close()
+            raise RuntimeError('scheduler is closed')
+        lease = self.acquire_work()
+        if lease is None:
+            close = getattr(coro, 'close', None)
+            if callable(close):
+                close()
+            raise RuntimeError('concurrency limit exceeded')
+        try:
+            task = self.dispatcher.spawn(coro)
+        except Exception:
+            lease.release()
+            raise
+        self.tasks.add(task)
+        self._owners[task] = owner
+        task.add_done_callback(self._owners.pop)
+        task.add_done_callback(lambda _task: lease.release())
+        return task
+
+    async def wait(self) -> None:
+        if not self.dispatcher.tasks:
+            return
+        await asyncio.gather(*list(self.dispatcher.tasks), return_exceptions=True)
+
+    async def drain(self, *, cancel_running: bool | None = None) -> None:
+        if self._closed:
+            return
+        self._draining = True
+        if cancel_running is None:
+            cancel_running = not self.policy.drain_on_shutdown
+        if cancel_running:
+            await self.tasks.cancel_all()
+        else:
+            await self.wait()
+        self._closed = True
+
+    async def close(self) -> None:
+        await self.drain()

tigrcorn_protocols/scheduler/tasks.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import asyncio
+from contextlib import suppress
+from dataclasses import dataclass, field
+
+from tigrcorn_protocols.scheduler.cancellation import CancellationResult, cancel_many_bounded
+
+
+@dataclass(slots=True)
+class TaskSet:
+    tasks: set[asyncio.Task] = field(default_factory=set)
+
+    def add(self, task: asyncio.Task) -> None:
+        self.tasks.add(task)
+        task.add_done_callback(self.tasks.discard)
+
+    async def cancel_all(self) -> None:
+        for task in list(self.tasks):
+            task.cancel()
+        for task in list(self.tasks):
+            with suppress(asyncio.CancelledError):
+                await task
+
+    async def cancel_all_bounded(self, *, timeout: float) -> CancellationResult:
+        return await cancel_many_bounded(list(self.tasks), timeout=timeout)
+
+
+async def cancel_tasks(tasks: list[asyncio.Task]) -> None:
+    taskset = TaskSet(set(tasks))
+    await taskset.cancel_all()

tigrcorn_protocols/sessions/__init__.py

@@ -0,0 +1 @@
+"""Session models."""

tigrcorn_protocols/sessions/base.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from time import monotonic
+
+
+@dataclass(slots=True)
+class BaseSession:
+    session_id: int
+    opened_at: float = field(default_factory=monotonic)
+    protocol: str = 'unknown'
+    closed_at: float | None = None
+
+    def close(self) -> None:
+        if self.closed_at is None:
+            self.closed_at = monotonic()

tigrcorn_protocols/sessions/connection.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .base import BaseSession
+
+
+@dataclass(slots=True)
+class ConnectionSession(BaseSession):
+    protocol: str = 'tcp'
+    peer: tuple[str | None, int | None] | None = None
+    server: tuple[str | None, int | None] | None = None

tigrcorn_protocols/sessions/limits.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class SessionLimits:
+    max_streams: int = 128
+    max_inflight_bytes: int = 1_048_576
+
+    def allow_stream(self, current: int) -> bool:
+        return current < self.max_streams

tigrcorn_protocols/sessions/manager.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+
+from tigrcorn_core.utils.ids import next_id
+
+from .base import BaseSession
+
+
+@dataclass(slots=True)
+class SessionManager:
+    sessions: dict[int, BaseSession] = field(default_factory=dict)
+    counts: dict[str, int] = field(default_factory=lambda: defaultdict(int))
+
+    def open(self, session: BaseSession | None = None, *, protocol: str = 'unknown') -> BaseSession:
+        if session is None:
+            session = BaseSession(session_id=next_id(), protocol=protocol)
+        self.sessions[session.session_id] = session
+        self.counts[session.protocol] += 1
+        return session
+
+    def close(self, session_id: int) -> None:
+        session = self.sessions.pop(session_id, None)
+        if session is None:
+            return
+        session.close()
+        self.counts[session.protocol] = max(0, self.counts[session.protocol] - 1)
+
+    def snapshot(self) -> dict[str, int]:
+        return dict(self.counts)

tigrcorn_protocols/sessions/metadata.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class SessionMetadata:
+    listener_name: str = 'default'
+    transport: str = 'tcp'
+    label: str = ''

tigrcorn_protocols/sessions/quic.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .base import BaseSession
+
+
+@dataclass(slots=True)
+class QuicSession(BaseSession):
+    protocol: str = 'quic'
+    stream_count: int = 0
+
+    def opened_stream(self) -> None:
+        self.stream_count += 1

tigrcorn_protocols/streams/__init__.py

@@ -0,0 +1 @@
+"""Logical stream models."""

tigrcorn_protocols/streams/base.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class LogicalStream:
+    stream_id: int
+    multiplexed: bool = False
+    open: bool = True
+
+    def close(self) -> None:
+        self.open = False

tigrcorn_protocols/streams/ids.py

@@ -0,0 +1,5 @@
+from tigrcorn_core.utils.ids import next_id
+
+
+def next_stream_id() -> int:
+    return next_id()

tigrcorn_protocols/streams/multiplex.py

@@ -0,0 +1,6 @@
+from .base import LogicalStream
+
+
+class MultiplexStream(LogicalStream):
+    def __init__(self, stream_id: int) -> None:
+        super().__init__(stream_id=stream_id, multiplexed=True)

tigrcorn_protocols/streams/registry.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from .base import LogicalStream
+
+
+@dataclass(slots=True)
+class StreamRegistry:
+    streams: dict[int, LogicalStream] = field(default_factory=dict)
+
+    def add(self, stream: LogicalStream) -> LogicalStream:
+        self.streams[stream.stream_id] = stream
+        return stream
+
+    def get(self, stream_id: int) -> LogicalStream | None:
+        return self.streams.get(stream_id)
+
+    def close(self, stream_id: int) -> None:
+        stream = self.streams.pop(stream_id, None)
+        if stream is not None:
+            stream.close()

tigrcorn_protocols/streams/singleplex.py

@@ -0,0 +1,6 @@
+from .base import LogicalStream
+
+
+class SingleplexStream(LogicalStream):
+    def __init__(self, stream_id: int = 1) -> None:
+        super().__init__(stream_id=stream_id, multiplexed=False)

tigrcorn_protocols/websocket/__init__.py

@@ -0,0 +1 @@
+__all__ = []

tigrcorn_protocols/websocket/codec.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from tigrcorn_protocols.websocket.frames import (
+    OP_BINARY,
+    OP_CLOSE,
+    OP_PING,
+    OP_PONG,
+    OP_TEXT,
+    encode_close_payload,
+    serialize_frame,
+)
+
+
+def text_frame(text: str, *, rsv1: bool = False) -> bytes:
+    return serialize_frame(OP_TEXT, text.encode('utf-8'), rsv1=rsv1)
+
+
+def binary_frame(data: bytes, *, rsv1: bool = False) -> bytes:
+    return serialize_frame(OP_BINARY, data, rsv1=rsv1)
+
+
+def ping_frame(data: bytes = b'') -> bytes:
+    return serialize_frame(OP_PING, data)
+
+
+def pong_frame(data: bytes = b'') -> bytes:
+    return serialize_frame(OP_PONG, data)
+
+
+def close_frame(code: int = 1000, reason: str = '') -> bytes:
+    return serialize_frame(OP_CLOSE, encode_close_payload(code, reason))