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

tigrcorn_protocols/http3/websocket.py

@@ -0,0 +1,360 @@
+from __future__ import annotations
+
+import asyncio
+from contextlib import suppress
+from typing import Awaitable, Callable
+
+from tigrcorn_protocols.flow.keepalive import KeepAlivePolicy, KeepAliveRuntime
+from tigrcorn_observability.metrics import Metrics
+
+from tigrcorn_asgi.events.websocket import websocket_connect, websocket_disconnect, websocket_receive_bytes, websocket_receive_text
+from tigrcorn_asgi.receive import QueueReceive
+from tigrcorn_asgi.scopes.websocket import build_websocket_scope
+from tigrcorn_config.model import ServerConfig
+from tigrcorn_core.errors import ProtocolError
+from tigrcorn_protocols.http1.parser import ParsedRequest
+from tigrcorn_protocols.websocket.codec import binary_frame, close_frame, pong_frame, text_frame
+from tigrcorn_protocols.websocket.extensions import PerMessageDeflateRuntime, default_permessage_deflate_agreement, negotiate_permessage_deflate, parse_permessage_deflate_offers
+from tigrcorn_protocols.websocket.frames import OP_BINARY, OP_CLOSE, OP_CONT, OP_PING, OP_PONG, OP_TEXT, decode_close_payload, parse_frame_bytes, serialize_frame
+from tigrcorn_core.types import ASGIApp
+from tigrcorn_core.utils.headers import get_header
+
+
+class H3WebSocketSession:
+    def __init__(
+        self,
+        *,
+        app: ASGIApp,
+        config: ServerConfig,
+        request: ParsedRequest,
+        client: tuple[str, int] | None,
+        server: tuple[str, int] | tuple[str, None] | None,
+        scheme: str,
+        send_headers: Callable[[int, list[tuple[bytes, bytes]], bool], Awaitable[None]],
+        send_data: Callable[[bytes, bool], Awaitable[None]],
+        metrics: Metrics | None = None,
+        on_close: Callable[[], None] | None = None,
+    ) -> None:
+        self.app = app
+        self.config = config
+        self.request = request
+        self.client = client
+        self.server = server
+        self.scheme = 'wss' if scheme == 'https' else 'ws'
+        self.send_headers = send_headers
+        self.send_data = send_data
+        self.metrics = metrics
+        self.on_close = on_close
+        self.receive = QueueReceive(max_size=self.config.websocket.max_queue)
+        self.task: asyncio.Task[None] | None = None
+        self.accepted = False
+        self.closed = False
+        self.http_denied = False
+        self.http_denial_status = 403
+        self.http_denial_headers: list[tuple[bytes, bytes]] = []
+        self.http_denial_started = False
+        self.subprotocols = build_websocket_scope(request, client=client, server=server, scheme=self.scheme, root_path=self.config.proxy.root_path, proxy=self.config.proxy)['subprotocols']
+        self.buffer = bytearray()
+        self.peer_end_stream_pending = False
+        self.fragmented_opcode: int | None = None
+        self.fragments: list[bytes] = []
+        self.current_message_size = 0
+        self.fragmented_compressed = False
+        self.permessage_deflate_offers = parse_permessage_deflate_offers(request.headers)
+        self.permessage_deflate_runtime: PerMessageDeflateRuntime | None = None
+        self.keepalive_policy = KeepAlivePolicy(
+            idle_timeout=self.config.http.idle_timeout,
+            ping_interval=self.config.websocket.ping_interval,
+            ping_timeout=self.config.websocket.ping_timeout,
+        )
+        self.keepalive = KeepAliveRuntime(self.keepalive_policy) if self.keepalive_policy.enabled else None
+        self.keepalive_task: asyncio.Task[None] | None = None
+        version = get_header(request.headers, b'sec-websocket-version')
+        if version != b'13':
+            raise ProtocolError('unsupported websocket version')
+
+    async def start(self) -> None:
+        scope = build_websocket_scope(self.request, client=self.client, server=self.server, scheme=self.scheme, root_path=self.config.proxy.root_path, proxy=self.config.proxy)
+        await self.receive.put(websocket_connect())
+        self.task = asyncio.create_task(self._run_app(scope), name=f'tigrcorn-h3-ws-{self.request.path}')
+        if self.keepalive is not None:
+            self.keepalive_task = asyncio.create_task(self._keepalive_loop(), name=f'tigrcorn-h3-ws-keepalive-{self.request.path}')
+
+    def _record_activity(self) -> None:
+        if self.keepalive is not None:
+            self.keepalive.record_activity()
+
+    def _notify_closed(self) -> None:
+        if self.on_close is not None:
+            callback = self.on_close
+            self.on_close = None
+            callback()
+
+    async def _keepalive_loop(self) -> None:
+        while not self.closed:
+            await asyncio.sleep(0.05)
+            if self.keepalive is None or self.closed:
+                return
+            if self.keepalive.ping_timed_out():
+                if self.metrics is not None:
+                    self.metrics.websocket_ping_timeout()
+                if not self.closed:
+                    await self.send_data(close_frame(1011, 'ping timeout'), True)
+                self.closed = True
+                self._notify_closed()
+                await self.receive.put(websocket_disconnect(1011, 'ping timeout'))
+                return
+            payload = self.keepalive.next_ping_payload()
+            if payload is None:
+                continue
+            if self.metrics is not None:
+                self.metrics.websocket_ping_sent()
+            await self.send_data(serialize_frame(OP_PING, payload), False)
+
+    async def _run_app(self, scope: dict) -> None:
+        try:
+            await self.app(scope, self.receive, self._send)
+        except Exception:
+            if self.accepted and not self.closed:
+                with suppress(Exception):
+                    await self.send_data(close_frame(1011, 'internal error'), True)
+            raise
+        finally:
+            if self.http_denied and not self.closed:
+                if not self.http_denial_started:
+                    await self.send_headers(self.http_denial_status, self.http_denial_headers, True)
+                    self.http_denial_started = True
+                self.closed = True
+            elif not self.accepted and not self.closed:
+                await self.send_headers(403, [], True)
+                self.closed = True
+            elif self.accepted and not self.closed:
+                await self.send_data(close_frame(1000, ''), True)
+                self.closed = True
+            self._notify_closed()
+            if self.keepalive_task is not None:
+                self.keepalive_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await self.keepalive_task
+
+    async def _send(self, message: dict) -> None:
+        typ = message['type']
+        if typ == 'websocket.accept':
+            if self.accepted or self.http_denied:
+                raise RuntimeError('websocket.accept sent more than once')
+            subprotocol = message.get('subprotocol')
+            if subprotocol is not None and subprotocol not in self.subprotocols:
+                raise RuntimeError('websocket.accept selected a subprotocol not offered by the client')
+            headers = [(bytes(k).lower(), bytes(v)) for k, v in message.get('headers', [])]
+            if self.config.websocket.compression != 'permessage-deflate':
+                headers = [(k, v) for k, v in headers if k != b'sec-websocket-extensions']
+            elif self.permessage_deflate_offers and get_header(headers, b'sec-websocket-extensions') is None:
+                default_agreement = default_permessage_deflate_agreement(self.permessage_deflate_offers)
+                if default_agreement is not None:
+                    headers = headers + [(b'sec-websocket-extensions', default_agreement.as_header_value())]
+            response_headers = [(k, v) for k, v in headers if k not in {b'sec-websocket-extensions', b'sec-websocket-protocol'}]
+            agreement = negotiate_permessage_deflate(
+                request_headers=self.request.headers,
+                response_headers=headers,
+            )
+            if agreement is not None:
+                self.permessage_deflate_runtime = PerMessageDeflateRuntime(agreement)
+                response_headers.append((b'sec-websocket-extensions', agreement.as_header_value()))
+            if subprotocol is not None:
+                response_headers.append((b'sec-websocket-protocol', subprotocol.encode('ascii')))
+            await self.send_headers(200, response_headers, False)
+            self.accepted = True
+            self._record_activity()
+            if self.buffer or self.peer_end_stream_pending:
+                pending_end_stream = self.peer_end_stream_pending
+                self.peer_end_stream_pending = False
+                await self.feed_data(b'', end_stream=pending_end_stream)
+            return
+        if typ == 'websocket.send':
+            if not self.accepted:
+                raise RuntimeError('websocket.send before websocket.accept')
+            if self.closed:
+                return
+            text = message.get('text')
+            data = message.get('bytes')
+            if text is not None and data is not None:
+                raise RuntimeError('websocket.send cannot contain both text and bytes')
+            if text is not None:
+                payload = text.encode('utf-8')
+                if self.permessage_deflate_runtime is not None:
+                    await self.send_data(serialize_frame(OP_TEXT, self.permessage_deflate_runtime.compress_message(payload), rsv1=True), False)
+                else:
+                    await self.send_data(text_frame(text), False)
+                self._record_activity()
+            else:
+                raw = data or b''
+                if self.permessage_deflate_runtime is not None:
+                    await self.send_data(binary_frame(self.permessage_deflate_runtime.compress_message(raw), rsv1=True), False)
+                else:
+                    await self.send_data(binary_frame(raw), False)
+            self._record_activity()
+            return
+        if typ == 'websocket.close':
+            code = int(message.get('code', 1000))
+            reason = message.get('reason', '')
+            if not self.accepted:
+                self.http_denied = True
+                self.http_denial_status = 403
+                self.http_denial_headers = []
+                return
+            if not self.closed:
+                await self.send_data(close_frame(code, reason), True)
+                self.closed = True
+                self._notify_closed()
+            return
+        if typ == 'websocket.http.response.start':
+            if self.accepted:
+                raise RuntimeError('cannot deny websocket after accept')
+            self.http_denied = True
+            self.http_denial_status = int(message['status'])
+            self.http_denial_headers = list(message.get('headers', []))
+            return
+        if typ == 'websocket.http.response.body':
+            if not self.http_denied:
+                raise RuntimeError('websocket.http.response.body before denial start')
+            body = bytes(message.get('body', b''))
+            more = bool(message.get('more_body', False))
+            if not self.http_denial_started:
+                headers = list(self.http_denial_headers)
+                if not more:
+                    headers.append((b'content-length', str(len(body)).encode('ascii')))
+                end_stream = (not body) and (not more)
+                await self.send_headers(self.http_denial_status, headers, end_stream)
+                self.http_denial_started = True
+                if end_stream:
+                    self.closed = True
+                    return
+            if body or not more:
+                await self.send_data(body, not more)
+            if not more:
+                self.closed = True
+            return
+        raise RuntimeError(f'unexpected websocket send message: {typ!r}')
+
+    def _frame_length(self, data: bytes) -> int | None:
+        if len(data) < 2:
+            return None
+        masked = bool(data[1] & 0x80)
+        length = data[1] & 0x7F
+        pos = 2
+        if length == 126:
+            if len(data) < pos + 2:
+                return None
+            length = int.from_bytes(data[pos:pos + 2], 'big')
+            pos += 2
+        elif length == 127:
+            if len(data) < pos + 8:
+                return None
+            length = int.from_bytes(data[pos:pos + 8], 'big')
+            pos += 8
+        if masked:
+            pos += 4
+        if len(data) < pos + length:
+            return None
+        return pos + length
+
+    def _inflate_if_needed(self, frame_payload: bytes, rsv1: bool) -> bytes:
+        if not rsv1:
+            return frame_payload
+        if self.permessage_deflate_runtime is None:
+            raise ProtocolError('RSV1 is not negotiated')
+        return self.permessage_deflate_runtime.decompress_message(frame_payload)
+
+    async def feed_data(self, data: bytes, *, end_stream: bool = False) -> None:
+        if self.closed:
+            return
+        self.buffer.extend(data)
+        if end_stream:
+            self.peer_end_stream_pending = True
+        if not self.accepted and not self.http_denied:
+            return
+        while self.buffer:
+            frame_len = self._frame_length(self.buffer)
+            if frame_len is None:
+                break
+            raw = bytes(self.buffer[:frame_len])
+            del self.buffer[:frame_len]
+            frame = parse_frame_bytes(
+                raw,
+                expect_masked=True,
+                max_payload_size=self.config.websocket_max_message_size,
+                allow_rsv1=self.permessage_deflate_runtime is not None,
+            )
+            self._record_activity()
+            if frame.opcode == OP_PING:
+                await self.send_data(pong_frame(frame.payload), False)
+                continue
+            if frame.opcode == OP_PONG:
+                if self.keepalive is not None:
+                    self.keepalive.acknowledge_pong(frame.payload)
+                continue
+            if frame.opcode == OP_CLOSE:
+                code, reason = decode_close_payload(frame.payload)
+                if not self.closed:
+                    await self.send_data(close_frame(code, reason), True)
+                self.closed = True
+                self._notify_closed()
+                await self.receive.put(websocket_disconnect(code, reason))
+                break
+            if frame.opcode in {OP_TEXT, OP_BINARY}:
+                if self.fragmented_opcode is not None:
+                    raise ProtocolError('new data frame before fragmented message completion')
+                self.current_message_size = len(frame.payload)
+                if self.current_message_size > self.config.websocket_max_message_size:
+                    raise ProtocolError('message too big')
+                if frame.fin:
+                    payload = self._inflate_if_needed(frame.payload, frame.rsv1)
+                    if frame.opcode == OP_TEXT:
+                        await self.receive.put(websocket_receive_text(payload.decode('utf-8')))
+                    else:
+                        await self.receive.put(websocket_receive_bytes(payload))
+                    self.current_message_size = 0
+                else:
+                    self.fragmented_opcode = frame.opcode
+                    self.fragmented_compressed = frame.rsv1
+                    self.fragments = [frame.payload]
+                continue
+            if frame.opcode == OP_CONT:
+                if self.fragmented_opcode is None:
+                    raise ProtocolError('unexpected continuation frame')
+                if frame.rsv1:
+                    raise ProtocolError('RSV1 is only valid on the first frame of a compressed message')
+                self.current_message_size += len(frame.payload)
+                if self.current_message_size > self.config.websocket_max_message_size:
+                    raise ProtocolError('message too big')
+                self.fragments.append(frame.payload)
+                if frame.fin:
+                    message = b''.join(self.fragments)
+                    if self.fragmented_compressed:
+                        message = self._inflate_if_needed(message, True)
+                    opcode = self.fragmented_opcode
+                    self.fragmented_opcode = None
+                    self.fragmented_compressed = False
+                    self.fragments = []
+                    self.current_message_size = 0
+                    if opcode == OP_TEXT:
+                        await self.receive.put(websocket_receive_text(message.decode('utf-8')))
+                    else:
+                        await self.receive.put(websocket_receive_bytes(message))
+                continue
+            raise ProtocolError('unsupported websocket opcode')
+        if self.peer_end_stream_pending and not self.closed:
+            self.peer_end_stream_pending = False
+            self.closed = True
+            self._notify_closed()
+            await self.receive.put(websocket_disconnect(1000, ''))
+
+    async def abort(self) -> None:
+        if not self.closed:
+            self.closed = True
+            self._notify_closed()
+            await self.receive.put(websocket_disconnect(1006, ''))
+        if self.task is not None:
+            self.task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self.task

tigrcorn_protocols/lifespan/__init__.py

@@ -0,0 +1,3 @@
+from .driver import LifespanManager
+
+__all__ = ["LifespanManager"]

tigrcorn_protocols/lifespan/driver.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+
+from tigrcorn_asgi.events.lifespan import lifespan_shutdown, lifespan_startup
+from tigrcorn_asgi.receive import LifespanReceive
+from tigrcorn_asgi.scopes.lifespan import build_lifespan_scope
+from tigrcorn_asgi.send import LifespanSend
+from tigrcorn_core.errors import TigrCornError
+from tigrcorn_core.types import ASGIApp
+
+
+class LifespanError(TigrCornError):
+    pass
+
+
+@dataclass(slots=True)
+class LifespanManager:
+    app: ASGIApp
+    mode: str = "auto"
+    timeout: float = 10.0
+    started: bool = False
+    _receive: LifespanReceive | None = None
+    _send: LifespanSend | None = None
+    _task: asyncio.Task | None = None
+
+    async def startup(self) -> None:
+        if self.mode == "off":
+            return
+        receive = LifespanReceive()
+        send = LifespanSend()
+        scope = build_lifespan_scope()
+
+        async def runner() -> None:
+            await self.app(scope, receive, send)
+
+        task = asyncio.create_task(runner(), name="tigrcorn-lifespan")
+        self._receive = receive
+        self._send = send
+        self._task = task
+        await receive.put(lifespan_startup())
+        try:
+            message = await asyncio.wait_for(send.get(), timeout=self.timeout)
+        except Exception as exc:
+            task.cancel()
+            if self.mode == "auto":
+                self._task = None
+                self._receive = None
+                self._send = None
+                return
+            raise LifespanError("lifespan startup did not complete") from exc
+
+        if message["type"] == "lifespan.startup.complete":
+            self.started = True
+            return
+        if message["type"] == "lifespan.startup.failed":
+            task.cancel()
+            raise LifespanError(str(message.get("message", "lifespan startup failed")))
+        if self.mode == "auto":
+            task.cancel()
+            self._task = None
+            self._receive = None
+            self._send = None
+            return
+        raise LifespanError(f"unexpected lifespan startup message: {message!r}")
+
+    async def shutdown(self) -> None:
+        if self.mode == "off":
+            return
+        if not self.started:
+            if self._task is not None:
+                self._task.cancel()
+            return
+        assert self._receive is not None and self._send is not None
+        await self._receive.put(lifespan_shutdown())
+        message = await asyncio.wait_for(self._send.get(), timeout=self.timeout)
+        if message["type"] == "lifespan.shutdown.failed":
+            raise LifespanError(str(message.get("message", "lifespan shutdown failed")))
+        if message["type"] != "lifespan.shutdown.complete":
+            raise LifespanError(f"unexpected lifespan shutdown message: {message!r}")
+        if self._task is not None:
+            await asyncio.wait({self._task}, timeout=self.timeout)

tigrcorn_protocols/py.typed

@@ -0,0 +1 @@
+

tigrcorn_protocols/rawframed/__init__.py

@@ -0,0 +1,5 @@
+from .codec import RawFrame, encode_frame, read_frame, try_decode_frame
+from .handler import RawFramedApplicationHandler
+from .state import RawFramedState
+
+__all__ = ['RawFrame', 'encode_frame', 'read_frame', 'try_decode_frame', 'RawFramedState', 'RawFramedApplicationHandler']

tigrcorn_protocols/rawframed/codec.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from tigrcorn_core.errors import ProtocolError
+from tigrcorn_protocols.rawframed.frames import RawFrame, encode_frame, try_decode_frame
+from tigrcorn_core.types import StreamReaderLike
+
+
+async def read_frame(reader: StreamReaderLike, *, max_frame_size: int = 16 * 1024 * 1024) -> RawFrame:
+    import struct
+
+    prefix = await reader.readexactly(4)
+    size = struct.unpack("!I", prefix)[0]
+    if size > max_frame_size:
+        raise ProtocolError("raw frame exceeds configured max_frame_size")
+    return RawFrame(await reader.readexactly(size))
+
+
+__all__ = ["RawFrame", "encode_frame", "read_frame", "try_decode_frame"]

tigrcorn_protocols/rawframed/frames.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class RawFrame:
+    payload: bytes
+
+    @property
+    def length(self) -> int:
+        return len(self.payload)
+
+
+def encode_frame(payload: bytes) -> bytes:
+    return struct.pack("!I", len(payload)) + payload
+
+
+def try_decode_frame(buffer: bytearray) -> RawFrame | None:
+    if len(buffer) < 4:
+        return None
+    size = struct.unpack("!I", buffer[:4])[0]
+    if len(buffer) < 4 + size:
+        return None
+    payload = bytes(buffer[4 : 4 + size])
+    del buffer[: 4 + size]
+    return RawFrame(payload)

tigrcorn_protocols/rawframed/handler.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Protocol
+
+from tigrcorn_asgi.events.custom import stream_receive
+from tigrcorn_asgi.receive import QueueReceive
+from tigrcorn_asgi.scopes.custom import build_custom_scope
+from tigrcorn_config.model import ListenerConfig, ServerConfig
+from tigrcorn_observability.logging import AccessLogger
+from tigrcorn_protocols.custom.adapters import adapt_scope
+from tigrcorn_protocols.rawframed.frames import encode_frame, try_decode_frame
+from tigrcorn_protocols.rawframed.state import RawFramedState
+from tigrcorn_core.types import ASGIApp
+
+
+class _Writable(Protocol):
+    def write(self, data: bytes) -> int: ...
+
+
+@dataclass(slots=True)
+class _RawAppSend:
+    connection: _Writable
+    outbound_frames: int = 0
+
+    async def __call__(self, message: dict) -> None:
+        typ = message.get('type')
+        if typ != 'tigrcorn.stream.send':
+            raise RuntimeError(f'unexpected raw framed send event: {typ!r}')
+        payload = bytes(message.get('data', b''))
+        self.connection.write(encode_frame(payload))
+        self.outbound_frames += 1
+
+
+@dataclass(slots=True)
+class RawFramedApplicationHandler:
+    app: ASGIApp
+    config: ServerConfig
+    listener: ListenerConfig
+    access_logger: AccessLogger
+    buffers: dict[int, bytearray] = field(default_factory=dict)
+    states: dict[int, RawFramedState] = field(default_factory=dict)
+
+    async def feed_bytes(self, connection: _Writable, data: bytes, *, path: str | None = None) -> int:
+        key = id(connection)
+        buffer = self.buffers.setdefault(key, bytearray())
+        state = self.states.setdefault(key, RawFramedState())
+        buffer.extend(data)
+        handled = 0
+        while True:
+            frame = try_decode_frame(buffer)
+            if frame is None:
+                return handled
+            state.frames_received += 1
+            handled += 1
+            await self._dispatch_frame(connection, frame.payload, state, path=path)
+
+    async def _dispatch_frame(self, connection: _Writable, payload: bytes, state: RawFramedState, *, path: str | None = None) -> None:
+        scope = adapt_scope(
+            build_custom_scope(
+                'tigrcorn.rawframed',
+                scheme=self.listener.scheme or 'tigrcorn+raw',
+                path=path or self.listener.path or '',
+                headers=[],
+                extensions={'tigrcorn.custom': {'transport': self.listener.kind}},
+            )
+        )
+        receive = QueueReceive()
+        await receive.put(stream_receive(payload, more_data=False))
+        send = _RawAppSend(connection=connection)
+        await self.app(scope, receive, send)
+        state.frames_sent += send.outbound_frames

tigrcorn_protocols/rawframed/state.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class RawFramedState:
+    frames_received: int = 0
+    frames_sent: int = 0

tigrcorn_protocols/registry.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class ProtocolDescriptor:
+    name: str
+    multiplexed: bool = False
+    asgi_scope_types: tuple[str, ...] = ()
+
+
+BUILTIN_PROTOCOLS = {
+    "http1": ProtocolDescriptor(name="http1", multiplexed=False, asgi_scope_types=("http",)),
+    "http2": ProtocolDescriptor(name="http2", multiplexed=True, asgi_scope_types=("http",)),
+    "http3": ProtocolDescriptor(name="http3", multiplexed=True, asgi_scope_types=("http",)),
+    "quic": ProtocolDescriptor(name="quic", multiplexed=True, asgi_scope_types=("tigrcorn.quic",)),
+    "websocket": ProtocolDescriptor(name="websocket", multiplexed=False, asgi_scope_types=("websocket",)),
+    "lifespan": ProtocolDescriptor(name="lifespan", multiplexed=False, asgi_scope_types=("lifespan",)),
+    "rawframed": ProtocolDescriptor(name="rawframed", multiplexed=False, asgi_scope_types=("tigrcorn.rawframed",)),
+    "custom": ProtocolDescriptor(name="custom", multiplexed=False, asgi_scope_types=("tigrcorn.stream",)),
+}

tigrcorn_protocols/scheduler/__init__.py

@@ -0,0 +1,17 @@
+"""Scheduling policy and runtime components."""
+
+from .dispatch import TaskDispatcher
+from .policy import SchedulerPolicy
+from .quotas import Quotas
+from .runtime import ConnectionLease, ProductionScheduler, WorkLease
+from .tasks import TaskSet
+
+__all__ = [
+    'ConnectionLease',
+    'ProductionScheduler',
+    'WorkLease',
+    'Quotas',
+    'SchedulerPolicy',
+    'TaskDispatcher',
+    'TaskSet',
+]

tigrcorn_protocols/scheduler/cancellation.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import asyncio
+from contextlib import suppress
+from dataclasses import dataclass
+from typing import Iterable
+
+
+@dataclass(slots=True)
+class CancellationResult:
+    completed: int
+    pending: int
+    timed_out: bool
+
+
+async def cancel(task: asyncio.Task | None) -> None:
+    if task is None:
+        return
+    task.cancel()
+    with suppress(asyncio.CancelledError):
+        await task
+
+
+async def cancel_many(tasks: Iterable[asyncio.Task]) -> None:
+    for task in tasks:
+        task.cancel()
+    for task in tasks:
+        with suppress(asyncio.CancelledError):
+            await task
+
+
+async def cancel_many_bounded(tasks: Iterable[asyncio.Task], *, timeout: float) -> CancellationResult:
+    task_list = list(tasks)
+    for task in task_list:
+        task.cancel()
+    done, pending = await asyncio.wait(task_list, timeout=timeout)
+    for task in done:
+        with suppress(asyncio.CancelledError):
+            task.result()
+    return CancellationResult(completed=len(done), pending=len(pending), timed_out=bool(pending))