sientia-dataops-opc-ingestor/ingestor/connection/pool.py

import asyncio

from asyncua import Client
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring

import ingestor.metrics as metrics
from ingestor.connection.certs import CertProvider
from ingestor.connection.clock import MonotonicClock
from ingestor.connection.security import apply_security


class OpcConnectionPool(SientiaMonitoring):
    """Per-server OPC UA client cache (F0 capstone): reuses live connections,
    evicts idle ones, and reconnects after a failed attempt (never caches it)."""

    def __init__(
        self,
        cert_provider: CertProvider,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
    ):
        self.cert_provider = cert_provider
        self.clock = MonotonicClock()
        self._connections: dict[str, tuple[Client, float]] = {}
        # Active subscriptions per server. A subscribed connection only receives
        # callbacks and never calls get(), so without a pin release_idle() would
        # evict it as idle and silently stop the reads (SIENTIAPDE-1974 G4).
        self._pins: dict[str, int] = {}
        # ponytail: one Lock per server_id, never cleaned up — server set is
        # small and bounded. Add eviction of stale locks only if it grows.
        self._locks: dict[str, asyncio.Lock] = {}

        SientiaMonitoring.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

    async def get(self, server_id: str, server_cfg: dict) -> Client:
        """Return the cached client for `server_id`, connecting one if needed.

        Double-checked under a per-server lock so concurrent callers for the
        same new server share one connection (no leaked, untracked sessions)."""
        client = self._cached(server_id)
        if client is not None:
            return client
        async with self._locks.setdefault(server_id, asyncio.Lock()):
            client = self._cached(server_id)  # re-check: a racer may have connected
            if client is not None:
                return client
            return await self._open(server_id, server_cfg)

    def _cached(self, server_id: str) -> Client | None:
        entry = self._connections.get(server_id)
        if entry is None:
            return None
        client, _ = entry
        self._connections[server_id] = (client, self.clock.elapsed())
        return client

    async def _open(self, server_id: str, server_cfg: dict) -> Client:
        try:
            client = await self._connect(server_cfg)
        except Exception:
            await self.emit_metric(
                metric_object=metrics.OPC_CONNECTIONS_FAILED,
                tags={'pod_id': self.pod_id, 'server_name': server_id},
            )
            raise

        self._connections[server_id] = (client, self.clock.elapsed())
        await self._record_connected(server_id, server_cfg['url'])
        return client

    async def _connect(self, server_cfg: dict) -> Client:
        client = Client(server_cfg['url'], timeout=10)
        server_uri = server_cfg.get('server_uri')
        if server_uri:
            client.application_uri = server_uri
        if server_cfg.get('use_security'):
            cert, key = await self.cert_provider.fetch()
            await apply_security(client, server_cfg.get('server_uri'), cert, key, None)
        await client.connect()
        return client

    async def _record_connected(self, server_id: str, url: str):
        await self.emit_metric(
            metric_object=metrics.OPC_CONNECTIONS_TOTAL,
            tags={'pod_id': self.pod_id, 'server_name': server_id},
        )
        await self.emit_metric(
            metric_object=metrics.OPC_CONNECTION_STATUS,
            method='set',
            value=1,
            tags={'pod_id': self.pod_id, 'server_name': server_id, 'server_url': url},
        )

    def pin(self, server_id: str) -> None:
        """Mark one more active subscription on `server_id` (never evicted as idle)."""
        self._pins[server_id] = self._pins.get(server_id, 0) + 1

    def unpin(self, server_id: str) -> None:
        """Drop one active subscription mark from `server_id`."""
        count = self._pins.get(server_id, 0) - 1
        if count > 0:
            self._pins[server_id] = count
        else:
            self._pins.pop(server_id, None)

    async def release_idle(self, max_idle_s: float):
        """Disconnect and drop clients unused for longer than `max_idle_s`,
        except the ones pinned by an active subscription."""
        now = self.clock.elapsed()
        idle_ids = [
            sid
            for sid, (_, last_used) in self._connections.items()
            if now - last_used > max_idle_s and not self._pins.get(sid)
        ]
        for server_id in idle_ids:
            await self._evict(server_id)

    async def evict(self, server_id: str):
        """Disconnect `server_id` now (lost server or changed config); the next
        get() reconnects. No-op when the server has no pooled client."""
        if server_id in self._connections:
            await self._evict(server_id)

    async def close_all(self):
        """Disconnect every pooled client (e.g. on shutdown)."""
        # list() snapshots the keys: _evict() pops from self._connections,
        # so iterating the dict directly would raise "changed size during iteration".
        for server_id in list(self._connections):  # NOSONAR
            await self._evict(server_id)

    async def _evict(self, server_id: str):
        client, _ = self._connections.pop(server_id)
        await self._disconnect_with_retries(client, server_id)
        await self.emit_metric(
            metric_object=metrics.OPC_DISCONNECTIONS_TOTAL,
            tags={'pod_id': self.pod_id, 'server_name': server_id},
        )

    async def _disconnect_with_retries(self, client: Client, server_id: str) -> None:
        """Tries 5 times to disconnect, with a delay of 100ms x try (mirrors
        the legacy OpcManager.disconnection_fallback)."""
        for i in range(5):
            try:
                await client.disconnect()
                return
            except Exception as e:
                self.logger.error(
                    f'Failed to disconnect from {server_id}, attempt {i + 1} of 5: {e}'
                )
                await asyncio.sleep(0.1 * i)