omlish 0.0.0.dev217__py3-none-any.whl → 0.0.0.dev218__py3-none-any.whl

omlish/sockets/bind.py

@@ -0,0 +1,332 @@
+# ruff: noqa: UP006 UP007
+# @omlish-lite
+"""
+TODO:
+ - DupSocketBinder
+"""
+import abc
+import dataclasses as dc
+import errno
+import os
+import socket as socket_
+import stat
+import typing as ta
+
+from omlish.lite.check import check
+from omlish.lite.dataclasses import dataclass_maybe_post_init
+from omlish.sockets.addresses import SocketAddress
+from omlish.sockets.addresses import SocketAndAddress
+
+
+SocketBinderT = ta.TypeVar('SocketBinderT', bound='SocketBinder')
+SocketBinderConfigT = ta.TypeVar('SocketBinderConfigT', bound='SocketBinder.Config')
+
+
+##
+
+
+class SocketBinder(abc.ABC, ta.Generic[SocketBinderConfigT]):
+    @dc.dataclass(frozen=True)
+    class Config:
+        listen_backlog: int = 5
+
+        allow_reuse_address: bool = True
+        allow_reuse_port: bool = True
+
+        set_inheritable: bool = False
+
+        #
+
+        @classmethod
+        def new(
+                cls,
+                target: ta.Union[
+                    int,
+                    ta.Tuple[str, int],
+                    str,
+                ],
+        ) -> 'SocketBinder.Config':
+            if isinstance(target, int):
+                return TcpSocketBinder.Config(
+                    port=target,
+                )
+
+            elif isinstance(target, tuple):
+                host, port = target
+                return TcpSocketBinder.Config(
+                    host=host,
+                    port=port,
+                )
+
+            elif isinstance(target, str):
+                return UnixSocketBinder.Config(
+                    file=target,
+                )
+
+            else:
+                raise TypeError(target)
+
+    #
+
+    def __init__(self, config: SocketBinderConfigT) -> None:
+        super().__init__()
+
+        self._config = config
+
+    #
+
+    @classmethod
+    def new(cls, target: ta.Any) -> 'SocketBinder':
+        config: SocketBinder.Config
+        if isinstance(target, SocketBinder.Config):
+            config = target
+
+        else:
+            config = SocketBinder.Config.new(target)
+
+        if isinstance(config, TcpSocketBinder.Config):
+            return TcpSocketBinder(config)
+
+        elif isinstance(config, UnixSocketBinder.Config):
+            return UnixSocketBinder(config)
+
+        else:
+            raise TypeError(config)
+
+    #
+
+    class Error(RuntimeError):
+        pass
+
+    class NotBoundError(Error):
+        pass
+
+    class AlreadyBoundError(Error):
+        pass
+
+    #
+
+    @property
+    @abc.abstractmethod
+    def address_family(self) -> int:
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def address(self) -> SocketAddress:
+        raise NotImplementedError
+
+    #
+
+    _socket: socket_.socket
+
+    @property
+    def is_bound(self) -> bool:
+        return hasattr(self, '_socket')
+
+    @property
+    def socket(self) -> socket_.socket:
+        try:
+            return self._socket
+        except AttributeError:
+            raise self.NotBoundError from None
+
+    _name: str
+
+    @property
+    def name(self) -> str:
+        try:
+            return self._name
+        except AttributeError:
+            raise self.NotBoundError from None
+
+    _port: ta.Optional[int]
+
+    @property
+    def port(self) -> ta.Optional[int]:
+        try:
+            return self._port
+        except AttributeError:
+            raise self.NotBoundError from None
+
+    #
+
+    def fileno(self) -> int:
+        return self.socket.fileno()
+
+    #
+
+    def __enter__(self: SocketBinderT) -> SocketBinderT:
+        self.bind()
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()
+
+    #
+
+    def _init_socket(self) -> None:
+        if hasattr(self, '_socket'):
+            raise self.AlreadyBoundError
+
+        socket = socket_.socket(self.address_family, socket_.SOCK_STREAM)
+        self._socket = socket
+
+        if self._config.allow_reuse_address and hasattr(socket_, 'SO_REUSEADDR'):
+            socket.setsockopt(socket_.SOL_SOCKET, socket_.SO_REUSEADDR, 1)
+
+        # Since Linux 6.12.9, SO_REUSEPORT is not allowed on other address families than AF_INET/AF_INET6.
+        if (
+                self._config.allow_reuse_port and hasattr(socket_, 'SO_REUSEPORT') and
+                self.address_family in (socket_.AF_INET, socket_.AF_INET6)
+        ):
+            try:
+                socket.setsockopt(socket_.SOL_SOCKET, socket_.SO_REUSEPORT, 1)
+            except OSError as err:
+                if err.errno not in (errno.ENOPROTOOPT, errno.EINVAL):
+                    raise
+
+        if self._config.set_inheritable and hasattr(socket, 'set_inheritable'):
+            socket.set_inheritable(True)
+
+    def _pre_bind(self) -> None:
+        pass
+
+    def _post_bind(self) -> None:
+        pass
+
+    def bind(self) -> None:
+        self._init_socket()
+
+        self._pre_bind()
+
+        self.socket.bind(self.address)
+
+        self._post_bind()
+
+        check.state(all(hasattr(self, a) for a in ('_socket', '_name', '_port')))
+
+    #
+
+    def close(self) -> None:
+        if hasattr(self, '_socket'):
+            self._socket.close()
+
+    #
+
+    def listen(self) -> None:
+        self.socket.listen(self._config.listen_backlog)
+
+    @abc.abstractmethod
+    def accept(self, socket: ta.Optional[socket_.socket] = None) -> SocketAndAddress:
+        raise NotImplementedError
+
+
+##
+
+
+class TcpSocketBinder(SocketBinder):
+    @dc.dataclass(frozen=True)
+    class Config(SocketBinder.Config):
+        DEFAULT_HOST: ta.ClassVar[str] = 'localhost'
+        host: str = DEFAULT_HOST
+
+        port: int = 0
+
+        def __post_init__(self) -> None:
+            dataclass_maybe_post_init(super())
+            check.non_empty_str(self.host)
+            check.isinstance(self.port, int)
+            check.arg(self.port > 0)
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(check.isinstance(config, self.Config))
+
+        self._address = (config.host, config.port)
+
+    #
+
+    address_family = socket_.AF_INET
+
+    @property
+    def address(self) -> SocketAddress:
+        return self._address
+
+    #
+
+    def _post_bind(self) -> None:
+        super()._post_bind()
+
+        host, port, *_ = self.socket.getsockname()
+
+        self._name = socket_.getfqdn(host)
+        self._port = port
+
+    #
+
+    def accept(self, socket: ta.Optional[socket_.socket] = None) -> SocketAndAddress:
+        if socket is None:
+            socket = self.socket
+
+        conn, client_address = socket.accept()
+        return SocketAndAddress(conn, client_address)
+
+
+##
+
+
+class UnixSocketBinder(SocketBinder):
+    @dc.dataclass(frozen=True)
+    class Config(SocketBinder.Config):
+        file: str = ''
+
+        unlink: bool = False
+
+        def __post_init__(self) -> None:
+            dataclass_maybe_post_init(super())
+            check.non_empty_str(self.file)
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(check.isinstance(config, self.Config))
+
+        self._address = config.file
+
+    #
+
+    address_family = socket_.AF_UNIX
+
+    @property
+    def address(self) -> SocketAddress:
+        return self._address
+
+    #
+
+    def _pre_bind(self) -> None:
+        super()._pre_bind()
+
+        if self._config.unlink:
+            try:
+                os.unlink(self._config.file)
+            except FileNotFoundError:
+                pass
+
+    def _post_bind(self) -> None:
+        super()._post_bind()
+
+        name = self.socket.getsockname()
+
+        os.chmod(name, stat.S_IRWXU | stat.S_IRWXG)  # noqa
+
+        self._name = name
+        self._port = None
+
+    #
+
+    def accept(self, sock: ta.Optional[socket_.socket] = None) -> SocketAndAddress:
+        if sock is None:
+            sock = self.socket
+
+        conn, _ = sock.accept()
+        client_address = ('', 0)
+        return SocketAndAddress(conn, client_address)

omlish/sockets/handlers.py

@@ -1,30 +1,12 @@
 # ruff: noqa: UP006 UP007
 # @omlish-lite
-import abc
 import typing as ta
 
 from .addresses import SocketAddress
+from .io import SocketIoPair  # noqa
 
 
-SocketHandlerFactory = ta.Callable[[SocketAddress, ta.BinaryIO, ta.BinaryIO], 'SocketHandler']
+SocketHandler = ta.Callable[[SocketAddress, 'SocketIoPair'], None]  # ta.TypeAlias
 
 
 ##
-
-
-class SocketHandler(abc.ABC):
-    def __init__(
-            self,
-            client_address: SocketAddress,
-            rfile: ta.BinaryIO,
-            wfile: ta.BinaryIO,
-    ) -> None:
-        super().__init__()
-
-        self._client_address = client_address
-        self._rfile = rfile
-        self._wfile = wfile
-
-    @abc.abstractmethod
-    def handle(self) -> None:
-        raise NotImplementedError

omlish/sockets/io.py

@@ -0,0 +1,69 @@
+# ruff: noqa: UP006 UP007
+# @omlish-lite
+import io
+import socket
+import typing as ta
+
+
+##
+
+
+class SocketWriter(io.BufferedIOBase):
+    """
+    Simple writable BufferedIOBase implementation for a socket
+
+    Does not hold data in a buffer, avoiding any need to call flush().
+    """
+
+    def __init__(self, sock):
+        super().__init__()
+
+        self._sock = sock
+
+    def writable(self):
+        return True
+
+    def write(self, b):
+        self._sock.sendall(b)
+        with memoryview(b) as view:
+            return view.nbytes
+
+    def fileno(self):
+        return self._sock.fileno()
+
+
+class SocketIoPair(ta.NamedTuple):
+    r: ta.BinaryIO
+    w: ta.BinaryIO
+
+    @classmethod
+    def from_socket(
+            cls,
+            sock: socket.socket,
+            *,
+            r_buf_size: int = -1,
+            w_buf_size: int = 0,
+    ) -> 'SocketIoPair':
+        rf: ta.Any = sock.makefile('rb', r_buf_size)
+
+        if w_buf_size:
+            wf: ta.Any = SocketWriter(sock)
+        else:
+            wf = sock.makefile('wb', w_buf_size)
+
+        return cls(rf, wf)
+
+
+##
+
+
+def close_socket_immediately(sock: socket.socket) -> None:
+    try:
+        # Explicitly shutdown. socket.close() merely releases the socket and waits for GC to perform the actual close.
+        sock.shutdown(socket.SHUT_WR)
+
+    except OSError:
+        # Some platforms may raise ENOTCONN here
+        pass
+
+    sock.close()

omlish/sockets/server/handlers.py

@@ -0,0 +1,99 @@
+# @omlish-lite
+# ruff: noqa: UP006 UP007
+import dataclasses as dc
+import socket
+import typing as ta
+
+from ..addresses import SocketAndAddress
+from ..handlers import SocketHandler
+from ..io import SocketIoPair
+from ..io import close_socket_immediately
+
+
+SocketServerHandler = ta.Callable[[SocketAndAddress], None]  # ta.TypeAlias
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class StandardSocketServerHandler:
+    handler: SocketServerHandler
+
+    timeout: ta.Optional[float] = None
+
+    # http://bugs.python.org/issue6192
+    # TODO: https://eklitzke.org/the-caveats-of-tcp-nodelay
+    disable_nagle_algorithm: bool = False
+
+    no_close: bool = False
+
+    def __call__(self, conn: SocketAndAddress) -> None:
+        try:
+            if self.timeout is not None:
+                conn.socket.settimeout(self.timeout)
+
+            if self.disable_nagle_algorithm:
+                conn.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True)
+
+            self.handler(conn)
+
+        finally:
+            close_socket_immediately(conn.socket)
+
+
+#
+
+
+@dc.dataclass(frozen=True)
+class CallbackWrappedSocketServerHandler:
+    handler: SocketServerHandler
+
+    before_handle: ta.Optional[SocketServerHandler] = None
+    after_handle: ta.Optional[SocketServerHandler] = None
+
+    # Return True if suppress like __exit__
+    on_error: ta.Optional[ta.Callable[[SocketAndAddress, Exception], bool]] = None
+
+    finally_: ta.Optional[SocketServerHandler] = None
+
+    def __call__(self, conn: SocketAndAddress) -> None:
+        try:
+            if (before_handle := self.before_handle) is not None:
+                before_handle(conn)
+
+            self.handler(conn)
+
+        except Exception as e:
+            if (on_error := self.on_error) is not None and on_error(conn, e):
+                pass
+            else:
+                raise
+
+        else:
+            if (after_handle := self.after_handle) is not None:
+                after_handle(conn)
+
+        finally:
+            if (finally_ := self.finally_) is not None:
+                finally_(conn)
+
+
+#
+
+
+@dc.dataclass(frozen=True)
+class SocketHandlerSocketServerHandler:
+    handler: SocketHandler
+
+    r_buf_size: int = -1
+    w_buf_size: int = 0
+
+    def __call__(self, conn: SocketAndAddress) -> None:
+        fp = SocketIoPair.from_socket(
+            conn.socket,
+            r_buf_size=self.r_buf_size,
+            w_buf_size=self.w_buf_size,
+        )
+
+        self.handler(conn.address, fp)

omlish/sockets/server/server.py

@@ -0,0 +1,144 @@
+# @omlish-lite
+# ruff: noqa: UP006 UP007
+import abc
+import contextlib
+import selectors
+import threading
+import typing as ta
+
+from ..bind import SocketBinder
+from .handlers import SocketServerHandler
+
+
+##
+
+
+class SocketServer(abc.ABC):
+    def __init__(
+            self,
+            binder: SocketBinder,
+            handler: SocketServerHandler,
+            *,
+            on_error: ta.Optional[ta.Callable[[BaseException], None]] = None,
+            poll_interval: float = .5,
+            shutdown_timeout: ta.Optional[float] = None,
+    ) -> None:
+        super().__init__()
+
+        self._binder = binder
+        self._handler = handler
+        self._on_error = on_error
+        self._poll_interval = poll_interval
+        self._shutdown_timeout = shutdown_timeout
+
+        self._lock = threading.RLock()
+        self._is_shutdown = threading.Event()
+        self._should_shutdown = False
+
+    @property
+    def binder(self) -> SocketBinder:
+        return self._binder
+
+    @property
+    def handler(self) -> SocketServerHandler:
+        return self._handler
+
+    #
+
+    class SelectorProtocol(ta.Protocol):
+        def register(self, *args, **kwargs) -> None:
+            raise NotImplementedError
+
+        def select(self, *args, **kwargs) -> bool:
+            raise NotImplementedError
+
+    Selector: ta.ClassVar[ta.Any]
+    if hasattr(selectors, 'PollSelector'):
+        Selector = selectors.PollSelector
+    else:
+        Selector = selectors.SelectSelector
+
+    #
+
+    @contextlib.contextmanager
+    def _listen_context(self) -> ta.Iterator[SelectorProtocol]:
+        with contextlib.ExitStack() as es:
+            es.enter_context(self._lock)
+            es.enter_context(self._binder)
+
+            self._binder.listen()
+
+            self._is_shutdown.clear()
+            try:
+                # XXX: Consider using another file descriptor or connecting to the socket to wake this up instead of
+                # polling. Polling reduces our responsiveness to a shutdown request and wastes cpu at all other times.
+                with self.Selector() as selector:
+                    selector.register(self._binder.fileno(), selectors.EVENT_READ)
+
+                    yield selector
+
+            finally:
+                self._is_shutdown.set()
+
+    @contextlib.contextmanager
+    def loop_context(self, poll_interval: ta.Optional[float] = None) -> ta.Iterator[ta.Iterator[bool]]:
+        if poll_interval is None:
+            poll_interval = self._poll_interval
+
+        with self._listen_context() as selector:
+            def loop():
+                while not self._should_shutdown:
+                    ready = selector.select(poll_interval)
+
+                    # bpo-35017: shutdown() called during select(), exit immediately.
+                    if self._should_shutdown:
+                        break  # type: ignore[unreachable]
+
+                    if ready:
+                        try:
+                            conn = self._binder.accept()
+
+                        except OSError as exc:
+                            if (on_error := self._on_error) is not None:
+                                on_error(exc)
+
+                            return
+
+                        self._handler(conn)
+
+                    yield bool(ready)
+
+            yield loop()
+
+    def run(self, poll_interval: ta.Optional[float] = None) -> None:
+        with self.loop_context(poll_interval=poll_interval) as loop:
+            for _ in loop:
+                pass
+
+    #
+
+    class _NOT_SET:  # noqa
+        def __new__(cls, *args, **kwargs):  # noqa
+            raise TypeError
+
+    def shutdown(
+            self,
+            block: bool = False,
+            timeout: ta.Union[float, None, ta.Type[_NOT_SET]] = _NOT_SET,
+    ) -> None:
+        self._should_shutdown = True
+
+        if block:
+            if timeout is self._NOT_SET:
+                timeout = self._shutdown_timeout
+
+            if not self._is_shutdown.wait(timeout=timeout):  # type: ignore
+                raise TimeoutError
+
+    #
+
+    def __enter__(self) -> 'SocketServer':
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.shutdown()