Skip to content

fix: add per-device execute lock to prevent I/O interleaving #902

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ChristopherJHart wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: add per-device execute lock to prevent I/O interleaving #902

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 63 additions & 46 deletions nac_test/pyats_core/broker/connection_broker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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)
Expand Down Expand Up @@ -359,33 +373,32 @@ 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
rule out the session dying between the probe and the command anyway.
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."""
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion tests/integration/test_connection_broker_failure_modes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
22 changes: 11 additions & 11 deletions tests/unit/pyats_core/broker/test_connection_broker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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()
Expand All @@ -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")


# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -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,
Expand All @@ -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."""
Expand All @@ -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,
Expand All @@ -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(
Expand All @@ -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.
Expand All @@ -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(
Expand All @@ -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."""
Expand Down
Loading