perf(broker): replace SSH liveness probe with reconnect-and-retry (#909)

[ChristopherJHart]

Summary

Removes the proactive SSH liveness probe from the connection broker and replaces it with reconnect-and-retry in _execute_command(), per the target design in #909.

This supersedes the TTL-cached probe originally proposed on this branch (see the discussion below). The cache preserved the probe's cost model and added problems of its own — cache identity via id(connection), no invalidation on disconnect, and 30s of deliberate blindness — so the root cause is addressed instead.

Problem

_get_connection() probed every cached connection before handing it out by evaluating the pyATS device.connected property, which performs a live SSH round-trip (~0.37s) synchronously on the broker's event loop:

• Double evaluation — hasattr(connection, "connected") triggered the property getter, then connection.connected triggered it again (~0.74s per call).
• Event loop blocking — the probe blocked all device traffic fleet-wide while one device was probed.
• Anti-scaling — cost grew as O(tests × devices): ~216s of blocking probes per device job at 169 tests.
• TOCTOU — the session can die between the probe passing and connection.execute(), so the check never guaranteed correctness anyway.

The old recovery path in _execute_command() disconnected and re-raised, which cleaned up for the next caller but failed the current request.

Solution

1. _get_connection() returns the cached connection as-is, or creates one. _is_connection_healthy() is removed.
2. _execute_command() retries once on transport failure. The failed attempt tears the connection down; the retry runs on a fresh one. This heals the current request rather than only the next.
3. Retry is scoped to failures a fresh session can fix — Unicon ConnectionError, TimeoutError, StateMachineError, plus OSError/EOFError. Classification lives in one small helper (_is_transport_failure) so it's easy to adjust in review.
4. Connection establishment errors are not retried, so an unreachable device still costs one connect attempt rather than two 60s timeouts.
5. Command rejections are untouched by this PR. SubCommandFailure is simply absent from the retry allowlist, so a rejected command is not re-run. Whether a rejection should also skip the disconnect is perf: don't disconnect + wipe cache on command rejection #900's change, left wholly to that PR.

Behavior changes worth a look in review

• connect() no longer validates liveness. _ensure_connection() returns success for a cached session without probing it; the first command execution handles recovery. Called out as acceptable in perf: proactive SSH health check in _get_connection() blocks event loop and anti-scales with fleet size #909.
• Retry assumes broker commands are safe to re-run. Every current caller issues read-only show commands, including the Genie supplementary calls routed through the broker by feat(broker): support Genie-driven command execution in broker mode #863. If non-idempotent commands are ever sent through the broker, the retry needs to become opt-in.
• Nothing here changes rejection handling. An earlier revision of this PR also made SubCommandFailure skip the disconnect, duplicating perf: don't disconnect + wipe cache on command rejection #900; that has been removed, so disconnect-on-any-exception behaves exactly as it does on main today.

Measured impact

Re-measured on a 21-device lab fleet (19 Nexus 9k + 2 Firepower 4215), SUITE=minimal
= 914 test instances per run, one variable changed between arms:

Arm	Broker behavior	NX-OS	FTD	Wall clock
A	pristine probe (double hasattr)	15m48s	1m24s	17m42s
B	probe with 30s TTL cache (this PR's first version)	2m32s	1m07s	4m09s
C	probe removed + retry once (this PR)	2m21s	1m06s	3m57s

A → C is 4.5x. The honest reading, though, is that the TTL cache already
captured almost all of it (A → B is 4.3x); removing the probe outright adds ~5%
more (B → C, 249s → 237s).

So this is not primarily a performance change over the cache it replaces. The
reason to prefer removal is the correctness argument in #909 — no id()-keyed
cache that can alias a recycled connection, no invalidation to forget on
disconnect, no window of deliberate blindness, and recovery that heals the
current request instead of only the next caller's. The 4.5x is what either
approach buys over what is on main today.

All three arms produced identical results: 914 tests, 726 passed, 188 failed,
0 skipped/errored, aggregated from every device's pyATS results.json, with the
per-test outcome sets matching name-for-name across arms (0 changed, 0 added, 0
missing). The retry never fired during arm C — correct for a healthy fleet.

Test plan

• tests/unit — 1085 passed. Broker tests reworked: the two _is_connection_healthy tests are gone; TestGetConnection now asserts the cached connection is handed out without touching device.connected; TestConnectionHealthRecovery exercises the full dead-session → disconnect → reconnect → retry path against the real _get_connection/_disconnect_device/_create_connection.

• New TestExecuteCommandRetry and TestFailureClassification: retry succeeds, retry result lands in a live cache (regression for the cache dropped by the intervening disconnect), retry exhaustion raises with both attempts torn down, SubCommandFailure is not retried, unclassified errors disconnect but don't retry, connect errors aren't retried.

• tests/integration — 77 passed, 2 skipped, 1 failed. The failure (test_connection_broker_pooling_and_caching) is pre-existing and environmental: this machine has no python on PATH, so the mock device fails with SpawnInitError. Verified identical (assert 6 == 2) on the branch point with the change stashed.

• ruff check, ruff format --check, mypy clean on both touched files.

• Live lab validation — the three-arm run above, plus a recovery check on a
  real Nexus switch, which is the part timing cannot cover. Two scenarios:

  • Killing the Unicon spawn out from under the broker never reaches the broker:
    Unicon's own execute() catches the EOF and reconnects
    (%UNICON-WARNING: +++ Reconnecting +++). Worth knowing — the broker's retry
    is a second line of defence, and part of why the probe bought so little.
  • Injecting a unicon ConnectionError at the boundary where the broker does
    see failures exercises the real path: disconnect → reconnect over SSH →
    command re-run → 872 chars of real show module output → follow-up command
    served by the healed session. Recovery cost ~14s, ~11s of which is Unicon's
    graceful disconnect teardown. That is the true price of a retry, paid only
    when a session is actually dead — against a probe that ran on every command.

  The error in the second scenario is injected, not spontaneous; it reproduces on
  demand what Unicon surfaces when it cannot self-heal.

• One reporting discrepancy, surfaced by the runs but not caused by this change:
  nac-test's console summary split NX-OS as 175 failed / 97 skipped in arm B
  against 176 / 96 in arms A and C, while the underlying pyATS data and every
  per-test result line are identical across all three arms. That looks like the
  summary layer classifying one result differently and deserves its own issue.

Files changed

File	Change
nac_test/pyats_core/broker/connection_broker.py	Remove _is_connection_healthy(), simplify _get_connection(), add retry-once + _run_and_cache()/_get_command_cache()/_is_transport_failure()
tests/unit/pyats_core/broker/test_connection_broker.py	Rework health-check tests, add retry and classification coverage
dev-docs/PRD_AND_ARCHITECTURE.md	Refresh the three broker sections that quoted the removed probe
CHANGELOG.md	Unreleased entry

Closes #909

---

🤖 AI Generation Metadata

• AI Generated: Yes
• AI Tool: claude-code
• AI Model: opus-5
• AI Contribution: ~71% (lines added / total changed)
• AI Reason: implement perf: proactive SSH health check in _get_connection() blocks event loop and anti-scales with fleet size #909 reactive reconnect-and-retry, replacing the TTL health cache
• Human Oversight: Direction set by @oboehmer in perf: proactive SSH health check in _get_connection() blocks event loop and anti-scales with fleet size #909; pending review

[ChristopherJHart]

Design Question: Should _is_connection_healthy() be removed entirely?

While working on this TTL cache fix, I wanted to surface a broader question for the maintainers.

Current behavior

_is_connection_healthy() is called on every _get_connection() invocation — i.e., before every command execution. It probes the live SSH session via device.connected (~0.37s round-trip) to detect stale connections before use.

The execute path already recovers from dead connections

async def _execute_command(self, hostname, cmd):
    connection = await self._get_connection(hostname)  # ← health check here
    try:
        output = await loop.run_in_executor(None, connection.execute, cmd)
        return output
    except Exception as e:
        if isinstance(e, SubCommandFailure):
            raise  # session is fine, command was rejected
        await self._disconnect_device(hostname)  # ← recovery here
        raise

If the connection is dead, connection.execute(cmd) raises a transport exception, which triggers disconnect. The next test's _get_connection() finds no cached connection and creates a fresh one. Recovery is already handled.

Cost-benefit of the proactive check

	Pro	Con
Proactive (current)	Avoids one failed execute attempt per dead connection	0.74s per test, blocks entire event loop, TOCTOU race (connection can die after check passes)
Reactive (remove check)	Zero overhead on every healthy call	One extra failed execute + reconnect when connection actually dies (rare)

The proactive check costs 0.74s × N tests × M devices to prevent a scenario (connection dies between tests) that happens at most once per device per session. It's also subject to a time-of-check/time-of-use race — the connection can pass the health check and die before the command runs.

Suggestion

Given the "retry on failure" model in _execute_command, would it be worth removing _is_connection_healthy() entirely and simplifying _get_connection() to:

async def _get_connection(self, hostname):
    async with self.connection_locks[hostname]:
        if hostname in self.connected_devices:
            return self.connected_devices[hostname]
        return await self._create_connection(hostname)

The TTL cache in this PR is a conservative middle ground — it preserves the check's intent while reducing the cost from 0.74s/call to 0.004s/call. But if the maintainers agree the check adds no value given the execute-path recovery, the simpler approach is to just remove it.

Happy to adjust this PR in either direction based on feedback.

[oboehmer]

@ChristopherJHart Thanks for the thorough analysis — both the PR and especially your design question comment. You're right that the proactive health check is problematic: the double evaluation via hasattr(), the event loop blocking, and the TOCTOU race are all real issues, and the scaling numbers speak for themselves.

I've been thinking along the same lines as your "remove it entirely" suggestion. The TTL cache is a reasonable middle ground, but it introduces its own complexity (cache identity via id(connection), missing invalidation on disconnect, 30s of intentional blindness). I'd rather address the root cause.

My preferred direction is: remove the proactive probe and add retry-once logic to _execute_command(). The key observation is that today's recovery path only cleans up for the next caller — it doesn't heal the current request. If we add a single retry after reconnect (for transport failures, not SubCommandFailure), we get the best of both worlds: zero overhead on the happy path, and graceful recovery for the rare dead-session case.

I've opened #909 to track this as the target design. Would you be interested in updating this PR in that direction? The scope would be:

1. Remove _is_connection_healthy() from _get_connection() — just return the cached connection or create a new one.
2. In _execute_command(), on transport failure: disconnect, reconnect, retry once, then raise if it fails again.

Happy to discuss further if you have concerns about retry safety or edge cases.

[ChristopherJHart]

@oboehmer Thanks for the direction — I've repointed this PR at #909 and dropped
the TTL cache entirely.

You were right on all three counts: the id(connection) keying can alias a
recycled address, _disconnect_device_internal() never cleared the cache, and
the "stale-positive is harmless" claim in my original description was wrong — the
old handler disconnected and re-raised, so a stale positive converted a
recoverable dead session into a test failure. Your "only cleans up for the next
caller" framing is the accurate one.

What this PR does now

1. _get_connection() returns the cached connection as-is, or creates one.
   _is_connection_healthy() is removed.
2. _execute_command() retries once on transport failure: the failed attempt
   tears the connection down, the retry runs on a fresh one.
3. Retry is an allowlist (_is_transport_failure): Unicon ConnectionError,
   TimeoutError, StateMachineError, plus OSError/EOFError.
4. Connection establishment errors are outside the retry, so an unreachable
   device costs one connect attempt, not two.

What it does NOT do

Command rejection handling (SubCommandFailure skip-disconnect) is left to #900,
which has its own justification and benchmark. This PR preserves today's
disconnect-on-any-exception semantics — SubCommandFailure is simply absent from
the retry allowlist, so it's not retried, but it still disconnects as before.

Lab validation (21-device fleet, 914 test instances per arm)

Arm	Broker behavior	NX-OS	FTD	Wall
A	pristine probe (double hasattr, what's on main)	15m48s	1m24s	17m42s
B	probe + 30s TTL cache (this PR's first version)	2m32s	1m07s	4m09s
C	probe removed + retry once (this PR)	2m21s	1m06s	3m57s

A → C is 4.5x — but the TTL cache already captured almost all of it (A → B is
4.3x). Removing the probe adds ~5% more. So this is not primarily a performance
win over the cache; it's the correctness and simplicity case you made in #909.

All three arms produced identical results (726 passed, 188 failed, 0
skipped/errored), matching name-for-name across every device's pyATS
results.json. The retry never fired in arm C — correct for a healthy fleet.

Recovery proof (the part timing can't check): Against a live Nexus switch,
injected a unicon ConnectionError at the boundary where the broker sees
failures. The broker disconnected, reconnected over SSH, re-ran the command, and
returned 872 chars of real show module output. Follow-up command worked on the
healed session. Recovery cost ~14s (mostly Unicon's graceful teardown). The error
was injected, not spontaneous — it reproduces what Unicon surfaces when it can't
self-heal.

One thing the recovery check revealed: killing the Unicon spawn (instead of
injecting at the boundary) never reaches the broker at all — Unicon's own
execute() catches the EOF and reconnects internally. The broker's retry is a
second line of defence, which is part of why the probe was buying so little.

Two things worth your eyes

1. feat(broker): support Genie-driven command execution in broker mode #863 interaction. Genie supplementary device.execute() calls now route
   through the broker, so the retry can fire mid-parse. Fine for show commands,
   but it widens the idempotency assumption perf: proactive SSH health check in _get_connection() blocks event loop and anti-scales with fleet size #909 flagged — flagging it as a
   conscious decision rather than an accident.

2. A retry-specific trap I found and fixed. _disconnect_device() deletes
   command_cache[hostname], so the naive retry would write its output to an
   orphaned CommandCache and silently lose it. _get_command_cache() re-resolves
   the cache after the retry; test_retry_result_is_cached covers it.

[oboehmer]

Direction is right and the diff is small. Removing the probe is the correct call, the separation from #900 (ab3e476) is clean, and the PR body is refreshingly honest that A→B captured 4.3x of the 4.5x — this is a correctness change, not a performance one, and it says so. CI green across all 13 checks.

Taking your two closing points first, since one of them shapes what I'd ask for in this PR.

On the #863 interaction

Right to flag, and I'd push it one step further. Before #863 the command set reaching the broker was whatever test authors wrote — in-repo, reviewable, demonstrably all show. After #863 it also includes whatever supplementary commands the installed Genie parsers decide to issue (your show running-config | section router ospf 1 example). That set isn't controlled by this repo and can change on a genie bump with no diff here. So the assumption isn't just wider, it's partly owned by a dependency.

That argues for making it cheap to enforce rather than documenting it. _is_transport_failure already answers "is this failure worth retrying"; add the symmetric "is this command safe to retry" next to it:

if self._is_transport_failure(e) and self._is_retryable_command(cmd):

with cmd.strip().lower().startswith("show") as the predicate. It's a no-op on every current caller, so it won't disturb the three-arm validation you've already banked, and it fails in the safe direction — a command that doesn't match simply loses the retry and behaves exactly as main does today. I'd rather see this here than as a tracking issue: the hazard only materialises when someone adds a non-show caller, which is precisely the moment a tracking issue goes unread. Three lines plus a test, and the conscious decision lives in the code rather than only in the PR description.

Smaller related note: broker_execute (ssh_base_test.py:334) keeps a client-side CommandCache that a broker-side disconnect doesn't touch, so the two diverge after any recovery. Not a correctness problem, but worth knowing when #900 gets benchmarked — the cache wipe costs less than its framing implies.

On the orphaned-cache trap

Verified, and it's the strongest defensive detail in the PR. The old binding (cache = self.command_cache[hostname] at the top) against del self.command_cache[hostname] in _disconnect_device_internal would have had the retry populate a cache nothing references — correct output returned, then silently re-executed on every subsequent request. A slow leak of exactly the property the broker exists to provide.

_get_command_cache() re-resolving inside _run_and_cache is the right shape, and the comment explaining why it re-resolves is worth keeping — it's the line that stops someone folding it back into a bug. test_retry_result_is_cached has good fidelity too: the side effect actually pops the entry rather than asserting against a mock, so it fails if the re-resolution goes away. No changes suggested.

The one other thing I'd fix in this PR

Retry fires even when the connection was created in the same call. The rationale throughout is "a session that died since its last use", but _execute_command doesn't distinguish a cache hit from a miss:

connection = await self._get_connection(hostname)   # may be brand new
try:
    return await self._run_and_cache(hostname, connection, cmd)
except Exception as e:
    if not self._is_transport_failure(e):
        raise
# → disconnect + full reconnect + re-execute

Unicon TimeoutError is in the allowlist, so a device that is reachable but hangs now costs 2× execute plus a full reconnect on a connection that was alive moments earlier. The "unreachable device still costs one connect attempt" argument holds for establishment errors but not execute-time ones. Having _get_connection report cached-vs-fresh, and gating the retry on cached, makes the code enforce the invariant the docstrings already describe.

Deferrable

Retry budget vs. the client timeout. ssh_base_test.py:345 uses future.result(timeout=DEVICE_EXECUTE_TIMEOUT) (120s). Failed execute + graceful Unicon disconnect (~11s of your measured ~14s) + reconnect + re-execute can exceed that, at which point the client abandons the future while the broker keeps working — run_coroutine_threadsafe won't cancel it. Happy path is fine; a docstring note that the retry can outlive its caller would do for now.

Teardown race — suggest its own issue. Not introduced here. The per-device lock is held only inside _get_connection, so A can fail → reconnect → be executing on a fresh connection when B's late _disconnect_device tears it down. Possible on main already, but retry widens the window. I checked whether draft #902 covers it — it doesn't: #902 takes _execute_locks[hostname] while _disconnect_device guards connection_locks[hostname], two locks over two halves of one critical section. Cheapest fix, if we open an issue: _disconnect_device(hostname, expect=connection) that only tears down when connected_devices.get(hostname) is connection. Separately — #899 and #902 both rewrite _execute_command and #899 moves the dispatch into _run_and_cache(), so whichever lands second will conflict.

Unicon import guard. The lazy import matches house style; the guard doesn't — the four equivalents in ssh/connection_manager.py (:127-134, :196, :277-281) are bare. As the # pragma: no cover acknowledges, the branch can't be reached: by the time _is_transport_failure runs, an execute has happened, so unicon is installed. Dropping it removes unreachable code. Unrelated to the guard, declaring "unicon>=26.5; sys_platform != 'win32'" alongside pyats/genie would make the import contract explicit.

Minor. stats_connection_cache_hits now increments for possibly-dead connections, so the ratio no longer proxies reuse quality. And the same probe pattern survives on the non-broker path (ssh/connection_manager.py:100, :247, :461) — out of scope for #909, but worth an issue so the two paths don't rot asymmetrically.

Verified

Classification checked against installed unicon 26.5: TimeoutError subclasses builtin OSError (doubly covered); ConnectionError/StateMachineError are plain Exception, so explicit listing is needed; SubCommandFailure, CredentialsExhaustedError, UniconAuthenticationError correctly excluded — not retrying auth failure is an important call. Builtin ConnectionError from _create_connection sits outside the try, matching your claim.

Tests assert the right things: test_does_not_probe_cached_connection verifies the property was never evaluated rather than mocking a deleted helper, and TestConnectionHealthRecovery now exercises the full dead-session path against the real methods. CHANGELOG lands correctly in a new # unreleased above 2.1.0b1, and the three PRD sections match the shipped code.

---

Approving — nothing in the diff is wrong, and main is worse in every dimension, so this shouldn't be gated on hardening for hazards that are all forward-looking. The retry-command guard would ideally land here since it's ~3 lines and a no-op today; the cached-vs-fresh gate here or immediately after. Teardown race, retry budget, import guard and the non-broker probe are all follow-up issues.

[oboehmer]

@ChristopherJHart , I did not open any issues for the deferrable points, please let me know which one you can/want to address in this PR (I would definitely remove the try/except around the unicon imports) and I will let my agent open the ones which are done outside..

[oboehmer]

Approval stands — one follow-up on _is_transport_failure that I only found after checking the actual exception hierarchy.

Two of the failures this retry is meant to catch aren't matched by the current allowlist:

>>> [c.__name__ for c in unicon.core.errors.SessionConnectionError.__mro__]
['SessionConnectionError', 'Exception', 'object']
>>> [c.__name__ for c in unicon.core.errors.EOF.__mro__]
['EOF', 'Exception', 'object']
>>> [c.__name__ for c in unicon.core.errors.TimeoutError.__mro__]
['TimeoutError', 'TimeoutError', 'OSError', 'Exception', 'BaseException']

1. SessionConnectionError is not a subclass of unicon's ConnectionError — it's a bare Exception, so isinstance() misses it. It's also the most literally-named "the SSH session died" error unicon has, which makes it an odd one to skip.
2. unicon.core.errors.EOF is not the builtin EOFError. The allowlist has the builtin; unicon's EOF is unrelated to it and won't match.
3. Minor: UniconTimeoutError is redundant — it already subclasses builtin TimeoutError → OSError, which is in the base list. Harmless, just noise.

Suggested:

from unicon.core.errors import EOF as UniconEOF
from unicon.core.errors import SessionConnectionError

retryable += [
    UniconConnectionError,
    SessionConnectionError,
    UniconEOF,
    StateMachineError,
]  # UniconTimeoutError already covered by OSError

---

Separately, for the record on the classification question, since it affects how much of that allowlist is reachable:

unicon.plugins.generic.service_implementation.Execute.call_service ends with

except StateMachineError:
    raise
except UniconBackendDecodeError:
    pass
except Exception as err:
    raise SubCommandFailure("Command execution failed", err) from err

so transport errors raised inside dialog.process() come back out as SubCommandFailure and won't reach the allowlist regardless of what's in it. I checked whether the platform plugins avoid this and they don't — nxos has no Execute override (runtime: NxosExecute, call_service resolving from generic), iosxe subclasses generic without overriding call_service, and iosxr overrides it only to set detect_state=False before delegating to super().

That's not an argument against this PR — con.sendline(), chatty_term_wait() and the trailing sm.go_to(end_state) all sit outside that try block, and StateMachineError is re-raised unwrapped, so the allowlist covers real paths. It just means the retry is narrower than the PR description implies, and the two additions above widen it where it's cheap to do so.

Worth noting the recovery here is also belt-and-braces: pyATS/unicon re-establishes a dead session on the next execute() by itself. I confirmed with a mock device — SIGKILL the spawn and its child, then execute again, and the call succeeds transparently on a new pid and fd (spawn object replaced, no exception surfaced). Which is consistent with your arm C observation that the retry never fired on a healthy fleet.

None of this is blocking. Happy for the allowlist tweak to land as a follow-up if you'd rather not disturb the validated branch.

[ChristopherJHart]

@oboehmer — thanks for the deep dive into the exception hierarchy; that was a gap I wouldn't have caught from the docs alone.

Addressed in this PR (0a98581):

The import guard and allowlist are fixed in one commit:

• Dropped the try/except around unicon imports — bare imports, matching ssh/connection_manager.py.
• Added SessionConnectionError and unicon.core.errors.EOF to the allowlist.
• Removed UniconTimeoutError (redundant via OSError) and builtin EOFError (not what unicon raises).

Tests updated: SessionConnectionError and UniconEOF in transport failures, builtin EOFError moved to non-transport.

Filed as follow-up issues for the deferrable items:

Issue	Item
#938	Guard retry on show-only commands
#939	Skip retry when connection was just created (cached vs. fresh)
#940	Retry can outlive the client's 120s timeout
#941	Teardown race between concurrent callers (_disconnect_device vs. fresh connection)
#942	Remove proactive probe from the non-broker SSH path

Each issue references back to this PR and the relevant review context. #941 also notes the conflict surface with #902 so whoever picks it up coordinates.

The stats_connection_cache_hits semantic drift you flagged is real but minor — happy to fold it into whichever of these touches the stats path, or leave it as-is if the counter is only used for logging.