pyfect-dev 0.1.0__py3-none-any.whl

pyfect/socket/_connection.py

@@ -0,0 +1,384 @@
+"""A stream connection on the socket loop: an asyncio protocol whose state
+fibers read under a lock.
+
+The loop thread fills the read buffer and records how the connection ended;
+fibers take from the buffer in their own step and park on ``Signal`` until
+there is something to take, so an interrupt between a wake-up and the take
+cannot lose bytes. Reading pauses while the buffer holds ``HIGH_WATER_MARK``
+bytes.
+
+How the read side ends follows Node's sockets: the peer ending the stream
+delivers the buffered bytes, then a close with code 1000; an error or a
+teardown discards them and fails with that error.
+
+``start_tls`` switches the connection to TLS (``TlsSession``): bytes are
+decrypted before they reach the buffer and encrypted on the way out.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import ssl
+import threading
+from collections import deque
+from collections.abc import Callable
+
+from pyfect._core import Effect
+from pyfect._os_wait import Resume, os_wait
+from pyfect.effect import fail, succeed, suspend, void
+from pyfect.socket._errors import SocketError, close_error, read_error, write_error
+from pyfect.socket._loop import call_soon
+from pyfect.socket._tls_session import Received, TlsSession
+
+HIGH_WATER_MARK = 64 * 1024
+
+
+class Signal:
+    """Fibers parked until a condition holds, woken from any thread to
+    retry. Conditions are checked under ``lock``, the lock of the state
+    they read."""
+
+    __slots__ = ("_waiters", "lock")
+
+    def __init__(self, lock: threading.Lock) -> None:
+        self.lock = lock
+        self._waiters: list[Callable[[], None]] = []
+
+    def wait(self, ready: Callable[[], bool]) -> Effect[None]:
+        def register(resume: Resume[None, object]) -> Callable[[], None] | None:
+            def wake() -> None:
+                resume(void)
+
+            with self.lock:
+                parked = not ready()
+                if parked:
+                    self._waiters.append(wake)
+            if not parked:
+                wake()
+                return None
+            return lambda: self._forget(wake)
+
+        return os_wait(register, waits_for="event")  # pyright: ignore[reportReturnType]
+
+    def wait_until(self, ready: Callable[[], bool]) -> Effect[None]:
+        """Wait, retrying after each wake, until ``ready`` holds (``wait``
+        returns at the first wake)."""
+        return self.wait(ready).flat_map(lambda _: void if ready() else self.wait_until(ready))
+
+    def notify_locked(self) -> list[Callable[[], None]]:
+        """Take the waiters (holding ``lock``); call them after releasing
+        it."""
+        waiters, self._waiters = self._waiters, []
+        return waiters
+
+    def _forget(self, wake: Callable[[], None]) -> None:
+        with self.lock:
+            if wake in self._waiters:
+                self._waiters.remove(wake)
+
+
+class Connection(asyncio.Protocol):
+    """One stream connection. Created on a fiber, driven by the loop.
+    ``on_made`` runs (on the loop) once it is connected: after the
+    handshake when it is accepted with ``server_tls``."""
+
+    def __init__(
+        self,
+        on_made: Callable[[Connection], None] | None = None,
+        *,
+        server_tls: ssl.SSLContext | None = None,
+    ) -> None:
+        self._on_made = on_made
+        self._server_tls = server_tls
+        self.lock = threading.Lock()
+        self.signal = Signal(self.lock)
+        self.transport: asyncio.Transport | None = None
+        self._tls: TlsSession | None = None
+        self._chunks: deque[bytes] = deque()
+        self._size = 0
+        self._error: SocketError | None = None
+        self._discard = False
+        self._reading_paused = False
+        self._writing_paused = False
+        self._close_requested: bool | None = None
+        self.lost = False
+        self.open = False
+
+    # --- loop side ----------------------------------------------------------
+
+    def connection_made(self, transport: asyncio.BaseTransport) -> None:
+        assert isinstance(transport, asyncio.Transport)
+        self.transport = transport
+        if self._close_requested is not None:
+            self._apply_close(self._close_requested)
+        elif self._server_tls is not None:
+            handshake = self.start_tls(self._server_tls, server_side=True, hostname=None)
+            handshake.add_done_callback(self._server_handshake_done)
+        elif self._on_made is not None:
+            self._on_made(self)
+
+    def _server_handshake_done(self, handshake: asyncio.Future[None]) -> None:
+        if handshake.cancelled() or handshake.exception() is not None:
+            self._apply_close(graceful=False)
+        elif self._on_made is not None:
+            self._on_made(self)
+
+    def data_received(self, data: bytes) -> None:
+        self._deliver(Received((data,)) if self._tls is None else self._tls.feed(data))
+
+    def _deliver(self, received: Received) -> None:
+        if received.data:
+            with self.lock:
+                if not self._discard:
+                    self._chunks.extend(received.data)
+                    self._size += sum(map(len, received.data))
+                    self._pause_if_full()
+                waiters = self.signal.notify_locked()
+            for wake in waiters:
+                wake()
+        if received.error is not None:
+            self.fail(read_error(received.error))
+        elif received.eof:
+            self.end(close_error(1000))
+
+    def _pause_if_full(self) -> None:
+        transport = self.transport
+        if self._size >= HIGH_WATER_MARK and not self._reading_paused and transport is not None:
+            self._reading_paused = True
+            transport.pause_reading()
+
+    def eof_received(self) -> bool:
+        self.end(close_error(1000))
+        # Stay half-open: the write side ends when the writer is released.
+        return True
+
+    def connection_lost(self, exc: Exception | None) -> None:
+        with self.lock:
+            self.lost = True
+            waiters = self.signal.notify_locked()
+        for wake in waiters:
+            wake()
+        if self._tls is not None and not self._tls.handshake.done():
+            self._tls.handshake.set_exception(exc or ConnectionError("socket closed"))
+        if exc is None:
+            self.end(close_error(1000))
+        else:
+            self.fail(read_error(exc))
+
+    def pause_writing(self) -> None:
+        with self.lock:
+            self._writing_paused = True
+
+    def resume_writing(self) -> None:
+        with self.lock:
+            self._writing_paused = False
+            waiters = self.signal.notify_locked()
+        for wake in waiters:
+            wake()
+
+    def start_tls(
+        self, context: ssl.SSLContext, *, server_side: bool, hostname: str | None
+    ) -> asyncio.Future[None]:
+        """Begin a TLS handshake (on the loop); bytes already buffered are
+        the peer's first handshake bytes. The future settles with it."""
+        transport = self.transport
+        assert transport is not None
+        session = TlsSession(transport, context, server_side=server_side, hostname=hostname)
+        self._tls = session
+        with self.lock:
+            buffered = tuple(self._chunks)
+            self._chunks.clear()
+            self._size = 0
+        self._deliver(session.start(buffered))
+        return session.handshake
+
+    def upgrade(
+        self, context: ssl.SSLContext, *, server_side: bool, hostname: str | None
+    ) -> Effect[None, BaseException]:
+        """``start_tls`` from a fiber, waiting for the handshake. It starts
+        in the order of the fiber's other loop calls, so a write made after
+        it is encrypted (queued until the handshake is done)."""
+
+        def register(resume: Resume[None, BaseException]) -> None:
+            def begin() -> None:
+                transport = self.transport
+                if transport is None or transport.is_closing():
+                    resume(fail(ConnectionError("socket closed during TLS upgrade")))
+                    return
+                handshake = self.start_tls(context, server_side=server_side, hostname=hostname)
+
+                def done(finished: asyncio.Future[None]) -> None:
+                    error = finished.exception()
+                    resume(void if error is None else fail(error))
+
+                handshake.add_done_callback(done)
+
+            call_soon(begin)
+
+        return os_wait(register, waits_for="work")
+
+    # --- either side ----------------------------------------------------------
+
+    def end(self, error: SocketError) -> None:
+        """The read side ended normally: buffered bytes are still taken
+        before ``error``."""
+        self._finish(error, discard=False)
+
+    def fail(self, error: SocketError) -> None:
+        """The connection failed or was torn down: buffered bytes are
+        dropped."""
+        self._finish(error, discard=True)
+
+    def _finish(self, error: SocketError, *, discard: bool) -> None:
+        with self.lock:
+            if self._error is None:
+                self._error = error
+            if discard:
+                self._discard = True
+                self._chunks.clear()
+                self._size = 0
+            waiters = self.signal.notify_locked()
+        for wake in waiters:
+            wake()
+
+    # --- fiber side -------------------------------------------------------------
+
+    def _take(self) -> tuple[bytes, ...] | SocketError | None:
+        with self.lock:
+            if self._chunks:
+                chunks = tuple(self._chunks)
+                self._chunks.clear()
+                self._size = 0
+                if self._reading_paused and self._error is None:
+                    self._reading_paused = False
+                    transport = self.transport
+                    if transport is not None:
+                        call_soon(transport.resume_reading)
+                return chunks
+            return self._error
+
+    def _readable(self) -> bool:
+        return bool(self._chunks) or self._error is not None
+
+    def pull(self) -> Effect[tuple[bytes, ...], SocketError]:
+        """The buffered chunks, waiting for some; the read side's error once
+        the buffer is empty."""
+        return self.pull_into(lambda chunks: chunks)
+
+    def pull_into[A](self, take: Callable[[tuple[bytes, ...]], A]) -> Effect[A, SocketError]:
+        """``take`` of the buffered chunks, called in the step that takes
+        them, so an interrupt cannot lose them."""
+
+        def step() -> Effect[A, SocketError]:
+            taken = self._take()
+            if isinstance(taken, tuple):
+                return succeed(take(taken))
+            if taken is not None:
+                return fail(taken)
+            return self.signal.wait(self._readable).flat_map(lambda _: suspend(step))
+
+        return suspend(step)
+
+    def closed(self) -> Effect[None]:
+        """Wait until the transport is lost."""
+        return self.signal.wait_until(lambda: self.lost)
+
+    def read_ended(self) -> Effect[None]:
+        """Wait until the read side has ended or failed."""
+        return self.signal.wait_until(lambda: self._error is not None)
+
+    def discard_input(self) -> None:
+        """Drop what was received and everything received from now on,
+        reading on (a server ignoring the rest of a request it rejected)."""
+
+        def run() -> None:
+            with self.lock:
+                self._discard = True
+                self._chunks.clear()
+                self._size = 0
+                resume = self._reading_paused
+                self._reading_paused = False
+            transport = self.transport
+            if resume and transport is not None and not transport.is_closing():
+                transport.resume_reading()
+
+        call_soon(run)
+
+    def peer_name(self) -> object:
+        """The transport's ``peername``, ``None`` before it is made."""
+        transport = self.transport
+        return None if transport is None else transport.get_extra_info("peername")
+
+    @property
+    def encrypted(self) -> bool:
+        return self._tls is not None or self._server_tls is not None
+
+    def write(self, chunks: tuple[bytes, ...]) -> Effect[None, SocketError]:
+        """Write ``chunks`` on the loop, then wait while the transport's
+        buffer is over its high-water mark."""
+
+        def register(resume: Resume[None, SocketError]) -> None:
+            def run() -> None:
+                transport = self.transport
+                if transport is None or transport.is_closing():
+                    resume(fail(write_error(ConnectionError("socket closed"))))
+                    return
+                try:
+                    if self._tls is None:
+                        transport.writelines(chunks)
+                    else:
+                        self._tls.write(chunks)
+                except Exception as cause:
+                    resume(fail(write_error(cause)))
+                    return
+                resume(void)
+
+            call_soon(run)
+
+        return os_wait(register, waits_for="work").flat_map(lambda _: self._drained())
+
+    def _drained(self) -> Effect[None, SocketError]:
+        def ready() -> bool:
+            return not self._writing_paused or self.lost
+
+        def check(_: None) -> Effect[None, SocketError]:
+            if self.lost and self._writing_paused:
+                return fail(write_error(ConnectionError("socket closed")))
+            return void
+
+        return self.signal.wait(ready).flat_map(check)
+
+    def end_writing(self) -> None:
+        """Half-close once buffered writes are flushed: close_notify on a
+        TLS connection, then end-of-stream."""
+
+        def run() -> None:
+            transport = self.transport
+            if transport is None or transport.is_closing():
+                return
+            if self._tls is not None:
+                self._tls.shutdown()
+            if transport.can_write_eof():
+                transport.write_eof()
+
+        call_soon(run)
+
+    def close(self, *, graceful: bool) -> None:
+        """Close the connection: flushing buffered writes first when
+        ``graceful``, else at once; a connection still being made is closed
+        when it is."""
+
+        def run() -> None:
+            self._close_requested = graceful
+            self._apply_close(graceful)
+
+        call_soon(run)
+
+    def _apply_close(self, graceful: bool) -> None:
+        transport = self.transport
+        if transport is None:
+            return
+        if graceful:
+            transport.close()
+        else:
+            transport.abort()

pyfect/socket/_errors.py

@@ -0,0 +1,135 @@
+"""Socket errors (Effect's ``SocketError`` and its reasons) and the close
+event a writer sends to close the socket."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Literal
+
+from pyfect.errors import TaggedError
+
+
+@dataclass(frozen=True, slots=True)
+class CloseEvent:
+    """Written to a socket, closes it with ``code`` (1000 is a normal
+    close)."""
+
+    code: int = 1000
+    reason: str | None = None
+
+    def __str__(self) -> str:
+        return f"{self.code}: {self.reason}" if self.reason else f"{self.code}"
+
+
+@dataclass(eq=False)
+class SocketReadError(TaggedError):
+    cause: object
+
+    @property
+    def message(self) -> str:
+        return "An error occurred during Read"
+
+
+@dataclass(eq=False)
+class SocketWriteError(TaggedError):
+    cause: object
+
+    @property
+    def message(self) -> str:
+        return "An error occurred during Write"
+
+
+type OpenErrorKind = Literal["Unknown", "Timeout"]
+
+
+@dataclass(eq=False)
+class SocketOpenError(TaggedError):
+    """Opening the socket failed; ``target`` is where it was going (a URL,
+    ``host:port`` or a path) when known."""
+
+    kind: OpenErrorKind
+    cause: object
+    target: str | None = None
+
+    @property
+    def message(self) -> str:
+        if self.kind == "Timeout":
+            return 'timeout waiting for "open"'
+        return "An error occurred during Open"
+
+    def __str__(self) -> str:
+        return self.message
+
+
+@dataclass(eq=False)
+class SocketUpgradeError(TaggedError):
+    """A TLS upgrade failed; without a ``cause``, the socket does not
+    support one."""
+
+    cause: object = None
+
+    @property
+    def message(self) -> str:
+        if self.cause is None:
+            return "Socket does not support TLS upgrade"
+        return "An error occurred during TLS upgrade"
+
+
+@dataclass(eq=False)
+class SocketCloseError(TaggedError):
+    """The socket closed: ``code`` 1000 is a normal close (the peer ended
+    the stream), 1006 an abnormal one."""
+
+    code: int
+    close_reason: str | None = None
+
+    @property
+    def message(self) -> str:
+        return f"{self.code}: {self.close_reason}" if self.close_reason else f"{self.code}"
+
+
+type SocketErrorReason = (
+    SocketReadError | SocketWriteError | SocketOpenError | SocketUpgradeError | SocketCloseError
+)
+
+
+@dataclass(eq=False)
+class SocketError(TaggedError):
+    """A socket operation failed; ``reason`` says how."""
+
+    reason: SocketErrorReason
+    cause: object = field(init=False, compare=False, repr=False)
+
+    def __post_init__(self) -> None:
+        self.cause = None if isinstance(self.reason, SocketCloseError) else self.reason.cause
+
+    @property
+    def message(self) -> str:
+        return self.reason.message
+
+
+def close_error(code: int, close_reason: str | None = None) -> SocketError:
+    return SocketError(SocketCloseError(code, close_reason))
+
+
+def is_clean_close(error: SocketError) -> bool:
+    """Whether ``error`` is a normal close (code 1000)."""
+    return isinstance(error.reason, SocketCloseError) and error.reason.code == 1000
+
+
+def read_error(cause: object) -> SocketError:
+    return SocketError(SocketReadError(cause))
+
+
+def write_error(cause: object) -> SocketError:
+    return SocketError(SocketWriteError(cause))
+
+
+def open_error(
+    cause: object, kind: OpenErrorKind = "Unknown", *, target: str | None = None
+) -> SocketError:
+    return SocketError(SocketOpenError(kind, cause, target))
+
+
+def upgrade_error(cause: object = None) -> SocketError:
+    return SocketError(SocketUpgradeError(cause))

pyfect/socket/_loop.py

@@ -0,0 +1,66 @@
+"""The asyncio loop the live sockets run on: one per process, on a daemon
+thread of its own, started on first use. Fibers never run on it: they ask
+it to act (``call_soon``, ``run``) and wait through ``os_wait``, which keeps
+a synchronous run (and its stall detection) aware of the outside work."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from collections.abc import Callable, Coroutine
+from concurrent.futures import Future
+from typing import Any
+
+from pyfect._core import Effect
+from pyfect._loop_thread import LoopThread
+from pyfect._os_wait import Resume, os_wait
+from pyfect.effect import fail, succeed
+
+
+class _Shared:
+    lock = threading.Lock()
+    loop: asyncio.AbstractEventLoop | None = None
+
+
+def _quiet(loop: asyncio.AbstractEventLoop, context: dict[str, Any]) -> None:
+    """Connection errors are reported to the socket that saw them; the
+    loop's own report of them (a failed TLS handshake, a reset during
+    close) would only be noise."""
+    if not isinstance(context.get("exception"), OSError):
+        loop.default_exception_handler(context)
+
+
+def get_loop() -> asyncio.AbstractEventLoop:
+    with _Shared.lock:
+        if _Shared.loop is None:
+            _Shared.loop = LoopThread("pyfect-net", exception_handler=_quiet).loop
+        return _Shared.loop
+
+
+def call_soon(f: Callable[[], object]) -> None:
+    """Run ``f`` on the loop thread."""
+    get_loop().call_soon_threadsafe(f)
+
+
+def run[A](make: Callable[[], Coroutine[Any, Any, A]]) -> Effect[A, BaseException]:
+    """Run the coroutine ``make()`` on the loop and wait for its result;
+    interrupting cancels it."""
+
+    def register(resume: Resume[A, BaseException]) -> Callable[[], None]:
+        future: Future[A] = asyncio.run_coroutine_threadsafe(make(), get_loop())
+
+        def done(finished: Future[A]) -> None:
+            if finished.cancelled():
+                resume(fail(asyncio.CancelledError()))
+                return
+            error = finished.exception()
+            resume(succeed(finished.result()) if error is None else fail(error))
+
+        future.add_done_callback(done)
+
+        def cancel() -> None:
+            future.cancel()
+
+        return cancel
+
+    return os_wait(register, waits_for="work")