mrok 0.4.6__py3-none-any.whl → 0.5.0__py3-none-any.whl

mrok/proxy/backend.py

@@ -0,0 +1,43 @@
+import asyncio
+from collections.abc import Iterable
+from pathlib import Path
+
+import openziti
+from httpcore import SOCKET_OPTION, AsyncNetworkBackend, AsyncNetworkStream
+from openziti.context import ZitiContext
+
+from mrok.proxy.exceptions import TargetUnavailableError
+from mrok.proxy.streams import AIONetworkStream
+
+
+class AIOZitiNetworkBackend(AsyncNetworkBackend):
+    def __init__(self, identity_file: str | Path) -> None:
+        self._identity_file = identity_file
+        self._ziti_ctx: ZitiContext | None = None
+
+    def _get_ziti_ctx(self) -> ZitiContext:
+        if self._ziti_ctx is None:
+            ctx, err = openziti.load(str(self._identity_file), timeout=10_000)
+            if err != 0:
+                raise Exception(f"Cannot create a Ziti context from the identity file: {err}")
+            self._ziti_ctx = ctx
+        return self._ziti_ctx
+
+    async def connect_tcp(
+        self,
+        host: str,
+        port: int,
+        timeout: float | None = None,
+        local_address: str | None = None,
+        socket_options: Iterable[SOCKET_OPTION] | None = None,
+    ) -> AsyncNetworkStream:
+        ctx = self._get_ziti_ctx()
+        try:
+            sock = ctx.connect(host)
+            reader, writer = await asyncio.open_connection(sock=sock)
+            return AIONetworkStream(reader, writer)
+        except Exception as e:
+            raise TargetUnavailableError() from e
+
+    async def sleep(self, seconds: float) -> None:
+        await asyncio.sleep(seconds)

mrok/{http → proxy}/config.py

@@ -8,12 +8,12 @@ from typing import Any
 import openziti
 from uvicorn import config
 
-from mrok.http.protocol import MrokHttpToolsProtocol
-from mrok.http.types import ASGIApp
+from mrok.proxy.protocol import MrokHttpToolsProtocol
+from mrok.proxy.types import ASGIApp
 
 logger = logging.getLogger("mrok.proxy")
 
-config.LIFESPAN["auto"] = "mrok.http.lifespan:MrokLifespan"
+config.LIFESPAN["auto"] = "mrok.proxy.lifespan:MrokLifespan"
 
 
 class MrokBackendConfig(config.Config):

mrok/{datastructures.py → proxy/datastructures.py}

@@ -1,11 +1,52 @@
 from __future__ import annotations
 
+import json
+from pathlib import Path
 from typing import Literal
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 from pydantic_core import core_schema
 
 
+class ZitiId(BaseModel):
+    key: str
+    cert: str
+    ca: str
+
+    @field_validator("key", "cert", "ca", mode="before")
+    @classmethod
+    def strip_pem_prefix(cls, value: str) -> str:
+        if isinstance(value, str) and value.startswith("pem:"):
+            return value[4:]
+        return value
+
+
+class ZitiMrokMeta(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    identity: str
+    extension: str
+    instance: str
+    domain: str | None = None
+    tags: dict[str, str | bool | None] | None = None
+
+
+class ZitiIdentity(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    zt_api: str = Field(validation_alias="ztAPI")
+    id: ZitiId
+    zt_apis: str | None = Field(default=None, validation_alias="ztAPIs")
+    config_types: str | None = Field(default=None, validation_alias="configTypes")
+    enable_ha: bool = Field(default=False, validation_alias="enableHa")
+    mrok: ZitiMrokMeta | None = None
+
+    @staticmethod
+    def load_from_file(path: str | Path) -> ZitiIdentity:
+        path = Path(path)
+        with path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        return ZitiIdentity.model_validate(data)
+
+
 class FixedSizeByteBuffer:
     def __init__(self, max_size: int):
         self._max_size = max_size
@@ -140,17 +181,9 @@ class WorkerMetrics(BaseModel):
     process: ProcessMetrics
 
 
-class Meta(BaseModel):
-    identity: str
-    extension: str
-    instance: str
-    domain: str
-    tags: dict[str, str] | None = None
-
-
 class Status(BaseModel):
     type: Literal["status"] = "status"
-    meta: Meta
+    meta: ZitiMrokMeta
     metrics: WorkerMetrics
 
 

mrok/proxy/exceptions.py

@@ -0,0 +1,22 @@
+from http import HTTPStatus
+
+from httpcore import ConnectError
+
+
+class ProxyError(Exception):
+    def __init__(self, http_status: HTTPStatus, message: str) -> None:
+        self.http_status: HTTPStatus = http_status
+        self.message: str = message
+
+
+class InvalidTargetError(ProxyError):
+    def __init__(self):
+        super().__init__(HTTPStatus.BAD_GATEWAY, "Bad Gateway: invalid target extension.")
+
+
+class TargetUnavailableError(ProxyError, ConnectError):
+    def __init__(self):
+        super().__init__(
+            HTTPStatus.SERVICE_UNAVAILABLE,
+            "Service Unavailable: the target extension is unavailable.",
+        )

mrok/proxy/lifespan.py

@@ -0,0 +1,10 @@
+import logging
+
+from uvicorn.config import Config
+from uvicorn.lifespan.on import LifespanOn
+
+
+class MrokLifespan(LifespanOn):
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+        self.logger = logging.getLogger("mrok.proxy")

mrok/{master.py → proxy/master.py}

@@ -1,6 +1,5 @@
 import asyncio
 import contextlib
-import json
 import logging
 import os
 import signal
@@ -18,14 +17,13 @@ from watchfiles.filters import PythonFilter
 from watchfiles.run import CombinedProcess, start_process
 
 from mrok.conf import get_settings
-from mrok.datastructures import Event, HTTPResponse, Meta, Status
-from mrok.http.config import MrokBackendConfig
-from mrok.http.lifespan import LifespanWrapper
-from mrok.http.middlewares import CaptureMiddleware, MetricsMiddleware
-from mrok.http.server import MrokServer
-from mrok.http.types import ASGIApp
 from mrok.logging import setup_logging
-from mrok.metrics import WorkerMetricsCollector
+from mrok.proxy.config import MrokBackendConfig
+from mrok.proxy.datastructures import Event, HTTPResponse, Status, ZitiIdentity
+from mrok.proxy.metrics import WorkerMetricsCollector
+from mrok.proxy.middlewares import CaptureMiddleware, LifespanMiddleware, MetricsMiddleware
+from mrok.proxy.server import MrokServer
+from mrok.proxy.types import ASGIApp
 
 logger = logging.getLogger("mrok.agent")
 
@@ -52,7 +50,7 @@ def start_events_router(events_pub_port: int, events_sub_port: int):
     try:
         logger.info(f"Events router process started: {os.getpid()}")
         zmq.proxy(frontend, backend)
-    except KeyboardInterrupt:
+    except KeyboardInterrupt:  # pragma: no cover
         pass
     finally:
         frontend.close()
@@ -62,7 +60,7 @@ def start_events_router(events_pub_port: int, events_sub_port: int):
 
 def start_uvicorn_worker(
     worker_id: str,
-    app: ASGIApp,
+    app: ASGIApp | str,
     identity_file: str,
     events_pub_port: int,
     metrics_interval: float = 5.0,
@@ -70,12 +68,10 @@ def start_uvicorn_worker(
     import sys
 
     sys.path.insert(0, os.getcwd())
-    if isinstance(app, str):
-        app = import_from_string(app)
+    asgi_app = app if not isinstance(app, str) else import_from_string(app)
 
     setup_logging(get_settings())
-    identity = json.load(open(identity_file))
-    meta = Meta(**identity["mrok"])
+    identity = ZitiIdentity.load_from_file(identity_file)
     ctx = zmq.asyncio.Context()
     pub = ctx.socket(zmq.PUB)
     pub.connect(f"tcp://localhost:{events_pub_port}")
@@ -83,30 +79,32 @@ def start_uvicorn_worker(
 
     task = None
 
-    async def status_sender():
+    async def status_sender():  # pragma: no cover
         while True:
             snap = await metrics.snapshot()
-            event = Event(type="status", data=Status(meta=meta, metrics=snap))
+            event = Event(type="status", data=Status(meta=identity.mrok, metrics=snap))
             await pub.send_string(event.model_dump_json())
             await asyncio.sleep(metrics_interval)
 
-    async def on_startup():  # noqa
+    async def on_startup():  # pragma: no cover
         nonlocal task
+        await asyncio.sleep(0)
         task = asyncio.create_task(status_sender())
 
-    async def on_shutdown():  # noqa
+    async def on_shutdown():  # pragma: no cover
+        await asyncio.sleep(0)
         if task:
             task.cancel()
 
-    async def on_response_complete(response: HTTPResponse):
+    async def on_response_complete(response: HTTPResponse):  # pragma: no cover
         event = Event(type="response", data=response)
         await pub.send_string(event.model_dump_json())
 
     config = MrokBackendConfig(
-        LifespanWrapper(
+        LifespanMiddleware(
             MetricsMiddleware(
                 CaptureMiddleware(
-                    app,
+                    asgi_app,
                     on_response_complete,
                 ),
                 metrics,
@@ -154,7 +152,7 @@ class MasterBase(ABC):
 
     @abstractmethod
     def get_asgi_app(self):
-        pass
+        raise NotImplementedError()
 
     def setup_signals_handler(self):
         for sig in (signal.SIGINT, signal.SIGTERM):
@@ -181,18 +179,6 @@ class MasterBase(ABC):
         logger.info(f"Worker {worker_id} [{p.pid}] started")
         return p
 
-    def start(self):
-        self.start_events_router()
-        self.start_workers()
-        self.monitor_thread.start()
-
-    def stop(self):
-        if self.monitor_thread.is_alive():
-            logger.debug("Wait for monitor worker to exit")
-            self.monitor_thread.join(timeout=MONITOR_THREAD_JOIN_TIMEOUT)
-        self.stop_workers()
-        self.stop_events_router()
-
     def start_events_router(self):
         self.zmq_pubsub_router_process = start_process(
             start_events_router,
@@ -204,21 +190,32 @@ class MasterBase(ABC):
             None,
         )
 
-    def stop_events_router(self):
-        self.zmq_pubsub_router_process.stop(sigint_timeout=5, sigkill_timeout=1)
-
     def start_workers(self):
         for i in range(self.workers):
             worker_id = self.worker_identifiers[i]
             p = self.start_worker(worker_id)
             self.worker_processes[worker_id] = p
 
+    def start(self):
+        self.start_events_router()
+        self.start_workers()
+        self.monitor_thread.start()
+
+    def stop_events_router(self):
+        self.zmq_pubsub_router_process.stop(sigint_timeout=5, sigkill_timeout=1)
+
     def stop_workers(self):
         for process in self.worker_processes.values():
-            if process.is_alive():
+            if process.is_alive():  # pragma: no branch
                 process.stop(sigint_timeout=5, sigkill_timeout=1)
         self.worker_processes.clear()
 
+    def stop(self):
+        if self.monitor_thread.is_alive():  # pragma: no branch
+            self.monitor_thread.join(timeout=MONITOR_THREAD_JOIN_TIMEOUT)
+        self.stop_workers()
+        self.stop_events_router()
+
     def restart(self):
         self.pause_event.set()
         self.stop_workers()

mrok/{metrics.py → proxy/metrics.py}

@@ -8,7 +8,7 @@ import time
 import psutil
 from hdrh.histogram import HdrHistogram
 
-from mrok.datastructures import (
+from mrok.proxy.datastructures import (
     DataTransferMetrics,
     ProcessMetrics,
     RequestsMetrics,
@@ -20,11 +20,7 @@ logger = logging.getLogger("mrok.proxy")
 
 
 def _collect_process_usage(interval: float) -> ProcessMetrics:
-    try:
-        proc = psutil.Process(os.getpid())
-    except psutil.NoSuchProcess:
-        return ProcessMetrics(cpu=0.0, mem=0.0)
-
+    proc = psutil.Process(os.getpid())
     total_cpu = 0.0
     total_mem = 0.0
 
@@ -33,29 +29,28 @@ def _collect_process_usage(interval: float) -> ProcessMetrics:
     except Exception:
         total_cpu = 0.0
 
-    if interval and interval > 0:
+    if interval and interval > 0:  # pragma: no branch
         time.sleep(interval)
 
     try:
         total_cpu = proc.cpu_percent(None)
-    except Exception:
+    except Exception:  # pragma: no cover
         total_cpu = 0.0
 
     try:
         total_mem = proc.memory_percent()
-    except Exception:
+    except Exception:  # pragma: no cover
         total_mem = 0.0
 
     return ProcessMetrics(cpu=total_cpu, mem=total_mem)
 
 
-async def get_process_and_children_usage(interval: float = 0.1) -> ProcessMetrics:
+async def get_process_metrics(interval: float = 0.1) -> ProcessMetrics:
     return await asyncio.to_thread(_collect_process_usage, interval)
 
 
 class WorkerMetricsCollector:
     def __init__(self, worker_id: str, lowest=1, highest=60000, sigfigs=3):
-        # Request-level counters
         self.worker_id = worker_id
         self.total_requests = 0
         self.successful_requests = 0
@@ -63,14 +58,11 @@ class WorkerMetricsCollector:
         self.bytes_in = 0
         self.bytes_out = 0
 
-        # RPS
         self._tick_last = time.time()
         self._tick_requests = 0
 
-        # latency histogram
         self.hist = HdrHistogram(lowest, highest, sigfigs)
 
-        # async lock
         self._lock = asyncio.Lock()
 
     async def on_request_start(self, scope):
@@ -102,38 +94,34 @@ class WorkerMetricsCollector:
             self.hist.record_value(elapsed_ms)
 
     async def snapshot(self) -> WorkerMetrics:
-        try:
-            async with self._lock:
-                now = time.time()
-                delta = now - self._tick_last
-                rps = int(self._tick_requests / delta) if delta > 0 else 0
-                data = WorkerMetrics(
-                    worker_id=self.worker_id,
-                    process=await get_process_and_children_usage(),
-                    requests=RequestsMetrics(
-                        rps=rps,
-                        total=self.total_requests,
-                        successful=self.successful_requests,
-                        failed=self.failed_requests,
-                    ),
-                    data_transfer=DataTransferMetrics(
-                        bytes_in=self.bytes_in,
-                        bytes_out=self.bytes_out,
-                    ),
-                    response_time=ResponseTimeMetrics(
-                        avg=self.hist.get_mean_value(),
-                        min=self.hist.get_min_value(),
-                        max=self.hist.get_max_value(),
-                        p50=self.hist.get_value_at_percentile(50),
-                        p90=self.hist.get_value_at_percentile(90),
-                        p99=self.hist.get_value_at_percentile(99),
-                    ),
-                )
-
-                self._tick_last = now
-                self._tick_requests = 0
-
-                return data
-        except Exception:
-            logger.exception("Exception calculating snapshot")
-            raise
+        async with self._lock:
+            now = time.time()
+            delta = now - self._tick_last
+            rps = int(self._tick_requests / delta) if delta > 0 else 0
+            data = WorkerMetrics(
+                worker_id=self.worker_id,
+                process=await get_process_metrics(),
+                requests=RequestsMetrics(
+                    rps=rps,
+                    total=self.total_requests,
+                    successful=self.successful_requests,
+                    failed=self.failed_requests,
+                ),
+                data_transfer=DataTransferMetrics(
+                    bytes_in=self.bytes_in,
+                    bytes_out=self.bytes_out,
+                ),
+                response_time=ResponseTimeMetrics(
+                    avg=self.hist.get_mean_value(),
+                    min=self.hist.get_min_value(),
+                    max=self.hist.get_max_value(),
+                    p50=self.hist.get_value_at_percentile(50),
+                    p90=self.hist.get_value_at_percentile(90),
+                    p99=self.hist.get_value_at_percentile(99),
+                ),
+            )
+
+            self._tick_last = now
+            self._tick_requests = 0
+
+            return data

mrok/{http → proxy}/middlewares.py

@@ -1,20 +1,23 @@
 import asyncio
-import inspect
 import logging
 import time
-from collections.abc import Callable, Coroutine
-from typing import Any
 
-from mrok.datastructures import FixedSizeByteBuffer, HTTPHeaders, HTTPRequest, HTTPResponse
-from mrok.http.constants import MAX_REQUEST_BODY_BYTES, MAX_RESPONSE_BODY_BYTES
-from mrok.http.types import ASGIApp, ASGIReceive, ASGISend, Message, Scope
-from mrok.http.utils import must_capture_request, must_capture_response
-from mrok.metrics import WorkerMetricsCollector
+from mrok.proxy.constants import MAX_REQUEST_BODY_BYTES, MAX_RESPONSE_BODY_BYTES
+from mrok.proxy.datastructures import FixedSizeByteBuffer, HTTPHeaders, HTTPRequest, HTTPResponse
+from mrok.proxy.metrics import WorkerMetricsCollector
+from mrok.proxy.types import (
+    ASGIApp,
+    ASGIReceive,
+    ASGISend,
+    LifespanCallback,
+    Message,
+    ResponseCompleteCallback,
+    Scope,
+)
+from mrok.proxy.utils import must_capture_request, must_capture_response
 
 logger = logging.getLogger("mrok.proxy")
 
-ResponseCompleteCallback = Callable[[HTTPResponse], Coroutine[Any, Any, None] | None]
-
 
 class CaptureMiddleware:
     def __init__(
@@ -92,22 +95,11 @@ class CaptureMiddleware:
             body=resp_buf.getvalue() if state["capture_resp_body"] else None,
             body_truncated=resp_buf.overflow if state["capture_resp_body"] else None,
         )
-        await asyncio.create_task(self.handle_callback(response))
-
-    async def handle_callback(self, response: HTTPResponse):
-        try:
-            if inspect.iscoroutinefunction(self._on_response_complete):
-                await self._on_response_complete(response)
-            else:
-                await asyncio.get_running_loop().run_in_executor(
-                    None, self._on_response_complete, response
-                )
-        except Exception:
-            logger.exception("Error invoking callback")
+        asyncio.create_task(self._on_response_complete(response))
 
 
 class MetricsMiddleware:
-    def __init__(self, app, metrics: WorkerMetricsCollector):
+    def __init__(self, app: ASGIApp, metrics: WorkerMetricsCollector):
         self.app = app
         self.metrics = metrics
 
@@ -116,11 +108,11 @@ class MetricsMiddleware:
             return await self.app(scope, receive, send)
 
         start_time = await self.metrics.on_request_start(scope)
-        status_code = 500  # default if app errors early
+        status_code = 500
 
         async def wrapped_receive():
             msg = await receive()
-            if msg["type"] == "http.request" and msg.get("body"):
+            if msg["type"] == "http.request" and msg.get("body"):  # pragma: no branch
                 await self.metrics.on_request_body(len(msg["body"]))
             return msg
 
@@ -131,7 +123,7 @@ class MetricsMiddleware:
                 status_code = msg["status"]
                 await self.metrics.on_response_start(status_code)
 
-            elif msg["type"] == "http.response.body":
+            elif msg["type"] == "http.response.body":  # pragma: no branch
                 body = msg.get("body", b"")
                 await self.metrics.on_response_chunk(len(body))
 
@@ -141,3 +133,32 @@ class MetricsMiddleware:
             await self.app(scope, wrapped_receive, wrapped_send)
         finally:
             await self.metrics.on_request_end(start_time, status_code)
+
+
+class LifespanMiddleware:
+    def __init__(
+        self,
+        app,
+        on_startup: LifespanCallback | None = None,
+        on_shutdown: LifespanCallback | None = None,
+    ):
+        self.app = app
+        self.on_startup = on_startup
+        self.on_shutdown = on_shutdown
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] == "lifespan":
+            while True:
+                event = await receive()
+                if event["type"] == "lifespan.startup":
+                    if self.on_startup:  # pragma: no branch
+                        await self.on_startup()
+                    await send({"type": "lifespan.startup.complete"})
+
+                elif event["type"] == "lifespan.shutdown":
+                    if self.on_shutdown:  # pragma: no branch
+                        await self.on_shutdown()
+                    await send({"type": "lifespan.shutdown.complete"})
+                    break
+        else:
+            await self.app(scope, receive, send)

mrok/proxy/streams.py

@@ -0,0 +1,45 @@
+import asyncio
+
+from httpcore import AsyncNetworkStream
+
+from mrok.proxy.types import ASGIReceive
+
+
+class AIONetworkStream(AsyncNetworkStream):
+    def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        self._reader = reader
+        self._writer = writer
+
+    async def read(self, n: int, timeout: float | None = None) -> bytes:
+        return await asyncio.wait_for(self._reader.read(n), timeout)
+
+    async def write(self, data: bytes, timeout: float | None = None) -> None:
+        self._writer.write(data)
+        await asyncio.wait_for(self._writer.drain(), timeout)
+
+    async def aclose(self) -> None:
+        self._writer.close()
+        await self._writer.wait_closed()
+
+
+class ASGIRequestBodyStream:
+    def __init__(self, receive: ASGIReceive):
+        self._receive = receive
+        self._more_body = True
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> bytes:
+        if not self._more_body:
+            raise StopAsyncIteration
+
+        msg = await self._receive()
+        if msg["type"] == "http.request":
+            chunk = msg.get("body", b"")
+            self._more_body = msg.get("more_body", False)
+            return chunk
+        elif msg["type"] == "http.disconnect":
+            raise Exception("Client disconnected.")
+
+        raise Exception("Unexpected asgi message.")

mrok/proxy/types.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Coroutine, MutableMapping
+from typing import Any, Never
+
+from mrok.proxy.datastructures import HTTPResponse
+
+Scope = MutableMapping[str, Any]
+Message = MutableMapping[str, Any]
+
+ASGIReceive = Callable[[], Awaitable[Message]]
+ASGISend = Callable[[Message], Awaitable[None]]
+ASGIApp = Callable[[Scope, ASGIReceive, ASGISend], Awaitable[None]]
+LifespanCallback = Callable[[], Awaitable[None]]
+ResponseCompleteCallback = Callable[[HTTPResponse], Coroutine[Any, Any, Never]]

mrok/{http → proxy}/utils.py

@@ -1,6 +1,6 @@
 from collections.abc import Mapping
 
-from mrok.http.constants import (
+from mrok.proxy.constants import (
     BINARY_CONTENT_TYPES,
     BINARY_PREFIXES,
     MAX_REQUEST_BODY_BYTES,