Coverage for bzfs_main/util/parallel_iterator.py: 100%
80 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"""Parallel execution utilities for I/O-bound operations, with configurable result ordering."""
17from __future__ import (
18 annotations,
19)
20import concurrent
21import itertools
22import os
23import shlex
24import sys
25from collections import (
26 deque,
27)
28from collections.abc import (
29 Iterable,
30 Iterator,
31)
32from concurrent.futures import (
33 FIRST_COMPLETED,
34 Executor,
35 Future,
36)
37from concurrent.futures.thread import (
38 ThreadPoolExecutor,
39)
40from typing import (
41 Callable,
42 Final,
43 TypeVar,
44)
46from bzfs_main.util.utils import (
47 SynchronousExecutor,
48)
50_T = TypeVar("_T")
53def parallel_iterator(
54 iterator_builder: Callable[[Executor], Iterable[Iterable[Future[_T]]]],
55 *,
56 max_workers: int = os.cpu_count() or 1,
57 ordered: bool = True,
58 is_terminated: Callable[[], bool] = lambda: False, # optional predicate to request early async termination
59) -> Iterator[_T]:
60 """Executes multiple iterators in parallel/concurrently, with explicit backpressure and configurable result ordering;
61 avoids pre-submitting the entire workload.
63 This function provides efficient parallel execution of iterator-based tasks using a shared thread pool, with precise
64 control over result delivery ordering and concurrency management through a bounded buffer (a sliding window of at most
65 ``max_workers`` in-flight futures).
67 Purpose:
68 --------
69 Enables parallel/concurrent execution of multiple iterator streams while providing either sequential (ordered) or
70 performance-optimized (unordered) result delivery. Primarily designed for I/O-bound operations like ZFS/SSH command
71 execution where parallelism significantly improves throughput.
73 Assumptions:
74 ------------
75 - The builder must submit tasks to the provided executor (e.g., via ``executor.submit(...)``) and yield an Iterator
76 over the corresponding Future[T] objects.
77 - Tasks are primarily I/O-bound and benefit from parallel execution
78 - Caller can handle potential exceptions propagated from ``Future.result()``
79 - The builder properly scopes any resources to the lifecycle of the provided ThreadPoolExecutor
81 Design Rationale:
82 -----------------
83 The design optimizes for bzfs's primary use case: executing similar ZFS/SSH commands across multiple remote systems
84 where I/O or ZFS overhead dominates and parallel execution provides substantial performance improvements over
85 sequential processing. The implementation addresses several key design challenges:
87 - Bounded Buffer: Maintains at most ``max_workers`` futures in flight, preventing resource exhaustion and bounding memory
88 consumption while maximizing thread utilization. New tasks are submitted as completed ones are consumed. This is
89 crucial when processing large numbers of datasets typical in ZFS operation.
91 - Ordered vs Unordered Execution:
93 - Ordered mode uses a FIFO queue (``deque.popleft()``) ensuring sequential delivery that preserves the order in
94 which the builder's iterators yield Futures (i.e., the chain order), regardless of completion order.
95 - Unordered mode uses ``concurrent.futures.wait(FIRST_COMPLETED)`` to yield results as soon as they complete for
96 minimum end-to-end latency and maximum throughput.
98 - Exception Propagation: ``Future.result()`` naturally propagates exceptions from worker threads, maintaining error
99 visibility for debugging.
101 Parameters:
102 -----------
103 iterator_builder : Callable[[ThreadPoolExecutor], Iterable[Iterable[Future[T]]]]
104 Factory function that is called once (and only once) with the ThreadPoolExecutor as parameter, returning a
105 corresponding series of iterators. Typically, each iterator is a lazy on-demand Python Generator of (a potentially
106 infinite number of) Future[T] objects representing the (future and eventually actual) result of tasks that have been
107 incrementally submitted to the thread pool, avoiding submitting all tasks at once. Typically, advancing the iterator
108 submits the next task to the executor and yields the corresponding Future.
110 max_workers : int, default=os.cpu_count() or 1
111 Maximum number of worker threads in the thread pool. Also determines the buffer size for the bounded-concurrency
112 execution model. Often higher than the number of available CPU cores for I/O-bound tasks.
114 ordered : bool, default=True
115 Controls result delivery mode:
116 - True: Results are yielded in the same order as produced by the builder's iterators (FIFO across the chained
117 iterators), not by task completion order.
118 - False: Results are yielded as soon as available (completion order) for minimum latency and maximum throughput.
120 Yields:
121 -------
122 Results from completed Future objects, either in iterator order (``ordered=True``) or completion order
123 (``ordered=False``).
125 Raises:
126 -------
127 Any exception raised by the submitted tasks will be propagated when their results are consumed via ``Future.result()``.
129 Example:
130 --------
131 # Parallel SSH command execution with ordered results
132 def build_ssh_commands(executor):
133 return [
134 (executor.submit(run_ssh_cmd, cmd) for cmd in commands) # lazy on-demand Python Generator of Future objects
135 ]
137 for result in parallel_iterator(build_ssh_commands, max_workers=4, ordered=True):
138 process_ssh_result(result)
139 """
140 with SynchronousExecutor.executor_for(max_workers=max_workers) as executor:
141 yield from parallel_iterator_results(
142 iterator=itertools.chain.from_iterable(iterator_builder(executor)),
143 max_workers=max_workers,
144 ordered=ordered,
145 is_terminated=is_terminated,
146 )
149def parallel_iterator_results(
150 iterator: Iterator[Future[_T]],
151 *,
152 max_workers: int,
153 ordered: bool,
154 is_terminated: Callable[[], bool] = lambda: False, # optional predicate to request early async termination
155) -> Iterator[_T]:
156 """Yield results from an iterator of Future[T] using bounded concurrency with optional ordered delivery."""
157 assert max_workers >= 0
158 max_workers = max(1, max_workers)
159 if is_terminated():
160 return
162 # Materialize the next N=max_workers futures into a buffer, causing submission + parallel execution of their CLI calls
163 fifo_buffer: deque[Future[_T]] = deque(itertools.islice(iterator, max_workers))
164 sentinel: Future[_T] = Future()
165 next_future: Future[_T]
167 if ordered:
168 while fifo_buffer: # submit the next CLI call whenever the current CLI call returns
169 if is_terminated():
170 for future in fifo_buffer:
171 future.cancel()
172 return
173 curr_future: Future[_T] = fifo_buffer.popleft()
174 next_future = next(iterator, sentinel) # keep the buffer full; causes the next CLI call to be submitted
175 if next_future is not sentinel:
176 fifo_buffer.append(next_future)
177 yield curr_future.result() # blocks until CLI returns
178 else:
179 todo_futures: set[Future[_T]] = set(fifo_buffer)
180 del fifo_buffer # help gc
181 done_futures: set[Future[_T]]
182 while todo_futures:
183 done_futures, todo_futures = concurrent.futures.wait(todo_futures, return_when=FIRST_COMPLETED) # blocks
184 while done_futures: # submit the next CLI call whenever a CLI call returns
185 if is_terminated():
186 for future in todo_futures:
187 future.cancel()
188 return
189 next_future = next(iterator, sentinel) # keep the buffer full; causes the next CLI call to be submitted
190 if next_future is not sentinel:
191 todo_futures.add(next_future)
192 yield done_futures.pop().result() # does not block as processing has already completed
193 assert next(iterator, sentinel) is sentinel
196_K = TypeVar("_K")
197_V = TypeVar("_V")
200def run_in_parallel(fn1: Callable[[], _K], fn2: Callable[[], _V]) -> tuple[_K, _V]:
201 """perf: Runs both I/O functions in parallel/concurrently."""
202 with ThreadPoolExecutor(max_workers=1) as executor:
203 future: Future[_V] = executor.submit(fn2) # async fn2
204 result1: _K = fn1() # blocks until fn1 call returns
205 result2: _V = future.result() # blocks until fn2 call returns
206 return result1, result2
209def batch_cmd_iterator(
210 cmd_args: Iterable[str], # list of arguments to be split across one or more commands
211 fn: Callable[[list[str]], _T], # callback that runs a CLI command with a single batch
212 *,
213 max_batch_items: int = 2**29, # max number of args per batch
214 max_batch_bytes: int = 127 * 1024, # max number of bytes per batch
215 sep: str = " ", # separator between batch args
216 quote_args: bool = False, # calculate byte size assuming each arg will be quoted
217) -> Iterator[_T]:
218 """Returns an iterator that runs fn(cmd_args) in sequential batches, without creating a cmdline that's too big for the OS
219 to handle; Can be seen as a Pythonic xargs -n / -s with OS-aware safety margin.
221 Except for the max_batch_bytes logic, this is essentially the same as:
222 >>>
223 while batch := itertools.batched(cmd_args, max_batch_items): # doctest: +SKIP
224 yield fn(batch)
225 """
226 assert isinstance(sep, str)
227 fsenc: str = sys.getfilesystemencoding()
228 seplen: int = len(sep.encode(fsenc))
229 batch: list[str] = []
230 total_bytes, total_items = 0, 0
231 for cmd_arg in cmd_args:
232 # Include a conservative eight-byte argv pointer
233 arg_bytes: int = 8 + seplen + len((shlex.quote(cmd_arg) if quote_args else cmd_arg).encode(fsenc))
234 if (total_items >= max_batch_items or total_bytes + arg_bytes > max_batch_bytes) and len(batch) > 0:
235 yield fn(batch)
236 batch = []
237 total_bytes, total_items = 0, 0
238 batch.append(cmd_arg)
239 total_bytes += arg_bytes
240 total_items += 1
241 if len(batch) > 0:
242 yield fn(batch)
245def get_max_command_line_bytes(os_name: str) -> int:
246 """Remote flavor of os.sysconf("SC_ARG_MAX") - size(os.environb) - safety margin"""
247 arg_max = _MAX_CMDLINE_BYTES.get(os_name, 256 * 1024)
248 environ_size = 4 * 1024 # typically is 1-4 KB
249 safety_margin = (8 * 2 * 4 + 4) * 1024 if arg_max >= 1 * 1024 * 1024 else 8 * 1024
250 max_bytes = max(4 * 1024, arg_max - environ_size - safety_margin)
251 return max_bytes
254# constants:
255_MAX_CMDLINE_BYTES: Final[dict[str, int]] = {
256 "Linux": 2 * 1024 * 1024,
257 "FreeBSD": 256 * 1024,
258 "Darwin": 1 * 1024 * 1024,
259 "Windows": 32 * 1024,
260}