Coverage for bzfs_main/util/parallel

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

7# http://www.apache.org/licenses/LICENSE-2.0

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."""

102

103

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.

107

108Purpose:

109- Decide follow-up scheduling after a ``process_dataset()`` task finished.

110

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"""

119

120

121#############################################################################

122@final

123class ParallelTaskTree:

124 """Main class for dependency-aware workflow scheduling of dataset jobs with optional barriers and priority ordering."""

125

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 enable_barriers: bool | None = None, # for testing only; None means 'auto-detect'

140 barrier_name: str = BARRIER_CHAR,

141 is_test_mode: bool = False,

142 ) -> None:

143 """Prepares to process datasets in parallel with dependency-aware workflow scheduling and fault tolerance.

144

145 This class orchestrates parallel execution of dataset operations while maintaining strict hierarchical dependencies.

146 Processing of a dataset only starts after processing of all its ancestor datasets has completed, ensuring data

147 consistency during operations like ZFS replication or snapshot deletion.

148

149 Purpose:

150 --------

151 - Process hierarchical datasets in parallel while respecting parent-child dependencies (parent-before-child order)

152 - Provide dependency-aware workflow scheduling; error handling and retries are implemented by callers via

153 ``CompletionCallback`` or thin wrappers

154 - Maximize throughput by processing independent dataset subtrees in parallel

155 - Support complex job scheduling patterns via optional barrier synchronization

156

157 Assumptions:

158 -----------------

159 - Input `datasets` list is sorted in lexicographical order (enforced in test mode)

160 - Input `datasets` list contains no duplicate entries (enforced in test mode)

161 - Input `datasets` list contains no empty dataset names and none that start with '/'

162 - Dataset hierarchy is determined by slash-separated path components

163 - The `process_dataset` callable is thread-safe and can be executed in parallel

164

165 Design Rationale:

166 -----------------

167 - The implementation uses a priority queue-based scheduler that maintains two key invariants:

168

169 - Dependency Ordering: Children are only made available for start of processing after their parent completes,

170 preventing inconsistent dataset states (parent-before-child order).

171

172 - Priority: Among the datasets available for start of processing, the "smallest" is always processed next, according

173 to a customizable priority callback function, which by default sorts by lexicographical order (not dataset size),

174 ensuring more deterministic execution order.

175

176 Algorithm Selection:

177 --------------------

178 - Simple Algorithm (default): Used when no barriers ('~') are detected in dataset names. Provides efficient

179 scheduling for standard parent-child dependencies via recursive child enqueueing after parent completion.

180

181 - Barrier Algorithm (advanced): Activated when barriers are detected or explicitly enabled. Supports complex

182 synchronization scenarios where jobs must wait for completion of entire subtrees before proceeding. Essential

183 for advanced job scheduling patterns like "complete all parallel replications before starting pruning phase."

184

185 - Both algorithms are CPU and memory efficient. They require main memory proportional to the number of datasets

186 (~400 bytes per dataset), and easily scale to millions of datasets. Time complexity is O(N log N), where

187 N is the number of datasets.

188

189 Concurrency Design:

190 -------------------

191 By default uses ThreadPoolExecutor with configurable worker limits to balance parallelism against resource

192 consumption. Optionally, plug in a custom Executor to submit tasks to scale-out clusters via frameworks like

193 Ray Core or Dask, etc.

194 The single-threaded coordination loop prevents race conditions while worker threads execute dataset operations in

195 parallel.

196

197 Params:

198 -------

199 - datasets: Sorted list of dataset names to process (must not contain duplicates)

200 - process_dataset: Thread-safe Callback function to execute on each dataset; returns a CompletionCallback determining

201 if to fail or skip subtree on error; CompletionCallback runs in the (single) main thread as part of the

202 coordination loop.

203 - priority: Callback function to determine dataset processing order; defaults to lexicographical order.

204 - interval_nanos: Callback that returns a non-negative delay (ns) to add to ``next_update_nanos`` for

205 jitter/back-pressure control; arguments are ``(last_update_nanos, dataset, submit_count)``

206 - max_workers: Maximum number of parallel worker threads

207 - executors: Factory returning an Executor to submit tasks to; None means 'auto-choose'

208 - enable_barriers: Force enable/disable barrier algorithm (None = auto-detect)

209 - barrier_name: Directory name that denotes a barrier within dataset/job strings (default '~'); must be non-empty and

210 not contain '/'

211 - timing: Optionally request early async termination; stops new submissions and cancels in-flight tasks

212 """

213 assert log is not None

214 assert (not is_test_mode) or datasets == sorted(datasets), "List is not sorted"

215 assert (not is_test_mode) or not has_duplicates(datasets), "List contains duplicates"

216 if COMPONENT_SEPARATOR in barrier_name or not barrier_name:

217 raise ValueError(f"Invalid barrier_name: {barrier_name}")

218 for dataset in datasets:

219 if dataset.startswith(COMPONENT_SEPARATOR) or not dataset:

220 raise ValueError(f"Invalid dataset name: {dataset}")

221 assert callable(process_dataset)

222 assert callable(priority)

223 assert max_workers > 0

224 assert callable(interval_nanos)

225 has_barrier: Final[bool] = any(barrier_name in dataset.split(COMPONENT_SEPARATOR) for dataset in datasets)

226 assert (enable_barriers is not False) or not has_barrier, "Barrier seen in datasets but barriers explicitly disabled"

227

228 self._barriers_enabled: Final[bool] = has_barrier or bool(enable_barriers)

229 self._barrier_name: Final[str] = barrier_name

230 self._log: Final[logging.Logger] = log

231 self._datasets: Final[list[str]] = datasets

232 self._process_dataset: Final[Callable[[str, int], CompletionCallback]] = process_dataset

233 self._priority: Final[Callable[[str], Comparable]] = priority

234 self._max_workers: Final[int] = max_workers

235 self._interval_nanos: Final[Callable[[int, str, int], int]] = interval_nanos

236 self._timing: Final[TaskTiming] = timing

237 self._is_test_mode: Final[bool] = is_test_mode

238 self._priority_queue: Final[list[_TreeNode]] = []

239 tree, has_siblings = _build_dataset_tree(datasets) # tree consists of nested dictionaries and is immutable

240 self._tree: Final[_Tree] = tree

241 self._has_siblings: Final[bool] = has_siblings

242 self._empty_barrier: Final[_TreeNode] = _make_tree_node("empty_barrier", "empty_barrier", {}) # immutable!

243 self._datasets_set: Final[SortedInterner[str]] = SortedInterner(datasets) # reduces memory footprint

244 if executors is None:

245 is_parallel: bool = max_workers > 1 and has_siblings # siblings can run in parallel

246

247 def _default_executor_factory() -> Executor:

248 return ThreadPoolExecutor(max_workers) if is_parallel else SynchronousExecutor()

249

250 executors = _default_executor_factory

251 self._executors: Final[Callable[[], Executor]] = executors

252 assert callable(executors)

253

254 def process_datasets_in_parallel(self) -> bool:

255 """Executes the configured tasks and returns True if any dataset processing failed, False if all succeeded."""

256 self._build_priority_queue()

257 executor: Executor = self._executors()

258 with executor:

259 todo_futures: set[Future[CompletionCallback]] = set()

260 future_to_node: dict[Future[CompletionCallback], _TreeNode] = {}

261 submit_count: int = 0

262 timing: TaskTiming = self._timing

263 next_update_nanos: int = timing.monotonic_ns()

264 wait_timeout: float | None = None

265 failed: bool = False

266

267 def submit_datasets() -> bool:

268 """Submits available datasets to worker threads and returns False if all tasks have been completed."""

269 nonlocal wait_timeout

270 wait_timeout = None # indicates to use blocking flavor of concurrent.futures.wait()

271 while len(self._priority_queue) > 0 and len(todo_futures) < self._max_workers:

272 # pick "smallest" dataset (wrt. sort order) available for start of processing; submit to thread pool

273 nonlocal next_update_nanos

274 sleep_nanos: int = next_update_nanos - timing.monotonic_ns()

275 if sleep_nanos > 0:

276 timing.sleep(sleep_nanos) # allow early wakeup on async termination

277 if timing.is_terminated():

278 break

279 if sleep_nanos > 0 and len(todo_futures) > 0:

280 wait_timeout = 0 # indicates to use non-blocking flavor of concurrent.futures.wait()

281 # It's possible an even "smaller" dataset (wrt. sort order) has become available while we slept.

282 # If so it's preferable to submit to the thread pool the smaller one first.

283 break # break out of loop to check if that's the case via non-blocking concurrent.futures.wait()

284 node: _TreeNode = heapq.heappop(self._priority_queue) # pick "smallest" dataset (wrt. sort order)

285 nonlocal submit_count

286 submit_count += 1

287 next_update_nanos += max(0, self._interval_nanos(next_update_nanos, node.dataset, submit_count))

288 future: Future[CompletionCallback] = executor.submit(self._process_dataset, node.dataset, submit_count)

289 future_to_node[future] = node

290 todo_futures.add(future)

291 return len(todo_futures) > 0 and not timing.is_terminated()

292

293 def complete_datasets() -> None:

294 """Waits for completed futures, processes results and errors, then enqueues follow-up tasks per policy."""

295 nonlocal failed

296 nonlocal todo_futures

297 done_futures: set[Future[CompletionCallback]]

298 done_futures, todo_futures = concurrent.futures.wait(todo_futures, wait_timeout, return_when=FIRST_COMPLETED)

299 for done_future in sorted(done_futures, key=lambda future: future_to_node[future].dataset):

300 done_node: _TreeNode = future_to_node.pop(done_future)

301 c_callback: CompletionCallback = done_future.result() # does not block as processing already completed

302 c_callback_result: CompletionCallbackResult = c_callback(todo_futures)

303 no_skip: bool = c_callback_result.no_skip

304 fail: bool = c_callback_result.fail

305 failed = failed or fail

306 self._complete_dataset(done_node, no_skip=no_skip)

307

308 # coordination loop; runs in the (single) main thread; submits tasks to worker threads and handles their results

309 while submit_datasets():

310 complete_datasets()

311

312 if timing.is_terminated():

313 for todo_future in todo_futures:

314 todo_future.cancel()

315 failed = failed or len(self._priority_queue) > 0 or len(todo_futures) > 0

316 self._priority_queue.clear()

317 todo_futures.clear()

318 future_to_node.clear()

319 assert len(self._priority_queue) == 0

320 assert len(todo_futures) == 0

321 assert len(future_to_node) == 0

322 return failed

323

324 def _build_priority_queue(self) -> None:

325 """Builds and fills initial priority queue of available root nodes for this task tree, ensuring the scheduler starts

326 from a synthetic root node while honoring barriers; the synthetic root simplifies enqueueing logic."""

327 self._priority_queue.clear()

328 root_node: _TreeNode = _make_tree_node(priority="", dataset="", children=self._tree)

329 self._complete_dataset(root_node, no_skip=True)

330

331 def _complete_dataset(self, node: _TreeNode, no_skip: bool) -> None:

332 """Enqueues child nodes for start of processing, using the appropriate algorithm."""

333 if self._barriers_enabled: # This barrier-based algorithm is for more general job scheduling, as in bzfs_jobrunner

334 self._complete_dataset_with_barriers(node, no_skip=no_skip)

335 elif no_skip: # This simple algorithm is sufficient for most uses

336 self._simple_enqueue_children(node)

337

338 def _simple_enqueue_children(self, node: _TreeNode) -> None:

339 """Enqueues child nodes for start of processing (using iteration to avoid potentially hitting recursion limits)."""

340 stack: list[_TreeNode] = [node]

341 while stack:

342 child_node: _TreeNode = stack.pop()

343 if child_node is not node and child_node.dataset in self._datasets_set:

344 heapq.heappush(

345 self._priority_queue,

346 _make_tree_node(self._priority(child_node.dataset), child_node.dataset, child_node.children),

347 )

348 else:

349 # it's an intermediate node with no job attached, or an initial node; pass enqueue recursively down the tree.

350 # child_node stores children in reverse lexicographic order; LIFO pop restores lexicogr. priority() traversal

351 for child, grandchildren in child_node.children.items():

352 child_abs_dataset: str = self._join_dataset(child_node.dataset, child)

353 stack.append(_make_tree_node(child_abs_dataset, child_abs_dataset, grandchildren)) # without priority()

354

355 def _complete_dataset_with_barriers(self, node: _TreeNode, no_skip: bool) -> None:

356 """After successful completion, enqueues children, opens barriers, and propagates completion upwards.

357

358 The (more complex) algorithm below is for more general job scheduling, as in bzfs_jobrunner. Here, a "dataset" string

359 is treated as an identifier for any kind of job rather than a reference to a concrete ZFS object. An example

360 "dataset" job string is "src_host1/createsnapshots/replicate_to_hostA". Jobs can depend on another job via a

361 parent/child relationship formed by '/' directory separators within the dataset string, and multiple "datasets" form

362 a job dependency tree by way of common dataset directory prefixes. Jobs that do not depend on each other can be

363 executed in parallel, and jobs can be told to first wait for other jobs to complete successfully. The algorithm is

364 based on a barrier primitive and is typically disabled. It is only required for rare jobrunner configs.

365

366 For example, a job scheduler can specify that all parallel replications jobs to multiple destinations must succeed

367 before the jobs of the pruning phase can start. More generally, with this algo, a job scheduler can specify that all

368 jobs within a given job subtree (containing any nested combination of sequential and/or parallel jobs) must

369 successfully complete before a certain other job within the job tree is started. This is specified via the barrier

370 directory named by ``barrier_name`` (default '~'). An example is "src_host1/createsnapshots/~/prune".

371

372 Note that the default '~' is unambiguous as it is not a valid ZFS dataset name component per the naming rules

373 enforced by the 'zfs create', 'zfs snapshot' and 'zfs bookmark' CLIs. Custom barrier names should avoid colliding

374 with real dataset/job components.

375 """

376

377 def enqueue_children(node: _TreeNode) -> int:

378 """Returns number of jobs that were added to priority_queue for immediate start of processing."""

379 n: int = 0

380 children: _Tree = node.children

381 for child, grandchildren in reversed(children.items()):

382 abs_dataset: str = self._join_dataset(node.dataset, child)

383 child_node: _TreeNode = _make_tree_node(abs_dataset, abs_dataset, grandchildren, parent=node)

384 k: int

385 if child != self._barrier_name:

386 if abs_dataset in self._datasets_set:

387 # it's not a barrier; make job available for immediate start of processing

388 child_node = _make_tree_node(self._priority(abs_dataset), abs_dataset, grandchildren, parent=node)

389 heapq.heappush(self._priority_queue, child_node)

390 k = 1

391 else: # it's an intermediate node that has no job attached; pass the enqueue operation

392 k = enqueue_children(child_node) # ... recursively down the tree

393 elif len(children) == 1: # if the only child is a barrier then pass the enqueue operation

394 k = enqueue_children(child_node) # ... recursively down the tree

395 else: # park the barrier node within the (still closed) barrier for the time being

396 assert node.mut.barrier is None

397 node.mut.barrier = child_node

398 k = 0

399 node.mut.pending += min(1, k)

400 n += k

401 assert n >= 0

402 return n

403

404 if no_skip:

405 enqueue_children(node) # make child datasets available for start of processing

406 else: # job completed without success

407 # ... thus, opening the barrier shall always do nothing in node and its ancestors.

408 # perf: Irrevocably mark (exactly once) barriers of this node and all its ancestors as cleared due to subtree

409 # skip, via barriers_cleared=True. This enables to avoid redundant re-walking the ancestor chain on subsequent

410 # skip.

411 tmp: _TreeNode | None = node

412 while (tmp is not None) and not tmp.mut.barriers_cleared:

413 tmp.mut.barriers_cleared = True

414 tmp.mut.barrier = self._empty_barrier

415 tmp = tmp.parent

416 assert node.mut.pending >= 0

417 while node.mut.pending == 0: # have all jobs in subtree of current node completed?

418 if no_skip: # ... if so open the barrier, if it exists, and enqueue jobs waiting on it

419 if not (node.mut.barrier is None or node.mut.barrier is self._empty_barrier):

420 node.mut.pending += min(1, enqueue_children(node.mut.barrier))

421 node.mut.barrier = self._empty_barrier

422 if node.mut.pending > 0: # did opening of barrier cause jobs to be enqueued in subtree?

423 break # ... if so we have not yet completed the subtree, so don't mark the subtree as completed yet

424 if node.parent is None:

425 break # we've reached the root node

426 node = node.parent # recurse up the tree to propagate completion upward

427 node.mut.pending -= 1 # mark subtree as completed

428 assert node.mut.pending >= 0

429

430 def _join_dataset(self, parent: str, child: str) -> str:

431 """Concatenates parent and child dataset names; accommodates synthetic root node; interns for memory footprint."""

432 return self._datasets_set.interned(f"{parent}{COMPONENT_SEPARATOR}{child}" if parent else child)

433

434

435#############################################################################

436@final

437class _TreeNodeMutableAttributes:

438 """Container for mutable attributes, stored space efficiently."""

439

440 __slots__ = ("barrier", "barriers_cleared", "pending") # uses more compact memory layout than __dict__

441

442 def __init__(self) -> None:

443 self.barrier: _TreeNode | None = None # zero or one barrier TreeNode waiting for this node to complete

444 self.pending: int = 0 # number of children added to priority queue that haven't completed their work yet

445 self.barriers_cleared: bool = False # irrevocably mark barriers of this node and all its ancestors as cleared?

446

447

448#############################################################################

449@final

450class _TreeNode(NamedTuple):

451 """Node in dataset dependency tree used by the scheduler; _TreeNodes are ordered by priority and dataset name within a

452 priority queue, via __lt__ comparisons."""

453

454 priority: Comparable # determines the processing order once this dataset has become available for start of processing

455 dataset: str # each dataset name is unique; attribs other than `priority` and `dataset` are never used for comparisons

456 children: _Tree # dataset "directory" tree consists of nested dicts; aka dict[str, dict]

457 parent: _TreeNode | None

458 mut: _TreeNodeMutableAttributes

459

460 def __repr__(self) -> str:

461 priority, dataset, pending, barrier = self.priority, self.dataset, self.mut.pending, self.mut.barrier

462 return str({"priority": priority, "dataset": dataset, "pending": pending, "barrier": barrier is not None})

463

464

465def _make_tree_node(priority: Comparable, dataset: str, children: _Tree, parent: _TreeNode | None = None) -> _TreeNode:

466 """Creates a TreeNode with mutable state container."""

467 return _TreeNode(priority, dataset, children, parent, _TreeNodeMutableAttributes())

468

469

470#############################################################################

471_Tree = dict[str, "_Tree"] # Type alias

472

473

474def _build_dataset_tree(sorted_datasets: list[str]) -> tuple[_Tree, bool]:

475 """Takes as input a sorted list of datasets and returns a (reverse) sorted directory tree containing the same dataset

476 names, in the form of nested dicts; This converts the dataset list into a dependency tree."""

477 tree: _Tree = {}

478 has_siblings: bool = False

479 interner: HashedInterner[str] = HashedInterner() # reduces memory footprint

480 shared_empty_leaf: _Tree = {} # tree with shared empty leafs has ~30% lower memory footprint than non-compacted version

481

482 for dataset in reversed(sorted_datasets):

483 current: _Tree = tree

484 components: list[str] = dataset.split(COMPONENT_SEPARATOR)

485 k: int = len(components) - 1

486 for i, component in enumerate(components):

487 child: _Tree | None = current.get(component)

488 if child is None:

489 child = {} if i < k else shared_empty_leaf # sharing is safe as the tree is treated as immutable henceforth

490 assert current is not shared_empty_leaf

491 has_siblings = has_siblings or len(current) > 0

492 component = interner.intern(component)

493 current[component] = child

494 current = child

495 return tree, has_siblings

Coverage for bzfs_main/util/parallel_tasktree.py: 100%

212 statements