mrok 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

mrok/master.py

@@ -12,6 +12,7 @@ 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
@@ -69,6 +70,9 @@ def start_uvicorn_worker(
     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"])
@@ -82,7 +86,6 @@ def start_uvicorn_worker(
     async def status_sender():
         while True:
             snap = await metrics.snapshot()
-            logger.info(f"New metrics snapshot taken: {snap}")
             event = Event(type="status", data=Status(meta=meta, metrics=snap))
             await pub.send_string(event.model_dump_json())
             await asyncio.sleep(metrics_interval)

mrok/proxy/app.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from pathlib import Path
 
 from mrok.conf import get_settings
@@ -35,33 +34,30 @@ class ProxyApp(ForwardAppBase):
         self._conn_manager = ZitiConnectionManager(
             identity_file,
             ttl_seconds=ziti_connection_ttl_seconds,
-            purge_interval=ziti_conn_cache_purge_interval_seconds,
+            cleanup_interval=ziti_conn_cache_purge_interval_seconds,
         )
 
-    def get_target_name(self, headers: dict[str, str]) -> str:
-        header_value = headers.get("x-forwarded-for", headers.get("host"))
-        if not header_value:
-            raise ProxyError(
-                "Cannot determine the target OpenZiti service/terminator name, "
-                "neither Host nor X-Forwarded-For headers have been sent in the request.",
-            )
-        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 Host or X-Forwarded-For header: `{header_value}`."
-            )
+    def get_target_from_header(self, headers: dict[str, str], name: str) -> str | None:
+        header_value = headers.get(name, "")
+        if self._proxy_wildcard_domain in header_value:
+            if ":" in header_value:
+                header_value, _ = header_value.split(":", 1)
+            return header_value[: -len(self._proxy_wildcard_domain)]
 
-        return header_value[: -len(self._proxy_wildcard_domain)]
+    def get_target_name(self, headers: dict[str, str]) -> str:
+        target = self.get_target_from_header(headers, "x-forwarded-host")
+        if not target:
+            target = self.get_target_from_header(headers, "host")
+        if not target:
+            raise ProxyError("Neither Host nor X-Forwarded-Host contain a valid target name")
+        return target
 
     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,
@@ -70,4 +66,4 @@ class ProxyApp(ForwardAppBase):
     ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
         target_name = self.get_target_name(headers)
 
-        return await self._conn_manager.get(target_name)
+        return await self._conn_manager.get_or_create(target_name)

mrok/proxy/streams.py

@@ -1,13 +1,13 @@
 import asyncio
 
-from mrok.proxy.types import ConnectionCache, ConnectionKey
+from mrok.proxy.types import ConnectionCache
 
 
 class CachedStreamReader:
     def __init__(
         self,
         reader: asyncio.StreamReader,
-        key: ConnectionKey,
+        key: str,
         manager: ConnectionCache,
     ):
         self._reader = reader
@@ -77,7 +77,7 @@ class CachedStreamWriter:
     def __init__(
         self,
         writer: asyncio.StreamWriter,
-        key: ConnectionKey,
+        key: str,
         manager: ConnectionCache,
     ):
         self._writer = writer

mrok/proxy/types.py

@@ -4,9 +4,8 @@ from typing import Protocol
 
 from mrok.http.types import StreamReader, StreamWriter
 
-ConnectionKey = tuple[str, str | None]
-CachedStream = tuple[StreamReader, StreamWriter]
+StreamPair = tuple[StreamReader, StreamWriter]
 
 
 class ConnectionCache(Protocol):
-    async def invalidate(self, key: ConnectionKey) -> None: ...
+    async def invalidate(self, key: str) -> None: ...

mrok/proxy/ziti.py

@@ -1,22 +1,13 @@
-"""Ziti-backed connection manager for the proxy.
-
-This manager owns creation of connections via an OpenZiti context, wraps
-streams to observe IO errors, evicts idle entries, and serializes creation
-per-key.
-"""
-
 import asyncio
+import contextlib
 import logging
-import time
 from pathlib import Path
 
-# typing imports intentionally minimized
 import openziti
+from aiocache import Cache
 
-from mrok.http.types import StreamReader, StreamWriter
-from mrok.proxy.dataclasses import CachedStreamEntry
 from mrok.proxy.streams import CachedStreamReader, CachedStreamWriter
-from mrok.proxy.types import CachedStream, ConnectionKey
+from mrok.proxy.types import StreamPair
 
 logger = logging.getLogger("mrok.proxy")
 
@@ -27,147 +18,100 @@ class ZitiConnectionManager:
         identity_file: str | Path,
         ziti_timeout_ms: int = 10000,
         ttl_seconds: float = 60.0,
-        purge_interval: float = 10.0,
+        cleanup_interval: float = 10.0,
     ):
-        self._identity_file = identity_file
-        self._ziti_ctx = None
-        self._ziti_timeout_ms = ziti_timeout_ms
-        self._ttl = float(ttl_seconds)
-        self._purge_interval = float(purge_interval)
-        self._cache: dict[ConnectionKey, CachedStreamEntry] = {}
-        self._lock = asyncio.Lock()
-        self._in_progress: dict[ConnectionKey, asyncio.Lock] = {}
-        self._purge_task: asyncio.Task | None = None
-
-    async def get(self, target: str) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
-        head, _, tail = target.partition(".")
-        terminator = target if head and tail else ""
-        service = tail if tail else head
-        r, w = await self._get_or_create_key((service, terminator))
-        return r, w
-
-    async def invalidate(self, key: ConnectionKey) -> None:
-        async with self._lock:
-            item = self._cache.pop(key, None)
-        if item is None:
-            return
-        await self._close_writer(item.writer)
+        self.identity_file = identity_file
+        self.ziti_timeout_ms = ziti_timeout_ms
+        self.ttl_seconds = ttl_seconds
+        self.cleanup_interval = cleanup_interval
+
+        self.cache = Cache(Cache.MEMORY)
+
+        self._active_pairs: dict[str, StreamPair] = {}
+
+        self._cleanup_task: asyncio.Task | None = None
+        self._ziti_ctx: openziti.context.ZitiContext | None = None
+
+    async def create_stream_pair(self, key: str) -> StreamPair:
+        if not self._ziti_ctx:
+            raise Exception("ZitiConnectionManager is not started")
+        sock = self._ziti_ctx.connect(key)
+        orig_reader, orig_writer = await asyncio.open_connection(sock=sock)
+
+        reader = CachedStreamReader(orig_reader, key, self)
+        writer = CachedStreamWriter(orig_writer, key, self)
+        return (reader, writer)
+
+    async def get_or_create(self, key: str) -> StreamPair:
+        pair = await self.cache.get(key)
+
+        if pair:
+            logger.info(f"return cached connection for {key}")
+            await self.cache.set(key, pair, ttl=self.ttl_seconds)
+            self._active_pairs[key] = pair
+            return pair
+
+        pair = await self.create_stream_pair(key)
+        await self.cache.set(key, pair, ttl=self.ttl_seconds)
+        self._active_pairs[key] = pair
+        logger.info(f"return new connection for {key}")
+        return pair
+
+    async def invalidate(self, key: str) -> None:
+        logger.info(f"invalidating connection for {key}")
+        pair = await self.cache.get(key)
+        if pair:
+            await self._close_pair(pair)
+
+        await self.cache.delete(key)
+        self._active_pairs.pop(key, None)
 
     async def start(self) -> None:
+        if self._cleanup_task is None:
+            self._cleanup_task = asyncio.create_task(self._periodic_cleanup())
         if self._ziti_ctx is None:
-            ctx, err = openziti.load(str(self._identity_file), timeout=self._ziti_timeout_ms)
+            ctx, err = openziti.load(str(self.identity_file), timeout=self.ziti_timeout_ms)
             if err != 0:
                 raise Exception(f"Cannot create a Ziti context from the identity file: {err}")
             self._ziti_ctx = ctx
-        if self._purge_task is None:
-            self._purge_task = asyncio.create_task(self._purge_loop())
-            logger.info("Ziti connection manager started")
 
     async def stop(self) -> None:
-        if self._purge_task is not None:
-            self._purge_task.cancel()
-            try:
-                await self._purge_task
-            except asyncio.CancelledError:
-                logger.debug("Purge task was cancelled")
-            except Exception as e:
-                logger.warning(f"An error occurred stopping the purge task: {e}")
-            self._purge_task = None
-            logger.info("Ziti connection manager stopped")
-
-        async with self._lock:
-            items = list(self._cache.items())
-            self._cache.clear()
-
-        for _, item in items:
-            await self._close_writer(item.writer)
-
-    async def _purge_loop(self) -> None:
+        if self._cleanup_task:
+            self._cleanup_task.cancel()
+            with contextlib.suppress(Exception):
+                await self._cleanup_task
+
+            for pair in list(self._active_pairs.values()):
+                await self._close_pair(pair)
+
+            self._active_pairs.clear()
+            await self.cache.clear()
+        openziti.shutdown()
+
+    @staticmethod
+    async def _close_pair(pair: StreamPair) -> None:
+        reader, writer = pair
+        writer.close()
+        with contextlib.suppress(Exception):
+            await writer.wait_closed()
+
+    async def _periodic_cleanup(self) -> None:
         try:
             while True:
-                await asyncio.sleep(self._purge_interval)
-                await self._purge_once()
+                await asyncio.sleep(self.cleanup_interval)
+                await self._cleanup_once()
         except asyncio.CancelledError:
             return
 
-    async def _purge_once(self) -> None:
-        to_close: list[tuple[StreamReader, StreamWriter]] = []
-        async with self._lock:
-            now = time.time()
-            for key, item in list(self._cache.items()):
-                if now - item.last_access > self._ttl:
-                    to_close.append((item.reader, item.writer))
-                    del self._cache[key]
-
-        for _, writer in to_close:
-            writer.close()
-            await self._close_writer(writer)
+    async def _cleanup_once(self) -> None:
+        # Keys currently stored in aiocache
+        keys_in_cache = set(await self.cache.keys())
+        # Keys we think are alive
+        known_keys = set(self._active_pairs.keys())
 
-    def _is_writer_closed(self, writer: StreamWriter) -> bool:
-        return writer.transport.is_closing()
+        expired_keys = known_keys - keys_in_cache
 
-    async def _close_writer(self, writer: StreamWriter) -> None:
-        writer.close()
-        try:
-            await writer.wait_closed()
-        except Exception as e:
-            logger.debug(f"Error closing writer: {e}")
-
-    async def _get_or_create_key(self, key: ConnectionKey) -> CachedStream:
-        """Internal: create or return a cached wrapped pair for the concrete key."""
-        await self._purge_once()
-        to_close = None
-        async with self._lock:
-            if key in self._cache:
-                now = time.time()
-                item = self._cache[key]
-                reader, writer = item.reader, item.writer
-                if not self._is_writer_closed(writer) and not reader.at_eof():
-                    self._cache[key] = CachedStreamEntry(reader, writer, now)
-                    return reader, writer
-                to_close = writer
-                del self._cache[key]
-
-            lock = self._in_progress.get(key)
-            if lock is None:
-                lock = asyncio.Lock()
-                self._in_progress[key] = lock
-
-        if to_close:
-            await self._close_writer(to_close)
-
-        async with lock:
-            try:
-                # # double-check cache after acquiring the per-key lock
-                # async with self._lock:
-                #     now = time.time()
-                #     if key in self._cache:
-                #         r, w, _ = self._cache[key]
-                #         if not self._is_writer_closed(w) and not r.at_eof():
-                #             self._cache[key] = (r, w, now)
-                #             return r, w
-
-                # perform creation via ziti context
-                extension, instance = key
-                logger.info(f"Create connection to {extension}: {instance}")
-                # loop = asyncio.get_running_loop()
-                # sock = await loop.run_in_executor(None, self._ziti_ctx.connect,
-                # extension, instance)
-                if instance:
-                    sock = self._ziti_ctx.connect(
-                        extension, terminator=instance
-                    )  # , terminator=instance)
-                else:
-                    sock = self._ziti_ctx.connect(extension)
-                orig_reader, orig_writer = await asyncio.open_connection(sock=sock)
-
-                reader = CachedStreamReader(orig_reader, key, self)
-                writer = CachedStreamWriter(orig_writer, key, self)
-
-                async with self._lock:
-                    self._cache[key] = CachedStreamEntry(reader, writer, time.time())
-
-                return reader, writer
-            finally:
-                async with self._lock:
-                    self._in_progress.pop(key, None)
+        for key in expired_keys:
+            pair = self._active_pairs.pop(key, None)
+            if pair:
+                await self._close_pair(pair)

{mrok-0.4.0.dist-info → mrok-0.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mrok
-Version: 0.4.0
+Version: 0.4.2
 Summary: MPT Extensions OpenZiti Orchestrator
 Author: SoftwareOne AG
 License:                                  Apache License
@@ -206,6 +206,7 @@ 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: cryptography<46.0.0,>=45.0.7
 Requires-Dist: dynaconf<4.0.0,>=3.2.11

{mrok-0.4.0.dist-info → mrok-0.4.2.dist-info}/RECORD

@@ -3,7 +3,7 @@ mrok/conf.py,sha256=_5Z-A5LyojQeY8J7W8C0QidsmrPl99r9qKYEoMf4kcI,840
 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
-mrok/master.py,sha256=6_bUic39mXE37lt0PBf_BtY2D2mBR7KzPj65JP7HWPc,8287
+mrok/master.py,sha256=XuketJZuB1YWdbTs819pjLum7Qfv232F9ZCxdwRztCQ,8340
 mrok/metrics.py,sha256=asweK_7xiV5MtkDkvbEm9Tktqrl2KHM8VflF0AkNGI0,4036
 mrok/agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mrok/agent/ziticorn.py,sha256=marXGcr6CbDdiNi8V3CZJhg8YCUbKw6ySuO3-0-zf8g,900
@@ -71,12 +71,12 @@ mrok/http/server.py,sha256=Mj7C85fc-DXp-WTBWaOd7ag808oliLmFBH5bf-G2FHg,370
 mrok/http/types.py,sha256=XpNrvbfpANKvmjOBYtLF1FmDHoJF3z_MIMQHXoJlvmE,1302
 mrok/http/utils.py,sha256=sOixYu3R9-nNoMFYdifrreYvcFRIHYVtb6AAmtVzaLE,2125
 mrok/proxy/__init__.py,sha256=vWXyImroqM1Eq8e_oFPBup8VJ3reyp8SVjFTbLzRkI8,51
-mrok/proxy/app.py,sha256=MwJZ91MQ1oWkbmsX7-NooZSV9BkLVZPFmUKpNPC4HHQ,2490
+mrok/proxy/app.py,sha256=xNgT-lqXRe53HWzZFz0ceyttEbz__1PD8J07nme8L2s,2339
 mrok/proxy/dataclasses.py,sha256=DtX-Yuma-uOECOPefJnoQJhZMEtT6Za_27cd-lJE9Iw,237
 mrok/proxy/main.py,sha256=ZXpticE6J4FABaslDB_8J5qklPsf3e7xIFSZmcPAAjQ,1588
-mrok/proxy/streams.py,sha256=6TMZwrQPbSyQqpqavsoTeyUmS2O026pJfiCnxLopPqg,3425
-mrok/proxy/types.py,sha256=dgWqAj6dFGVH_Q8-k8sU5h18yoUF_fTn-SRPIfEs_gA,308
-mrok/proxy/ziti.py,sha256=dKd6UzmEAFu9-gey871sPEDUZTkt4YVPyCYRzeA5mlA,6539
+mrok/proxy/streams.py,sha256=a7EMKn3R7JB3iHKdmbs8QiEHd1xlT4N-vnrzuaiZSTU,3390
+mrok/proxy/types.py,sha256=XpAfTklmJfcQilyKVTkYbaFHvWZSTcr_6Rg_feiq9Mw,257
+mrok/proxy/ziti.py,sha256=kWnX1d-BaZcc0tdk_xwSp8rmQ3joZIxs7MlLScHPvMg,3879
 mrok/ziti/__init__.py,sha256=20OWMiexRhOovZOX19zlX87-V78QyWnEnSZfyAftUdE,263
 mrok/ziti/api.py,sha256=KvGiT9d4oSgC3JbFWLDQyuHcLX2HuZJoJ8nHmWtCDkY,16154
 mrok/ziti/bootstrap.py,sha256=QIDhlkIxPW2QRuumFq2D1WDbD003P5f3z24pAUsyeBI,2696
@@ -85,8 +85,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.0.dist-info/METADATA,sha256=39Su41kvqC4oAsQl1_ZG7b5v1xRr6c2YK6x84XgURi8,15796
-mrok-0.4.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mrok-0.4.0.dist-info/entry_points.txt,sha256=tloXwvU1uJicBJR2h-8HoVclPgwJWDwuREMHN8Zq-nU,38
-mrok-0.4.0.dist-info/licenses/LICENSE.txt,sha256=6PaICaoA3yNsZKLv5G6OKqSfLSoX7MakYqTDgJoTCBs,11346
-mrok-0.4.0.dist-info/RECORD,,
+mrok-0.4.2.dist-info/METADATA,sha256=b2BYs3KtydCUqoBQxxet4p7KJFJKECthQ6Vn_0M-uvM,15836
+mrok-0.4.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mrok-0.4.2.dist-info/entry_points.txt,sha256=tloXwvU1uJicBJR2h-8HoVclPgwJWDwuREMHN8Zq-nU,38
+mrok-0.4.2.dist-info/licenses/LICENSE.txt,sha256=6PaICaoA3yNsZKLv5G6OKqSfLSoX7MakYqTDgJoTCBs,11346
+mrok-0.4.2.dist-info/RECORD,,