mrok 0.3.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

mrok/master.py

@@ -0,0 +1,272 @@
+import asyncio
+import contextlib
+import json
+import logging
+import os
+import signal
+import threading
+import time
+from abc import ABC, abstractmethod
+from pathlib import Path
+from uuid import uuid4
+
+import zmq
+import zmq.asyncio
+from uvicorn.importer import import_from_string
+from watchfiles import watch
+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
+
+logger = logging.getLogger("mrok.agent")
+
+MONITOR_THREAD_JOIN_TIMEOUT = 5
+MONITOR_THREAD_CHECK_DELAY = 1
+MONITOR_THREAD_ERROR_DELAY = 3
+
+
+def print_path(path):
+    try:
+        return f'"{path.relative_to(Path.cwd())}"'
+    except ValueError:
+        return f'"{path}"'
+
+
+def start_events_router(events_pub_port: int, events_sub_port: int):
+    setup_logging(get_settings())
+    context = zmq.Context()
+    frontend = context.socket(zmq.XSUB)
+    frontend.bind(f"tcp://localhost:{events_pub_port}")
+    backend = context.socket(zmq.XPUB)
+    backend.bind(f"tcp://localhost:{events_sub_port}")
+
+    try:
+        logger.info(f"Events router process started: {os.getpid()}")
+        zmq.proxy(frontend, backend)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        frontend.close()
+        backend.close()
+        context.term()
+
+
+def start_uvicorn_worker(
+    worker_id: str,
+    app: ASGIApp,
+    identity_file: str,
+    events_pub_port: int,
+    metrics_interval: float = 5.0,
+):
+    import sys
+
+    sys.path.insert(0, os.getcwd())
+    if isinstance(app, str):
+        app = import_from_string(app)
+
+    setup_logging(get_settings())
+    identity = json.load(open(identity_file))
+    meta = Meta(**identity["mrok"])
+    ctx = zmq.asyncio.Context()
+    pub = ctx.socket(zmq.PUB)
+    pub.connect(f"tcp://localhost:{events_pub_port}")
+    metrics = WorkerMetricsCollector(worker_id)
+
+    task = None
+
+    async def status_sender():
+        while True:
+            snap = await metrics.snapshot()
+            event = Event(type="status", data=Status(meta=meta, metrics=snap))
+            await pub.send_string(event.model_dump_json())
+            await asyncio.sleep(metrics_interval)
+
+    async def on_startup():  # noqa
+        nonlocal task
+        task = asyncio.create_task(status_sender())
+
+    async def on_shutdown():  # noqa
+        if task:
+            task.cancel()
+
+    async def on_response_complete(response: HTTPResponse):
+        event = Event(type="response", data=response)
+        await pub.send_string(event.model_dump_json())
+
+    config = MrokBackendConfig(
+        LifespanWrapper(
+            MetricsMiddleware(
+                CaptureMiddleware(
+                    app,
+                    on_response_complete,
+                ),
+                metrics,
+            ),
+            on_startup=on_startup,
+            on_shutdown=on_shutdown,
+        ),
+        identity_file,
+    )
+    server = MrokServer(config)
+    with contextlib.suppress(KeyboardInterrupt, asyncio.CancelledError):
+        server.run()
+
+
+class MasterBase(ABC):
+    def __init__(
+        self,
+        identity_file: str,
+        workers: int,
+        reload: bool,
+        events_pub_port: int,
+        events_sub_port: int,
+        metrics_interval: float = 5.0,
+    ):
+        self.identity_file = identity_file
+        self.workers = workers
+        self.reload = reload
+        self.events_pub_port = events_pub_port
+        self.events_sub_port = events_sub_port
+        self.metrics_interval = metrics_interval
+        self.worker_identifiers = [str(uuid4()) for _ in range(workers)]
+        self.worker_processes: dict[str, CombinedProcess] = {}
+        self.zmq_pubsub_router_process = None
+        self.monitor_thread = threading.Thread(target=self.monitor_workers, daemon=True)
+        self.stop_event = threading.Event()
+        self.pause_event = threading.Event()
+        self.watch_filter = PythonFilter(ignore_paths=None)
+        self.watcher = watch(
+            Path.cwd(),
+            watch_filter=self.watch_filter,
+            stop_event=self.stop_event,
+            yield_on_timeout=True,
+        )
+        self.setup_signals_handler()
+
+    @abstractmethod
+    def get_asgi_app(self):
+        pass
+
+    def setup_signals_handler(self):
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            signal.signal(sig, self.handle_signal)
+
+    def handle_signal(self, *args, **kwargs):
+        self.stop_event.set()
+
+    def start_worker(self, worker_id: str):
+        """Start a single worker process"""
+
+        p = start_process(
+            start_uvicorn_worker,
+            "function",
+            (
+                worker_id,
+                self.get_asgi_app(),
+                self.identity_file,
+                self.events_pub_port,
+                self.metrics_interval,
+            ),
+            None,
+        )
+        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,
+            "function",
+            (
+                self.events_pub_port,
+                self.events_sub_port,
+            ),
+            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 stop_workers(self):
+        for process in self.worker_processes.values():
+            if process.is_alive():
+                process.stop(sigint_timeout=5, sigkill_timeout=1)
+        self.worker_processes.clear()
+
+    def restart(self):
+        self.pause_event.set()
+        self.stop_workers()
+        self.start_workers()
+        self.pause_event.clear()
+
+    def monitor_workers(self):
+        while not self.stop_event.is_set():
+            try:
+                self.pause_event.wait()
+                for worker_id, process in self.worker_processes.items():
+                    if not process.is_alive():
+                        logger.warning(f"Worker {worker_id} [{process.pid}] died unexpectedly")
+                        process.stop(sigint_timeout=1, sigkill_timeout=1)
+                        new_process = self.start_worker(worker_id)
+                        self.worker_processes[worker_id] = new_process
+                        logger.info(
+                            f"Restarted worker {worker_id} [{process.pid}] -> [{new_process.pid}]"
+                        )
+
+                time.sleep(MONITOR_THREAD_CHECK_DELAY)
+
+            except Exception as e:
+                logger.error(f"Error in worker monitoring: {e}")
+                time.sleep(MONITOR_THREAD_ERROR_DELAY)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        changes = next(self.watcher)
+        if changes:
+            return list({Path(change[1]) for change in changes})
+        return None
+
+    def run(self):
+        setup_logging(get_settings())
+        logger.info(f"Master process started: {os.getpid()}")
+        self.start()
+        try:
+            if self.reload:
+                for files_changed in self:
+                    if files_changed:
+                        logger.warning(
+                            f"{', '.join(map(print_path, files_changed))} changed, reloading...",
+                        )
+                        self.restart()
+            else:
+                self.stop_event.wait()
+        finally:
+            self.stop()

mrok/metrics.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import time
+
+import psutil
+from hdrh.histogram import HdrHistogram
+
+from mrok.datastructures import (
+    DataTransferMetrics,
+    ProcessMetrics,
+    RequestsMetrics,
+    ResponseTimeMetrics,
+    WorkerMetrics,
+)
+
+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)
+
+    total_cpu = 0.0
+    total_mem = 0.0
+
+    try:
+        total_cpu = proc.cpu_percent(None)
+    except Exception:
+        total_cpu = 0.0
+
+    if interval and interval > 0:
+        time.sleep(interval)
+
+    try:
+        total_cpu = proc.cpu_percent(None)
+    except Exception:
+        total_cpu = 0.0
+
+    try:
+        total_mem = proc.memory_percent()
+    except Exception:
+        total_mem = 0.0
+
+    return ProcessMetrics(cpu=total_cpu, mem=total_mem)
+
+
+async def get_process_and_children_usage(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
+        self.failed_requests = 0
+        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):
+        return time.perf_counter()
+
+    async def on_request_body(self, length):
+        async with self._lock:
+            self.bytes_in += length
+
+    async def on_response_start(self, status_code):
+        pass  # reserved
+
+    async def on_response_chunk(self, length):
+        async with self._lock:
+            self.bytes_out += length
+
+    async def on_request_end(self, start_time, status_code):
+        elapsed_ms = (time.perf_counter() - start_time) * 1000
+
+        async with self._lock:
+            self.total_requests += 1
+            self._tick_requests += 1
+
+            if status_code < 500:
+                self.successful_requests += 1
+            else:
+                self.failed_requests += 1
+
+            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

mrok/proxy/__init__.py

@@ -0,0 +1,3 @@
+from mrok.proxy.main import run
+
+__all__ = ["run"]

mrok/proxy/app.py

@@ -0,0 +1,77 @@
+import logging
+import os
+from pathlib import Path
+
+from mrok.conf import get_settings
+from mrok.http.forwarder import ForwardAppBase
+from mrok.http.types import Scope, StreamReader, StreamWriter
+from mrok.logging import setup_logging
+from mrok.proxy.ziti import ZitiConnectionManager
+
+logger = logging.getLogger("mrok.proxy")
+
+
+class ProxyError(Exception):
+    pass
+
+
+class ProxyApp(ForwardAppBase):
+    def __init__(
+        self,
+        identity_file: str | Path,
+        *,
+        read_chunk_size: int = 65536,
+        ziti_connection_ttl_seconds: float = 60,
+        ziti_conn_cache_purge_interval_seconds: float = 10,
+    ) -> None:
+        super().__init__(read_chunk_size=read_chunk_size)
+        self._identity_file = identity_file
+        settings = get_settings()
+        self._proxy_wildcard_domain = (
+            settings.proxy.domain
+            if settings.proxy.domain[0] == "."
+            else f".{settings.proxy.domain}"
+        )
+        self._conn_manager = ZitiConnectionManager(
+            identity_file,
+            ttl_seconds=ziti_connection_ttl_seconds,
+            purge_interval=ziti_conn_cache_purge_interval_seconds,
+        )
+
+    def get_target_from_header(self, name: str, headers: dict[str, str]) -> str:
+        header_value = headers.get(name)
+        if not header_value:
+            raise ProxyError(
+                f"Header {name} not found!",
+            )
+        if ":" in header_value:
+            header_value, _ = header_value.split(":", 1)
+        if not header_value.endswith(self._proxy_wildcard_domain):
+            raise ProxyError(f"Unexpected value for {name} header: `{header_value}`.")
+
+        return header_value[: -len(self._proxy_wildcard_domain)]
+
+    def get_target_name(self, headers: dict[str, str]) -> str:
+        try:
+            return self.get_target_from_header("x-forwared-for", headers)
+        except ProxyError as pe:
+            logger.warning(pe)
+            return self.get_target_from_header("host", headers)
+
+    async def startup(self):
+        setup_logging(get_settings())
+        await self._conn_manager.start()
+        logger.info(f"Proxy app startup completed: {os.getpid()}")
+
+    async def shutdown(self):
+        await self._conn_manager.stop()
+        logger.info(f"Proxy app shutdown completed: {os.getpid()}")
+
+    async def select_backend(
+        self,
+        scope: Scope,
+        headers: dict[str, str],
+    ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
+        target_name = self.get_target_name(headers)
+
+        return await self._conn_manager.get(target_name)

mrok/proxy/dataclasses.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from mrok.http.types import StreamReader, StreamWriter
+
+
+@dataclass
+class CachedStreamEntry:
+    reader: StreamReader
+    writer: StreamWriter
+    last_access: float

mrok/proxy/main.py

@@ -0,0 +1,58 @@
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from gunicorn.app.base import BaseApplication
+from uvicorn_worker import UvicornWorker
+
+from mrok.conf import get_settings
+from mrok.http.lifespan import LifespanWrapper
+from mrok.logging import get_logging_config
+from mrok.proxy.app import ProxyApp
+
+
+class MrokUvicornWorker(UvicornWorker):
+    CONFIG_KWARGS: dict[str, Any] = {"loop": "asyncio", "http": "auto", "lifespan": "on"}
+
+
+class StandaloneApplication(BaseApplication):  # pragma: no cover
+    def __init__(self, application: Callable, options: dict[str, Any] | None = None):
+        self.options = options or {}
+        self.application = application
+        super().__init__()
+
+    def load_config(self):
+        config = {
+            key: value
+            for key, value in self.options.items()
+            if key in self.cfg.settings and value is not None
+        }
+        for key, value in config.items():
+            self.cfg.set(key.lower(), value)
+
+    def load(self):
+        return self.application
+
+
+def run(
+    identity_file: str | Path,
+    host: str,
+    port: int,
+    workers: int,
+):
+    proxy_app = ProxyApp(identity_file)
+
+    asgi_app = LifespanWrapper(
+        proxy_app,
+        proxy_app.startup,
+        proxy_app.shutdown,
+    )
+    options = {
+        "bind": f"{host}:{port}",
+        "workers": workers,
+        "worker_class": "mrok.proxy.main.MrokUvicornWorker",
+        "logconfig_dict": get_logging_config(get_settings()),
+        "reload": False,
+    }
+
+    StandaloneApplication(asgi_app, options).run()

mrok/proxy/streams.py

@@ -0,0 +1,124 @@
+import asyncio
+
+from mrok.proxy.types import ConnectionCache, ConnectionKey
+
+
+class CachedStreamReader:
+    def __init__(
+        self,
+        reader: asyncio.StreamReader,
+        key: ConnectionKey,
+        manager: ConnectionCache,
+    ):
+        self._reader = reader
+        self._key = key
+        self._manager = manager
+
+    async def read(self, n: int = -1) -> bytes:
+        try:
+            return await self._reader.read(n)
+        except (
+            asyncio.CancelledError,
+            asyncio.IncompleteReadError,
+            asyncio.LimitOverrunError,
+            BrokenPipeError,
+            ConnectionAbortedError,
+            ConnectionResetError,
+            RuntimeError,
+            TimeoutError,
+            UnicodeDecodeError,
+        ):
+            asyncio.create_task(self._manager.invalidate(self._key))
+            raise
+
+    async def readexactly(self, n: int) -> bytes:
+        try:
+            return await self._reader.readexactly(n)
+        except (
+            asyncio.CancelledError,
+            asyncio.IncompleteReadError,
+            asyncio.LimitOverrunError,
+            BrokenPipeError,
+            ConnectionAbortedError,
+            ConnectionResetError,
+            RuntimeError,
+            TimeoutError,
+            UnicodeDecodeError,
+        ):
+            asyncio.create_task(self._manager.invalidate(self._key))
+            raise
+
+    async def readline(self) -> bytes:
+        try:
+            return await self._reader.readline()
+        except (
+            asyncio.CancelledError,
+            asyncio.IncompleteReadError,
+            asyncio.LimitOverrunError,
+            BrokenPipeError,
+            ConnectionAbortedError,
+            ConnectionResetError,
+            RuntimeError,
+            TimeoutError,
+            UnicodeDecodeError,
+        ):
+            asyncio.create_task(self._manager.invalidate(self._key))
+            raise
+
+    def at_eof(self) -> bool:
+        return self._reader.at_eof()
+
+    @property
+    def underlying(self) -> asyncio.StreamReader:
+        return self._reader
+
+
+class CachedStreamWriter:
+    def __init__(
+        self,
+        writer: asyncio.StreamWriter,
+        key: ConnectionKey,
+        manager: ConnectionCache,
+    ):
+        self._writer = writer
+        self._key = key
+        self._manager = manager
+
+    def write(self, data: bytes) -> None:
+        try:
+            return self._writer.write(data)
+        except (RuntimeError, TypeError):
+            asyncio.create_task(self._manager.invalidate(self._key))
+            raise
+
+    async def drain(self) -> None:
+        try:
+            return await self._writer.drain()
+        except (
+            asyncio.CancelledError,
+            BrokenPipeError,
+            ConnectionAbortedError,
+            ConnectionResetError,
+            RuntimeError,
+            TimeoutError,
+        ):
+            asyncio.create_task(self._manager.invalidate(self._key))
+            raise
+
+    def close(self) -> None:
+        return self._writer.close()
+
+    async def wait_closed(self) -> None:
+        try:
+            return await self._writer.wait_closed()
+        except (ConnectionResetError, BrokenPipeError):
+            asyncio.create_task(self._manager.invalidate(self._key))
+            raise
+
+    @property
+    def transport(self):
+        return self._writer.transport
+
+    @property
+    def underlying(self) -> asyncio.StreamWriter:
+        return self._writer

mrok/proxy/types.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+from mrok.http.types import StreamReader, StreamWriter
+
+ConnectionKey = tuple[str, str | None]
+CachedStream = tuple[StreamReader, StreamWriter]
+
+
+class ConnectionCache(Protocol):
+    async def invalidate(self, key: ConnectionKey) -> None: ...