feat: thread-safe SSH/NETCONF connection pooling by h-network · Pull Request #27 · Juniper/junos-mcp-server

h-network · 2026-04-02T08:13:27Z

feat: thread-safe SSH/NETCONF connection pooling

Summary

Reuse SSH/NETCONF sessions across tool invocations instead of opening
and tearing down a connection per call.

Thread-safe ConnectionPool with per-router locking
Auto-detection and reconnection of stale sessions
Background cleanup of idle connections (configurable timeout)
Graceful shutdown on server exit

Performance

Tested on 5 Junos routers, 7 commands each (35 operations):

Mode	Time	Speedup
Original	24.5s	baseline
Pooled + sequential	4.0s	6.2x
Pooled + parallel	1.2s	20.7x

Benchmark methodology and reproduction steps: https://github.com/h-network/junos-mcp-server/tree/h-dev/benchmark

Closes #26

Wrap int() parsing of JMCP_POOL_IDLE_TIMEOUT in try/except ValueError, log a warning, and fall back to the 300s default — mirrors the pattern in get_timeout_with_fallback. Addresses review feedback on Juniper#27.

nileshsimaria · 2026-05-01T15:11:35Z

Looks good.

nileshsimaria · 2026-05-01T15:46:34Z

@h-network once you resolve conflict, it will merge. Thanks for your contributions. Appreciate it.

Addresses the second round of code review on PR Juniper#27. - Config-lock leak (completeness): last round's close-on-failed-cleanup only ran for an allowlist of exception types, but unlock() raises UnlockError (a direct RpcError subclass) and rollback() can raise a bare RpcError -- neither in the list -- so the realistic failure mode still returned a config-locked session to the pool. Widen the inner except to 'except Exception:' before closing the transport; classifying the failure is the outer handler's job. - Idle-cleanup gap: get_connection only refreshed last_used after a successful op, so a connection whose operation raised while still connected kept last_used=0.0 and the idle-cleanup loop (which requires last_used > 0) never reaped it. Refresh last_used in finally. - reload_devices race: close_all() ran before the device-map swap, so a concurrent tool call could open a pooled connection with the OLD config during the reset and survive it. Swap devices first; entry creation is serialized with close_all()'s snapshot via the pool lock and config is read only after entry creation, so connections opened during the reset use the NEW config. Tests: lock-leak regression test now exercises the real UnlockError path and asserts the pool evicts the session (rebuilds Device on next borrow); adds a last_used-refreshed-after-failure test. Both fail on pre-fix code.

Addresses the third review round on PR Juniper#27 — concurrency issues our pool introduced (scoped to our changes; pre-existing items left alone per the reviewers' own classification). - close_all() no longer clears _connections: clearing detached an entry an in-flight get_connection() borrow already held, orphaning the connection it then opened (never reaped/closed). Entries are kept (device=None is cheap) and each is closed under its own per-router lock. - _cleanup_idle() collects idle entries under _pool_lock and closes them OUTSIDE it. Previously a black-holed device.close() held the global lock and froze every router's borrows. Mirrors close_all(). - get_connection() closes a partially-opened Device if open() raises (it was never stored in entry[device], so it leaked), and evicts a session on RpcTimeoutError (which leaves connected=True but poisons reuse). - close_all(shutdown=True) uses non-blocking per-router locks so SIGTERM during a long command doesn't hang teardown. - execute_junos_command_batch dedupes router_names (pooling made repeated routers serialize on one lock instead of running in parallel). - Drop unused ConnectionPool.stats; type-hint idle_timeout as int | None; drop vestigial __enter__ mock; update CLAUDE.md (pool + idle timeout). Tests: add open-failure, RpcTimeoutError-eviction and idle-reaper regression tests (all fail on pre-round-3 code); update close_all tests for the keep-entry / acquire-release behavior.

Reuse SSH sessions across tool invocations instead of connect/teardown per call. Thread-safe ConnectionPool with per-router locking, stale detection, idle cleanup, and graceful shutdown. Benchmarked: 6.2x sequential, 20.7x parallel (5 routers, 35 ops).

Wrap int() parsing of JMCP_POOL_IDLE_TIMEOUT in try/except ValueError, log a warning, and fall back to the 300s default — mirrors the pattern in get_timeout_with_fallback. Addresses review feedback on Juniper#27.

Addresses the event-loop-blocking concern raised on this PR. Four handlers patched to dispatch the sync ConnectionPool.get_connection() call through anyio.to_thread.run_sync — same pattern already used in handle_execute_junos_command_batch: - handle_execute_junos_command (jmcp.py:1163) - handle_gather_device_facts (jmcp.py:1810, _gather_facts_sync) - handle_execute_junos_pfe_command (jmcp.py:1963) - handle_load_and_commit_config (jmcp.py:2023, _load_and_commit_sync) handle_render_and_apply_j2_template (line 1632) intentionally not included — has interleaved `await context.info(...)` inside the connection block; deferred to a separate PR. Validated against a 5-router vMX lab. Detail in the PR comment.

handle_reload_devices reused close_all(), which set _running=False and permanently stopped the idle-cleanup thread for the rest of the process, and ran blocking device.close() calls directly on the async event loop. Add a 'shutdown' flag to close_all(): reload now calls close_all(shutdown=False) to drop stale connections while keeping the pool (and its cleanup thread) running, dispatched via anyio.to_thread.run_sync so the event loop is not blocked. Add ConnectionPool lifecycle tests covering the reload vs shutdown paths.

The connection pool is a module-level singleton, so _run_junos_pfe_command now reuses cached connections across test cases: a device mocked in one test leaked into the next, and one test still mocked the obsolete 'with Device()' __enter__ path the pool no longer uses. Reset the pool in setUp/tearDown and simulate the generic failure on the RPC call itself instead of __enter__.

Addresses code-review Juniper#1 and Juniper#2 (both pooling regressions vs the old fresh-connection-per-call behavior). Juniper#1 close_all() held only _pool_lock and closed every cached Device without taking entry["lock"], so a reload or SIGTERM during an in-flight RPC could tear down the transport mid-operation. Snapshot+clear entries under the global lock, then close each under its own per-router lock (waits for in-flight ops; stops blocking unrelated routers). Juniper#2 In _load_and_commit_sync, when commit() fails and the rollback()/ unlock() in the except-block also fail (swallowed), the still-config- locked session went back into the pool, so every later commit failed with "database is locked". Best-effort close the device so the pool evicts it via the not-device.connected check. Adds ConnectionPool teardown and lock-leak regression tests.

Addresses code-review Juniper#3. These two handlers called the now-pool-backed _run_junos_cli_command directly on the event loop, so a contended per-router lock wait would stall the whole server. Wrap them in anyio.to_thread.run_sync like the other four handlers.

Addresses the second round of code review on PR Juniper#27. - Config-lock leak (completeness): last round's close-on-failed-cleanup only ran for an allowlist of exception types, but unlock() raises UnlockError (a direct RpcError subclass) and rollback() can raise a bare RpcError -- neither in the list -- so the realistic failure mode still returned a config-locked session to the pool. Widen the inner except to 'except Exception:' before closing the transport; classifying the failure is the outer handler's job. - Idle-cleanup gap: get_connection only refreshed last_used after a successful op, so a connection whose operation raised while still connected kept last_used=0.0 and the idle-cleanup loop (which requires last_used > 0) never reaped it. Refresh last_used in finally. - reload_devices race: close_all() ran before the device-map swap, so a concurrent tool call could open a pooled connection with the OLD config during the reset and survive it. Swap devices first; entry creation is serialized with close_all()'s snapshot via the pool lock and config is read only after entry creation, so connections opened during the reset use the NEW config. Tests: lock-leak regression test now exercises the real UnlockError path and asserts the pool evicts the session (rebuilds Device on next borrow); adds a last_used-refreshed-after-failure test. Both fail on pre-fix code.

The original feat commit split this log.info() across three lines, but it fits on one and 'black --check --fast .' (the upstream CI gate in .github/workflows/ci.yml) flags it. Pre-existing formatting, surfaced by running their CI locally.

Concurrency issues our pool introduced (scoped to our changes): - close_all() no longer clears _connections: clearing detached an entry an in-flight get_connection() borrow already held, orphaning the connection it then opened. Entries are kept (device=None is cheap) and each is closed under its own per-router lock. - _cleanup_idle() collects idle entries under _pool_lock and closes them OUTSIDE it, so a black-holed device.close() can't freeze every router's borrow. Mirrors close_all(). - get_connection() closes a partially-opened Device if open() raises, and evicts a session on RpcTimeoutError (which leaves connected=True but poisons reuse). - close_all(shutdown=True) uses non-blocking per-router locks so SIGTERM during a long command doesn't hang teardown. - execute_junos_command_batch dedupes router_names (pooling made repeated routers serialize on one lock). - Drop unused ConnectionPool.stats; type-hint idle_timeout as int | None; drop vestigial mock. Tests: add open-failure, RpcTimeoutError-eviction and idle-reaper regression tests (all fail on pre-round-3 code).

- note device connections now flow through the thread-safe pool (connection_pool.get_connection) and document JMCP_POOL_IDLE_TIMEOUT - correct the tool count (6 -> 11) and list all tools - replace 'no tests / no linting' with how to run the unittest suite and Black (both enforced in CI) - fix 'adding new tools': handlers register in TOOL_HANDLERS and are routed by a single @app.call_tool() dispatcher, not per-tool decorators

h-network · 2026-06-21T13:53:51Z

@nileshsimaria
While rebasing I did a concurrency self-review of the pool and tightened a few things the first cut missed (all scoped to the pooling change):

Reload safety: reload_devices no longer stops the idle-cleanup thread, and its connection teardown runs off the event loop.
Teardown races: close_all and _cleanup_idle no longer hold the global pool lock during a blocking device.close() (a black-holed device could otherwise freeze all routers), and no longer orphan a connection that races a borrow.
Session hygiene: a session is now evicted on a failed commit+rollback (prevents a leaked config lock), on RpcTimeoutError (which leaves connected true), and on a failed Device.open().
Event loop: get_junos_config / junos_config_diff now dispatch their pooled call via anyio.to_thread.run_sync, matching the other handlers.
Tests/CI: restored test isolation the pool broke, added ConnectionPool regression tests; black --check and the unit tests pass.

Deferred (flagging explicitly): render_and_apply_j2_template still opens its own session rather than borrowing from the pool. Pooling it needs a small refactor (it interleaves awaits inside the connection block), so it's left for a follow-up. Consequence worth noting: render_and_apply_j2_template (configure exclusive) and load_and_commit_config (lock-configuration) on the same router are two independent sessions that both lock the configuration database, so they can contend for that lock rather than being serialized by the pool.

Ready for re-review.

- guidance doc: most handlers borrow from the pool and offload PyEZ calls via anyio.to_thread.run_sync, but render_and_apply_j2_template and the add_device connectivity check do not -- say so instead of implying all do. - get_connection: note that pooled calls share anyio's process-wide default thread limiter (40 tokens), the ceiling for large fleets / same-router fan-out. - handle_reload_devices: the guaranteed property is 'no old-config connection is leaked', not 'every concurrent borrow uses the new config'; reword to match, and note the caught KeyError if a router is dropped mid-reload. Comments/docs only; no behavior change.

h-network · 2026-07-21T20:31:40Z

@nileshsimaria Do you have time soon to take a look please?

nileshsimaria · 2026-07-29T01:51:51Z

Looks great. Thanks a lot @h-network for working on this.

… pool The last unpooled config path opened a fresh SSH session per router and ran every blocking PyEZ call (open/lock/load/diff/commit_check/commit/rollback) directly on the async event loop, stalling all other tool calls. - Per-router work now borrows from connection_pool.get_connection and runs via anyio.to_thread.run_sync; progress messages are collected in the worker and replayed through context.info/error after it returns (this interleaving was why PR #27 deferred pooling this handler). - Config-format detection depends only on the rendered config, so it runs once before the router loop instead of per router. - Lock-leak protection mirrors _load_and_commit_sync: ConfigLoadError/ CommitError/LockError mean Config.__exit__ already unlocked (a failed unlock raises UnlockError instead), so the session stays pooled; any other failure closes the transport so the pool evicts the possibly- still-locked session. - Effective RPC timeout is unchanged: the pool default (360s) matches what prepare_connection_params set on each fresh Device. Docs: document JMCP_POOL_IDLE_TIMEOUT in the README environment-variable section (default 300s, reaper runs every 60s, invalid values fall back with a warning); update CLAUDE.md now that this handler is pooled. Tests: pooled-commit session reuse, dry-run rollback + retention, and UnlockError-eviction regression (tests/test_render_j2_pool.py). Full suite passes (53 tests) and black --check is clean. Validated end-to-end against a 6-router vMX lab: warm batch across all six routers took 0.25s vs 2.03s cold (~8x from session reuse); j2 dry-run (commit-check + auto-rollback) and a real commit + revert both succeeded through the pool, with the session staying usable after every exclusive-lock cycle.

nileshsimaria reviewed Apr 2, 2026

View reviewed changes

Comment thread jmcp.py Outdated

shigechika mentioned this pull request Apr 16, 2026

feat: per-host SSH/NETCONF connection pool with idle timeout shigechika/junos-mcp#5

Closed

6 tasks

H-Network added 11 commits June 21, 2026 13:33

fix: handle invalid JMCP_POOL_IDLE_TIMEOUT gracefully

acb820e

Wrap int() parsing of JMCP_POOL_IDLE_TIMEOUT in try/except ValueError, log a warning, and fall back to the 300s default — mirrors the pattern in get_timeout_with_fallback. Addresses review feedback on Juniper#27.

h-network force-pushed the connection-pooling branch from 41b29c5 to 60fe1c7 Compare June 21, 2026 13:34

h-network mentioned this pull request Jul 1, 2026

Connection pooling: persist/reuse ConnectHandler across tool calls ktbyers/netmiko_mcp#17

Open

nileshsimaria merged commit 4fec562 into Juniper:main Jul 29, 2026

Navigation Menu

Uh oh!

feat: thread-safe SSH/NETCONF connection pooling - #27

feat: thread-safe SSH/NETCONF connection pooling#27
nileshsimaria merged 12 commits into
Juniper:mainfrom
h-network:connection-pooling

h-network commented Apr 2, 2026

Uh oh!

Uh oh!

nileshsimaria commented May 1, 2026

Uh oh!

nileshsimaria commented May 1, 2026

Uh oh!

h-network commented Jun 21, 2026 •

edited

Loading

Uh oh!

h-network commented Jul 21, 2026

Uh oh!

nileshsimaria commented Jul 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

h-network commented Apr 2, 2026

Summary

Performance

Uh oh!

Uh oh!

nileshsimaria commented May 1, 2026

Uh oh!

nileshsimaria commented May 1, 2026

Uh oh!

h-network commented Jun 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

h-network commented Jul 21, 2026

Uh oh!

nileshsimaria commented Jul 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

h-network commented Jun 21, 2026 •

edited

Loading