mrok 0.4.2__tar.gz → 0.4.4__tar.gz

{mrok-0.4.2 → mrok-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mrok
-Version: 0.4.2
+Version: 0.4.4
 Summary: MPT Extensions OpenZiti Orchestrator
 Author: SoftwareOne AG
 License:                                  Apache License

{mrok-0.4.2 → mrok-0.4.4}/mrok/proxy/app.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from pathlib import Path
 
@@ -5,7 +6,7 @@ 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
+from mrok.proxy.ziti import ZitiSocketCache
 
 logger = logging.getLogger("mrok.proxy")
 
@@ -20,8 +21,6 @@ class ProxyApp(ForwardAppBase):
         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
@@ -31,11 +30,7 @@ class ProxyApp(ForwardAppBase):
             if settings.proxy.domain[0] == "."
             else f".{settings.proxy.domain}"
         )
-        self._conn_manager = ZitiConnectionManager(
-            identity_file,
-            ttl_seconds=ziti_connection_ttl_seconds,
-            cleanup_interval=ziti_conn_cache_purge_interval_seconds,
-        )
+        self._ziti_socket_cache = ZitiSocketCache(self._identity_file)
 
     def get_target_from_header(self, headers: dict[str, str], name: str) -> str | None:
         header_value = headers.get(name, "")
@@ -54,10 +49,9 @@ class ProxyApp(ForwardAppBase):
 
     async def startup(self):
         setup_logging(get_settings())
-        await self._conn_manager.start()
 
     async def shutdown(self):
-        await self._conn_manager.stop()
+        await self._ziti_socket_cache.stop()
 
     async def select_backend(
         self,
@@ -65,5 +59,6 @@ class ProxyApp(ForwardAppBase):
         headers: dict[str, str],
     ) -> tuple[StreamReader, StreamWriter] | tuple[None, None]:
         target_name = self.get_target_name(headers)
-
-        return await self._conn_manager.get_or_create(target_name)
+        sock = await self._ziti_socket_cache.get_or_create(target_name)
+        reader, writer = await asyncio.open_connection(sock=sock)
+        return reader, writer

mrok-0.4.4/mrok/proxy/ziti.py

@@ -0,0 +1,102 @@
+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)

{mrok-0.4.2 → mrok-0.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mrok"
-version = "0.4.2"
+version = "0.4.4"
 description = "MPT Extensions OpenZiti Orchestrator"
 readme = { file = "README.md", content-type = "text/markdown" }
 authors = [

mrok-0.4.2/mrok/proxy/dataclasses.py

@@ -1,12 +0,0 @@
-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-0.4.2/mrok/proxy/streams.py

@@ -1,124 +0,0 @@
-import asyncio
-
-from mrok.proxy.types import ConnectionCache
-
-
-class CachedStreamReader:
-    def __init__(
-        self,
-        reader: asyncio.StreamReader,
-        key: str,
-        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: str,
-        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-0.4.2/mrok/proxy/types.py

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

mrok-0.4.2/mrok/proxy/ziti.py

@@ -1,117 +0,0 @@
-import asyncio
-import contextlib
-import logging
-from pathlib import Path
-
-import openziti
-from aiocache import Cache
-
-from mrok.proxy.streams import CachedStreamReader, CachedStreamWriter
-from mrok.proxy.types import StreamPair
-
-logger = logging.getLogger("mrok.proxy")
-
-
-class ZitiConnectionManager:
-    def __init__(
-        self,
-        identity_file: str | Path,
-        ziti_timeout_ms: int = 10000,
-        ttl_seconds: float = 60.0,
-        cleanup_interval: float = 10.0,
-    ):
-        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)
-            if err != 0:
-                raise Exception(f"Cannot create a Ziti context from the identity file: {err}")
-            self._ziti_ctx = ctx
-
-    async def stop(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.cleanup_interval)
-                await self._cleanup_once()
-        except asyncio.CancelledError:
-            return
-
-    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())
-
-        expired_keys = known_keys - keys_in_cache
-
-        for key in expired_keys:
-            pair = self._active_pairs.pop(key, None)
-            if pair:
-                await self._close_pair(pair)