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

mrok/agent/sidecar/app.py

@@ -1,9 +1,12 @@
 import asyncio
 import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
 from pathlib import Path
 
 from mrok.http.forwarder import ForwardAppBase
-from mrok.http.types import Scope, StreamReader, StreamWriter
+from mrok.http.pool import ConnectionPool
+from mrok.http.types import Scope, StreamPair
 
 logger = logging.getLogger("mrok.agent")
 
@@ -18,12 +21,31 @@ class ForwardApp(ForwardAppBase):
             read_chunk_size=read_chunk_size,
         )
         self._target_address = target_address
+        self._pool = ConnectionPool(
+            pool_name=str(self._target_address),
+            factory=self.connect,
+            initial_connections=5,
+            max_size=100,
+            idle_timeout=20.0,
+            reaper_interval=5.0,
+        )
+
+    async def connect(self) -> StreamPair:
+        if isinstance(self._target_address, tuple):
+            return await asyncio.open_connection(*self._target_address)
+        return await asyncio.open_unix_connection(str(self._target_address))
 
+    async def startup(self):
+        await self._pool.start()
+
+    async def shutdown(self):
+        await self._pool.stop()
+
+    @asynccontextmanager
     async def select_backend(
         self,
         scope: Scope,
         headers: dict[str, str],
-    ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
-        if isinstance(self._target_address, tuple):
-            return await asyncio.open_connection(*self._target_address)
-        return await asyncio.open_unix_connection(str(self._target_address))
+    ) -> AsyncGenerator[StreamPair, None]:
+        async with self._pool.acquire() as (reader, writer):
+            yield reader, writer

mrok/cli/commands/admin/register/extensions.py

@@ -1,5 +1,4 @@
 import asyncio
-import re
 from typing import Annotated
 
 import typer
@@ -7,11 +6,10 @@ from rich import print
 
 from mrok.cli.commands.admin.utils import parse_tags
 from mrok.conf import Settings
+from mrok.constants import RE_EXTENSION_ID
 from mrok.ziti.api import ZitiManagementAPI
 from mrok.ziti.services import register_service
 
-RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
-
 
 async def do_register(settings: Settings, extension_id: str, tags: list[str] | None):
     async with ZitiManagementAPI(settings) as api:
@@ -20,7 +18,7 @@ async def do_register(settings: Settings, extension_id: str, tags: list[str] | N
 
 def validate_extension_id(extension_id: str) -> str:
     if not RE_EXTENSION_ID.fullmatch(extension_id):
-        raise typer.BadParameter("ext_id must match EXT-xxxx-yyyy (case-insensitive)")
+        raise typer.BadParameter("it must match EXT-xxxx-yyyy (case-insensitive)")
     return extension_id
 
 

mrok/cli/commands/admin/register/instances.py

@@ -1,6 +1,5 @@
 import asyncio
 import json
-import re
 from pathlib import Path
 from typing import Annotated
 
@@ -8,11 +7,10 @@ import typer
 
 from mrok.cli.commands.admin.utils import parse_tags
 from mrok.conf import Settings
+from mrok.constants import RE_EXTENSION_ID, RE_INSTANCE_ID
 from mrok.ziti.api import ZitiClientAPI, ZitiManagementAPI
 from mrok.ziti.identities import register_identity
 
-RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
-
 
 async def do_register(
     settings: Settings, extension_id: str, instance_id: str, tags: list[str] | None
@@ -25,10 +23,16 @@ async def do_register(
 
 def validate_extension_id(extension_id: str):
     if not RE_EXTENSION_ID.fullmatch(extension_id):
-        raise typer.BadParameter("ext_id must match EXT-xxxx-yyyy (case-insensitive)")
+        raise typer.BadParameter("it must match EXT-xxxx-yyyy (case-insensitive)")
     return extension_id
 
 
+def validate_instance_id(instance_id: str):
+    if not RE_INSTANCE_ID.fullmatch(instance_id):
+        raise typer.BadParameter("it must match INS-xxxx-yyyy-zzzz (case-insensitive)")
+    return instance_id
+
+
 def register(app: typer.Typer) -> None:
     @app.command("instance")
     def register_instance(
@@ -36,7 +40,9 @@ def register(app: typer.Typer) -> None:
         extension_id: str = typer.Argument(
             ..., callback=validate_extension_id, help="Extension ID in format EXT-xxxx-yyyy"
         ),
-        instance_id: str = typer.Argument(..., help="Instance ID"),
+        instance_id: str = typer.Argument(
+            ..., callback=validate_instance_id, help="Instance ID in format INS-xxxx-yyyy-zzzz"
+        ),
         output: Path = typer.Argument(
             ...,
             file_okay=True,

mrok/constants.py

@@ -0,0 +1,6 @@
+import re
+
+RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
+RE_INSTANCE_ID = re.compile(r"(?i)INS-\d{4}-\d{4}-\d{4}")
+
+RE_SUBDOMAIN = re.compile(r"(?i)^(?:EXT-\d{4}-\d{4}|INS-\d{4}-\d{4}-\d{4})$")

mrok/http/forwarder.py

@@ -1,14 +1,27 @@
 import abc
 import asyncio
 import logging
+from contextlib import AbstractAsyncContextManager
 
-from mrok.http.types import ASGIReceive, ASGISend, Scope, StreamReader, StreamWriter
+from mrok.http.types import ASGIReceive, ASGISend, Scope, StreamPair
 
 logger = logging.getLogger("mrok.proxy")
 
 
-class BackendNotFoundError(Exception):
-    pass
+class BackendSelectionError(Exception):
+    def __init__(self, status: int = 500, message: str = "Internal Server Error"):
+        self.status = status
+        self.message = message
+
+
+class InvalidBackendError(BackendSelectionError):
+    def __init__(self):
+        super().__init__(status=502, message="Bad Gateway")
+
+
+class BackendUnavailableError(BackendSelectionError):
+    def __init__(self):
+        super().__init__(status=503, message="Service Unavailable")
 
 
 class ForwardAppBase(abc.ABC):
@@ -60,20 +73,18 @@ class ForwardAppBase(abc.ABC):
                 await send({"type": "lifespan.shutdown.complete"})
                 return
 
-    @abc.abstractmethod
-    async def select_backend(
-        self,
-        scope: Scope,
-        headers: dict[str, str],
-    ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
-        """Return (reader, writer) connected to the target backend."""
-
     async def startup(self):
         return
 
     async def shutdown(self):
         return
 
+    @abc.abstractmethod
+    def select_backend(
+        self, scope: Scope, headers: dict[str, str]
+    ) -> AbstractAsyncContextManager[StreamPair]:
+        raise NotImplementedError()
+
     async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend) -> None:
         """ASGI callable entry point.
 
@@ -95,37 +106,40 @@ class ForwardAppBase(abc.ABC):
 
         headers = list(scope.get("headers", []))
         headers = self.ensure_host_header(headers, scope)
-        reader, writer = await self.select_backend(
-            scope, {k[0].decode().lower(): k[1].decode() for k in headers}
-        )
-
-        if not (reader and writer):
-            await send({"type": "http.response.start", "status": 502, "headers": []})
-            await send({"type": "http.response.body", "body": b"Bad Gateway"})
-            return
-
-        use_chunked = self.ensure_chunked_if_needed(headers)
+        try:
+            async with self.select_backend(
+                scope, {k[0].decode().lower(): k[1].decode() for k in headers}
+            ) as (reader, writer):
+                if not (reader and writer):
+                    await send({"type": "http.response.start", "status": 502, "headers": []})
+                    await send({"type": "http.response.body", "body": b"Bad Gateway"})
+                    return
 
-        await self.write_request_line_and_headers(writer, method, path_qs, headers)
+                use_chunked = self.ensure_chunked_if_needed(headers)
 
-        await self.stream_request_body(receive, writer, use_chunked)
+                await self.write_request_line_and_headers(writer, method, path_qs, headers)
 
-        status_line = await reader.readline()
-        if not status_line:
-            await send({"type": "http.response.start", "status": 502, "headers": []})
-            await send({"type": "http.response.body", "body": b"Bad Gateway"})
-            writer.close()
-            await writer.wait_closed()
-            return
+                await self.stream_request_body(receive, writer, use_chunked)
 
-        status, headers_out, raw_headers = await self.read_status_and_headers(reader, status_line)
+                status_line = await reader.readline()
+                if not status_line:
+                    await send({"type": "http.response.start", "status": 502, "headers": []})
+                    await send({"type": "http.response.body", "body": b"Bad Gateway"})
+                    return
 
-        await send({"type": "http.response.start", "status": status, "headers": headers_out})
+                status, headers_out, raw_headers = await self.read_status_and_headers(
+                    reader, status_line
+                )
 
-        await self.stream_response_body(reader, send, raw_headers)
+                await send(
+                    {"type": "http.response.start", "status": status, "headers": headers_out}
+                )
 
-        writer.close()
-        await writer.wait_closed()
+                await self.stream_response_body(reader, send, raw_headers)
+        except BackendSelectionError as bse:
+            await send({"type": "http.response.start", "status": bse.status, "headers": []})
+            await send({"type": "http.response.body", "body": bse.message.encode()})
+            return
 
     def format_path(self, scope: Scope) -> str:
         raw_path = scope.get("raw_path")
@@ -159,7 +173,7 @@ class ForwardAppBase(abc.ABC):
 
     async def write_request_line_and_headers(
         self,
-        writer: StreamWriter,
+        writer: asyncio.StreamWriter,
         method: str,
         path_qs: str,
         headers: list[tuple[bytes, bytes]],
@@ -173,7 +187,7 @@ class ForwardAppBase(abc.ABC):
         await writer.drain()
 
     async def stream_request_body(
-        self, receive: ASGIReceive, writer: StreamWriter, use_chunked: bool
+        self, receive: ASGIReceive, writer: asyncio.StreamWriter, use_chunked: bool
     ) -> None:
         if use_chunked:
             await self.stream_request_chunked(receive, writer)
@@ -181,7 +195,9 @@ class ForwardAppBase(abc.ABC):
 
         await self.stream_request_until_end(receive, writer)
 
-    async def stream_request_chunked(self, receive: ASGIReceive, writer: StreamWriter) -> None:
+    async def stream_request_chunked(
+        self, receive: ASGIReceive, writer: asyncio.StreamWriter
+    ) -> None:
         """Send request body to backend using HTTP/1.1 chunked encoding."""
         while True:
             event = await receive()
@@ -195,13 +211,14 @@ class ForwardAppBase(abc.ABC):
                 if not event.get("more_body", False):
                     break
             elif event["type"] == "http.disconnect":
-                writer.close()
                 return
 
         writer.write(b"0\r\n\r\n")
         await writer.drain()
 
-    async def stream_request_until_end(self, receive: ASGIReceive, writer: StreamWriter) -> None:
+    async def stream_request_until_end(
+        self, receive: ASGIReceive, writer: asyncio.StreamWriter
+    ) -> None:
         """Send request body to backend when content length/transfer-encoding
         already provided (no chunking).
         """
@@ -215,11 +232,10 @@ class ForwardAppBase(abc.ABC):
                 if not event.get("more_body", False):
                     break
             elif event["type"] == "http.disconnect":
-                writer.close()
                 return
 
     async def read_status_and_headers(
-        self, reader: StreamReader, first_line: bytes
+        self, reader: asyncio.StreamReader, first_line: bytes
     ) -> tuple[int, list[tuple[bytes, bytes]], dict[bytes, bytes]]:
         parts = first_line.decode(errors="ignore").split(" ", 2)
         status = int(parts[1]) if len(parts) >= 2 and parts[1].isdigit() else 502
@@ -256,14 +272,14 @@ class ForwardAppBase(abc.ABC):
         except Exception:
             return None
 
-    async def drain_trailers(self, reader: StreamReader) -> None:
+    async def drain_trailers(self, reader: asyncio.StreamReader) -> None:
         """Consume trailer header lines until an empty line is encountered."""
         while True:
             trailer = await reader.readline()
             if trailer in (b"\r\n", b"\n", b""):
                 break
 
-    async def stream_response_chunked(self, reader: StreamReader, send: ASGISend) -> None:
+    async def stream_response_chunked(self, reader: asyncio.StreamReader, send: ASGISend) -> None:
         """Read chunked-encoded response from reader, decode and forward to ASGI send."""
         while True:
             size_line = await reader.readline()
@@ -292,7 +308,7 @@ class ForwardAppBase(abc.ABC):
         await send({"type": "http.response.body", "body": b"", "more_body": False})
 
     async def stream_response_with_content_length(
-        self, reader: StreamReader, send: ASGISend, content_length: int
+        self, reader: asyncio.StreamReader, send: ASGISend, content_length: int
     ) -> None:
         """Read exactly content_length bytes and forward to ASGI send events."""
         remaining = content_length
@@ -311,7 +327,7 @@ class ForwardAppBase(abc.ABC):
         if not sent_final:
             await send({"type": "http.response.body", "body": b"", "more_body": False})
 
-    async def stream_response_until_eof(self, reader: StreamReader, send: ASGISend) -> None:
+    async def stream_response_until_eof(self, reader: asyncio.StreamReader, send: ASGISend) -> None:
         """Read from reader until EOF and forward chunks to ASGI send events."""
         while True:
             chunk = await reader.read(self._read_chunk_size)
@@ -321,7 +337,7 @@ class ForwardAppBase(abc.ABC):
         await send({"type": "http.response.body", "body": b"", "more_body": False})
 
     async def stream_response_body(
-        self, reader: StreamReader, send: ASGISend, raw_headers: dict[bytes, bytes]
+        self, reader: asyncio.StreamReader, send: ASGISend, raw_headers: dict[bytes, bytes]
     ) -> None:
         te = raw_headers.get(b"transfer-encoding", b"").lower()
         cl = raw_headers.get(b"content-length")

mrok/http/pool.py

@@ -0,0 +1,239 @@
+import asyncio
+import contextlib
+import logging
+import time
+from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
+from contextlib import asynccontextmanager
+from typing import Any
+
+from cachetools import TTLCache
+
+from mrok.http.types import StreamPair
+
+PoolItem = tuple[asyncio.StreamReader, asyncio.StreamWriter, float]
+
+logger = logging.getLogger("mrok.proxy")
+
+
+class ConnectionPool:
+    def __init__(
+        self,
+        pool_name: str,
+        factory: Callable[[], Awaitable[StreamPair]],
+        *,
+        initial_connections: int = 0,
+        max_size: int = 10,
+        idle_timeout: float = 30.0,
+        reaper_interval: float = 5.0,
+    ) -> None:
+        if initial_connections < 0:
+            raise ValueError("initial_connections must be >= 0")
+        if max_size < 1:
+            raise ValueError("max_size must be >= 1")
+        if initial_connections > max_size:
+            raise ValueError("initial_connections cannot exceed max_size")
+        self.pool_name = pool_name
+        self.factory = factory
+        self.initial_connections = initial_connections
+        self.max_size = max_size
+        self.idle_timeout = idle_timeout
+        self.reaper_interval = reaper_interval
+
+        self._pool: list[PoolItem] = []
+        self._in_use = 0
+        self._lock = asyncio.Lock()
+        self._cond = asyncio.Condition()
+        self._stop_event = asyncio.Event()
+
+        self._started = False
+        self._reaper_task: asyncio.Task | None = None
+
+    async def start(self) -> None:
+        if self._started:
+            return
+        self._reaper_task = asyncio.create_task(self._reaper())
+        await self._prewarm()
+        self._started = True
+
+    async def stop(self) -> None:
+        self._stop_event.set()
+        if self._reaper_task is not None:
+            self._reaper_task.cancel()
+            with contextlib.suppress(Exception):
+                await self._reaper_task
+
+        to_close: list[asyncio.StreamWriter] = []
+        async with self._lock:
+            to_close = [writer for _, writer, _ in self._pool]
+            self._pool.clear()
+        for w in to_close:
+            with contextlib.suppress(Exception):
+                w.close()
+                await w.wait_closed()
+
+        async with self._cond:
+            self._cond.notify_all()
+
+    @asynccontextmanager
+    async def acquire(self) -> AsyncGenerator[StreamPair]:
+        if not self._started:
+            await self.start()
+        reader, writer = await self._acquire()
+        logger.info(
+            f"Acquire stats for pool {self.pool_name}: "
+            f"in_use={self._in_use}, size={len(self._pool)}"
+        )
+        try:
+            yield (reader, writer)
+        finally:
+            await self._release(reader, writer)
+
+    async def _prewarm(self) -> None:
+        conns: list[PoolItem] = []
+        needed = max(0, self.initial_connections - (self._in_use + len(self._pool)))
+        for _ in range(needed):
+            reader, writer = await self.factory()
+            conns.append((reader, writer, time.time()))
+        if conns:
+            async with self._lock:
+                self._pool.extend(conns)
+            # notify any waiters
+            async with self._cond:
+                self._cond.notify_all()
+
+    async def _acquire(self) -> StreamPair:  # type: ignore
+        to_close: list[asyncio.StreamWriter] = []
+        create_new = False
+        while True:
+            need_prewarm = False
+            async with self._cond:
+                now = time.time()
+                if not self._pool:
+                    need_prewarm = True
+                while self._pool:
+                    reader, writer, ts = self._pool.pop()
+                    if now - ts <= self.idle_timeout and not writer.is_closing():
+                        self._in_use += 1
+                        return reader, writer
+                    to_close.append(writer)
+
+                total = self._in_use + len(self._pool)
+                if total < self.max_size:
+                    self._in_use += 1
+                    create_new = True
+                    break
+                await self._cond.wait()
+
+            if need_prewarm:
+                await self._prewarm()
+                continue
+
+        for w in to_close:
+            with contextlib.suppress(Exception):
+                w.close()
+                await w.wait_closed()
+
+        if create_new:
+            try:
+                reader, writer = await self.factory()
+            except Exception:
+                async with self._cond:
+                    if self._in_use > 0:
+                        self._in_use -= 1
+                    self._cond.notify()
+                raise
+            return reader, writer
+
+    async def _release(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        async with self._cond:
+            if self._in_use > 0:
+                self._in_use -= 1
+
+            if not writer.is_closing():
+                self._pool.append((reader, writer, time.time()))
+
+            self._cond.notify()
+        logger.info(
+            f"Release stats for pool {self.pool_name}: "
+            f"in_use={self._in_use}, size={len(self._pool)}"
+        )
+
+    async def _reaper(self) -> None:
+        try:
+            while not self._stop_event.is_set():
+                await asyncio.sleep(self.reaper_interval)
+                to_close: list[asyncio.StreamWriter] = []
+                now = time.time()
+                async with self._lock:
+                    new_pool: list[PoolItem] = []
+                    for reader, writer, ts in self._pool:
+                        if now - ts > self.idle_timeout or writer.is_closing():
+                            to_close.append(writer)
+                        else:
+                            new_pool.append((reader, writer, ts))
+                    self._pool = new_pool
+
+                for w in to_close:
+                    with contextlib.suppress(Exception):
+                        w.close()
+                        await w.wait_closed()
+
+                async with self._cond:
+                    self._cond.notify_all()
+        except asyncio.CancelledError:
+            pass
+
+
+class SlidingTTLCache(TTLCache):
+    def __init__(
+        self,
+        *,
+        maxsize: float,
+        ttl: float,
+        on_evict: Callable[[Any], Coroutine[Any, Any, None]] | None,
+    ) -> None:
+        super().__init__(maxsize=maxsize, ttl=ttl)
+        self.on_evict = on_evict
+
+    def __getitem__(self, key: Any) -> Any:
+        value = super().__getitem__(key)
+        super().__setitem__(key, value)
+        return value
+
+    def popitem(self) -> Any:
+        key, value = super().popitem()
+        if self.on_evict:
+            asyncio.create_task(self.on_evict(value))
+        return key, value
+
+
+class PoolManager:
+    def __init__(
+        self,
+        pool_factory: Callable[[Any], Awaitable[ConnectionPool]],
+        idle_timeout: int = 300,
+    ):
+        self.factory = pool_factory
+        self.cache = SlidingTTLCache(
+            maxsize=float("inf"),
+            ttl=idle_timeout,
+            on_evict=self._close_pool,
+        )
+
+    async def _close_pool(self, pool: ConnectionPool):
+        with contextlib.suppress(Exception):
+            await pool.stop()
+
+    async def get_pool(self, key) -> ConnectionPool:
+        try:
+            return self.cache[key]
+        except KeyError:
+            pool = await self.factory(key)
+            await pool.start()
+            self.cache[key] = pool
+            return pool
+
+    async def shutdown(self) -> None:
+        pools = list(self.cache.values())
+        for pool in pools:
+            await self._close_pool(pool)

mrok/http/types.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import Awaitable, Callable, MutableMapping
-from typing import Any, Protocol
+from typing import Any
 
 from mrok.datastructures import HTTPRequest, HTTPResponse
 
@@ -15,29 +15,4 @@ ASGIApp = Callable[[Scope, ASGIReceive, ASGISend], Awaitable[None]]
 RequestCompleteCallback = Callable[[HTTPRequest], Awaitable | None]
 ResponseCompleteCallback = Callable[[HTTPResponse], Awaitable | None]
 
-
-class StreamReaderWrapper(Protocol):
-    async def read(self, n: int = -1) -> bytes: ...
-    async def readexactly(self, n: int) -> bytes: ...
-    async def readline(self) -> bytes: ...
-    def at_eof(self) -> bool: ...
-
-    @property
-    def underlying(self) -> asyncio.StreamReader: ...
-
-
-class StreamWriterWrapper(Protocol):
-    def write(self, data: bytes) -> None: ...
-    async def drain(self) -> None: ...
-    def close(self) -> None: ...
-    async def wait_closed(self) -> None: ...
-
-    @property
-    def transport(self): ...
-
-    @property
-    def underlying(self) -> asyncio.StreamWriter: ...
-
-
-StreamReader = StreamReaderWrapper | asyncio.StreamReader
-StreamWriter = StreamWriterWrapper | asyncio.StreamWriter
+StreamPair = tuple[asyncio.StreamReader, asyncio.StreamWriter]

mrok/proxy/app.py

@@ -1,12 +1,18 @@
 import asyncio
 import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
 from pathlib import Path
 
+import openziti
+from openziti.context import ZitiContext
+
 from mrok.conf import get_settings
-from mrok.http.forwarder import ForwardAppBase
-from mrok.http.types import Scope, StreamReader, StreamWriter
+from mrok.constants import RE_SUBDOMAIN
+from mrok.http.forwarder import BackendUnavailableError, ForwardAppBase, InvalidBackendError
+from mrok.http.pool import ConnectionPool, PoolManager
+from mrok.http.types import Scope, StreamPair
 from mrok.logging import setup_logging
-from mrok.proxy.ziti import ZitiSocketCache
 
 logger = logging.getLogger("mrok.proxy")
 
@@ -30,7 +36,8 @@ class ProxyApp(ForwardAppBase):
             if settings.proxy.domain[0] == "."
             else f".{settings.proxy.domain}"
         )
-        self._ziti_socket_cache = ZitiSocketCache(self._identity_file)
+        self._ziti_ctx: ZitiContext | None = None
+        self._pool_manager = PoolManager(self.build_connection_pool)
 
     def get_target_from_header(self, headers: dict[str, str], name: str) -> str | None:
         header_value = headers.get(name, "")
@@ -47,18 +54,48 @@ class ProxyApp(ForwardAppBase):
             raise ProxyError("Neither Host nor X-Forwarded-Host contain a valid target name")
         return target
 
+    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 startup(self):
         setup_logging(get_settings())
+        self._get_ziti_ctx()
 
     async def shutdown(self):
-        await self._ziti_socket_cache.stop()
+        await self._pool_manager.shutdown()
+
+    async def build_connection_pool(self, key: str) -> ConnectionPool:
+        async def connect():
+            sock = self._get_ziti_ctx().connect(key)
+            reader, writer = await asyncio.open_connection(sock=sock)
+            return reader, writer
+
+        return ConnectionPool(
+            pool_name=key,
+            factory=connect,
+            initial_connections=5,
+            max_size=100,
+            idle_timeout=20.0,
+            reaper_interval=5.0,
+        )
 
+    @asynccontextmanager
     async def select_backend(
         self,
         scope: Scope,
         headers: dict[str, str],
-    ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
+    ) -> AsyncGenerator[StreamPair]:
         target_name = self.get_target_name(headers)
-        sock = await self._ziti_socket_cache.get_or_create(target_name)
-        reader, writer = await asyncio.open_connection(sock=sock)
-        return reader, writer
+        if not target_name or not RE_SUBDOMAIN.fullmatch(target_name):
+            raise InvalidBackendError()
+        pool = await self._pool_manager.get_pool(target_name)
+        try:
+            async with pool.acquire() as (reader, writer):
+                yield reader, writer
+        except Exception:
+            raise BackendUnavailableError()

{mrok-0.4.4.dist-info → mrok-0.4.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mrok
-Version: 0.4.4
+Version: 0.4.6
 Summary: MPT Extensions OpenZiti Orchestrator
 Author: SoftwareOne AG
 License:                                  Apache License
@@ -206,8 +206,8 @@ License:                                  Apache License
            limitations under the License.
 License-File: LICENSE.txt
 Requires-Python: <4,>=3.12
-Requires-Dist: aiocache<0.13.0,>=0.12.3
 Requires-Dist: asn1crypto<2.0.0,>=1.5.1
+Requires-Dist: cachetools<7.0.0,>=6.2.2
 Requires-Dist: cryptography<46.0.0,>=45.0.7
 Requires-Dist: dynaconf<4.0.0,>=3.2.11
 Requires-Dist: fastapi-pagination<0.15.0,>=0.14.1

{mrok-0.4.4.dist-info → mrok-0.4.6.dist-info}/RECORD

@@ -1,5 +1,6 @@
 mrok/__init__.py,sha256=D1PUs3KtMCqG4bFLceVNG62L3RN53NS95uSCNXpgvzs,181
 mrok/conf.py,sha256=_5Z-A5LyojQeY8J7W8C0QidsmrPl99r9qKYEoMf4kcI,840
+mrok/constants.py,sha256=65OlmploxfND686E4mt9LR9MqYn8I5k-L0H-R5KsLG8,201
 mrok/datastructures.py,sha256=gp8KF2JoNOxIRzYStVZLKL_XVDbcIVSIDnmpQo4FNt0,4067
 mrok/errors.py,sha256=ruNMDFr2_0ezCGXuCG1OswCEv-bHOIzMMd02J_0ABcs,37
 mrok/logging.py,sha256=ZMWn0w4fJ-F_g-L37H_GM14BSXAIF2mFF_ougX5S7mg,2856
@@ -14,7 +15,7 @@ mrok/agent/devtools/inspector/__main__.py,sha256=HeYcRf1bjXPji2LKMPCcTU61afrRH2P
 mrok/agent/devtools/inspector/app.py,sha256=_pzxemMqIunE5EdMq5amjqpOGsMWIOw17GgiCtRAi6Q,16464
 mrok/agent/devtools/inspector/server.py,sha256=C4uD6_1psSHMjJLUDCMPGvKdQYKaEwYTw27NAbwuuA0,636
 mrok/agent/sidecar/__init__.py,sha256=DrjJGhqFyxsVODW06KI20Wpr6HsD2lD6qFCKUXc7GIE,59
-mrok/agent/sidecar/app.py,sha256=1p6qWkXVq78zcJ2dhCYlw8CqfwPsgEtu07Lp5csK3Iw,874
+mrok/agent/sidecar/app.py,sha256=YOQLwPPqcElbF2kU15bcw-ePzZM09eVJQZ6Z5NYg9u8,1509
 mrok/agent/sidecar/main.py,sha256=h31wynUCcFmRckvqLHtH97w1QgMv4fzcmYjhRPUobxY,1076
 mrok/cli/__init__.py,sha256=mtFEa8IeS1x6Gm4dUYoSnAxyEzNqbUVSmWxtuZUMR84,61
 mrok/cli/main.py,sha256=DFcYPwDskXi8SKAgEsuP4GMFzaniIf_6bZaSDWvYKDk,2724
@@ -28,8 +29,8 @@ mrok/cli/commands/admin/list/__init__.py,sha256=kjCMcpn1gopcrQaaHxfFh8Kyngldepnl
 mrok/cli/commands/admin/list/extensions.py,sha256=16fhDB5ucL8su2WQnSaQ1E6MhgC4vkP9-nuHAcPpzyE,4405
 mrok/cli/commands/admin/list/instances.py,sha256=kaqeyidwUxgYqfaHXqp2m76rm5h2ErBsYyZcNeaBRwY,5912
 mrok/cli/commands/admin/register/__init__.py,sha256=5Jb_bc2L47MEpQIrOcquzduTFWQ01Jd1U1MpqaR-Ekw,209
-mrok/cli/commands/admin/register/extensions.py,sha256=p1qX5gSQX1IGpOQjO2MJzbc09v1ebdFuPo94QzJErKk,1485
-mrok/cli/commands/admin/register/instances.py,sha256=XB6uAchc7Rm8uAu7o3-oHaN_rS8CCIBf0QKWZGW86fI,1940
+mrok/cli/commands/admin/register/extensions.py,sha256=dxciVA_S31rZSm0A7lkecn2mI9TMlWDhcJTgwgNXbM4,1460
+mrok/cli/commands/admin/register/instances.py,sha256=raF57jPUTryWdvNqGCosth1C-8jjv9IbA0UuNbDel3A,2220
 mrok/cli/commands/admin/unregister/__init__.py,sha256=-GjjCPX1pISbWmJK6GpKO3ijGsDQb21URjU1hNu99O4,215
 mrok/cli/commands/admin/unregister/extensions.py,sha256=GR3Iwzeksk_R0GkgmCSG7iHRcUrI7ABqDi25Gbes64Y,1016
 mrok/cli/commands/admin/unregister/instances.py,sha256=-28wL8pTXTWHVHtw93y8-dqi-Dlf0OZOnlBCKOyGo80,1138
@@ -63,17 +64,17 @@ mrok/controller/routes/instances.py,sha256=v-fn_F6JHbDZ4YUNCIZzClgHp6aC1Eu5HB7k7
 mrok/http/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mrok/http/config.py,sha256=k73-4hBo6jag1RpyZagJLtTCL6EQoebZaX8Vv-CMN_k,2050
 mrok/http/constants.py,sha256=ao5gI2HFBWmrdd2Yc6XFK_RGaHk-omxI4AqvfIiGes8,409
-mrok/http/forwarder.py,sha256=DakD9hrrCWAzB1B_4SgQaQaEHcDYLLI9WaYs5F0O36I,12977
+mrok/http/forwarder.py,sha256=vAf2nh6Fmr07JdRJkK4dPHKJilP9PnsYZcroqsnilB8,13751
 mrok/http/lifespan.py,sha256=UdbOqjWZsHzJJjX0CTd2hY96Jpk5QWtdHJEzPG6Z4hQ,1288
 mrok/http/middlewares.py,sha256=SGo4EwhTId2uJx1aMuqGbNy7MXgZlDEdZI0buzBYVv0,5011
+mrok/http/pool.py,sha256=Q-pRwgYPusqEKQCwZsRQ2mnGaDfyWknWpvydUu5KtEU,7696
 mrok/http/protocol.py,sha256=ap8jbLUvgbAH81ZJZCBkQiYR7mkV_eL3rpfwEkoE8sU,392
 mrok/http/server.py,sha256=Mj7C85fc-DXp-WTBWaOd7ag808oliLmFBH5bf-G2FHg,370
-mrok/http/types.py,sha256=XpNrvbfpANKvmjOBYtLF1FmDHoJF3z_MIMQHXoJlvmE,1302
+mrok/http/types.py,sha256=A82zloEqW8KdKahdNrbW5fhlJNUo2enLNRVMWIJTatA,632
 mrok/http/utils.py,sha256=sOixYu3R9-nNoMFYdifrreYvcFRIHYVtb6AAmtVzaLE,2125
 mrok/proxy/__init__.py,sha256=vWXyImroqM1Eq8e_oFPBup8VJ3reyp8SVjFTbLzRkI8,51
-mrok/proxy/app.py,sha256=yulfBdTdxesVxF1h2lli_5zjd5wP-jTx17FRdbkaV7A,2163
+mrok/proxy/app.py,sha256=VvMRmYLwsItjCcecy6ccrkk564LnArIermHTRVDxh9U,3469
 mrok/proxy/main.py,sha256=ZXpticE6J4FABaslDB_8J5qklPsf3e7xIFSZmcPAAjQ,1588
-mrok/proxy/ziti.py,sha256=rKgIXpOvtBeVopZkQlNUZa3Fdci9jgiog_i6egb17ps,3318
 mrok/ziti/__init__.py,sha256=20OWMiexRhOovZOX19zlX87-V78QyWnEnSZfyAftUdE,263
 mrok/ziti/api.py,sha256=KvGiT9d4oSgC3JbFWLDQyuHcLX2HuZJoJ8nHmWtCDkY,16154
 mrok/ziti/bootstrap.py,sha256=QIDhlkIxPW2QRuumFq2D1WDbD003P5f3z24pAUsyeBI,2696
@@ -82,8 +83,8 @@ mrok/ziti/errors.py,sha256=yYCbVDwktnR0AYduqtynIjo73K3HOhIrwA_vQimvEd4,368
 mrok/ziti/identities.py,sha256=1BcwfqAJHMBhc3vRaf0aLaIkoHskj5Xe2Lsq2lO9Vs8,6735
 mrok/ziti/pki.py,sha256=o2tySqHC8-7bvFuI2Tqxg9vX6H6ZSxWxfP_9x29e19M,1954
 mrok/ziti/services.py,sha256=zR1PEBYwXVou20iJK4euh0ZZFAo9UB8PZk8f6SDmiUE,3194
-mrok-0.4.4.dist-info/METADATA,sha256=nzjalRGet1yhkJf1L4t022A-NTDG-xQ9a5cWZfbDkdg,15836
-mrok-0.4.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mrok-0.4.4.dist-info/entry_points.txt,sha256=tloXwvU1uJicBJR2h-8HoVclPgwJWDwuREMHN8Zq-nU,38
-mrok-0.4.4.dist-info/licenses/LICENSE.txt,sha256=6PaICaoA3yNsZKLv5G6OKqSfLSoX7MakYqTDgJoTCBs,11346
-mrok-0.4.4.dist-info/RECORD,,
+mrok-0.4.6.dist-info/METADATA,sha256=Io64noW9WGLw9asC4xjeuLS7Wh8bFefufJmTjUK8Syo,15836
+mrok-0.4.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mrok-0.4.6.dist-info/entry_points.txt,sha256=tloXwvU1uJicBJR2h-8HoVclPgwJWDwuREMHN8Zq-nU,38
+mrok-0.4.6.dist-info/licenses/LICENSE.txt,sha256=6PaICaoA3yNsZKLv5G6OKqSfLSoX7MakYqTDgJoTCBs,11346
+mrok-0.4.6.dist-info/RECORD,,

mrok/proxy/ziti.py

@@ -1,102 +0,0 @@
-import asyncio
-import contextlib
-import logging
-from asyncio import Task
-from pathlib import Path
-
-import openziti
-from aiocache import Cache
-from openziti.context import ZitiContext
-from openziti.zitisock import ZitiSocket
-
-logger = logging.getLogger("mrok.proxy")
-
-
-class ZitiSocketCache:
-    def __init__(
-        self,
-        identity_file: str | Path,
-        ziti_ctx_timeout_ms: int = 10_000,
-        ttl_seconds: float = 60.0,
-        cleanup_interval: float = 10.0,
-    ) -> None:
-        self._identity_file = identity_file
-        self._ziti_ctx_timeout_ms = ziti_ctx_timeout_ms
-        self._ttl_seconds = ttl_seconds
-        self._cleanup_interval = cleanup_interval
-
-        self._ziti_ctx: ZitiContext | None = None
-        self._cache = Cache(Cache.MEMORY)
-        self._active_sockets: dict[str, ZitiSocket] = {}
-        self._cleanup_task: Task | None = None
-
-    def _get_ziti_ctx(self) -> ZitiContext:
-        if self._ziti_ctx is None:
-            ctx, err = openziti.load(str(self._identity_file), timeout=self._ziti_ctx_timeout_ms)
-            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 _create_socket(self, key: str):
-        return self._get_ziti_ctx().connect(key)
-
-    async def get_or_create(self, key: str):
-        sock = await self._cache.get(key)
-
-        if sock:
-            await self._cache.expire(key, self._ttl_seconds)
-            self._active_sockets[key] = sock
-            logger.debug(f"Ziti socket found for service {key}")
-            return sock
-
-        sock = await self._create_socket(key)
-        await self._cache.set(key, sock, self._ttl_seconds)
-        self._active_sockets[key] = sock
-        logger.info(f"New Ziti socket created for service {key}")
-        return sock
-
-    # async def invalidate(self, key: str):
-    #     sock = await self._cache.get(key)
-    #     if sock:
-    #         await self._close_socket(sock)
-
-    #     await self._cache.delete(key)
-    #     self._active_sockets.pop(key, None)
-
-    async def start(self):
-        self._cleanup_task = asyncio.create_task(self._periodic_cleanup())
-        # Warmup ziti context
-        self._get_ziti_ctx()
-
-    async def stop(self):
-        self._cleanup_task.cancel()
-        with contextlib.suppress(Exception):
-            await self._cleanup_task
-
-        for sock in list(self._active_sockets.values()):
-            await self._close_socket(sock)
-
-        self._active_sockets.clear()
-        await self._cache.clear()
-
-    @staticmethod
-    async def _close_socket(sock: ZitiSocket):
-        with contextlib.suppress(Exception):
-            sock.close()
-
-    async def _periodic_cleanup(self):
-        try:
-            while True:
-                await asyncio.sleep(self._cleanup_interval)
-                await self._cleanup_once()
-        except asyncio.CancelledError:
-            return
-
-    async def _cleanup_once(self):
-        expired = {key for key in self._active_sockets.keys() if not self._cache.exists(key)}
-        for key in expired:
-            logger.debug(f"Cleaning up expired socket connection {key}")
-            sock = self._active_sockets.pop(key, None)
-            if sock:
-                await self._close_socket(sock)