From 5d30a5c5472b171a9ccdebaa0e3d3a05d088cd6a Mon Sep 17 00:00:00 2001 From: PedroHMCosme Date: Tue, 14 Jul 2026 09:35:53 -0300 Subject: [PATCH] feat(gateway): OpcConnectionPool (per-server, reuse, idle-evict) --- ingestor/connection/pool.py | 120 ++++++++++++++++++++++++++ tests/unit/connection/test_pool.py | 133 +++++++++++++++++++++++++++++ 2 files changed, 253 insertions(+) create mode 100644 ingestor/connection/pool.py create mode 100644 tests/unit/connection/test_pool.py diff --git a/ingestor/connection/pool.py b/ingestor/connection/pool.py new file mode 100644 index 0000000..5378c28 --- /dev/null +++ b/ingestor/connection/pool.py @@ -0,0 +1,120 @@ +import asyncio +import traceback + +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, + metadata: dict | None = None, + ): + self.cert_provider = cert_provider + self.metadata = metadata or {} + self.clock = MonotonicClock() + self._connections: dict[str, tuple[Client, float]] = {} + + 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.""" + cached = self._connections.get(server_id) + if cached is not None: + client, _ = cached + self._connections[server_id] = (client, self.clock.elapsed()) + return 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) + client.application_uri = server_cfg.get('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}, + ) + + async def release_idle(self, max_idle_s: float): + """Disconnect and drop clients unused for longer than `max_idle_s`.""" + now = self.clock.elapsed() + idle_ids = [ + sid for sid, (_, last_used) in self._connections.items() if now - last_used > max_idle_s + ] + for server_id in idle_ids: + await self._evict(server_id) + + async def close_all(self): + """Disconnect every pooled client (e.g. on shutdown).""" + for server_id in list(self._connections): + 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) -> list: + """Tries 5 times to disconnect, with a delay of 100ms x try (mirrors + OpcManager.disconnection_fallback).""" + error_stack = [] + 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}' + ) + error_stack.append( + {'attempt': i + 1, 'error': str(e), 'traceback': traceback.format_exc()} + ) + await asyncio.sleep(0.1 * i) + return error_stack diff --git a/tests/unit/connection/test_pool.py b/tests/unit/connection/test_pool.py new file mode 100644 index 0000000..7e048bc --- /dev/null +++ b/tests/unit/connection/test_pool.py @@ -0,0 +1,133 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import ingestor.metrics as metrics +from ingestor.connection.pool import OpcConnectionPool + + +def make_pool(cert_provider=None): + pool = OpcConnectionPool( + cert_provider=cert_provider or AsyncMock(), + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + metadata={}, + ) + pool.emit_metric = AsyncMock() + return pool + + +@pytest.mark.asyncio +async def test_get_reuses_connection_per_server(): + pool = make_pool() + fake_client = AsyncMock() + with patch('ingestor.connection.pool.Client', return_value=fake_client): + c1 = await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False}) + c2 = await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False}) + + assert c1 is c2 + assert fake_client.connect.await_count == 1 + + +@pytest.mark.asyncio +async def test_get_failure_records_metric_and_raises(): + pool = make_pool() + fake_client = AsyncMock() + fake_client.connect.side_effect = RuntimeError('boom') + with patch('ingestor.connection.pool.Client', return_value=fake_client): + with pytest.raises(RuntimeError): + await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False}) + + pool.emit_metric.assert_any_await( + metric_object=metrics.OPC_CONNECTIONS_FAILED, + tags={'pod_id': pool.pod_id, 'server_name': 'srv'}, + ) + assert 'srv' not in pool._connections + + # a later successful get must reconnect (failed attempt was not cached) + fake_client2 = AsyncMock() + with patch('ingestor.connection.pool.Client', return_value=fake_client2): + client = await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False}) + assert client is fake_client2 + assert fake_client2.connect.await_count == 1 + + +@pytest.mark.asyncio +async def test_get_applies_security_when_requested(): + cert_provider = AsyncMock() + cert_provider.fetch.return_value = ('cert.pem', 'key.pem') + pool = make_pool(cert_provider=cert_provider) + fake_client = AsyncMock() + with ( + patch('ingestor.connection.pool.Client', return_value=fake_client), + patch('ingestor.connection.pool.apply_security', new=AsyncMock()) as mock_apply, + ): + await pool.get( + 'srv', + {'url': 'opc.tcp://h:4840', 'server_uri': 'urn:test', 'use_security': True}, + ) + + cert_provider.fetch.assert_awaited_once() + mock_apply.assert_awaited_once_with(fake_client, 'urn:test', 'cert.pem', 'key.pem', None) + + +@pytest.mark.asyncio +async def test_release_idle_evicts_and_disconnects(): + pool = make_pool() + fake_client = AsyncMock() + with patch('ingestor.connection.pool.Client', return_value=fake_client): + await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False}) + + client, _ = pool._connections['srv'] + pool._connections['srv'] = (client, -1000.0) # force idle + + await pool.release_idle(max_idle_s=0) + + fake_client.disconnect.assert_awaited_once() + assert 'srv' not in pool._connections + + +@pytest.mark.asyncio +async def test_release_idle_keeps_recently_used(): + pool = make_pool() + fake_client = AsyncMock() + with patch('ingestor.connection.pool.Client', return_value=fake_client): + await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False}) + + await pool.release_idle(max_idle_s=3600) + + fake_client.disconnect.assert_not_awaited() + assert 'srv' in pool._connections + + +@pytest.mark.asyncio +async def test_close_all_disconnects_everything(): + pool = make_pool() + fake_a, fake_b = AsyncMock(), AsyncMock() + with patch('ingestor.connection.pool.Client', side_effect=[fake_a, fake_b]): + await pool.get('a', {'url': 'opc.tcp://h:4840', 'use_security': False}) + await pool.get('b', {'url': 'opc.tcp://h:4840', 'use_security': False}) + + await pool.close_all() + + fake_a.disconnect.assert_awaited_once() + fake_b.disconnect.assert_awaited_once() + assert pool._connections == {} + + +@pytest.mark.asyncio +async def test_release_idle_disconnect_retries_then_gives_up(): + pool = make_pool() + fake_client = AsyncMock() + fake_client.disconnect.side_effect = RuntimeError('nope') + with patch('ingestor.connection.pool.Client', return_value=fake_client): + await pool.get('srv', {'url': 'opc.tcp://h:4840', 'use_security': False}) + + client, _ = pool._connections['srv'] + pool._connections['srv'] = (client, -1000.0) + + await pool.release_idle(max_idle_s=0) + + assert fake_client.disconnect.await_count == 5 + assert 'srv' not in pool._connections