diff --git a/nac_test/pyats_core/broker/connection_broker.py b/nac_test/pyats_core/broker/connection_broker.py index a68618c7..d1898dbf 100644 --- a/nac_test/pyats_core/broker/connection_broker.py +++ b/nac_test/pyats_core/broker/connection_broker.py @@ -55,7 +55,7 @@ def __init__( # Connection management self.testbed: Any | None = None self.connected_devices: dict[str, Any] = {} # hostname -> device connection - self.connection_locks: dict[str, asyncio.Lock] = {} + self._device_locks: dict[str, asyncio.Lock] = {} self.connection_semaphore = asyncio.Semaphore(max_connections) # Command caching - shared across all clients @@ -106,9 +106,9 @@ async def _load_testbed(self) -> None: logger.info(f"Loaded testbed with {len(self.testbed.devices)} devices") # type: ignore[attr-defined] - # Initialize connection locks for all devices + # Initialize per-device locks for all devices for hostname in self.testbed.devices: # type: ignore[attr-defined] - self.connection_locks[hostname] = asyncio.Lock() + self._device_locks[hostname] = asyncio.Lock() except Exception as e: logger.error(f"Failed to load testbed: {e}", exc_info=True) @@ -256,10 +256,17 @@ async def _execute_command(self, hostname: str, cmd: str) -> str: disconnecting, reconnecting and retrying the command exactly once, which recovers the current request instead of only cleaning up for the next one. + + A single per-device lock serialises the entire get-connection → + execute → failure-handling → retry cycle. This prevents two hazards: + (1) interleaved PTY I/O from concurrent execute() calls on Unicon's + non-thread-safe spawn, and (2) a stale caller tearing down a + successor's connection during the reconnect-and-retry window. + Cross-device parallelism is unaffected. """ cache = self._get_command_cache(hostname) - # Check cache first + # Check cache first (no lock needed — cache is per-device and read-only here) cached_output = cache.get(cmd) if cached_output is not None: self.stats_command_cache_hits += 1 @@ -270,27 +277,38 @@ async def _execute_command(self, hostname: str, cmd: str) -> str: self.stats_command_cache_misses += 1 logger.debug(f"Broker cache miss for '{cmd}' on {hostname}, executing...") - # Connection establishment errors are never retried here - only the - # execution itself is, and only when the session looks dead. - connection = await self._get_connection(hostname) - try: - return await self._run_and_cache(hostname, connection, cmd) - except Exception as e: - if not self._is_transport_failure(e): - raise - logger.warning( - f"Transport failure executing '{cmd}' on {hostname} ({e}); " - f"reconnecting and retrying once" - ) + if hostname not in self._device_locks: + self._device_locks[hostname] = asyncio.Lock() + + async with self._device_locks[hostname]: + # Re-check cache under lock — another caller may have populated it + cached_output = cache.get(cmd) + if cached_output is not None: + self.stats_command_cache_hits += 1 + logger.debug(f"Broker cache hit (under lock) for '{cmd}' on {hostname}") + return cached_output + + connection = await self._get_connection(hostname) + try: + return await self._run_and_cache(hostname, connection, cmd) + except Exception as e: + if not self._is_transport_failure(e): + raise + logger.warning( + f"Transport failure executing '{cmd}' on {hostname} ({e}); " + f"reconnecting and retrying once" + ) - # Final attempt. The failed attempt above disconnected the device, so - # this creates a fresh connection. Failures propagate to the caller. - connection = await self._get_connection(hostname) - return await self._run_and_cache(hostname, connection, cmd) + # Final attempt — the failed attempt disconnected the device, so + # this creates a fresh connection. Failures propagate to the caller. + connection = await self._get_connection(hostname) + return await self._run_and_cache(hostname, connection, cmd) async def _run_and_cache(self, hostname: str, connection: Any, cmd: str) -> str: """Run a command on a connection once and cache its output. + The caller must hold ``_device_locks[hostname]``. + Raises: SubCommandFailure: Re-raised as-is. The device answered and rejected the command, so the session and its cache are left intact. @@ -307,15 +325,11 @@ async def _run_and_cache(self, hostname: str, connection: Any, cmd: str) -> str: try: output = await loop.run_in_executor(None, connection.execute, cmd) except Exception as e: - # SubCommandFailure is a fast path, not a health guarantee - unicon wraps - # transport errors in it too. Those get neither a disconnect here nor a retry - # in _execute_command; recovery relies on unicon re-establishing the session - # on the next execute(), which it does transparently. if SubCommandFailure is not None and isinstance(e, SubCommandFailure): logger.warning(f"Command rejected by {hostname} (session intact): {e}") raise logger.error(f"Command execution failed on {hostname}: {e}") - await self._disconnect_device(hostname) + await self._disconnect_device_internal(hostname) raise output_str = str(output) @@ -359,6 +373,9 @@ def _is_transport_failure(error: Exception) -> bool: async def _get_connection(self, hostname: str) -> Any: """Get or create connection to device. + When called from ``_execute_command`` the caller already holds + ``_device_locks[hostname]``, so this method must not re-acquire it. + A cached connection is returned as-is, without probing it. Evaluating ``device.connected`` performs a live SSH round-trip (~0.37s) on the broker's event loop, blocking traffic for every device, and it cannot @@ -366,26 +383,22 @@ async def _get_connection(self, hostname: str) -> Any: Dead sessions are detected and healed by ``_execute_command``'s reconnect-and-retry instead. """ - if hostname not in self.connection_locks: - self.connection_locks[hostname] = asyncio.Lock() - - async with self.connection_locks[hostname]: - # Return existing connection - if hostname in self.connected_devices: - self.stats_connection_cache_hits += 1 - logger.info( - f"[BROKER] Reusing existing connection for {hostname} " - f"(total connections: {len(self.connected_devices)})" - ) - return self.connected_devices[hostname] - - # Create new connection - self.stats_connection_cache_misses += 1 + # Return existing connection + if hostname in self.connected_devices: + self.stats_connection_cache_hits += 1 logger.info( - f"[BROKER] Creating NEW connection for {hostname} " - f"(current connections: {len(self.connected_devices)})" + f"[BROKER] Reusing existing connection for {hostname} " + f"(total connections: {len(self.connected_devices)})" ) - return await self._create_connection(hostname) + return self.connected_devices[hostname] + + # Create new connection + self.stats_connection_cache_misses += 1 + logger.info( + f"[BROKER] Creating NEW connection for {hostname} " + f"(current connections: {len(self.connected_devices)})" + ) + return await self._create_connection(hostname) async def _create_connection(self, hostname: str) -> Any: """Create new connection to device using testbed.""" @@ -435,16 +448,20 @@ async def _create_connection(self, hostname: str) -> Any: async def _ensure_connection(self, hostname: str) -> tuple[bool, str]: """Ensure device is connected, return (success, error_message).""" + if hostname not in self._device_locks: + self._device_locks[hostname] = asyncio.Lock() + try: - await self._get_connection(hostname) + async with self._device_locks[hostname]: + await self._get_connection(hostname) return True, "" except Exception as e: return False, str(e) async def _disconnect_device(self, hostname: str) -> None: """Disconnect from device and clean up.""" - if hostname in self.connection_locks: - async with self.connection_locks[hostname]: + if hostname in self._device_locks: + async with self._device_locks[hostname]: await self._disconnect_device_internal(hostname) async def _disconnect_device_internal(self, hostname: str) -> None: diff --git a/tests/integration/test_connection_broker_failure_modes.py b/tests/integration/test_connection_broker_failure_modes.py index 57c79a26..25dc2446 100644 --- a/tests/integration/test_connection_broker_failure_modes.py +++ b/tests/integration/test_connection_broker_failure_modes.py @@ -72,7 +72,7 @@ def _factory(devices: dict[str, MagicMock]) -> ConnectionBroker: broker.testbed = MagicMock() broker.testbed.devices = devices for hostname in devices: - broker.connection_locks[hostname] = asyncio.Lock() + broker._device_locks[hostname] = asyncio.Lock() return broker return _factory diff --git a/tests/unit/pyats_core/broker/test_connection_broker.py b/tests/unit/pyats_core/broker/test_connection_broker.py index 79aef654..57960d0e 100644 --- a/tests/unit/pyats_core/broker/test_connection_broker.py +++ b/tests/unit/pyats_core/broker/test_connection_broker.py @@ -28,7 +28,7 @@ def broker(tmp_path: Path) -> ConnectionBroker: b.testbed = MagicMock() b.testbed.devices = {"router-1": MagicMock(), "router-2": MagicMock()} for hostname in b.testbed.devices: - b.connection_locks[hostname] = asyncio.Lock() + b._device_locks[hostname] = asyncio.Lock() return b @@ -307,7 +307,7 @@ def test_execute_command_disconnects_on_execution_failure( cache.get.return_value = None broker.command_cache["router-1"] = cache broker._get_connection = AsyncMock(return_value=MagicMock()) # type: ignore[method-assign] - broker._disconnect_device = AsyncMock() # type: ignore[method-assign] + broker._disconnect_device_internal = AsyncMock() # type: ignore[method-assign] async def _run() -> None: loop = asyncio.get_event_loop() @@ -326,7 +326,7 @@ async def _run() -> None: with pytest.raises(Exception, match="timeout"): asyncio.run(_run()) - broker._disconnect_device.assert_called_once_with("router-1") + broker._disconnect_device_internal.assert_called_once_with("router-1") # --------------------------------------------------------------------------- @@ -365,7 +365,7 @@ def test_retries_once_after_transport_failure( from unicon.core.errors import ConnectionError as UniconConnectionError broker._get_connection = AsyncMock(side_effect=[MagicMock(), MagicMock()]) # type: ignore[method-assign] - broker._disconnect_device = AsyncMock() # type: ignore[method-assign] + broker._disconnect_device_internal = AsyncMock() # type: ignore[method-assign] result = _run_execute( broker, @@ -376,7 +376,7 @@ def test_retries_once_after_transport_failure( assert result == "live output" assert broker._get_connection.await_count == 2 - broker._disconnect_device.assert_awaited_once_with("router-1") + broker._disconnect_device_internal.assert_awaited_once_with("router-1") def test_retry_result_is_cached(self, broker: ConnectionBroker) -> None: """The retry's output lands in a live cache, not the one disconnect dropped.""" @@ -388,7 +388,7 @@ def test_retry_result_is_cached(self, broker: ConnectionBroker) -> None: async def _disconnect(hostname: str) -> None: broker.command_cache.pop(hostname, None) - broker._disconnect_device = AsyncMock(side_effect=_disconnect) # type: ignore[method-assign] + broker._disconnect_device_internal = AsyncMock(side_effect=_disconnect) # type: ignore[method-assign] result = _run_execute( broker, @@ -404,7 +404,7 @@ def test_raises_when_retry_also_fails(self, broker: ConnectionBroker) -> None: from unicon.core.errors import ConnectionError as UniconConnectionError broker._get_connection = AsyncMock(side_effect=[MagicMock(), MagicMock()]) # type: ignore[method-assign] - broker._disconnect_device = AsyncMock() # type: ignore[method-assign] + broker._disconnect_device_internal = AsyncMock() # type: ignore[method-assign] with pytest.raises(UniconConnectionError): _run_execute( @@ -419,7 +419,7 @@ def test_raises_when_retry_also_fails(self, broker: ConnectionBroker) -> None: assert broker._get_connection.await_count == 2 # Both attempts tear the connection down; nothing dead is left cached - assert broker._disconnect_device.await_count == 2 + assert broker._disconnect_device_internal.await_count == 2 def test_does_not_retry_command_rejection(self, broker: ConnectionBroker) -> None: """A rejected command is not retried - a fresh session rejects it too. @@ -430,7 +430,7 @@ def test_does_not_retry_command_rejection(self, broker: ConnectionBroker) -> Non from unicon.core.errors import SubCommandFailure broker._get_connection = AsyncMock(return_value=MagicMock()) # type: ignore[method-assign] - broker._disconnect_device = AsyncMock() # type: ignore[method-assign] + broker._disconnect_device_internal = AsyncMock() # type: ignore[method-assign] with pytest.raises(SubCommandFailure): _run_execute( @@ -444,13 +444,13 @@ def test_does_not_retry_unclassified_failure( ) -> None: """An error that is neither rejection nor transport disconnects but is not retried.""" broker._get_connection = AsyncMock(return_value=MagicMock()) # type: ignore[method-assign] - broker._disconnect_device = AsyncMock() # type: ignore[method-assign] + broker._disconnect_device_internal = AsyncMock() # type: ignore[method-assign] with pytest.raises(RuntimeError, match="broker bug"): _run_execute(broker, "router-1", "show version", RuntimeError("broker bug")) assert broker._get_connection.await_count == 1 - broker._disconnect_device.assert_awaited_once_with("router-1") + broker._disconnect_device_internal.assert_awaited_once_with("router-1") def test_connection_errors_are_not_retried(self, broker: ConnectionBroker) -> None: """Failure to establish a connection must not double the connect attempts."""