From 5b471407008ec6946e2c2f40926a2f4548493fbb Mon Sep 17 00:00:00 2001 From: Christopher Hart Date: Fri, 4 Sep 2026 17:05:43 +0000 Subject: [PATCH 1/5] fix: unify per-device locking to prevent I/O interleaving and teardown races MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidate connection_locks and the proposed _execute_locks into a single _device_locks dict. The lock in _execute_command now spans the full get-connection → execute → failure-handling → retry cycle, which closes 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. _get_connection and _run_and_cache are now lock-free — callers hold the lock. External entry points (_ensure_connection, _disconnect_device) acquire _device_locks themselves. Addresses review feedback on PR #902 and overlap with #899. Co-Authored-By: Claude Opus 4.6 (1M context) AI-Generated: yes AI-Tool: claude-code AI-Model: opus-4.6 AI-Percent: 56 AI-Reason: per-device lock consolidation addressing PR #902 review feedback --- .../pyats_core/broker/connection_broker.py | 109 ++++++++++-------- .../test_connection_broker_failure_modes.py | 2 +- .../broker/test_connection_broker.py | 22 ++-- 3 files changed, 75 insertions(+), 58 deletions(-) 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.""" From cbf114424a3071677c5d7c1c997fa28c06e5fda2 Mon Sep 17 00:00:00 2001 From: Christopher Hart Date: Mon, 7 Sep 2026 10:50:32 -0400 Subject: [PATCH 2/5] fix(broker): bound shutdown disconnect with asyncio.wait_for MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit shutdown() → _disconnect_device() now acquires _device_locks, which is held across the full execute + reconnect + retry cycle. If a client abandons a request but the broker's run_in_executor is still blocked in the thread pool, shutdown waits on that command — potentially twice over given the retry. Wrap each disconnect in asyncio.wait_for(timeout=30) so the run doesn't hang at teardown on flaky devices. Addresses review feedback item #1 on PR #902. Co-Authored-By: Claude Opus 4.6 (1M context) --- nac_test/pyats_core/broker/connection_broker.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/nac_test/pyats_core/broker/connection_broker.py b/nac_test/pyats_core/broker/connection_broker.py index d1898dbf..72edb90e 100644 --- a/nac_test/pyats_core/broker/connection_broker.py +++ b/nac_test/pyats_core/broker/connection_broker.py @@ -522,9 +522,15 @@ async def shutdown(self) -> None: writer.close() await writer.wait_closed() - # Disconnect all devices + # Disconnect all devices — bounded so an in-flight execute holding + # the device lock cannot block shutdown indefinitely. for hostname in list(self.connected_devices.keys()): - await self._disconnect_device(hostname) + try: + await asyncio.wait_for(self._disconnect_device(hostname), timeout=30.0) + except asyncio.TimeoutError: + logger.warning( + f"Timed out waiting for device lock on {hostname} during shutdown" + ) # Stop socket server if self.server: From 4605072d421b985f442c9de8603c4e70d822c20b Mon Sep 17 00:00:00 2001 From: Christopher Hart Date: Mon, 7 Sep 2026 10:51:02 -0400 Subject: [PATCH 3/5] fix(broker): replace SubCommandFailure import guard with placeholder class MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The local try/except in _run_and_cache set SubCommandFailure = None on ImportError, which silently skipped the isinstance check and routed every command rejection through disconnect-and-wipe — reverting #900. Hoist the import to module level with a placeholder Exception subclass on ImportError (Windows). isinstance() now works correctly on all platforms — command rejections are always caught. Addresses review feedback item #2 on PR #902. Co-Authored-By: Claude Opus 4.6 (1M context) --- nac_test/pyats_core/broker/connection_broker.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/nac_test/pyats_core/broker/connection_broker.py b/nac_test/pyats_core/broker/connection_broker.py index 72edb90e..285777fc 100644 --- a/nac_test/pyats_core/broker/connection_broker.py +++ b/nac_test/pyats_core/broker/connection_broker.py @@ -20,6 +20,14 @@ from pathlib import Path from typing import Any +try: + from unicon.core.errors import SubCommandFailure +except ImportError: + + class SubCommandFailure(Exception): # type: ignore[no-redef] + """Placeholder when unicon is not installed (Windows).""" + + from nac_test.pyats_core.constants import MAX_BROKER_MESSAGE_BYTES from nac_test.pyats_core.ssh.command_cache import CommandCache from nac_test.utils import get_or_create_event_loop @@ -317,15 +325,10 @@ async def _run_and_cache(self, hostname: str, connection: Any, cmd: str) -> str: """ # Execute command in thread pool (since Unicon is synchronous) loop = get_or_create_event_loop() - try: - from unicon.core.errors import SubCommandFailure - except ImportError: - SubCommandFailure = None # type: ignore[assignment,misc] - try: output = await loop.run_in_executor(None, connection.execute, cmd) except Exception as e: - if SubCommandFailure is not None and isinstance(e, SubCommandFailure): + if isinstance(e, SubCommandFailure): logger.warning(f"Command rejected by {hostname} (session intact): {e}") raise logger.error(f"Command execution failed on {hostname}: {e}") From 8b59b2f28b6766b9997201c1c8b64c404c4a34ed Mon Sep 17 00:00:00 2001 From: Christopher Hart Date: Mon, 7 Sep 2026 10:51:33 -0400 Subject: [PATCH 4/5] test(broker): pin lock contract on the retry path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assert that _device_locks[hostname].locked() is True when _get_connection is called on both the initial and retry attempts. This pins the contract that _get_connection and _run_and_cache are lock-free — if a future change narrows or drops the lock in _execute_command, this test catches it before the teardown race reappears. Addresses review feedback item #3 on PR #902. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../broker/test_connection_broker.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/unit/pyats_core/broker/test_connection_broker.py b/tests/unit/pyats_core/broker/test_connection_broker.py index 57960d0e..40799689 100644 --- a/tests/unit/pyats_core/broker/test_connection_broker.py +++ b/tests/unit/pyats_core/broker/test_connection_broker.py @@ -452,6 +452,33 @@ def test_does_not_retry_unclassified_failure( assert broker._get_connection.await_count == 1 broker._disconnect_device_internal.assert_awaited_once_with("router-1") + def test_retry_path_stays_under_device_lock(self, broker: ConnectionBroker) -> None: + """_get_connection on both attempts must run under _device_locks[hostname]. + + This pins the contract that _get_connection and _run_and_cache are + lock-free — callers hold the lock. If a future change narrows or + drops the lock in _execute_command, this test catches it. + """ + from unicon.core.errors import ConnectionError as UniconConnectionError + + observed: list[bool] = [] + broker._disconnect_device_internal = AsyncMock() # type: ignore[method-assign] + + async def _spy(hostname: str) -> MagicMock: + observed.append(broker._device_locks[hostname].locked()) + return MagicMock() + + broker._get_connection = _spy # type: ignore[method-assign] + + _run_execute( + broker, + "router-1", + "show version", + [UniconConnectionError("socket closed"), "live output"], + ) + + assert observed == [True, True] + def test_connection_errors_are_not_retried(self, broker: ConnectionBroker) -> None: """Failure to establish a connection must not double the connect attempts.""" broker._get_connection = AsyncMock( # type: ignore[method-assign] From 99bbf9e2d9f2d671a5d7a27450bd0ea40a957734 Mon Sep 17 00:00:00 2001 From: Christopher Hart Date: Mon, 7 Sep 2026 10:52:10 -0400 Subject: [PATCH 5/5] docs: add CHANGELOG entry and document broker concurrency model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Bug Fixes entry for the unified per-device locking. Document the Connection Broker concurrency invariant in PRD_AND_ARCHITECTURE.md: the per-device lock is uncontended by design (one subprocess per device, _request_lock serialisation, per-device testbed), per-device sequential execution is intentional, and lock ordering is _device_locks → connection_semaphore. Addresses review feedback item #4 on PR #902. Co-Authored-By: Claude Opus 4.6 (1M context) --- CHANGELOG.md | 4 ++++ dev-docs/PRD_AND_ARCHITECTURE.md | 17 +++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c3be54f7..f0420c32 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,10 @@ - pyats broker: removed the per-command SSH liveness probe from the connection broker; dead sessions are now recovered by reconnecting and retrying the command once. Removes ~0.77s of event-loop blocking per test and the fleet-wide slowdown it caused. - pyats broker: command rejections (SubCommandFailure) no longer trigger a full SSH disconnect and cache wipe; only transport-level failures do. Eliminates a 4.7x per-failure penalty on devices with unsupported commands. +## Bug Fixes + +- pyats broker: unified per-device locking to prevent a stale caller from tearing down a successor's connection during the reconnect-and-retry window. The execute and disconnect paths previously used separate locks over two halves of one critical section. + # 2.1.0b1 ## Features diff --git a/dev-docs/PRD_AND_ARCHITECTURE.md b/dev-docs/PRD_AND_ARCHITECTURE.md index 200ab634..865acc8d 100644 --- a/dev-docs/PRD_AND_ARCHITECTURE.md +++ b/dev-docs/PRD_AND_ARCHITECTURE.md @@ -2409,6 +2409,23 @@ archive_paths = await asyncio.gather(*tasks, return_exceptions=True) Return d2d_aggregated_archive path ``` +##### Connection Broker Concurrency Model + +The Connection Broker uses a single `_device_locks[hostname]` asyncio.Lock per device to serialise the full get-connection → execute → failure-handling → retry cycle inside `_execute_command`. This lock prevents two hazards: + +1. **Interleaved PTY I/O** — Unicon's spawn is not thread-safe; concurrent `execute()` calls on the same device interleave I/O on the PTY, corrupting command output. +2. **Teardown race** — without a unified lock, a stale caller could tear down a successor's freshly-created connection mid-execute during the reconnect-and-retry window. + +**The per-device lock is uncontended by design.** Three mechanisms guarantee this at runtime: + +- One subprocess per device with unique hostnames (`orchestrator.py` device-centric execution) +- `BrokerClient._request_lock` serialises every request within each subprocess +- Per-device testbeds mean a test cannot address a peer hostname + +Per-device sequential execution is a deliberate design choice to avoid overloading devices, not an implementation artefact. The lock is therefore a runtime no-op today — it exists as a correctness invariant against the teardown race, not as active contention management. + +Lock ordering is `_device_locks[hostname]` → `connection_semaphore`, never the reverse. + ##### Development Mode Flow (--pyats Flag) ```