Coverage for bzfs_main/parallel_batch_cmd.py: 96%
43 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"""Helpers for running CLI commands in (sequential or parallel) batches, without exceeding operating system limits.
17The batch size aka max_batch_items splits one CLI command into one or more CLI commands. The resulting commands are executed
18sequentially (via functions *_batched()), or in parallel across max_workers threads (via functions *_parallel()).
20The degree of parallelism (max_workers) is specified by the job (via --threads).
21Batch size is a trade-off between resource consumption, latency, bandwidth and throughput.
23Example:
24--------
26- max_batch_items=1 (seq or par):
27```
28zfs list -t snapshot d1
29zfs list -t snapshot d2
30zfs list -t snapshot d3
31zfs list -t snapshot d4
32zfs list -t snapshot d5
33```
35- max_batch_items=2 (seq or par):
36```
37zfs list -t snapshot d1 d2
38zfs list -t snapshot d3 d4
39zfs list -t snapshot d5
40```
42- max_batch_items=N (seq or par):
43```
44zfs list -t snapshot d1 d2 d3 d4 d5
45```
46"""
48from __future__ import (
49 annotations,
50)
51import shlex
52import sys
53from collections.abc import (
54 Iterable,
55 Iterator,
56 Sequence,
57)
58from typing import (
59 TYPE_CHECKING,
60 Callable,
61 TypeVar,
62)
64from bzfs_main.util.connection import (
65 SHARED,
66 ConnectionPool,
67 MiniRemote,
68)
69from bzfs_main.util.parallel_iterator import (
70 batch_cmd_iterator,
71 get_max_command_line_bytes,
72 parallel_iterator,
73)
74from bzfs_main.util.utils import (
75 LOG_TRACE,
76 drain,
77)
79if TYPE_CHECKING: # pragma: no cover - for type hints only
80 from bzfs_main.bzfs import (
81 Job,
82 )
84_T = TypeVar("_T")
87def run_ssh_cmd_batched(
88 job: Job,
89 r: MiniRemote,
90 cmd: list[str],
91 cmd_args: Iterable[str],
92 fn: Callable[[list[str]], object],
93 *,
94 max_batch_items: int = 2**29,
95 sep: str = " ",
96) -> None:
97 """Runs ssh command for each sequential batch of args, without creating a cmdline that's too big for the OS to handle."""
98 drain(itr_ssh_cmd_batched(job, r, cmd, cmd_args, fn, max_batch_items=max_batch_items, sep=sep))
101def itr_ssh_cmd_batched(
102 job: Job,
103 r: MiniRemote,
104 cmd: list[str],
105 cmd_args: Iterable[str],
106 fn: Callable[[list[str]], _T],
107 *,
108 max_batch_items: int = 2**29,
109 sep: str = " ",
110) -> Iterator[_T]:
111 """Runs fn(cmd_args) in sequential batches w/ cmd, without creating a cmdline that's too big for the OS to handle."""
112 max_bytes: int = _max_batch_bytes(job, r, cmd, sep)
113 return batch_cmd_iterator(
114 cmd_args=cmd_args,
115 fn=fn,
116 max_batch_items=max_batch_items,
117 max_batch_bytes=max_bytes,
118 sep=sep,
119 quote_args=bool(r.ssh_user_host),
120 )
123def run_ssh_cmd_parallel(
124 job: Job,
125 r: MiniRemote,
126 cmd_args_list: Iterable[tuple[list[str], Iterable[str]]],
127 fn: Callable[[list[str], list[str]], object],
128 *,
129 max_batch_items: int = 2**29,
130) -> None:
131 """Runs multiple ssh commands in parallel, batching each set of args."""
132 drain(itr_ssh_cmd_parallel(job, r, cmd_args_list, fn=fn, max_batch_items=max_batch_items, ordered=False))
135def itr_ssh_cmd_parallel(
136 job: Job,
137 r: MiniRemote,
138 cmd_args_list: Iterable[tuple[list[str], Iterable[str]]],
139 fn: Callable[[list[str], list[str]], _T],
140 *,
141 max_batch_items: int = 2**29,
142 ordered: bool = True,
143) -> Iterator[_T]:
144 """Streams results from multiple parallel (batched) SSH commands.
146 When ordered=True, preserves the order of the batches as provided by cmd_args_list (i.e. yields results in the same order
147 as the input), not in completion order. When ordered=False, yields results as they complete for minimum latency.
148 """
149 return parallel_iterator(
150 iterator_builder=lambda executr: (
151 itr_ssh_cmd_batched( # advancing the Generator submits the next task and yields the corresponding Future
152 job, r, cmd, cmd_args, lambda batch, cmd=cmd: executr.submit(fn, cmd, batch), max_batch_items=max_batch_items # type: ignore[misc]
153 )
154 for cmd, cmd_args in cmd_args_list # lazy on-demand Python Generator
155 ),
156 max_workers=job.max_workers[r.location],
157 ordered=ordered,
158 is_terminated=job.termination_event.is_set,
159 )
162def zfs_list_snapshots_in_parallel(
163 job: Job,
164 r: MiniRemote,
165 cmd: list[str],
166 datasets: list[str],
167 *,
168 ordered: bool = True,
169) -> Iterator[list[str]]:
170 """Runs 'zfs list -t snapshot' on multiple datasets at the same time.
172 Implemented with a time and space efficient streaming algorithm; easily scales to millions of datasets and any number of
173 snapshots. Attempts to use at least 4 datasets per remote cmd mini batch to reflect increased communication latency.
174 """
175 max_workers: int = job.max_workers[r.location]
176 max_batch_items: int = min(
177 job.max_datasets_per_minibatch_on_list_snaps[r.location],
178 max(
179 len(datasets) // (max_workers * 8 if max_workers > 1 else 1),
180 4 if r.ssh_user_host else 1,
181 ),
182 )
183 return itr_ssh_cmd_parallel(
184 job,
185 r,
186 [(cmd, datasets)],
187 fn=lambda cmd, batch: (job.try_ssh_command_with_retries(r, LOG_TRACE, cmd=cmd + batch) or "").splitlines(),
188 max_batch_items=max_batch_items,
189 ordered=ordered,
190 )
193def _max_batch_bytes(job: Job, r: MiniRemote, cmd: Sequence[str], sep: str) -> int:
194 """Avoids creating a cmdline that's too big for the OS or SSH transport to handle.
196 The calculation subtracts 'header_bytes', which accounts for the full SSH invocation (including control socket/options)
197 plus the fixed subcommand prefix, so that the remaining budget is reserved exclusively for the batched arguments.
198 """
199 assert isinstance(sep, str)
200 max_bytes: int = min(_get_max_command_line_bytes(job, "local"), _get_max_command_line_bytes(job, r.location))
201 if r.ssh_user_host or sep != " ": # e.g. 'zfs destroy foo@s1,s2,...,sN', or sshd executes `sh -c <large_blurb>`
202 # Max size of a single argument is 128KB on Linux - https://lists.gnu.org/archive/html/bug-bash/2020-09/msg00095.html
203 max_bytes = min(max_bytes, 128 * 1024 - 1)
204 if r.ssh_user_host and job.params.ssh_program == "hpnssh": 204 ↛ 207line 204 didn't jump to line 207 because the condition on line 204 was never true
205 # HPN-SSH caps authenticated packets at 33 KiB; leave room for SSH packet framing.
206 # https://github.com/rapier1/hpn-ssh/blob/hpn-18.9.0/packet.c#L2466-L2478
207 max_bytes = min(max_bytes, 32 * 1024 - 1)
208 conn_pool: ConnectionPool = job.params.connection_pools[r.location].pool(SHARED)
209 with conn_pool.connection() as conn:
210 ssh_cmd: tuple[str, ...] = conn.ssh_cmd
211 if r.ssh_user_host:
212 cmd = tuple(shlex.quote(arg) for arg in cmd)
213 header_bytes: int = len(" ".join(ssh_cmd + tuple(cmd)).encode(sys.getfilesystemencoding()))
214 return max_bytes - header_bytes
217def _get_max_command_line_bytes(job: Job, location: str, os_name: str | None = None) -> int:
218 """Remote flavor of os.sysconf("SC_ARG_MAX") - size(os.environb) - safety margin"""
219 os_name = os_name if os_name else job.params.available_programs[location].get("os")
220 os_name = os_name if os_name else "n/a"
221 max_bytes = get_max_command_line_bytes(os_name)
222 if job.max_command_line_bytes is not None:
223 return job.max_command_line_bytes # for testing only
224 else:
225 return max_bytes