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

ominfra/scripts/supervisor.py

@@ -142,10 +142,6 @@ ConfigDataT = ta.TypeVar('ConfigDataT', bound='ConfigData')
 # ../../omlish/http/parsing.py
 HttpHeaders = http.client.HTTPMessage  # ta.TypeAlias
 
-# ../../omlish/lite/contextmanagers.py
-ExitStackedT = ta.TypeVar('ExitStackedT', bound='ExitStacked')
-AsyncExitStackedT = ta.TypeVar('AsyncExitStackedT', bound='AsyncExitStacked')
-
 # ../../omlish/lite/inject.py
 U = ta.TypeVar('U')
 InjectorKeyCls = ta.Union[type, ta.NewType]
@@ -154,10 +150,11 @@ InjectorProviderFnMap = ta.Mapping['InjectorKey', 'InjectorProviderFn']
 InjectorBindingOrBindings = ta.Union['InjectorBinding', 'InjectorBindings']
 
 # ../../omlish/sockets/handlers.py
-SocketHandlerFactory = ta.Callable[[SocketAddress, ta.BinaryIO, ta.BinaryIO], 'SocketHandler']
+SocketHandler = ta.Callable[[SocketAddress, 'SocketIoPair'], None]  # ta.TypeAlias
 
 # ../../omlish/http/handlers.py
 HttpHandler = ta.Callable[['HttpHandlerRequest'], 'HttpHandlerResponse']  # ta.TypeAlias
+HttpHandlerResponseData = ta.Union[bytes, 'HttpHandlerResponseStreamedData']  # ta.TypeAlias  # noqa
 
 # ../../omlish/http/coro/server.py
 CoroHttpServerFactory = ta.Callable[[SocketAddress], 'CoroHttpServer']
@@ -2812,8 +2809,7 @@ class ProxyLogHandler(ProxyLogFilterer, logging.Handler):
 # ../../../omlish/sockets/addresses.py
 """
 TODO:
- - SocketClientAddress family / tuple pairs
-  + codification of https://docs.python.org/3/library/socket.html#socket-families
+ - codification of https://docs.python.org/3/library/socket.html#socket-families
 """
 
 
@@ -2839,11 +2835,16 @@ class SocketAddressInfo:
     sockaddr: SocketAddress
 
 
+class SocketFamilyAndAddress(ta.NamedTuple):
+    family: socket.AddressFamily
+    address: SocketAddress
+
+
 def get_best_socket_family(
         host: ta.Optional[str],
         port: ta.Union[str, int, None],
         family: ta.Union[int, socket.AddressFamily] = socket.AddressFamily.AF_UNSPEC,
-) -> ta.Tuple[socket.AddressFamily, SocketAddress]:
+) -> SocketFamilyAndAddress:
     """https://github.com/python/cpython/commit/f289084c83190cc72db4a70c58f007ec62e75247"""
 
     infos = socket.getaddrinfo(
@@ -2854,7 +2855,80 @@ def get_best_socket_family(
         flags=socket.AI_PASSIVE,
     )
     ai = SocketAddressInfo(*next(iter(infos)))
-    return ai.family, ai.sockaddr
+    return SocketFamilyAndAddress(ai.family, ai.sockaddr)
+
+
+class SocketAndAddress(ta.NamedTuple):
+    socket: socket.socket
+    address: SocketAddress
+
+
+########################################
+# ../../../omlish/sockets/io.py
+
+
+##
+
+
+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()
 
 
 ########################################
@@ -4426,116 +4500,6 @@ else:
     KqueueFdioPoller = None
 
 
-########################################
-# ../../../omlish/lite/contextmanagers.py
-
-
-##
-
-
-class ExitStacked:
-    _exit_stack: ta.Optional[contextlib.ExitStack] = None
-
-    def __enter__(self: ExitStackedT) -> ExitStackedT:
-        check.state(self._exit_stack is None)
-        es = self._exit_stack = contextlib.ExitStack()
-        es.__enter__()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if (es := self._exit_stack) is None:
-            return None
-        self._exit_contexts()
-        return es.__exit__(exc_type, exc_val, exc_tb)
-
-    def _exit_contexts(self) -> None:
-        pass
-
-    def _enter_context(self, cm: ta.ContextManager[T]) -> T:
-        es = check.not_none(self._exit_stack)
-        return es.enter_context(cm)
-
-
-class AsyncExitStacked:
-    _exit_stack: ta.Optional[contextlib.AsyncExitStack] = None
-
-    async def __aenter__(self: AsyncExitStackedT) -> AsyncExitStackedT:
-        check.state(self._exit_stack is None)
-        es = self._exit_stack = contextlib.AsyncExitStack()
-        await es.__aenter__()
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        if (es := self._exit_stack) is None:
-            return None
-        await self._async_exit_contexts()
-        return await es.__aexit__(exc_type, exc_val, exc_tb)
-
-    async def _async_exit_contexts(self) -> None:
-        pass
-
-    def _enter_context(self, cm: ta.ContextManager[T]) -> T:
-        es = check.not_none(self._exit_stack)
-        return es.enter_context(cm)
-
-    async def _enter_async_context(self, cm: ta.AsyncContextManager[T]) -> T:
-        es = check.not_none(self._exit_stack)
-        return await es.enter_async_context(cm)
-
-
-##
-
-
-@contextlib.contextmanager
-def defer(fn: ta.Callable) -> ta.Generator[ta.Callable, None, None]:
-    try:
-        yield fn
-    finally:
-        fn()
-
-
-@contextlib.asynccontextmanager
-async def adefer(fn: ta.Callable) -> ta.AsyncGenerator[ta.Callable, None]:
-    try:
-        yield fn
-    finally:
-        await fn()
-
-
-##
-
-
-@contextlib.contextmanager
-def attr_setting(obj, attr, val, *, default=None):  # noqa
-    not_set = object()
-    orig = getattr(obj, attr, not_set)
-    try:
-        setattr(obj, attr, val)
-        if orig is not not_set:
-            yield orig
-        else:
-            yield default
-    finally:
-        if orig is not_set:
-            delattr(obj, attr)
-        else:
-            setattr(obj, attr, orig)
-
-
-##
-
-
-class aclosing(contextlib.AbstractAsyncContextManager):  # noqa
-    def __init__(self, thing):
-        self.thing = thing
-
-    async def __aenter__(self):
-        return self.thing
-
-    async def __aexit__(self, *exc_info):
-        await self.thing.aclose()
-
-
 ########################################
 # ../../../omlish/lite/inject.py
 
@@ -6427,24 +6391,6 @@ def journald_log_handler_factory(
 ##
 
 
-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
-
-
 ########################################
 # ../configs.py
 
@@ -6865,6 +6811,9 @@ class SupervisorSetup(abc.ABC):
 # ../../../omlish/http/handlers.py
 
 
+##
+
+
 @dc.dataclass(frozen=True)
 class HttpHandlerRequest:
     client_address: SocketAddress
@@ -6879,10 +6828,16 @@ class HttpHandlerResponse:
     status: ta.Union[http.HTTPStatus, int]
 
     headers: ta.Optional[ta.Mapping[str, str]] = None
-    data: ta.Optional[bytes] = None
+    data: ta.Optional[HttpHandlerResponseData] = None
     close_connection: ta.Optional[bool] = None
 
 
+@dc.dataclass(frozen=True)
+class HttpHandlerResponseStreamedData:
+    iter: ta.Iterable[bytes]
+    length: ta.Optional[int] = None
+
+
 class HttpHandlerError(Exception):
     pass
 
@@ -7394,7 +7349,7 @@ class CoroHttpServer:
 
         message: ta.Optional[str] = None
         headers: ta.Optional[ta.Sequence['CoroHttpServer._Header']] = None
-        data: ta.Optional[bytes] = None
+        data: ta.Optional[HttpHandlerResponseData] = None
         close_connection: ta.Optional[bool] = False
 
         def get_header(self, key: str) -> ta.Optional['CoroHttpServer._Header']:
@@ -7405,7 +7360,7 @@ class CoroHttpServer:
 
     #
 
-    def _build_response_bytes(self, a: _Response) -> bytes:
+    def _build_response_head_bytes(self, a: _Response) -> bytes:
         out = io.BytesIO()
 
         if a.version >= HttpProtocolVersions.HTTP_1_0:
@@ -7420,11 +7375,22 @@ class CoroHttpServer:
 
             out.write(b'\r\n')
 
-        if a.data is not None:
-            out.write(a.data)
-
         return out.getvalue()
 
+    def _yield_response_data(self, a: _Response) -> ta.Iterator[bytes]:
+        if a.data is None:
+            return
+
+        elif isinstance(a.data, bytes):
+            yield a.data
+            return
+
+        elif isinstance(a.data, HttpHandlerResponseStreamedData):
+            yield from a.data.iter
+
+        else:
+            raise TypeError(a.data)
+
     #
 
     DEFAULT_CONTENT_TYPE = 'text/plain'
@@ -7435,8 +7401,17 @@ class CoroHttpServer:
 
         if resp.get_header('Content-Type') is None:
             nh.append(self._Header('Content-Type', self._default_content_type))
+
         if resp.data is not None and resp.get_header('Content-Length') is None:
-            nh.append(self._Header('Content-Length', str(len(resp.data))))
+            cl: ta.Optional[int]
+            if isinstance(resp.data, bytes):
+                cl = len(resp.data)
+            elif isinstance(resp.data, HttpHandlerResponseStreamedData):
+                cl = resp.data.length
+            else:
+                raise TypeError(resp.data)
+            if cl is not None:
+                nh.append(self._Header('Content-Length', str(cl)))
 
         if nh:
             kw.update(headers=[*(resp.headers or []), *nh])
@@ -7610,9 +7585,13 @@ class CoroHttpServer:
 
                 elif isinstance(o, self._Response):
                     i = None
+
                     r = self._preprocess_response(o)
-                    b = self._build_response_bytes(r)
-                    check.none((yield self.WriteIo(b)))
+                    hb = self._build_response_head_bytes(r)
+                    check.none((yield self.WriteIo(hb)))
+
+                    for b in self._yield_response_data(r):
+                        yield self.WriteIo(b)
 
                 else:
                     raise TypeError(o)
@@ -7731,27 +7710,20 @@ class CoroHttpServer:
 ##
 
 
-class CoroHttpServerSocketHandler(SocketHandler):
+class CoroHttpServerSocketHandler:  # SocketHandler
     def __init__(
             self,
-            client_address: SocketAddress,
-            rfile: ta.BinaryIO,
-            wfile: ta.BinaryIO,
-            *,
             server_factory: CoroHttpServerFactory,
+            *,
             log_handler: ta.Optional[ta.Callable[[CoroHttpServer, CoroHttpServer.AnyLogIo], None]] = None,
     ) -> None:
-        super().__init__(
-            client_address,
-            rfile,
-            wfile,
-        )
+        super().__init__()
 
         self._server_factory = server_factory
         self._log_handler = log_handler
 
-    def handle(self) -> None:
-        server = self._server_factory(self._client_address)
+    def __call__(self, client_address: SocketAddress, fp: SocketIoPair) -> None:
+        server = self._server_factory(client_address)
 
         gen = server.coro_handle()
 
@@ -7763,15 +7735,15 @@ class CoroHttpServerSocketHandler(SocketHandler):
                     self._log_handler(server, o)
 
             elif isinstance(o, CoroHttpServer.ReadIo):
-                i = self._rfile.read(o.sz)
+                i = fp.r.read(o.sz)
 
             elif isinstance(o, CoroHttpServer.ReadLineIo):
-                i = self._rfile.readline(o.sz)
+                i = fp.r.readline(o.sz)
 
             elif isinstance(o, CoroHttpServer.WriteIo):
                 i = None
-                self._wfile.write(o.data)
-                self._wfile.flush()
+                fp.w.write(o.data)
+                fp.w.flush()
 
             else:
                 raise TypeError(o)
@@ -8877,7 +8849,7 @@ class HttpServer(HasDispatchers):
 
         self._conns: ta.List[CoroHttpServerConnectionFdioHandler] = []
 
-        exit_stack.enter_context(defer(self._server.close))  # noqa
+        exit_stack.callback(self._server.close)
 
     def get_dispatchers(self) -> Dispatchers:
         l = []

ominfra/supervisor/http.py

@@ -10,7 +10,6 @@ from omlish.http.handlers import HttpHandlerRequest
 from omlish.http.handlers import HttpHandlerResponse
 from omlish.io.fdio.handlers import SocketFdioHandler
 from omlish.lite.check import check
-from omlish.lite.contextmanagers import defer
 from omlish.lite.json import JSON_PRETTY_KWARGS
 from omlish.sockets.addresses import SocketAddress
 
@@ -73,7 +72,7 @@ class HttpServer(HasDispatchers):
 
         self._conns: ta.List[CoroHttpServerConnectionFdioHandler] = []
 
-        exit_stack.enter_context(defer(self._server.close))  # noqa
+        exit_stack.callback(self._server.close)
 
     def get_dispatchers(self) -> Dispatchers:
         l = []

{ominfra-0.0.0.dev217.dist-info → ominfra-0.0.0.dev218.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ominfra
-Version: 0.0.0.dev217
+Version: 0.0.0.dev218
 Summary: ominfra
 Author: wrmsr
 License: BSD-3-Clause
@@ -12,8 +12,8 @@ Classifier: Operating System :: OS Independent
 Classifier: Operating System :: POSIX
 Requires-Python: >=3.12
 License-File: LICENSE
-Requires-Dist: omdev==0.0.0.dev217
-Requires-Dist: omlish==0.0.0.dev217
+Requires-Dist: omdev==0.0.0.dev218
+Requires-Dist: omlish==0.0.0.dev218
 Provides-Extra: all
 Requires-Dist: paramiko~=3.5; extra == "all"
 Requires-Dist: asyncssh~=2.18; extra == "all"

{ominfra-0.0.0.dev217.dist-info → ominfra-0.0.0.dev218.dist-info}/RECORD

@@ -114,7 +114,7 @@ ominfra/manage/targets/targets.py,sha256=7GP6UAZyJFEhpkJN6UQdpr_WN3p7C76v-s445y-
 ominfra/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ominfra/scripts/journald2aws.py,sha256=CM1nZMpecXJpcg7icJUCJ-_DX-7CduQurxZQ7EmhQmA,171450
 ominfra/scripts/manage.py,sha256=MaUur1_dDdMHt5YnE6Waxa_v70kydHzKCBHXAuTmKd0,362827
-ominfra/scripts/supervisor.py,sha256=Laa7EM9tsLzjzb2Pkg2DVeUJX7OpsWLHuwlA9WP0cNM,295534
+ominfra/scripts/supervisor.py,sha256=ejPO7r3Ta0RQbP28MyL8MJNS1jop6g8d9qdv-vJ0OYA,294602
 ominfra/supervisor/LICENSE.txt,sha256=ZrHY15PVR98y26Yg6iQfa-SXnUaYTDhrUsPVcEO5OKM,1874
 ominfra/supervisor/__init__.py,sha256=Y3l4WY4JRi2uLG6kgbGp93fuGfkxkKwZDvhsa0Rwgtk,15
 ominfra/supervisor/__main__.py,sha256=I0yFw-C08OOiZ3BF6lF1Oiv789EQXu-_j6whDhQUTEA,66
@@ -125,7 +125,7 @@ ominfra/supervisor/events.py,sha256=XGrtzHr1xm0dwjz329fn9eR0_Ap-LQL6Sk8LJ8eVDEo,
 ominfra/supervisor/exceptions.py,sha256=Qbu211H3CLlSmi9LsSikOwrcL5HgJP9ugvcKWlGTAoI,750
 ominfra/supervisor/groups.py,sha256=MBbsbt8Zh_WEYkGOr1KXa82gnPVw9wPB2lAnqhugXSc,2457
 ominfra/supervisor/groupsimpl.py,sha256=PCDyc_Wc-pnvIj56_aEyiA5exCpK6n9iErqnJzO2rZk,2281
-ominfra/supervisor/http.py,sha256=Tl5eT6dyf4KuQXKbhGo1B6CFEzUDtPl_DZAuUFVD9zI,3453
+ominfra/supervisor/http.py,sha256=P7afN223jeR_g36bPK-cWulnJo1wAeFUBWAaMKPycmw,3387
 ominfra/supervisor/inject.py,sha256=IeR-WKvK1sGNxMe6G2OBT5kSP7EUP5ipkDgcUFmDPCA,4838
 ominfra/supervisor/io.py,sha256=moaGNaPuYXEAUzLg8Qjo05DEIcOUNYUj8SSr8eT0d24,3198
 ominfra/supervisor/main.py,sha256=zCVuHZG2kGIPHwTLH4EUr5xpN4vJYessmKO2P0NE5RU,4381
@@ -156,9 +156,9 @@ ominfra/tailscale/api.py,sha256=C5-t_b6jZXUWcy5k8bXm7CFnk73pSdrlMOgGDeGVrpw,1370
 ominfra/tailscale/cli.py,sha256=3FnJbgpLw6gInTfhERd1mDy9ijjMUGxkdYVo43Tnxx4,3555
 ominfra/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ominfra/tools/listresources.py,sha256=4qVg5txsb10EHhvqXXeM6gJ2jx9LbroEnPydDv1uXs0,6176
-ominfra-0.0.0.dev217.dist-info/LICENSE,sha256=B_hVtavaA8zCYDW99DYdcpDLKz1n3BBRjZrcbv8uG8c,1451
-ominfra-0.0.0.dev217.dist-info/METADATA,sha256=iTe6sIvgKblZxtFkb4vMZGGDzVFpVoV0l2bXB-3Jf5k,731
-ominfra-0.0.0.dev217.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-ominfra-0.0.0.dev217.dist-info/entry_points.txt,sha256=kgecQ2MgGrM9qK744BoKS3tMesaC3yjLnl9pa5CRczg,37
-ominfra-0.0.0.dev217.dist-info/top_level.txt,sha256=E-b2OHkk_AOBLXHYZQ2EOFKl-_6uOGd8EjeG-Zy6h_w,8
-ominfra-0.0.0.dev217.dist-info/RECORD,,
+ominfra-0.0.0.dev218.dist-info/LICENSE,sha256=B_hVtavaA8zCYDW99DYdcpDLKz1n3BBRjZrcbv8uG8c,1451
+ominfra-0.0.0.dev218.dist-info/METADATA,sha256=CIoVjeFWhR47YXXqQ65hW7lHeJK6XpLd0jTGOcP6RA4,731
+ominfra-0.0.0.dev218.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+ominfra-0.0.0.dev218.dist-info/entry_points.txt,sha256=kgecQ2MgGrM9qK744BoKS3tMesaC3yjLnl9pa5CRczg,37
+ominfra-0.0.0.dev218.dist-info/top_level.txt,sha256=E-b2OHkk_AOBLXHYZQ2EOFKl-_6uOGd8EjeG-Zy6h_w,8
+ominfra-0.0.0.dev218.dist-info/RECORD,,