diff --git a/SmallPackage/Kernel.py b/SmallPackage/Kernel.py index e01bed6..7342f34 100644 --- a/SmallPackage/Kernel.py +++ b/SmallPackage/Kernel.py @@ -24,9 +24,21 @@ if TYPE_CHECKING: from collections.abc import Iterable, Mapping, Sequence from typing import Any, cast + from ._types import SocketBuffer, SocketOperation, SocketRetryMode _UNSET = object() +_SOCKET_OPERATIONS = ('accept', 'recv', 'send', 'handshake') + + +def _validate_socket_operation(operation): + """Reject retry classifications that do not name a supported operation.""" + if operation not in _SOCKET_OPERATIONS: + raise ValueError( + 'Unknown socket operation {!r}; expected one of: {}.'.format( + operation, ', '.join(_SOCKET_OPERATIONS) + ) + ) def _import_first(*module_names: str) -> Any | None: @@ -556,7 +568,8 @@ def socket_connect(self, sock: Any, sockaddr: Any) -> bool: def socket_connection_error(self, sock: Any) -> int: return 0 - def socket_send(self, sock: Any, data: bytes) -> int: + def socket_send(self, sock: Any, data: SocketBuffer) -> int: + """Send a bytes-like buffer without requiring an intermediate copy.""" return 0 def socket_recv(self, sock: Any, buffer_size: int) -> bytes: @@ -593,6 +606,17 @@ def socket_needs_read(self, exc: BaseException) -> bool: def socket_needs_write(self, exc: BaseException) -> bool: return False + def socket_retry_mode( + self, exc: BaseException, operation: SocketOperation + ) -> SocketRetryMode: + """Return the readiness direction needed to retry one socket operation.""" + _validate_socket_operation(operation) + if self.socket_needs_write(exc): + return 'write' + if self.socket_needs_read(exc): + return 'write' if operation == 'send' else 'read' + return None + def _poll_lookup_key(self, obj: Any) -> Any: """ Return the identity key used to map poll events back to registered objects. @@ -766,7 +790,7 @@ def socket_connect(self, sock, sockaddr): def socket_connection_error(self, sock): return sock.getsockopt(self._socket.SOL_SOCKET, self._socket.SO_ERROR) - def socket_send(self, sock, data): + def socket_send(self, sock: Any, data: SocketBuffer) -> int: return sock.send(data) def socket_recv(self, sock, buffer_size): @@ -811,6 +835,25 @@ def socket_needs_read(self, exc): def socket_needs_write(self, exc): return isinstance(exc, self._ssl.SSLWantWriteError) + def socket_retry_mode( + self, exc: BaseException, operation: SocketOperation + ) -> SocketRetryMode: + _validate_socket_operation(operation) + if isinstance(exc, self._ssl.SSLWantReadError): + return 'read' + if isinstance(exc, self._ssl.SSLWantWriteError): + return 'write' + would_block = { + value for value in ( + getattr(self._errno, 'EAGAIN', None), + getattr(self._errno, 'EWOULDBLOCK', None), + ) if value is not None + } + err = self._extract_errno(exc) + if err in would_block or (isinstance(exc, BlockingIOError) and err is None): + return 'write' if operation == 'send' else 'read' + return None + class MicroPythonKernel(Kernel): ''' @@ -848,6 +891,7 @@ def __init__( self._ssl = _import_first('ssl', 'ussl') self._sys = modules.get('sys') or _import_first('sys') self._os = modules.get('os') or _import_first('os') + self._errno = modules.get('errno') or _import_first('errno', 'uerrno') self._network = modules.get('network') or _import_first('network') self._rp2 = modules.get('rp2') or _import_first('rp2') self._machine = modules.get('machine') or _import_first('machine') @@ -957,7 +1001,7 @@ def socket_connect(self, sock, sockaddr): pending = { 115, 11, } - if err in pending or self.socket_needs_read(exc) or self.socket_needs_write(exc): + if err in pending or isinstance(exc, BlockingIOError): return False raise @@ -966,7 +1010,7 @@ def socket_connection_error(self, sock): so_error = getattr(self._socket, 'SO_ERROR', 4) return sock.getsockopt(sol_socket, so_error) - def socket_send(self, sock, data): + def socket_send(self, sock: Any, data: SocketBuffer) -> int: return sock.send(data) def socket_recv(self, sock, buffer_size): @@ -1050,6 +1094,27 @@ def socket_needs_write(self, exc): return True return False + def socket_retry_mode( + self, exc: BaseException, operation: SocketOperation + ) -> SocketRetryMode: + _validate_socket_operation(operation) + want_read_error = getattr(self._ssl, 'SSLWantReadError', None) + if isinstance(want_read_error, type) and isinstance(exc, want_read_error): + return 'read' + want_write_error = getattr(self._ssl, 'SSLWantWriteError', None) + if isinstance(want_write_error, type) and isinstance(exc, want_write_error): + return 'write' + err = self._extract_errno(exc) + pending = {11} + if self._errno is not None: + for name in ('EAGAIN', 'EWOULDBLOCK'): + value = getattr(self._errno, name, None) + if value is not None: + pending.add(value) + if err in pending or (isinstance(exc, BlockingIOError) and err is None): + return 'write' if operation == 'send' else 'read' + return None + def machine_name(self): ''' Return the current firmware/board descriptor when the platform exposes one. diff --git a/SmallPackage/_types.py b/SmallPackage/_types.py index 6589d5e..03b78bf 100644 --- a/SmallPackage/_types.py +++ b/SmallPackage/_types.py @@ -7,13 +7,16 @@ from __future__ import annotations from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Protocol, TypeAlias, TypedDict, TypeVar +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, TypedDict, TypeVar if TYPE_CHECKING: from .SmallOS import SmallOS T = TypeVar("T") +SocketBuffer: TypeAlias = bytes | bytearray | memoryview +SocketOperation: TypeAlias = Literal["accept", "recv", "send", "handshake"] +SocketRetryMode: TypeAlias = Literal["read", "write"] | None class TaskLike(Protocol): @@ -42,6 +45,15 @@ def io_wait( ) -> tuple[Sequence[Any], Sequence[Any]]: ... +class SocketKernelLike(Protocol): + """Optional kernel boundary for operation-aware byte-stream I/O.""" + + def socket_send(self, sock: Any, data: SocketBuffer) -> int: ... + def socket_retry_mode( + self, exc: BaseException, operation: SocketOperation + ) -> SocketRetryMode: ... + + class AdapterCompletionLike(Protocol): """Structural completion record consumed by the scheduler.""" diff --git a/SmallPackage/clients/SmallStream.py b/SmallPackage/clients/SmallStream.py index 2cbcd2f..3125f6c 100644 --- a/SmallPackage/clients/SmallStream.py +++ b/SmallPackage/clients/SmallStream.py @@ -100,10 +100,11 @@ async def connect(self): kernel.socket_do_handshake(sock) break except Exception as exc: - if kernel.socket_needs_read(exc): + retry_mode = kernel.socket_retry_mode(exc, "handshake") + if retry_mode == "read": await self.task.wait_readable(sock) continue - if kernel.socket_needs_write(exc): + if retry_mode == "write": await self.task.wait_writable(sock) continue raise @@ -135,7 +136,7 @@ async def send_all(self, data): if not self._connected or self.sock is None: await self.connect() - view = memoryview(bytes(data)) + view = data if isinstance(data, memoryview) else memoryview(data) while view: try: sent = self.kernel.socket_send(self.sock, view) @@ -143,10 +144,11 @@ async def send_all(self, data): raise StreamClosedError("socket closed while sending data") view = view[sent:] except Exception as exc: - if self.kernel.socket_needs_read(exc): + retry_mode = self.kernel.socket_retry_mode(exc, "send") + if retry_mode == "read": await self.task.wait_readable(self.sock) continue - if self.kernel.socket_needs_write(exc): + if retry_mode == "write": await self.task.wait_writable(self.sock) continue raise @@ -164,10 +166,11 @@ async def recv_some(self, size=4096): raise StreamClosedError("socket closed while reading data") return bytes(chunk) except Exception as exc: - if self.kernel.socket_needs_read(exc): + retry_mode = self.kernel.socket_retry_mode(exc, "recv") + if retry_mode == "read": await self.task.wait_readable(self.sock) continue - if self.kernel.socket_needs_write(exc): + if retry_mode == "write": await self.task.wait_writable(self.sock) continue raise diff --git a/demos/web_app_demo.py b/demos/web_app_demo.py index 6244b30..55481c6 100644 --- a/demos/web_app_demo.py +++ b/demos/web_app_demo.py @@ -184,16 +184,17 @@ def _home_page(): async def _send_all(task, sock, data): """Write all bytes to a non-blocking socket using smallOS waits.""" kernel = task.OS.kernel - remaining = memoryview(bytes(data)) + remaining = data if isinstance(data, memoryview) else memoryview(data) while remaining: try: sent = kernel.socket_send(sock, remaining) except Exception as exc: - if kernel.socket_needs_read(exc): + retry_mode = kernel.socket_retry_mode(exc, "send") + if retry_mode == "read": await task.wait_readable(sock) continue - if kernel.socket_needs_write(exc): + if retry_mode == "write": await task.wait_writable(sock) continue raise @@ -212,10 +213,11 @@ async def _read_request_head(task, sock): try: chunk = kernel.socket_recv(sock, 1024) except Exception as exc: - if kernel.socket_needs_read(exc): + retry_mode = kernel.socket_retry_mode(exc, "recv") + if retry_mode == "read": await task.wait_readable(sock) continue - if kernel.socket_needs_write(exc): + if retry_mode == "write": await task.wait_writable(sock) continue raise @@ -369,10 +371,11 @@ async def web_server_task(task, state): try: client_sock, client_addr = listener.accept() except Exception as exc: - if kernel.socket_needs_read(exc): + retry_mode = kernel.socket_retry_mode(exc, "accept") + if retry_mode == "read": await task.wait_readable(listener) continue - if kernel.socket_needs_write(exc): + if retry_mode == "write": await task.wait_writable(listener) continue raise diff --git a/tests/test_socket_retry.py b/tests/test_socket_retry.py new file mode 100644 index 0000000..e65eed2 --- /dev/null +++ b/tests/test_socket_retry.py @@ -0,0 +1,214 @@ +import asyncio +import errno +import os +import sys +import unittest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "demos")) + +from demos.web_app_demo import _read_request_head, _send_all +from SmallPackage.Kernel import Kernel, MicroPythonKernel, Unix +from SmallPackage.clients.SmallStream import SmallStream + + +class FakeSSLWantReadError(OSError): + pass + + +class FakeSSLWantWriteError(OSError): + pass + + +class FakeSSLModule: + SSLWantReadError = FakeSSLWantReadError + SSLWantWriteError = FakeSSLWantWriteError + + +class StreamSocket: + pass + + +class ConnectWouldBlockSocket: + def connect(self, _sockaddr): + raise BlockingIOError() + + +class StreamKernel(Kernel): + def __init__(self): + super().__init__() + self.send_attempts = [] + self.recv_attempts = 0 + self.handshake_attempts = 0 + + def resolve_address(self, host, port): + return (0, 0, 0, "", (host, port)) + + def socket_open(self, address_info): + return StreamSocket() + + def socket_setblocking(self, sock, flag): + return None + + def socket_connect(self, sock, sockaddr): + return True + + def socket_wrap_tls_client(self, sock, **kwargs): + return sock + + def socket_do_handshake(self, sock): + self.handshake_attempts += 1 + if self.handshake_attempts == 1: + raise BlockingIOError(errno.EAGAIN, "try again") + + def socket_send(self, sock, data): + self.send_attempts.append(data) + if len(self.send_attempts) == 1: + raise BlockingIOError(errno.EAGAIN, "try again") + return len(data) + + def socket_recv(self, sock, buffer_size): + self.recv_attempts += 1 + if self.recv_attempts == 1: + raise BlockingIOError(errno.EAGAIN, "try again") + return b"response" + + +class DemoStreamKernel(StreamKernel): + def socket_recv(self, sock, buffer_size): + self.recv_attempts += 1 + if self.recv_attempts == 1: + raise BlockingIOError(errno.EAGAIN, "try again") + return b"GET / HTTP/1.1\r\n\r\n" + + +class RuntimeRef: + def __init__(self, kernel): + self.kernel = kernel + + +class StreamTask: + def __init__(self, kernel): + self.OS = RuntimeRef(kernel) + self.read_waits = [] + self.write_waits = [] + + async def wait_readable(self, sock): + self.read_waits.append(sock) + + async def wait_writable(self, sock): + self.write_waits.append(sock) + + +class TestSocketRetryClassification(unittest.TestCase): + def test_base_retry_direction_depends_on_operation(self): + kernel = Kernel() + error = BlockingIOError(errno.EAGAIN, "try again") + + self.assertEqual("read", kernel.socket_retry_mode(error, "accept")) + self.assertEqual("read", kernel.socket_retry_mode(error, "recv")) + self.assertEqual("read", kernel.socket_retry_mode(error, "handshake")) + self.assertEqual("write", kernel.socket_retry_mode(error, "send")) + self.assertIsNone(kernel.socket_retry_mode(OSError(errno.EINVAL, "bad"), "recv")) + + def test_unix_retry_mode_preserves_tls_direction(self): + kernel = Unix() + pending = BlockingIOError(errno.EAGAIN, "try again") + + self.assertEqual("read", kernel.socket_retry_mode(pending, "recv")) + self.assertEqual("write", kernel.socket_retry_mode(pending, "send")) + self.assertEqual( + "write", kernel.socket_retry_mode(OSError(errno.EAGAIN, "again"), "send") + ) + self.assertEqual( + "read", + kernel.socket_retry_mode(kernel._ssl.SSLWantReadError(), "send"), + ) + self.assertEqual( + "write", + kernel.socket_retry_mode(kernel._ssl.SSLWantWriteError(), "recv"), + ) + self.assertIsNone( + kernel.socket_retry_mode( + BlockingIOError(errno.EINPROGRESS, "pending connect"), "send" + ) + ) + + def test_micropython_matches_unix_and_preserves_tls_direction(self): + kernel = MicroPythonKernel(modules={"ssl": FakeSSLModule()}) + pending = OSError(errno.EAGAIN, "try again") + + self.assertEqual("read", kernel.socket_retry_mode(pending, "accept")) + self.assertEqual("read", kernel.socket_retry_mode(pending, "recv")) + self.assertEqual("write", kernel.socket_retry_mode(pending, "send")) + self.assertEqual( + "read", kernel.socket_retry_mode(FakeSSLWantReadError(), "send") + ) + self.assertEqual( + "write", kernel.socket_retry_mode(FakeSSLWantWriteError(), "recv") + ) + self.assertIsNone( + kernel.socket_retry_mode(OSError(115, "pending connect"), "recv") + ) + + def test_micropython_connect_preserves_errno_less_would_block(self): + kernel = MicroPythonKernel() + + self.assertFalse(kernel.socket_connect(ConnectWouldBlockSocket(), ("host", 80))) + + def test_unknown_operation_is_rejected_before_classification(self): + kernels = (Kernel(), Unix(), MicroPythonKernel()) + + for kernel in kernels: + with self.subTest(kernel=type(kernel).__name__): + with self.assertRaisesRegex(ValueError, "Unknown socket operation"): + kernel.socket_retry_mode(BlockingIOError(), "connect") + + def test_small_stream_forwards_memoryview_and_waits_by_operation(self): + kernel = StreamKernel() + task = StreamTask(kernel) + stream = SmallStream(task, "example.test", 80) + stream.sock = StreamSocket() + stream._connected = True + payload = memoryview(b"request") + + async def exercise_stream(): + await stream.send_all(payload) + return await stream.recv_some() + + result = asyncio.run(exercise_stream()) + + self.assertEqual(b"response", result) + self.assertIs(payload, kernel.send_attempts[0]) + self.assertEqual([stream.sock], task.write_waits) + self.assertEqual([stream.sock], task.read_waits) + + def test_small_stream_handshake_waits_for_readability(self): + kernel = StreamKernel() + task = StreamTask(kernel) + stream = SmallStream(task, "example.test", 443, use_tls=True) + + result = asyncio.run(stream.connect()) + + self.assertIs(stream, result) + self.assertEqual(2, kernel.handshake_attempts) + self.assertEqual([stream.sock], task.read_waits) + self.assertEqual([], task.write_waits) + + def test_web_demo_routes_send_and_recv_backpressure_by_operation(self): + kernel = DemoStreamKernel() + task = StreamTask(kernel) + sock = StreamSocket() + + async def exercise_helpers(): + await _send_all(task, sock, memoryview(b"response")) + return await _read_request_head(task, sock) + + result = asyncio.run(exercise_helpers()) + + self.assertEqual(b"GET / HTTP/1.1\r\n\r\n", result) + self.assertEqual([sock], task.write_waits) + self.assertEqual([sock], task.read_waits) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/typing/kernel_socket_contract.py b/tests/typing/kernel_socket_contract.py new file mode 100644 index 0000000..caa7e9d --- /dev/null +++ b/tests/typing/kernel_socket_contract.py @@ -0,0 +1,34 @@ +# pyright: strict, reportUnnecessaryTypeIgnoreComment=true + +"""Static contract for operation-aware retries and bytes-like socket sends.""" + +from typing import Any + +from SmallPackage.Kernel import Kernel +from SmallPackage._types import SocketBuffer, SocketKernelLike, SocketOperation + + +class SocketKernel(Kernel): + received: SocketBuffer | None = None + + def socket_send(self, sock: Any, data: SocketBuffer) -> int: + self.received = data + return len(data) + + +def accepts_kernel(kernel: SocketKernelLike) -> None: + """Require the concrete kernel to satisfy the public structural contract.""" + + +def verify_socket_contract() -> None: + kernel = SocketKernel() + accepts_kernel(kernel) + payload = memoryview(b"response") + sent: int = kernel.socket_send(object(), payload) + operation: SocketOperation = "send" + retry_mode = kernel.socket_retry_mode(BlockingIOError(), operation) + assert sent == len(payload) + assert kernel.received is payload + assert retry_mode == "write" + + kernel.socket_retry_mode(BlockingIOError(), "connect") # pyright: ignore[reportArgumentType]