Coverage for bzfs_main/util/parallel_tasktree.py: 100%
216 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"""Fault-tolerant, dependency-aware workflow scheduling and execution of parallel operations, ensuring that ancestor datasets
16finish before descendants start (parent-before-child order); The design maximizes throughput while preventing inconsistent
17dataset states during replication or snapshot deletion.
19This module contains only generic scheduling and coordination (the "algorithm"). Error handling, retries, and skip policies
20are customizable and implemented by callers via the CompletionCallback API or wrappers such as ``parallel_tasktree_policy``.
22Has zero dependencies beyond the Python standard library.
24Example Usage:
25--------------
26 import logging
27 from bzfs_main.util.parallel_tasktree import CompletionCallback, CompletionCallbackResult, ParallelTaskTree
29 datasets = ["a", "a/b", "a/b/c", "a/d", "e"]
31 def process_dataset(dataset: str, _submit_count: int) -> CompletionCallback:
32 print(dataset)
34 def completion_callback(_todo_futures) -> CompletionCallbackResult:
35 return CompletionCallbackResult(no_skip=True, fail=False)
37 return completion_callback
39 log = logging.getLogger(__name__)
40 tasktree = ParallelTaskTree(
41 log=log,
42 datasets=datasets,
43 process_dataset=process_dataset,
44 max_workers=2,
45 )
46 tasktree.process_datasets_in_parallel()
48 # Sample output:
49 # a
50 # e
51 # a/b
52 # a/d
53 # a/b/c
54"""
56from __future__ import (
57 annotations,
58)
59import concurrent
60import heapq
61import logging
62import os
63from concurrent.futures import (
64 FIRST_COMPLETED,
65 Executor,
66 Future,
67)
68from concurrent.futures.thread import (
69 ThreadPoolExecutor,
70)
71from typing import (
72 Callable,
73 Final,
74 NamedTuple,
75 final,
76)
78from bzfs_main.util.utils import (
79 Comparable,
80 HashedInterner,
81 SortedInterner,
82 SynchronousExecutor,
83 TaskTiming,
84 has_duplicates,
85)
87# constants:
88BARRIER_CHAR: Final[str] = "~"
89COMPONENT_SEPARATOR: Final[str] = "/" # ZFS dataset component separator
92#############################################################################
93@final
94class CompletionCallbackResult(NamedTuple):
95 """Result of a CompletionCallback invocation."""
97 no_skip: bool
98 """True enqueues children, False skips the subtree."""
100 fail: bool
101 """True marks overall run as failed."""
104#############################################################################
105CompletionCallback = Callable[[set[Future["CompletionCallback"]]], CompletionCallbackResult] # Type alias
106"""Callable that is run by the main coordination thread after a ``process_dataset()`` task finishes.
108Purpose:
109- Decide follow-up scheduling after a ``process_dataset()`` task finished.
111Assumptions:
112- Runs in the single coordination thread.
113- May inspect and cancel in-flight futures to implement fail-fast semantics.
114- If cancelling in-flight futures for tasks that spawn subprocesses (e.g. via subprocess.run()), callers should also
115consider terminating the corresponding process subtree to avoid child processes lingering longer than desired. Skipping
116termination will not hang the scheduler (workers will complete naturally), but those subprocesses may outlive cancellation
117until they exit or time out.
118"""
121#############################################################################
122@final
123class ParallelTaskTree:
124 """Main class for dependency-aware workflow scheduling of dataset jobs with optional barriers and priority ordering."""
126 def __init__(
127 self,
128 *,
129 log: logging.Logger,
130 datasets: list[str], # (sorted) list of datasets to process
131 process_dataset: Callable[[str, int], CompletionCallback], # lambda: dataset, tid; must be thread-safe
132 priority: Callable[[str], Comparable] = lambda dataset: dataset, # lexicographical order by default
133 max_workers: int = os.cpu_count() or 1,
134 executors: Callable[[], Executor] | None = None, # factory producing Executor; None means 'auto-choose'
135 interval_nanos: Callable[
136 [int, str, int], int
137 ] = lambda last_update_nanos, dataset, submit_count: 0, # optionally spread tasks out over time; e.g. for jitter
138 timing: TaskTiming = TaskTiming(), # noqa: B008
139 component_separator: str = COMPONENT_SEPARATOR,
140 enable_barriers: bool | None = None, # for testing only; None means 'auto-detect'
141 barrier_name: str = BARRIER_CHAR,
142 is_test_mode: bool = False,
143 ) -> None:
144 """Prepares to process datasets in parallel with dependency-aware workflow scheduling and fault tolerance.
146 This class orchestrates parallel execution of dataset operations while maintaining strict hierarchical dependencies.
147 Processing of a dataset only starts after processing of all its ancestor datasets has completed, ensuring data
148 consistency during operations like ZFS replication or snapshot deletion.
150 Purpose:
151 --------
152 - Process hierarchical datasets in parallel while respecting parent-child dependencies (parent-before-child order)
153 - Provide dependency-aware workflow scheduling; error handling and retries are implemented by callers via
154 ``CompletionCallback`` or thin wrappers
155 - Maximize throughput by processing independent dataset subtrees in parallel
156 - Support complex job scheduling patterns via optional barrier synchronization
158 Assumptions:
159 -----------------
160 - Input `datasets` list is sorted in lexicographical order (enforced in test mode)
161 - Input `datasets` list contains no duplicate entries (enforced in test mode)
162 - Input `datasets` list contains no empty dataset names and none that start with '/'
163 - Dataset hierarchy is determined by slash-separated path components
164 - The `process_dataset` callable is thread-safe and can be executed in parallel
166 Design Rationale:
167 -----------------
168 - The implementation uses a priority queue-based scheduler that maintains two key invariants:
170 - Dependency Ordering: Children are only made available for start of processing after their parent completes,
171 preventing inconsistent dataset states (parent-before-child order).
173 - Priority: Among the datasets available for start of processing, the "smallest" is always processed next, according
174 to a customizable priority callback function, which by default sorts by lexicographical order (not dataset size),
175 ensuring more deterministic execution order.
177 Algorithm Selection:
178 --------------------
179 - Simple Algorithm (default): Used when no barriers ('~') are detected in dataset names. Provides efficient
180 scheduling for standard parent-child dependencies via recursive child enqueueing after parent completion.
182 - Barrier Algorithm (advanced): Activated when barriers are detected or explicitly enabled. Supports complex
183 synchronization scenarios where jobs must wait for completion of entire subtrees before proceeding. Essential
184 for advanced job scheduling patterns like "complete all parallel replications before starting pruning phase."
186 - Both algorithms are CPU and memory efficient. They require main memory proportional to the number of datasets
187 (~400 bytes per dataset), and easily scale to millions of datasets. Time complexity is O(N log N), where
188 N is the number of datasets.
190 Concurrency Design:
191 -------------------
192 By default uses ThreadPoolExecutor with configurable worker limits to balance parallelism against resource
193 consumption. Optionally, plug in a custom Executor to submit tasks to scale-out clusters via frameworks like
194 Ray Core or Dask, etc.
195 The single-threaded coordination loop prevents race conditions while worker threads execute dataset operations in
196 parallel.
198 Params:
199 -------
200 - datasets: Sorted list of dataset names to process (must not contain duplicates)
201 - process_dataset: Thread-safe Callback function to execute on each dataset; returns a CompletionCallback determining
202 if to fail or skip subtree on error; CompletionCallback runs in the (single) main thread as part of the
203 coordination loop.
204 - priority: Callback function to determine dataset processing order; defaults to lexicographical order.
205 - interval_nanos: Callback that returns a non-negative delay (ns) to add to ``next_update_nanos`` for
206 jitter/back-pressure control; arguments are ``(last_update_nanos, dataset, submit_count)``
207 - max_workers: Maximum number of parallel worker threads
208 - executors: Factory returning an Executor to submit tasks to; None means 'auto-choose'
209 - component_separator: String that separates dataset components (default '/'); must be non-empty
210 - enable_barriers: Force enable/disable barrier algorithm (None = auto-detect)
211 - barrier_name: Directory name that denotes a barrier within dataset/job strings (default '~'); must be non-empty and
212 not contain the component_separator
213 - timing: Optionally request early async termination; stops new submissions and cancels in-flight tasks
214 """
215 assert log is not None
216 assert (not is_test_mode) or datasets == sorted(datasets), "List is not sorted"
217 assert (not is_test_mode) or not has_duplicates(datasets), "List contains duplicates"
218 if not component_separator:
219 raise ValueError(f"Invalid component_separator: {component_separator}")
220 if component_separator in barrier_name or not barrier_name:
221 raise ValueError(f"Invalid barrier_name: {barrier_name}")
222 for dataset in datasets:
223 if dataset.startswith(component_separator) or not dataset:
224 raise ValueError(f"Invalid dataset name: {dataset}")
225 assert callable(process_dataset)
226 assert callable(priority)
227 assert max_workers > 0
228 assert callable(interval_nanos)
229 has_barrier: Final[bool] = any(barrier_name in dataset.split(component_separator) for dataset in datasets)
230 assert (enable_barriers is not False) or not has_barrier, "Barrier seen in datasets but barriers explicitly disabled"
232 self._component_separator: Final[str] = component_separator
233 self._barriers_enabled: Final[bool] = has_barrier or bool(enable_barriers)
234 self._barrier_name: Final[str] = barrier_name
235 self._log: Final[logging.Logger] = log
236 self._datasets: Final[list[str]] = datasets
237 self._process_dataset: Final[Callable[[str, int], CompletionCallback]] = process_dataset
238 self._priority: Final[Callable[[str], Comparable]] = priority
239 self._max_workers: Final[int] = max_workers
240 self._interval_nanos: Final[Callable[[int, str, int], int]] = interval_nanos
241 self._timing: Final[TaskTiming] = timing
242 self._is_test_mode: Final[bool] = is_test_mode
243 self._priority_queue: Final[list[_TreeNode]] = []
244 tree, has_siblings = _build_dataset_tree(datasets, component_separator) # tree is nested dicts and is immutable
245 self._tree: Final[_Tree] = tree
246 self._has_siblings: Final[bool] = has_siblings
247 self._empty_barrier: Final[_TreeNode] = _make_tree_node("empty_barrier", "empty_barrier", {}) # immutable!
248 self._datasets_set: Final[SortedInterner[str]] = SortedInterner(datasets) # reduces memory footprint
249 if executors is None:
250 is_parallel: bool = max_workers > 1 and has_siblings # siblings can run in parallel
252 def _default_executor_factory() -> Executor:
253 return ThreadPoolExecutor(max_workers) if is_parallel else SynchronousExecutor()
255 executors = _default_executor_factory
256 self._executors: Final[Callable[[], Executor]] = executors
257 assert callable(executors)
259 def process_datasets_in_parallel(self) -> bool:
260 """Executes the configured tasks and returns True if any dataset processing failed, False if all succeeded."""
261 self._build_priority_queue()
262 executor: Executor = self._executors()
263 with executor:
264 todo_futures: set[Future[CompletionCallback]] = set()
265 future_to_node: dict[Future[CompletionCallback], _TreeNode] = {}
266 submit_count: int = 0
267 timing: TaskTiming = self._timing
268 next_update_nanos: int = timing.monotonic_ns()
269 wait_timeout: float | None = None
270 failed: bool = False
272 def submit_datasets() -> bool:
273 """Submits available datasets to worker threads and returns False if all tasks have been completed."""
274 nonlocal wait_timeout
275 wait_timeout = None # indicates to use blocking flavor of concurrent.futures.wait()
276 while len(self._priority_queue) > 0 and len(todo_futures) < self._max_workers:
277 # pick "smallest" dataset (wrt. sort order) available for start of processing; submit to thread pool
278 nonlocal next_update_nanos
279 sleep_nanos: int = next_update_nanos - timing.monotonic_ns()
280 if sleep_nanos > 0:
281 timing.sleep(sleep_nanos) # allow early wakeup on async termination
282 if timing.is_terminated():
283 break
284 if sleep_nanos > 0 and len(todo_futures) > 0:
285 wait_timeout = 0 # indicates to use non-blocking flavor of concurrent.futures.wait()
286 # It's possible an even "smaller" dataset (wrt. sort order) has become available while we slept.
287 # If so it's preferable to submit to the thread pool the smaller one first.
288 break # break out of loop to check if that's the case via non-blocking concurrent.futures.wait()
289 node: _TreeNode = heapq.heappop(self._priority_queue) # pick "smallest" dataset (wrt. sort order)
290 nonlocal submit_count
291 submit_count += 1
292 next_update_nanos += max(0, self._interval_nanos(next_update_nanos, node.dataset, submit_count))
293 future: Future[CompletionCallback] = executor.submit(self._process_dataset, node.dataset, submit_count)
294 future_to_node[future] = node
295 todo_futures.add(future)
296 return len(todo_futures) > 0 and not timing.is_terminated()
298 def complete_datasets() -> None:
299 """Waits for completed futures, processes results and errors, then enqueues follow-up tasks per policy."""
300 nonlocal failed
301 nonlocal todo_futures
302 done_futures: set[Future[CompletionCallback]]
303 done_futures, todo_futures = concurrent.futures.wait(todo_futures, wait_timeout, return_when=FIRST_COMPLETED)
304 for done_future in sorted(done_futures, key=lambda future: future_to_node[future].dataset):
305 done_node: _TreeNode = future_to_node.pop(done_future)
306 c_callback: CompletionCallback = done_future.result() # does not block as processing already completed
307 c_callback_result: CompletionCallbackResult = c_callback(todo_futures)
308 no_skip: bool = c_callback_result.no_skip
309 fail: bool = c_callback_result.fail
310 failed = failed or fail
311 self._complete_dataset(done_node, no_skip=no_skip)
313 # coordination loop; runs in the (single) main thread; submits tasks to worker threads and handles their results
314 while submit_datasets():
315 complete_datasets()
317 if timing.is_terminated():
318 for todo_future in todo_futures:
319 todo_future.cancel()
320 failed = failed or len(self._priority_queue) > 0 or len(todo_futures) > 0
321 self._priority_queue.clear()
322 todo_futures.clear()
323 future_to_node.clear()
324 assert len(self._priority_queue) == 0
325 assert len(todo_futures) == 0
326 assert len(future_to_node) == 0
327 return failed
329 def _build_priority_queue(self) -> None:
330 """Builds and fills initial priority queue of available root nodes for this task tree, ensuring the scheduler starts
331 from a synthetic root node while honoring barriers; the synthetic root simplifies enqueueing logic."""
332 self._priority_queue.clear()
333 root_node: _TreeNode = _make_tree_node(priority="", dataset="", children=self._tree)
334 self._complete_dataset(root_node, no_skip=True)
336 def _complete_dataset(self, node: _TreeNode, no_skip: bool) -> None:
337 """Enqueues child nodes for start of processing, using the appropriate algorithm."""
338 if self._barriers_enabled: # This barrier-based algorithm is for more general job scheduling, as in bzfs_jobrunner
339 self._complete_dataset_with_barriers(node, no_skip=no_skip)
340 elif no_skip: # This simple algorithm is sufficient for most uses
341 self._simple_enqueue_children(node)
343 def _simple_enqueue_children(self, node: _TreeNode) -> None:
344 """Enqueues child nodes for start of processing (using iteration to avoid potentially hitting recursion limits)."""
345 stack: list[_TreeNode] = [node]
346 while stack:
347 child_node: _TreeNode = stack.pop()
348 if child_node is not node and child_node.dataset in self._datasets_set:
349 heapq.heappush(
350 self._priority_queue,
351 _make_tree_node(self._priority(child_node.dataset), child_node.dataset, child_node.children),
352 )
353 else:
354 # it's an intermediate node with no job attached, or an initial node; pass enqueue recursively down the tree.
355 # child_node stores children in reverse lexicographic order; LIFO pop restores lexicogr. priority() traversal
356 for child, grandchildren in child_node.children.items():
357 child_abs_dataset: str = self._join_dataset(child_node.dataset, child)
358 stack.append(_make_tree_node(child_abs_dataset, child_abs_dataset, grandchildren)) # without priority()
360 def _complete_dataset_with_barriers(self, node: _TreeNode, no_skip: bool) -> None:
361 """After successful completion, enqueues children, opens barriers, and propagates completion upwards.
363 The (more complex) algorithm below is for more general job scheduling, as in bzfs_jobrunner. Here, a "dataset" string
364 is treated as an identifier for any kind of job rather than a reference to a concrete ZFS object. An example
365 "dataset" job string is "src_host1/createsnapshots/replicate_to_hostA". Jobs can depend on another job via a
366 parent/child relationship formed by '/' directory separators within the dataset string, and multiple "datasets" form
367 a job dependency tree by way of common dataset directory prefixes. Jobs that do not depend on each other can be
368 executed in parallel, and jobs can be told to first wait for other jobs to complete successfully. The algorithm is
369 based on a barrier primitive and is typically disabled. It is only required for rare jobrunner configs.
371 For example, a job scheduler can specify that all parallel replications jobs to multiple destinations must succeed
372 before the jobs of the pruning phase can start. More generally, with this algo, a job scheduler can specify that all
373 jobs within a given job subtree (containing any nested combination of sequential and/or parallel jobs) must
374 successfully complete before a certain other job within the job tree is started. This is specified via the barrier
375 directory named by ``barrier_name`` (default '~'). An example is "src_host1/createsnapshots/~/prune".
377 Note that the default '~' is unambiguous as it is not a valid ZFS dataset name component per the naming rules
378 enforced by the 'zfs create', 'zfs snapshot' and 'zfs bookmark' CLIs. Custom barrier names should avoid colliding
379 with real dataset/job components.
380 """
382 def enqueue_children(node: _TreeNode) -> int:
383 """Returns number of jobs that were added to priority_queue for immediate start of processing."""
384 n: int = 0
385 children: _Tree = node.children
386 for child, grandchildren in reversed(children.items()):
387 abs_dataset: str = self._join_dataset(node.dataset, child)
388 child_node: _TreeNode = _make_tree_node(abs_dataset, abs_dataset, grandchildren, parent=node)
389 k: int
390 if child != self._barrier_name:
391 if abs_dataset in self._datasets_set:
392 # it's not a barrier; make job available for immediate start of processing
393 child_node = _make_tree_node(self._priority(abs_dataset), abs_dataset, grandchildren, parent=node)
394 heapq.heappush(self._priority_queue, child_node)
395 k = 1
396 else: # it's an intermediate node that has no job attached; pass the enqueue operation
397 k = enqueue_children(child_node) # ... recursively down the tree
398 elif len(children) == 1: # if the only child is a barrier then pass the enqueue operation
399 k = enqueue_children(child_node) # ... recursively down the tree
400 else: # park the barrier node within the (still closed) barrier for the time being
401 assert node.mut.barrier is None
402 node.mut.barrier = child_node
403 k = 0
404 node.mut.pending += min(1, k)
405 n += k
406 assert n >= 0
407 return n
409 if no_skip:
410 enqueue_children(node) # make child datasets available for start of processing
411 else: # job completed without success
412 # ... thus, opening the barrier shall always do nothing in node and its ancestors.
413 # perf: Irrevocably mark (exactly once) barriers of this node and all its ancestors as cleared due to subtree
414 # skip, via barriers_cleared=True. This enables to avoid redundant re-walking the ancestor chain on subsequent
415 # skip.
416 tmp: _TreeNode | None = node
417 while (tmp is not None) and not tmp.mut.barriers_cleared:
418 tmp.mut.barriers_cleared = True
419 tmp.mut.barrier = self._empty_barrier
420 tmp = tmp.parent
421 assert node.mut.pending >= 0
422 while node.mut.pending == 0: # have all jobs in subtree of current node completed?
423 if no_skip: # ... if so open the barrier, if it exists, and enqueue jobs waiting on it
424 if not (node.mut.barrier is None or node.mut.barrier is self._empty_barrier):
425 node.mut.pending += min(1, enqueue_children(node.mut.barrier))
426 node.mut.barrier = self._empty_barrier
427 if node.mut.pending > 0: # did opening of barrier cause jobs to be enqueued in subtree?
428 break # ... if so we have not yet completed the subtree, so don't mark the subtree as completed yet
429 if node.parent is None:
430 break # we've reached the root node
431 node = node.parent # recurse up the tree to propagate completion upward
432 node.mut.pending -= 1 # mark subtree as completed
433 assert node.mut.pending >= 0
435 def _join_dataset(self, parent: str, child: str) -> str:
436 """Concatenates parent and child dataset names; accommodates synthetic root node; interns for memory footprint."""
437 return self._datasets_set.interned(f"{parent}{self._component_separator}{child}" if parent else child)
440#############################################################################
441@final
442class _TreeNodeMutableAttributes:
443 """Container for mutable attributes, stored space efficiently."""
445 __slots__ = ("barrier", "barriers_cleared", "pending") # uses more compact memory layout than __dict__
447 def __init__(self) -> None:
448 self.barrier: _TreeNode | None = None # zero or one barrier TreeNode waiting for this node to complete
449 self.pending: int = 0 # number of children added to priority queue that haven't completed their work yet
450 self.barriers_cleared: bool = False # irrevocably mark barriers of this node and all its ancestors as cleared?
453#############################################################################
454@final
455class _TreeNode(NamedTuple):
456 """Node in dataset dependency tree used by the scheduler; _TreeNodes are ordered by priority and dataset name within a
457 priority queue, via __lt__ comparisons."""
459 priority: Comparable # determines the processing order once this dataset has become available for start of processing
460 dataset: str # each dataset name is unique; attribs other than `priority` and `dataset` are never used for comparisons
461 children: _Tree # dataset "directory" tree consists of nested dicts; aka dict[str, dict]
462 parent: _TreeNode | None
463 mut: _TreeNodeMutableAttributes
465 def __repr__(self) -> str:
466 priority, dataset, pending, barrier = self.priority, self.dataset, self.mut.pending, self.mut.barrier
467 return str({"priority": priority, "dataset": dataset, "pending": pending, "barrier": barrier is not None})
470def _make_tree_node(priority: Comparable, dataset: str, children: _Tree, parent: _TreeNode | None = None) -> _TreeNode:
471 """Creates a TreeNode with mutable state container."""
472 return _TreeNode(priority, dataset, children, parent, _TreeNodeMutableAttributes())
475#############################################################################
476_Tree = dict[str, "_Tree"] # Type alias
479def _build_dataset_tree(sorted_datasets: list[str], component_separator: str) -> tuple[_Tree, bool]:
480 """Takes as input a sorted list of datasets and returns a (reverse) sorted directory tree containing the same dataset
481 names, in the form of nested dicts; This converts the dataset list into a dependency tree."""
482 assert component_separator
483 tree: _Tree = {}
484 has_siblings: bool = False
485 interner: HashedInterner[str] = HashedInterner() # reduces memory footprint
486 shared_empty_leaf: _Tree = {} # tree with shared empty leafs has ~30% lower memory footprint than non-compacted version
488 for dataset in reversed(sorted_datasets):
489 current: _Tree = tree
490 components: list[str] = dataset.split(component_separator)
491 k: int = len(components) - 1
492 for i, component in enumerate(components):
493 child: _Tree | None = current.get(component)
494 if child is None:
495 child = {} if i < k else shared_empty_leaf # sharing is safe as the tree is treated as immutable henceforth
496 assert current is not shared_empty_leaf
497 has_siblings = has_siblings or len(current) > 0
498 component = interner.intern(component)
499 current[component] = child
500 current = child
501 return tree, has_siblings