Coverage for bzfs_main/util/connection.py: 99%
339 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 18:44 +0000
1# Copyright 2024 Wolfgang Hoschek AT mac DOT com
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14#
15"""Efficient thread-safe SSH command client; See run_ssh_command() and refresh_ssh_connection_if_necessary() and class
16ConnectionPool and class ConnectionLease.
18Can be configured to reuse multiplexed SSH connections for low latency, even on fresh process startup, for example leading to
19ballpark 3-5ms total time for running `/bin/echo hello` end-to-end over SSH on LAN, which requires two (sequential) network
20round trips (one for CHANNEL_OPEN, plus a subsequent one for CHANNEL_REQUEST).
21Has zero dependencies beyond the standard OpenSSH client CLI (`ssh`); also works with `hpnssh`. The latter uses larger TCP
22window sizes for best throughput over high speed long distance networks, aka paths with large bandwidth-delay product. Also
23see https://youtu.be/fcHXOgl3dis?t=473 and https://gist.github.com/rapier1/325de17bbb85f1ce663ccb866ce22639
25Example usage:
27import logging
28from subprocess import DEVNULL, PIPE
29from bzfs_main.util.connection import ConnectionPool, create_simple_minijob, create_simple_miniremote
30from bzfs_main.util.retry import Retry, RetryPolicy, call_with_retries
32log = logging.getLogger(__name__)
33remote = create_simple_miniremote(log=log, ssh_user_host="alice@127.0.0.1")
34connection_pool = ConnectionPool(remote, connpool_name="example")
35try:
36 job = create_simple_minijob()
37 retry_policy = RetryPolicy(
38 max_retries=10,
39 min_sleep_secs=0,
40 initial_max_sleep_secs=0.125,
41 max_sleep_secs=10,
42 max_elapsed_secs=60,
43 )
45 def run_cmd(retry: Retry) -> str:
46 with connection_pool.connection() as conn:
47 stdout: str = conn.run_ssh_command(
48 cmd=["echo", "hello"], job=job, check=True, stdin=DEVNULL, stdout=PIPE, stderr=PIPE, text=True
49 ).stdout
50 return stdout
52 stdout = call_with_retries(fn=run_cmd, policy=retry_policy, log=log)
53 print(f"stdout: {stdout}")
54finally:
55 connection_pool.shutdown()
56"""
58from __future__ import (
59 annotations,
60)
61import contextlib
62import copy
63import functools
64import logging
65import os
66import shlex
67import socket
68import stat
69import subprocess
70import threading
71import time
72from collections.abc import (
73 Generator,
74 Sequence,
75)
76from dataclasses import (
77 dataclass,
78)
79from subprocess import (
80 DEVNULL,
81 PIPE,
82)
83from typing import (
84 Any,
85 Final,
86 Protocol,
87 final,
88 runtime_checkable,
89)
91from bzfs_main.util.connection_lease import (
92 ConnectionLease,
93 ConnectionLeaseManager,
94)
95from bzfs_main.util.retry import (
96 RetryableError,
97)
98from bzfs_main.util.utils import (
99 LOG_TRACE,
100 SHELL_CHARS_AND_SLASH,
101 Subprocesses,
102 UpdatablePriorityQueue,
103 die,
104 get_home_directory,
105 list_formatter,
106 sha256_urlsafe_base64,
107 stderr_to_str,
108 subprocess_run,
109)
111# constants:
112SHARED: Final[str] = "shared"
113DEDICATED: Final[str] = "dedicated"
116#############################################################################
117@runtime_checkable
118class MiniJob(Protocol):
119 """Minimal Job interface required by the connections module; for loose coupling."""
121 @property
122 def timeout_nanos(self) -> int | None:
123 """Timeout timestamp aka instant in time, in monotonic nanoseconds, or None when no timeout applies."""
125 @property
126 def timeout_duration_nanos(self) -> int | None:
127 """Timeout duration (not a timestamp), in nanoseconds; used only for logging."""
129 @property
130 def subprocesses(self) -> Subprocesses:
131 """Subprocess manager used to run and track connection commands."""
134#############################################################################
135@runtime_checkable
136class MiniParams(Protocol):
137 """Minimal Params interface used by the connections module; for loose coupling."""
139 @property
140 def log(self) -> logging.Logger:
141 """Logger used for connection diagnostics."""
143 @property
144 def ssh_program(self) -> str:
145 """Name or path of the SSH executable; 'hpnssh' is also valid."""
148#############################################################################
149@runtime_checkable
150class MiniRemote(Protocol):
151 """Minimal Remote interface used by the connections module; for loose coupling."""
153 @property
154 def params(self) -> MiniParams:
155 """Settings read by connection operations."""
157 @property
158 def location(self) -> str:
159 """Endpoint role: 'src' or 'dst'."""
161 @property
162 def ssh_user_host(self) -> str:
163 """SSH destination, or the empty string to indicate local mode without SSH."""
165 @property
166 def ssh_extra_opts(self) -> tuple[str, ...]:
167 """Additional SSH command-line options used by connections to this endpoint."""
169 @property
170 def reuse_ssh_connection(self) -> bool:
171 """Whether connections should reuse SSH multiplexing masters."""
173 @property
174 def ssh_control_persist_secs(self) -> int:
175 """Idle lifetime (in seconds) requested for SSH multiplexing masters."""
177 @property
178 def ssh_control_persist_margin_secs(self) -> int:
179 """Safety margin (in seconds) for refreshing a master before its idle timeout."""
181 @property
182 def ssh_exit_on_shutdown(self) -> bool:
183 """Whether shutdown should close this endpoint's SSH multiplexing masters."""
185 @property
186 def ssh_socket_dir(self) -> str:
187 """Directory containing SSH multiplexing sockets for this endpoint."""
189 def is_ssh_available(self) -> bool:
190 """Return True if the ssh client program required for this remote is available on the local host."""
192 def local_ssh_command(self, socket_file: str | None) -> tuple[tuple[str, ...], str | None]:
193 """Returns the ssh CLI command to run locally in order to talk to the remote host; This excludes the (trailing)
194 command to run on the remote host, which will be appended later; also returns the effective ControlPath used by the
195 ssh CLI command, or ``None`` when SSH multiplexing is not active."""
197 def cache_namespace(self) -> str:
198 """Returns cache namespace string which is a stable, unique directory component for caches that distinguishes
199 endpoints by username+host+port+ssh_config_file where applicable, and uses '-' when no user/host is present (local
200 mode)."""
203#############################################################################
204def create_simple_miniremote(
205 *,
206 log: logging.Logger,
207 ssh_user_host: str = "", # option passed to `ssh` CLI; empty string indicates local mode
208 ssh_port: int | None = None, # option passed to `ssh -p` CLI
209 ssh_extra_opts: list[str] | None = None, # optional args passed to `ssh` CLI
210 ssh_verbose: bool = False, # option passed to `ssh -v` CLI
211 ssh_config_file: str = "", # option passed to `ssh -F` CLI; path to ssh_config(5) file; e.g /path/to/homedir/.ssh/config
212 ssh_cipher: str = "^aes256-gcm@openssh.com", # option passed to `ssh -c` CLI
213 ssh_connect_timeout_secs: int | None = None, # option passed to `ssh -oConnectTimeout=N`; default is system TCP timeout
214 ssh_program: str = "ssh", # name or path of CLI executable; "hpnssh" is also valid
215 reuse_ssh_connection: bool = True,
216 ssh_control_persist_secs: int = 600,
217 ssh_control_persist_margin_secs: int = 2,
218 ssh_socket_dir: str = os.path.join(get_home_directory(), ".ssh", "bzfs"),
219 location: str = "dst",
220) -> MiniRemote:
221 """Factory that returns a simple implementation of the MiniRemote interface."""
223 @dataclass(frozen=True) # aka immutable
224 @final
225 class SimpleMiniParams:
226 log: logging.Logger
227 ssh_program: str
229 @dataclass(frozen=True) # aka immutable
230 @final
231 class SimpleMiniRemote:
232 params: MiniParams
233 location: str # "src" or "dst"
234 ssh_user_host: str
235 ssh_extra_opts: tuple[str, ...]
236 reuse_ssh_connection: bool
237 ssh_control_persist_secs: int
238 ssh_control_persist_margin_secs: int
239 ssh_exit_on_shutdown: bool
240 ssh_socket_dir: str
241 ssh_port: int | None
242 ssh_config_file: str
244 def is_ssh_available(self) -> bool:
245 return True
247 def local_ssh_command(self, socket_file: str | None) -> tuple[tuple[str, ...], str | None]:
248 if not self.ssh_user_host:
249 return (), None # local mode
250 ssh_cmd: list[str] = [self.params.ssh_program]
251 ssh_cmd += self.ssh_extra_opts
252 socket_path: str | None = None
253 if self.reuse_ssh_connection and socket_file:
254 ssh_cmd.append("-S")
255 ssh_cmd.append(socket_file)
256 socket_path = socket_file
257 ssh_cmd.append(self.ssh_user_host)
258 return tuple(ssh_cmd), socket_path
260 def cache_namespace(self) -> str:
261 return self._cache_namespace
263 @functools.cached_property
264 def _cache_namespace(self) -> str:
265 if not self.ssh_user_host:
266 return "-" # local mode
267 ssh_config_file_hash: str = get_ssh_config_file_hash(self, self.ssh_config_file)
268 return f"{self.ssh_user_host}#{self.ssh_port or ''}#{ssh_config_file_hash}"
270 def validate_userhost(userhost: str) -> None:
271 invalid_chars: str = SHELL_CHARS_AND_SLASH
272 uh: str = userhost.replace("@", "", 1)
273 if (not uh) or userhost.startswith("-") or ".." in userhost or any(c.isspace() or c in invalid_chars for c in uh):
274 raise ValueError(f"Invalid [user@]host: '{userhost}'")
276 if log is None:
277 raise ValueError("log must not be None")
278 if not ssh_program:
279 raise ValueError("ssh_program must be a non-empty string")
280 if location not in ("src", "dst"):
281 raise ValueError("location must be 'src' or 'dst'")
282 if ssh_user_host:
283 validate_userhost(ssh_user_host)
284 if ssh_control_persist_secs < 1:
285 raise ValueError("ssh_control_persist_secs must be >= 1")
286 params: MiniParams = SimpleMiniParams(log=log, ssh_program=ssh_program)
288 ssh_extra_opts = ( # disable interactive password prompts and X11 forwarding and pseudo-terminal allocation
289 ["-oBatchMode=yes", "-oServerAliveInterval=0", "-x", "-T"] if ssh_extra_opts is None else list(ssh_extra_opts)
290 )
291 ssh_extra_opts += ["-v"] if ssh_verbose else []
292 ssh_extra_opts += ["-F", ssh_config_file] if ssh_config_file else []
293 ssh_extra_opts += ["-c", ssh_cipher] if ssh_cipher else []
294 ssh_extra_opts += ["-p", str(ssh_port)] if ssh_port is not None else []
295 ssh_extra_opts += [] if ssh_connect_timeout_secs is None else [f"-oConnectTimeout={max(0, ssh_connect_timeout_secs)}s"]
296 return SimpleMiniRemote(
297 params=params,
298 location=location,
299 ssh_user_host=ssh_user_host,
300 ssh_extra_opts=tuple(ssh_extra_opts),
301 reuse_ssh_connection=reuse_ssh_connection,
302 ssh_control_persist_secs=ssh_control_persist_secs,
303 ssh_control_persist_margin_secs=ssh_control_persist_margin_secs,
304 ssh_exit_on_shutdown=False,
305 ssh_socket_dir=ssh_socket_dir,
306 ssh_port=ssh_port,
307 ssh_config_file=ssh_config_file,
308 )
311def get_ssh_config_file_hash(remote: MiniRemote, ssh_config_file: str | None) -> str:
312 """Hashes the explicit ssh config path and effective SSH settings, including nested and wildcard `includes`, using
313 `ssh -G` as fingerprint."""
314 p = remote.params
315 ssh_cmd, _ = remote.local_ssh_command("none")
316 if not ssh_cmd:
317 return "" # local mode
318 cmd: list[str] = [ssh_cmd[0], "-G", *ssh_cmd[1:], "exit"]
319 try:
320 config: str = subprocess_run(cmd, stdin=DEVNULL, stdout=PIPE, stderr=PIPE, text=True, check=True, log=p.log).stdout
321 except subprocess.CalledProcessError as e:
322 die(f"Cannot evaluate SSH configuration: {e}\n{stderr_to_str(e.stderr or '')}".rstrip())
323 except OSError as e:
324 die(f"Cannot evaluate SSH configuration: {e}")
325 path: str = os.path.abspath(ssh_config_file) if ssh_config_file else ""
326 return sha256_urlsafe_base64(path + "\x00" + config, padding=False) # "\x00" ASCII NUL char is unambiguous separator
329def create_simple_minijob(
330 *, timeout_duration_secs: float | None = None, subprocesses: Subprocesses | None = None
331) -> MiniJob:
332 """Factory that returns a simple implementation of the MiniJob interface."""
334 @dataclass(frozen=True) # aka immutable
335 @final
336 class SimpleMiniJob:
337 timeout_nanos: int | None # timestamp aka instant in time
338 timeout_duration_nanos: int | None # duration (not a timestamp); for logging only
339 subprocesses: Subprocesses
341 t_duration_nanos: int | None = None if timeout_duration_secs is None else int(timeout_duration_secs * 1_000_000_000)
342 timeout_nanos: int | None = None if t_duration_nanos is None else time.monotonic_ns() + t_duration_nanos
343 subprocesses = Subprocesses() if subprocesses is None else subprocesses
344 return SimpleMiniJob(timeout_nanos=timeout_nanos, timeout_duration_nanos=t_duration_nanos, subprocesses=subprocesses)
347#############################################################################
348def timeout(job: MiniJob) -> float | None:
349 """Raises TimeoutExpired if necessary, else returns the number of seconds left until timeout is to occur."""
350 timeout_nanos: int | None = job.timeout_nanos
351 if timeout_nanos is None:
352 return None # never raise a timeout
353 assert job.timeout_duration_nanos is not None
354 delta_nanos: int = timeout_nanos - time.monotonic_ns()
355 if delta_nanos <= 0:
356 raise subprocess.TimeoutExpired("_timeout", timeout=job.timeout_duration_nanos / 1_000_000_000)
357 return delta_nanos / 1_000_000_000 # seconds
360def squote(remote: MiniRemote, arg: str) -> str:
361 """Quotes an argument only when running remotely over ssh."""
362 assert arg is not None
363 return shlex.quote(arg) if remote.ssh_user_host else arg
366def dquote(arg: str) -> str:
367 """Shell-escapes backslash and double quotes and dollar and backticks, then surrounds with double quotes; For an example
368 how to safely construct and quote complex shell pipeline commands for use over SSH, see
369 replication.py:_prepare_zfs_send_receive()"""
370 arg = arg.replace("\\", "\\\\").replace('"', '\\"').replace("$", "\\$").replace("`", "\\`")
371 return '"' + arg + '"'
374#############################################################################
375@dataclass(order=True, repr=False)
376@final
377class Connection:
378 """Represents the ability to multiplex N=capacity concurrent SSH sessions over the same TCP connection."""
380 _free: int # sort order evens out the number of concurrent sessions among the TCP connections
381 _last_modified: int # LIFO: tiebreaker favors latest returned conn as that's most alive and hot; also ensures no dupes
383 def __init__(
384 self,
385 remote: MiniRemote,
386 max_concurrent_ssh_sessions_per_tcp_connection: int,
387 *,
388 lease: ConnectionLease | None = None,
389 ) -> None:
390 assert max_concurrent_ssh_sessions_per_tcp_connection > 0
391 self._remote: Final[MiniRemote] = remote
392 self._capacity: Final[int] = max_concurrent_ssh_sessions_per_tcp_connection
393 self._free: int = max_concurrent_ssh_sessions_per_tcp_connection
394 self._last_modified: int = 0 # monotonically increasing
395 self._last_refresh_time: int = 1 - (1 << 150) # negative infinity for all practical purposes
396 self._lock: Final[threading.Lock] = threading.Lock()
397 self._reuse_ssh_connection: Final[bool] = remote.reuse_ssh_connection
398 self._connection_lease: Final[ConnectionLease | None] = lease
399 ssh_cmd, ssh_socket_path = remote.local_ssh_command(
400 None if self._connection_lease is None else self._connection_lease.socket_path
401 )
402 self._ssh_socket_path: Final[str | None] = ssh_socket_path
403 self.ssh_cmd: Final[tuple[str, ...]] = ssh_cmd
404 self.ssh_cmd_quoted: Final[tuple[str, ...]] = tuple(shlex.quote(item) for item in self.ssh_cmd)
406 def __repr__(self) -> str:
407 return str({"free": self._free})
409 def run_ssh_command(
410 self,
411 cmd: Sequence[str],
412 *,
413 job: MiniJob,
414 loglevel: int = logging.INFO,
415 is_dry: bool = False,
416 **kwargs: Any, # optional low-level keyword args to be forwarded to subprocess.run()
417 ) -> subprocess.CompletedProcess:
418 """Runs the given CLI cmd via ssh on the given remote, and returns CompletedProcess including stdout and stderr.
420 The full command is the concatenation of both the command to run on the localhost in order to talk to the remote host
421 (``remote.local_ssh_command(...)[0]``) and the command to run on the given remote host (``cmd``).
423 Note: When executing on a remote host (remote.ssh_user_host is set), cmd arguments are pre-quoted with shlex.quote to
424 safely traverse the ssh "remote shell" boundary, as ssh concatenates argv into a single remote shell string. In local
425 mode (no remote.ssh_user_host) argv is executed directly without an intermediate shell.
426 """
427 cmd = tuple(cmd)
428 if not cmd:
429 raise ValueError("run_ssh_command requires a non-empty cmd sequence")
430 log: logging.Logger = self._remote.params.log
431 quoted_cmd: tuple[str, ...] = tuple(shlex.quote(arg) for arg in cmd)
432 if self._remote.ssh_user_host:
433 self.refresh_ssh_connection_if_necessary(job)
434 cmd = quoted_cmd
435 cmd = self.ssh_cmd + cmd
436 msg: str = "Would execute: %s" if is_dry else "Executing: %s"
437 log.log(loglevel, msg, list_formatter(self.ssh_cmd_quoted + quoted_cmd, lstrip=True))
438 if is_dry:
439 return subprocess.CompletedProcess(cmd, returncode=0, stdout=None, stderr=None)
440 else:
441 sp: Subprocesses = job.subprocesses
442 return sp.subprocess_run(cmd, timeout=timeout(job), log=log, **kwargs)
444 def refresh_ssh_connection_if_necessary(self, job: MiniJob) -> None:
445 """Maintain or create an ssh master connection for low latency reuse."""
446 remote: MiniRemote = self._remote
447 p: MiniParams = remote.params
448 log: logging.Logger = p.log
449 if not remote.ssh_user_host:
450 return # we're in local mode; no ssh required
451 if not remote.is_ssh_available():
452 die(f"{p.ssh_program} CLI is not available to talk to remote host. Install {p.ssh_program} first!")
453 if not remote.reuse_ssh_connection:
454 return
456 # Performance: reuse ssh connection for low latency startup of frequent ssh invocations via the 'ssh -S' and
457 # 'ssh -S -M -oControlPersist=90s' options. See https://en.wikibooks.org/wiki/OpenSSH/Cookbook/Multiplexing
458 # and https://chessman7.substack.com/p/how-ssh-multiplexing-reuses-master
459 control_limit_nanos: int = (remote.ssh_control_persist_secs - remote.ssh_control_persist_margin_secs) * 1_000_000_000
460 socket_path: str | None = self._ssh_socket_path
461 with self._lock:
462 if time.monotonic_ns() < self._last_refresh_time + control_limit_nanos:
463 if socket_path is None or self._is_ssh_control_socket_usable(socket_path):
464 return # ssh master is alive, reuse its TCP connection (this is the common case and the ultra-fast path)
465 ssh_sock_cmd: list[str] = list(self.ssh_cmd[0:-1]) # omit trailing ssh_user_host
466 ssh_sock_cmd += ["-O", "check", remote.ssh_user_host]
467 # extend lifetime of ssh master by $ssh_control_persist_secs via `ssh -O check` if master is still running.
468 # `ssh -S /path/to/socket -O check` doesn't talk over the network, hence is still a low latency fast path.
469 sp: Subprocesses = job.subprocesses
470 t: float | None = timeout(job)
471 if sp.subprocess_run(ssh_sock_cmd, stdin=DEVNULL, stdout=PIPE, stderr=PIPE, timeout=t, log=log).returncode == 0:
472 log.log(LOG_TRACE, "ssh connection is alive: %s", list_formatter(ssh_sock_cmd))
473 else: # ssh master is not alive; start a new master:
474 log.log(LOG_TRACE, "ssh connection is not yet alive: %s", list_formatter(ssh_sock_cmd))
475 if socket_path is not None and not self._is_ssh_control_socket_usable(socket_path):
476 with contextlib.suppress(OSError):
477 os.unlink(socket_path) # if present, remove stale ssh control socket path before master restart
478 ssh_control_persist_secs: int = max(1, remote.ssh_control_persist_secs)
479 if any(opt.startswith("-v") and all(char == "v" for char in opt[1:]) for opt in remote.ssh_extra_opts):
480 # Unfortunately, with `ssh -v` (debug mode), the ssh master won't background; instead it stays in the
481 # foreground and blocks until the ControlPersist timer expires (90 secs). To make progress earlier we ...
482 ssh_control_persist_secs = min(1, ssh_control_persist_secs) # tell ssh block as briefly as possible (1s)
483 ssh_sock_cmd = list(self.ssh_cmd[0:-1]) # omit trailing ssh_user_host
484 ssh_sock_cmd += ["-M", f"-oControlPersist={ssh_control_persist_secs}s", remote.ssh_user_host, "exit"]
485 log.log(LOG_TRACE, "Executing: %s", list_formatter(ssh_sock_cmd))
486 t = timeout(job)
487 try:
488 sp.subprocess_run(ssh_sock_cmd, stdin=DEVNULL, stdout=PIPE, stderr=PIPE, check=True, timeout=t, log=log)
489 except subprocess.CalledProcessError as e:
490 log.error("%s", stderr_to_str(e.stderr).rstrip())
491 raise RetryableError(
492 f"Cannot ssh into remote host via '{' '.join(ssh_sock_cmd)}'. Fix ssh configuration first, "
493 "considering diagnostic log file output from running with -v -v -v.",
494 display_msg="ssh connect",
495 ) from e
496 self._last_refresh_time = time.monotonic_ns()
497 if self._connection_lease is not None:
498 self._connection_lease.set_socket_mtime_to_now()
500 def _is_ssh_control_socket_usable(self, socket_path: str) -> bool:
501 """To improve ssh perf, check whether a control socket path is a live Unix-domain listener; this helps detect stale
502 socket files that still exist after master crashes.
504 _is_ssh_control_socket_usable() is ~300x faster than `ssh ... -O check`: ~5 microseconds vs ~1-2 milliseconds.
505 """
506 try:
507 st_mode: int = os.stat(socket_path, follow_symlinks=False).st_mode
508 if not stat.S_ISSOCK(st_mode):
509 return False
510 except OSError:
511 return False
512 try:
513 with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
514 return sock.connect_ex(socket_path) == 0
515 except OSError:
516 return False
518 def _increment_free(self, value: int) -> None:
519 """Adjusts the count of available SSH slots."""
520 self._free += value
521 assert self._free >= 0
522 assert self._free <= self._capacity
524 def _is_full(self) -> bool:
525 """Returns True if no more SSH sessions may be opened over this TCP connection."""
526 return self._free <= 0
528 def _update_last_modified(self, last_modified: int) -> None:
529 """Records when the connection was last used."""
530 self._last_modified = last_modified
532 def shutdown(self, msg_prefix: str) -> None:
533 """Closes the underlying SSH master connection and releases the corresponding connection lease."""
534 if self.ssh_cmd and self._reuse_ssh_connection:
535 if self._connection_lease is None:
536 ssh_sock_cmd: tuple[str, ...] = self.ssh_cmd[0:-1] + ("-O", "exit", self.ssh_cmd[-1])
537 log = self._remote.params.log
538 log.log(LOG_TRACE, f"Executing {msg_prefix}: %s", shlex.join(ssh_sock_cmd))
539 try:
540 proc: subprocess.CompletedProcess = subprocess.run(
541 ssh_sock_cmd, stdin=DEVNULL, stderr=PIPE, text=True, timeout=0.1, check=False
542 )
543 except subprocess.TimeoutExpired as e: # harmless as master auto-exits after ssh_control_persist_secs anyway
544 log.log(LOG_TRACE, "Harmless ssh master connection shutdown timeout: %s", e)
545 else:
546 if proc.returncode != 0: # harmless for the same reason
547 log.log(LOG_TRACE, "Harmless ssh master connection shutdown issue: %s", proc.stderr.rstrip())
548 else:
549 self._connection_lease.release()
552#############################################################################
553class ConnectionPool:
554 """Fetch a TCP connection for use in an SSH session, use it, finally return it back to the pool for future reuse;
555 Note that max_concurrent_ssh_sessions_per_tcp_connection must not be larger than the server-side sshd_config(5)
556 MaxSessions parameter (which defaults to 10, see https://manpages.ubuntu.com/manpages/man5/sshd_config.5.html)."""
558 def __init__(
559 self, remote: MiniRemote, connpool_name: str, max_concurrent_ssh_sessions_per_tcp_connection: int = 8
560 ) -> None:
561 assert max_concurrent_ssh_sessions_per_tcp_connection > 0
562 self._remote: Final[MiniRemote] = copy.copy(remote) # shallow copy for immutability (Remote is mutable)
563 self._capacity: Final[int] = max_concurrent_ssh_sessions_per_tcp_connection
564 self._connpool_name: Final[str] = connpool_name
565 self._priority_queue: Final[UpdatablePriorityQueue[Connection]] = UpdatablePriorityQueue(
566 reverse=True # sorted by #free slots and last_modified
567 )
568 self._last_modified: int = 0 # monotonically increasing sequence number
569 self._lock: Final[threading.Lock] = threading.Lock()
570 lease_mgr: ConnectionLeaseManager | None = None
571 if self._remote.ssh_user_host and self._remote.reuse_ssh_connection and not self._remote.ssh_exit_on_shutdown:
572 lease_mgr = ConnectionLeaseManager(
573 root_dir=self._remote.ssh_socket_dir,
574 namespace=f"{self._remote.location}#{self._remote.cache_namespace()}#{self._connpool_name}",
575 ssh_control_persist_secs=max(90 * 60, 2 * self._remote.ssh_control_persist_secs + 2),
576 log=self._remote.params.log,
577 )
578 self._lease_mgr: Final[ConnectionLeaseManager | None] = lease_mgr
580 @contextlib.contextmanager
581 def connection(self) -> Generator[Connection]:
582 """Context manager that yields a connection from the pool and automatically returns it on __exit__."""
583 conn: Connection = self.get_connection()
584 try:
585 yield conn
586 finally:
587 self.return_connection(conn)
589 def get_connection(self) -> Connection:
590 """Any Connection object returned on get_connection() also remains intentionally contained in the priority queue
591 while it is "checked out", and that identical Connection object is later, on return_connection(), temporarily removed
592 from the priority queue, updated with an incremented "free" slot count and then immediately reinserted into the
593 priority queue.
595 In effect, any Connection object remains intentionally contained in the priority queue at all times. This design
596 keeps ordering/fairness accurate while avoiding duplicate Connection instances.
597 """
598 with self._lock:
599 conn = self._priority_queue.pop() if len(self._priority_queue) > 0 else None
600 if conn is None or conn._is_full(): # noqa: SLF001 # pylint: disable=protected-access
601 if conn is not None:
602 self._priority_queue.push(conn)
603 conn = self._new_connection() # add a new connection
604 self._last_modified += 1
605 conn._update_last_modified(self._last_modified) # noqa: SLF001 # pylint: disable=protected-access
606 conn._increment_free(-1) # noqa: SLF001 # pylint: disable=protected-access
607 self._priority_queue.push(conn)
608 return conn
610 def _new_connection(self) -> Connection:
611 lease: ConnectionLease | None = None if self._lease_mgr is None else self._lease_mgr.acquire()
612 return Connection(self._remote, self._capacity, lease=lease)
614 def return_connection(self, conn: Connection) -> None:
615 """Returns the given connection to the pool and updates its priority."""
616 assert conn is not None
617 with self._lock:
618 # update priority = remove conn from queue, increment priority, finally reinsert updated conn into queue
619 if self._priority_queue.remove(conn): # conn is not contained only if ConnectionPool.shutdown() was called
620 conn._increment_free(1) # noqa: SLF001 # pylint: disable=protected-access
621 self._last_modified += 1
622 conn._update_last_modified(self._last_modified) # noqa: SLF001 # pylint: disable=protected-access
623 self._priority_queue.push(conn)
625 def shutdown(self, msg_prefix: str = "") -> None:
626 """Closes all SSH connections managed by this pool."""
627 with self._lock:
628 try:
629 if self._remote.reuse_ssh_connection:
630 msg_prefix = msg_prefix + "/" + self._connpool_name
631 for conn in self._priority_queue:
632 conn.shutdown(msg_prefix)
633 finally:
634 self._priority_queue.clear()
636 def __repr__(self) -> str:
637 with self._lock:
638 queue = self._priority_queue
639 return str({"capacity": self._capacity, "queue_len": len(queue), "queue": queue})
642#############################################################################
643@final
644class ConnectionPools:
645 """A bunch of named connection pools with various multiplexing capacities."""
647 def __init__(self, remote: MiniRemote, *, capacities: dict[str, int]) -> None:
648 """Creates one connection pool per name with the given capacities."""
649 self._pools: Final[dict[str, ConnectionPool]] = {
650 name: ConnectionPool(remote, name, capacity) for name, capacity in capacities.items()
651 }
653 def __repr__(self) -> str:
654 return str(self._pools)
656 def pool(self, name: str) -> ConnectionPool:
657 """Returns the pool associated with the given name."""
658 return self._pools[name]
660 def shutdown(self, msg_prefix: str = "") -> None:
661 """Shuts down every contained pool."""
662 for pool in self._pools.values():
663 pool.shutdown(msg_prefix)