From 0b29857089011337709c72e346218ccbe69110be Mon Sep 17 00:00:00 2001
From: PedroHMCosme <pedro.cosme@aignosi.com.br>
Date: Tue, 14 Jul 2026 09:50:14 -0300
Subject: [PATCH] fix(gateway): per-server lock in OpcConnectionPool to avoid
 duplicate connects

---
 ingestor/connection/pool.py        | 30 +++++++++++++----
 tests/unit/connection/test_pool.py | 52 ++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 6 deletions(-)

diff --git a/ingestor/connection/pool.py b/ingestor/connection/pool.py
index 5378c28..a530e4c 100644
--- a/ingestor/connection/pool.py
+++ b/ingestor/connection/pool.py
@@ -29,6 +29,9 @@ class OpcConnectionPool(SientiaMonitoring):
         self.metadata = metadata or {}
         self.clock = MonotonicClock()
         self._connections: dict[str, tuple[Client, float]] = {}
+        # 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,
@@ -38,13 +41,28 @@ class OpcConnectionPool(SientiaMonitoring):
         )
 
     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
+        """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:
diff --git a/tests/unit/connection/test_pool.py b/tests/unit/connection/test_pool.py
index 7e048bc..759c94a 100644
--- a/tests/unit/connection/test_pool.py
+++ b/tests/unit/connection/test_pool.py
@@ -1,3 +1,4 @@
+import asyncio
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -86,6 +87,10 @@ async def test_release_idle_evicts_and_disconnects():
 
     fake_client.disconnect.assert_awaited_once()
     assert 'srv' not in pool._connections
+    pool.emit_metric.assert_any_await(
+        metric_object=metrics.OPC_DISCONNECTIONS_TOTAL,
+        tags={'pod_id': pool.pod_id, 'server_name': 'srv'},
+    )
 
 
 @pytest.mark.asyncio
@@ -131,3 +136,50 @@ async def test_release_idle_disconnect_retries_then_gives_up():
 
     assert fake_client.disconnect.await_count == 5
     assert 'srv' not in pool._connections
+
+
+@pytest.mark.asyncio
+async def test_success_path_emits_connection_metrics():
+    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})
+
+    pool.emit_metric.assert_any_await(
+        metric_object=metrics.OPC_CONNECTIONS_TOTAL,
+        tags={'pod_id': pool.pod_id, 'server_name': 'srv'},
+    )
+    pool.emit_metric.assert_any_await(
+        metric_object=metrics.OPC_CONNECTION_STATUS,
+        method='set',
+        value=1,
+        tags={'pod_id': pool.pod_id, 'server_name': 'srv', 'server_url': 'opc.tcp://h:4840'},
+    )
+
+
+@pytest.mark.asyncio
+async def test_concurrent_get_same_server_connects_once():
+    pool = make_pool()
+    started = asyncio.Event()
+    proceed = asyncio.Event()
+
+    clients = [AsyncMock(), AsyncMock()]
+
+    async def slow_connect():
+        started.set()
+        await proceed.wait()
+
+    clients[0].connect.side_effect = slow_connect
+
+    cfg = {'url': 'opc.tcp://h:4840', 'use_security': False}
+    with patch('ingestor.connection.pool.Client', side_effect=clients) as ctor:
+        t1 = asyncio.create_task(pool.get('srv', cfg))
+        await started.wait()  # t1 is mid-connect, holding the per-server lock
+        t2 = asyncio.create_task(pool.get('srv', cfg))
+        await asyncio.sleep(0)  # let t2 block on the lock (not build a 2nd Client)
+        proceed.set()
+        c1, c2 = await asyncio.gather(t1, t2)
+
+    assert c1 is c2 is clients[0]
+    assert ctor.call_count == 1
+    assert clients[0].connect.await_count == 1