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

mrok/agent/devtools/inspector/app.py

@@ -27,7 +27,7 @@ from textual.widgets.data_table import ColumnKey
 from textual.worker import get_current_worker
 
 from mrok import __version__
-from mrok.proxy.datastructures import Event, HTTPHeaders, HTTPResponse, WorkerMetrics, ZitiMrokMeta
+from mrok.proxy.models import Event, HTTPHeaders, HTTPResponse, ServiceMetadata, WorkerMetrics
 
 
 def build_tree(node, data):
@@ -185,7 +185,7 @@ class InfoPanel(Static):
         #         mem=int(mean([m.process.mem for m in self.workers_metrics.values()])),
         #     )
 
-    def update_meta(self, meta: ZitiMrokMeta) -> None:
+    def update_meta(self, meta: ServiceMetadata) -> None:
         table = self.query_one(DataTable)
         if len(table.rows) == 0:
             table.add_row("URL", f"https://{meta.extension}.{meta.domain}")

mrok/agent/sidecar/app.py

@@ -5,7 +5,7 @@ from typing import Literal
 from httpcore import AsyncConnectionPool
 
 from mrok.proxy.app import ProxyAppBase
-from mrok.proxy.types import Scope
+from mrok.types.proxy import Scope
 
 logger = logging.getLogger("mrok.agent")
 
@@ -18,10 +18,10 @@ class SidecarProxyApp(ProxyAppBase):
         self,
         target: str | Path | tuple[str, int],
         *,
-        max_connections=1000,
-        max_keepalive_connections=10,
-        keepalive_expiry=120,
-        retries=0,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
+        retries: int = 0,
     ):
         self._target = target
         self._target_type, self._target_address = self._parse_target()
@@ -34,10 +34,10 @@ class SidecarProxyApp(ProxyAppBase):
 
     def setup_connection_pool(
         self,
-        max_connections: int | None = 1000,
-        max_keepalive_connections: int | None = 10,
-        keepalive_expiry: float | None = 120.0,
-        retries: int = 0,
+        max_connections: int | None,
+        max_keepalive_connections: int | None,
+        keepalive_expiry: float | None,
+        retries: int,
     ) -> AsyncConnectionPool:
         if self._target_type == "unix":
             return AsyncConnectionPool(

mrok/agent/sidecar/main.py

@@ -13,26 +13,47 @@ class SidecarAgent(MasterBase):
         identity_file: str,
         target: str | Path | tuple[str, int],
         workers: int = 4,
+        events_enabled: bool = True,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
+        retries: int = 0,
         publishers_port: int = 50000,
         subscribers_port: int = 50001,
     ):
         super().__init__(
             identity_file,
-            workers,
-            False,
-            publishers_port,
-            subscribers_port,
+            workers=workers,
+            reload=False,
+            events_enabled=events_enabled,
+            events_pub_port=publishers_port,
+            events_sub_port=subscribers_port,
         )
         self._target = target
+        self._max_connections = max_connections
+        self._max_keepalive_connections = max_keepalive_connections
+        self._keepalive_expiry = keepalive_expiry
+        self._retries = retries
 
     def get_asgi_app(self):
-        return SidecarProxyApp(self._target)
+        return SidecarProxyApp(
+            self._target,
+            max_connections=self._max_connections,
+            max_keepalive_connections=self._max_keepalive_connections,
+            keepalive_expiry=self._keepalive_expiry,
+            retries=self._retries,
+        )
 
 
 def run(
     identity_file: str,
     target_addr: str | Path | tuple[str, int],
     workers: int = 4,
+    events_enabled: bool = True,
+    max_connections: int | None = 10,
+    max_keepalive_connections: int | None = None,
+    keepalive_expiry: float | None = None,
+    retries: int = 0,
     publishers_port: int = 50000,
     subscribers_port: int = 50001,
 ):
@@ -40,6 +61,11 @@ def run(
         identity_file,
         target_addr,
         workers=workers,
+        events_enabled=events_enabled,
+        max_connections=max_connections,
+        max_keepalive_connections=max_keepalive_connections,
+        keepalive_expiry=keepalive_expiry,
+        retries=retries,
         publishers_port=publishers_port,
         subscribers_port=subscribers_port,
     )

mrok/agent/ziticorn.py

@@ -1,5 +1,5 @@
 from mrok.proxy.master import MasterBase
-from mrok.proxy.types import ASGIApp
+from mrok.types.proxy import ASGIApp
 
 
 class ZiticornAgent(MasterBase):
@@ -12,7 +12,13 @@ class ZiticornAgent(MasterBase):
         publishers_port: int = 50000,
         subscribers_port: int = 50001,
     ):
-        super().__init__(identity_file, workers, reload, publishers_port, subscribers_port)
+        super().__init__(
+            identity_file,
+            workers=workers,
+            reload=reload,
+            events_pub_port=publishers_port,
+            events_sub_port=subscribers_port,
+        )
         self.app = app
 
     def get_asgi_app(self):

mrok/cli/commands/admin/bootstrap.py

@@ -8,14 +8,15 @@ import typer
 
 from mrok.cli.commands.admin.utils import parse_tags
 from mrok.conf import Settings
-from mrok.ziti.api import TagsType, ZitiClientAPI, ZitiManagementAPI
+from mrok.types.ziti import Tags
+from mrok.ziti.api import ZitiClientAPI, ZitiManagementAPI
 from mrok.ziti.bootstrap import bootstrap_identity
 
 logger = logging.getLogger(__name__)
 
 
 async def bootstrap(
-    settings: Settings, forced: bool, tags: TagsType | None
+    settings: Settings, forced: bool, tags: Tags | None
 ) -> tuple[str, dict[str, Any] | None]:
     async with ZitiManagementAPI(settings) as mgmt_api, ZitiClientAPI(settings) as client_api:
         return await bootstrap_identity(

mrok/cli/commands/admin/utils.py

@@ -2,10 +2,10 @@ from datetime import datetime
 
 import typer
 
-from mrok.ziti.api import TagsType
+from mrok.types.ziti import Tags
 
 
-def parse_tags(pairs: list[str] | None) -> TagsType | None:
+def parse_tags(pairs: list[str] | None) -> Tags | None:
     if not pairs:
         return None
 

mrok/cli/commands/agent/run/sidecar.py

@@ -25,10 +25,55 @@ def register(app: typer.Typer) -> None:
             typer.Option(
                 "--workers",
                 "-w",
-                help=f"Number of workers. Default: {default_workers}",
+                help="Number of workers.",
                 show_default=True,
             ),
         ] = default_workers,
+        max_connections: Annotated[
+            int,
+            typer.Option(
+                "--max-pool-connections",
+                help=(
+                    "The maximum number of concurrent HTTP connections that "
+                    "the pool should allow. Any attempt to send a request on a pool that "
+                    "would exceed this amount will block until a connection is available."
+                ),
+                show_default=True,
+            ),
+        ] = 10,
+        max_keepalive_connections: Annotated[
+            int | None,
+            typer.Option(
+                "--max-pool-keepalive-connections",
+                help=(
+                    "The maximum number of idle HTTP connections "
+                    "that will be maintained in the pool."
+                ),
+                show_default=True,
+            ),
+        ] = None,
+        keepalive_expiry: Annotated[
+            float | None,
+            typer.Option(
+                "--max-pool-keepalive-expiry",
+                help=(
+                    "The duration in seconds that an idle HTTP connection "
+                    "may be maintained for before being expired from the pool."
+                ),
+                show_default=True,
+            ),
+        ] = None,
+        retries: Annotated[
+            int,
+            typer.Option(
+                "--max-pool-connect-retries",
+                help=(
+                    "The duration in seconds that an idle HTTP connection "
+                    "may be maintained for before being expired from the pool."
+                ),
+                show_default=True,
+            ),
+        ] = 0,
         publishers_port: Annotated[
             int,
             typer.Option(
@@ -53,6 +98,14 @@ def register(app: typer.Typer) -> None:
                 show_default=True,
             ),
         ] = 50001,
+        no_events: Annotated[
+            bool,
+            typer.Option(
+                "--no-events",
+                help="Disable events. Default: False",
+                show_default=True,
+            ),
+        ] = False,
     ):
         """Run a Sidecar Proxy to expose a web application through OpenZiti."""
         if ":" in str(target):
@@ -65,6 +118,11 @@ def register(app: typer.Typer) -> None:
             str(identity_file),
             target_addr,
             workers=workers,
+            events_enabled=not no_events,
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+            retries=retries,
             publishers_port=publishers_port,
             subscribers_port=subscribers_port,
         )

mrok/cli/commands/frontend/run.py

@@ -44,6 +44,48 @@ def register(app: typer.Typer) -> None:
                 show_default=True,
             ),
         ] = default_workers,
+        max_connections: Annotated[
+            int,
+            typer.Option(
+                "--max-pool-connections",
+                help=(
+                    "The maximum number of concurrent HTTP connections that "
+                    "the pool should allow. Any attempt to send a request on a pool that "
+                    "would exceed this amount will block until a connection is available."
+                ),
+                show_default=True,
+            ),
+        ] = 1000,
+        max_keepalive_connections: Annotated[
+            int | None,
+            typer.Option(
+                "--max-pool-keepalive-connections",
+                help=(
+                    "The maximum number of idle HTTP connections "
+                    "that will be maintained in the pool."
+                ),
+                show_default=True,
+            ),
+        ] = 100,
+        keepalive_expiry: Annotated[
+            float | None,
+            typer.Option(
+                "--max-pool-keepalive-expiry",
+                help=(
+                    "The duration in seconds that an idle HTTP connection "
+                    "may be maintained for before being expired from the pool."
+                ),
+                show_default=True,
+            ),
+        ] = 300,
     ):
         """Run the mrok frontend with Gunicorn and Uvicorn workers."""
-        frontend.run(identity_file, host, port, workers)
+        frontend.run(
+            identity_file,
+            host,
+            port,
+            workers,
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+        )

mrok/controller/schemas.py

@@ -9,12 +9,12 @@ from pydantic import (
     computed_field,
 )
 
-from mrok.ziti.api import TagsType
+from mrok.types.ziti import Tags
 
 
 class BaseSchema(BaseModel):
     model_config = ConfigDict(from_attributes=True, extra="ignore")
-    tags: TagsType | None = None
+    tags: Tags | None = None
 
 
 class IdSchema(BaseModel):

mrok/frontend/app.py

@@ -6,7 +6,7 @@ from mrok.conf import get_settings
 from mrok.proxy.app import ProxyAppBase
 from mrok.proxy.backend import AIOZitiNetworkBackend
 from mrok.proxy.exceptions import InvalidTargetError
-from mrok.proxy.types import Scope
+from mrok.types.proxy import Scope
 
 RE_SUBDOMAIN = re.compile(r"(?i)^(?:EXT-\d{4}-\d{4}|INS-\d{4}-\d{4}-\d{4})$")
 
@@ -16,9 +16,9 @@ class FrontendProxyApp(ProxyAppBase):
         self,
         identity_file: str,
         *,
-        max_connections: int = 1000,
-        max_keepalive_connections: int = 10,
-        keepalive_expiry: float = 120.0,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
         retries=0,
     ):
         self._identity_file = identity_file
@@ -32,10 +32,10 @@ class FrontendProxyApp(ProxyAppBase):
 
     def setup_connection_pool(
         self,
-        max_connections: int | None = 1000,
-        max_keepalive_connections: int | None = 100,
-        keepalive_expiry: float | None = 120.0,
-        retries: int = 0,
+        max_connections: int | None,
+        max_keepalive_connections: int | None,
+        keepalive_expiry: float | None,
+        retries: int,
     ) -> AsyncConnectionPool:
         return AsyncConnectionPool(
             max_connections=max_connections,

mrok/frontend/main.py

@@ -38,8 +38,16 @@ def run(
     host: str,
     port: int,
     workers: int,
+    max_connections: int | None,
+    max_keepalive_connections: int | None,
+    keepalive_expiry: float | None,
 ):
-    app = FrontendProxyApp(str(identity_file))
+    app = FrontendProxyApp(
+        str(identity_file),
+        max_connections=max_connections,
+        max_keepalive_connections=max_keepalive_connections,
+        keepalive_expiry=keepalive_expiry,
+    )
 
     options = {
         "bind": f"{host}:{port}",

mrok/proxy/app.py

@@ -4,8 +4,8 @@ import logging
 from httpcore import AsyncConnectionPool, Request
 
 from mrok.proxy.exceptions import ProxyError
-from mrok.proxy.streams import ASGIRequestBodyStream
-from mrok.proxy.types import ASGIReceive, ASGISend, Scope
+from mrok.proxy.stream import ASGIRequestBodyStream
+from mrok.types.proxy import ASGIReceive, ASGISend, Scope
 
 logger = logging.getLogger("mrok.proxy")
 
@@ -26,9 +26,9 @@ class ProxyAppBase(abc.ABC):
     def __init__(
         self,
         *,
-        max_connections: int | None = 1000,
-        max_keepalive_connections: int | None = 10,
-        keepalive_expiry: float | None = 120.0,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
         retries: int = 0,
     ) -> None:
         self._pool = self.setup_connection_pool(
@@ -41,10 +41,10 @@ class ProxyAppBase(abc.ABC):
     @abc.abstractmethod
     def setup_connection_pool(
         self,
-        max_connections: int | None = 1000,
-        max_keepalive_connections: int | None = 10,
-        keepalive_expiry: float | None = 120.0,
-        retries: int = 0,
+        max_connections: int | None,
+        max_keepalive_connections: int | None,
+        keepalive_expiry: float | None,
+        retries: int,
     ) -> AsyncConnectionPool:
         raise NotImplementedError()
 
@@ -78,6 +78,7 @@ class ProxyAppBase(abc.ABC):
                 content=body_stream,
             )
             response = await self._pool.handle_async_request(request)
+            logger.debug(f"connection pool status: {self._pool}")
             response_headers = []
             for k, v in response.headers:
                 if k.lower() not in HOP_BY_HOP_HEADERS:

mrok/proxy/asgi.py

@@ -0,0 +1,96 @@
+from collections.abc import Iterator
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import Any, ParamSpec, Protocol
+
+from mrok.types.proxy import ASGIApp, ASGIReceive, ASGISend, Lifespan, Scope
+
+P = ParamSpec("P")
+
+
+class ASGIMiddleware(Protocol[P]):
+    def __call__(
+        self, app: ASGIApp, /, *args: P.args, **kwargs: P.kwargs
+    ) -> ASGIApp: ...  # pragma: no cover
+
+
+class Middleware:
+    def __init__(self, cls: ASGIMiddleware[P], *args: P.args, **kwargs: P.kwargs) -> None:
+        self.cls = cls
+        self.args = args
+        self.kwargs = kwargs
+
+    def __iter__(self) -> Iterator[Any]:
+        as_tuple = (self.cls, self.args, self.kwargs)
+        return iter(as_tuple)
+
+
+class ASGIAppWrapper:
+    def __init__(
+        self,
+        app: ASGIApp,
+        lifespan: Lifespan | None = None,
+    ) -> None:
+        self.app = app
+        self.lifespan = lifespan
+        self.middlware: list[Middleware] = []
+        self.middleare_stack: ASGIApp | None = None
+
+    def add_middleware(self, cls: ASGIMiddleware[P], *args: P.args, **kwargs: P.kwargs):
+        self.middlware.insert(0, Middleware(cls, *args, **kwargs))
+
+    def build_middleware_stack(self):
+        app = self.app
+        for cls, args, kwargs in reversed(self.middlware):
+            app = cls(app, *args, **kwargs)
+        return app
+
+    def get_starlette_lifespan(self):
+        router = getattr(self.app, "router", None)
+        if router is None:
+            return None
+        return getattr(router, "lifespan_context", None)
+
+    @asynccontextmanager
+    async def merge_lifespan(self, app: ASGIApp):
+        async with AsyncExitStack() as stack:
+            state: dict[Any, Any] = {}
+            if self.lifespan is not None:
+                outer_state = await stack.enter_async_context(self.lifespan(app))
+                state.update(outer_state or {})
+            starlette_lifespan = self.get_starlette_lifespan()
+            if starlette_lifespan is not None:
+                inner_state = await stack.enter_async_context(starlette_lifespan(app))
+                state.update(inner_state or {})
+            yield state
+
+    async def handle_lifespan(self, scope: Scope, receive: ASGIReceive, send: ASGISend) -> None:
+        started = False
+        app: Any = scope.get("app")
+        await receive()
+        try:
+            async with self.merge_lifespan(app) as state:
+                if state:
+                    if "state" not in scope:
+                        raise RuntimeError('"state" is unsupported by the current ASGI Server.')
+                    scope["state"].update(state)
+                await send({"type": "lifespan.startup.complete"})
+                started = True
+                await receive()
+        except Exception as e:  # pragma: no cover
+            if started:
+                await send({"type": "lifespan.shutdown.failed", "message": str(e)})
+            else:
+                await send({"type": "lifespan.startup.failed", "message": str(e)})
+            raise
+        else:
+            await send({"type": "lifespan.shutdown.complete"})
+
+    async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend) -> None:
+        if self.middleare_stack is None:  # pragma: no branch
+            self.middleware_stack = self.build_middleware_stack()
+        if scope["type"] == "lifespan":
+            scope["app"] = self
+            await self.handle_lifespan(scope, receive, send)
+            return
+
+        await self.middleware_stack(scope, receive, send)

mrok/proxy/backend.py

@@ -6,8 +6,8 @@ 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
+from mrok.proxy.exceptions import InvalidTargetError, TargetUnavailableError
+from mrok.proxy.stream import AIONetworkStream
 
 
 class AIOZitiNetworkBackend(AsyncNetworkBackend):
@@ -37,7 +37,9 @@ class AIOZitiNetworkBackend(AsyncNetworkBackend):
             reader, writer = await asyncio.open_connection(sock=sock)
             return AIONetworkStream(reader, writer)
         except Exception as e:
-            raise TargetUnavailableError() from e
+            if e.args and e.args[0] == -24:  # the service exists but is not available
+                raise TargetUnavailableError() from e
+            raise InvalidTargetError() from e
 
     async def sleep(self, seconds: float) -> None:
         await asyncio.sleep(seconds)

mrok/proxy/event_publisher.py

@@ -0,0 +1,66 @@
+import asyncio
+import contextlib
+import logging
+
+import zmq
+import zmq.asyncio
+
+from mrok.proxy.asgi import ASGIAppWrapper
+from mrok.proxy.metrics import MetricsCollector
+from mrok.proxy.middleware import CaptureMiddleware, MetricsMiddleware
+from mrok.proxy.models import Event, HTTPResponse, ServiceMetadata, Status
+from mrok.types.proxy import ASGIApp
+
+logger = logging.getLogger("mrok.proxy")
+
+
+class EventPublisher:
+    def __init__(
+        self,
+        worker_id: str,
+        meta: ServiceMetadata | None = None,
+        event_publisher_port: int = 50000,
+        metrics_interval: float = 5.0,
+    ):
+        self._worker_id = worker_id
+        self._meta = meta
+        self._metrics_interval = metrics_interval
+        self.publisher_port = event_publisher_port
+        self._zmq_ctx = zmq.asyncio.Context()
+        self._publisher = self._zmq_ctx.socket(zmq.PUB)
+        self._metrics_collector = MetricsCollector(self._worker_id)
+        self._publish_task = None
+
+    async def on_startup(self):
+        self._publisher.connect(f"tcp://localhost:{self.publisher_port}")
+        self._publish_task = asyncio.create_task(self.publish_metrics_event())
+        logger.info(f"Events publishing for worker {self._worker_id} started")
+
+    async def on_shutdown(self):
+        self._publish_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self._publish_task
+        self._publisher.close()
+        self._zmq_ctx.term()
+        logger.info(f"Events publishing for worker {self._worker_id} stopped")
+
+    async def publish_metrics_event(self):
+        while True:
+            snap = await self._metrics_collector.snapshot()
+            event = Event(type="status", data=Status(meta=self._meta, metrics=snap))
+            await self._publisher.send_string(event.model_dump_json())
+            await asyncio.sleep(self._metrics_interval)
+
+    async def publish_response_event(self, response: HTTPResponse):
+        event = Event(type="response", data=response)
+        await self._publisher.send_string(event.model_dump_json())  # type: ignore[attr-defined]
+
+    def setup_middleware(self, app: ASGIAppWrapper):
+        app.add_middleware(CaptureMiddleware, self.publish_response_event)
+        app.add_middleware(MetricsMiddleware, self._metrics_collector)  # type: ignore
+
+    @contextlib.asynccontextmanager
+    async def lifespan(self, app: ASGIApp):
+        await self.on_startup()  # type: ignore
+        yield
+        await self.on_shutdown()  # type: ignore