Coverage for bzfs_main/progress_reporter.py: 99%
262 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"""Thread-based progress monitor for `pv` output during data transfers.
17It tails the log files produced by (parallel) ``pv`` processes and periodically prints a single status line showing
18aggregated throughput, ETA, etc. The reporter runs in a separate daemon thread to avoid blocking replication. All methods
19are designed for minimal synchronization overhead.
20"""
22from __future__ import (
23 annotations,
24)
25import argparse
26import glob
27import os
28import re
29import selectors
30import threading
31import time
32from collections import (
33 deque,
34)
35from dataclasses import (
36 dataclass,
37 field,
38)
39from datetime import (
40 datetime,
41)
42from enum import (
43 Enum,
44 auto,
45)
46from logging import (
47 Logger,
48)
49from pathlib import (
50 Path,
51)
52from typing import (
53 IO,
54 Any,
55 Final,
56 NamedTuple,
57 final,
58)
60from bzfs_main.util.utils import (
61 FILE_PERMISSIONS,
62 LOG_STDOUT,
63 InterruptibleSleep,
64 human_readable_bytes,
65 open_nofollow,
66)
68# constants
69PV_FILE_THREAD_SEPARATOR: Final[str] = "_"
70_ARABIC_DECIMAL_SEPARATOR: Final[str] = "\u066b" # "٫" # noqa: RUF003
71_PV_SIZE_TO_BYTES_REGEX: Final[re.Pattern[str]] = re.compile(
72 rf"(\d+[.,{_ARABIC_DECIMAL_SEPARATOR}]?\d*)\s*([KMGTPEZYRQ]?)(i?)([Bb])(.*)"
73)
76#############################################################################
77@final
78class State(Enum):
79 """Progress reporter lifecycle state transitions."""
81 IS_PAUSING = auto()
82 IS_RESETTING = auto()
85#############################################################################
86@final
87class ProgressReporter:
88 """Periodically prints progress updates to the same console status line, which is helpful if the program runs in an
89 interactive Unix terminal session.
91 Tails the 'pv' output log files that are being written to by (parallel) replication,
92 and extracts aggregate progress and throughput metrics from them, such as MB, MB/s, ETA, etc. Periodically prints these
93 metrics to the console status line, and in doing so "visually overwrites" the previous status line, via appending a \r
94 carriage return control char rather than a \n newline char.
95 Example console status line:
96 2025-01-17 01:23:04 [I] zfs sent 41.7 GiB 0:00:46 [963 MiB/s] [907 MiB/s] [==========> ] 80% ETA 0:00:04 ETA 01:23:08
97 """
99 def __init__(
100 self,
101 log: Logger,
102 pv_program_opts: list[str],
103 use_select: bool,
104 progress_update_intervals: tuple[float, float] | None,
105 fail: bool = False,
106 ) -> None:
107 """Creates a reporter configured for ``pv`` log parsing."""
108 # immutable variables:
109 self._log: Final[Logger] = log
110 self._pv_program_opts: Final[list[str]] = pv_program_opts
111 self._use_select: Final[bool] = use_select
112 self._progress_update_intervals: Final[tuple[float, float] | None] = progress_update_intervals
113 self._inject_error: bool = fail # for testing only
115 # mutable variables:
116 self._thread: threading.Thread | None = None
117 self._exception: BaseException | None = None
118 self._lock: Final[threading.Lock] = threading.Lock()
119 self._sleeper: Final[InterruptibleSleep] = InterruptibleSleep(self._lock) # sleeper shares lock with reporter
120 self._file_name_queue: set[str] = set()
121 self._file_name_set: Final[set[str]] = set()
122 self._is_stopping: bool = False
123 self._states: list[State] = [State.IS_RESETTING]
124 self._first_status_line_emitted: Final[threading.Event] = threading.Event() # for testing only
126 def start(self) -> None:
127 """Starts the monitoring thread and begins asynchronous parsing of ``pv`` log files."""
128 with self._lock:
129 assert self._thread is None
130 self._thread = threading.Thread(target=lambda: self._run(), name="progress_reporter", daemon=True)
131 self._thread.start()
133 def stop(self) -> None:
134 """Blocks until reporter is stopped, then reraises any exception that may have happened during log processing."""
135 with self._lock:
136 self._is_stopping = True
137 self._sleeper.interrupt()
138 t = self._thread
139 if t is not None:
140 t.join()
141 e = self._exception
142 if e is not None:
143 raise e # reraise exception in current thread
145 def pause(self) -> None:
146 """Temporarily suspends status logging."""
147 self._append_state(State.IS_PAUSING)
149 def reset(self) -> None:
150 """Clears metrics before processing a new batch of logs; the purpose is to discard previous totals to avoid mixing
151 unrelated transfers."""
152 self._append_state(State.IS_RESETTING)
154 def _append_state(self, state: State) -> None:
155 with self._lock:
156 states: list[State] = self._states
157 if len(states) > 0 and states[-1] is state:
158 return # same state twice in a row is a no-op
159 states.append(state)
160 if len(states) >= 3:
161 del states[0] # cap time and memory consumption by removing redundant state transitions
162 self._sleeper.interrupt()
164 def wait_for_first_status_line(self) -> None:
165 """Blocks until at least the first status line has been emitted; for testing only."""
166 if not self._first_status_line_emitted.is_set(): 166 ↛ exitline 166 didn't return from function 'wait_for_first_status_line' because the condition on line 166 was always true
167 self._first_status_line_emitted.wait()
169 def enqueue_pv_log_file(self, pv_log_file: str) -> None:
170 """Tells progress reporter thread to also monitor and tail the given pv log file."""
171 with self._lock:
172 if pv_log_file not in self._file_name_set:
173 self._file_name_queue.add(pv_log_file)
175 def _run(self) -> None:
176 """Thread entry point consuming pv logs and updating metrics."""
177 log = self._log
178 try:
179 fds: list[IO[Any]] = []
180 try:
181 selector = selectors.SelectSelector() if self._use_select else selectors.PollSelector()
182 try:
183 self._run_internal(fds, selector)
184 finally:
185 selector.close()
186 finally:
187 for fd in fds:
188 fd.close()
189 except BaseException as e:
190 if not isinstance(e, BrokenPipeError): 190 ↛ exitline 190 didn't return from function '_run' because the condition on line 190 was always true
191 self._exception = e # will be reraised in stop()
192 log.exception("%s", "ProgressReporter:")
194 @dataclass
195 @final
196 class TransferStat:
197 """Tracks per-file transfer state and ETA."""
199 @dataclass(order=True)
200 @final
201 class ETA:
202 """Estimated time of arrival."""
204 timestamp_nanos: int # sorted by future time at which current zfs send/recv transfer is estimated to complete
205 seq_nr: int # tiebreaker wrt. sort order
206 line_tail: str = field(compare=False) # trailing pv log line part w/ progress bar, duration ETA, timestamp ETA
208 bytes_in_flight: int
209 eta: ETA
211 def _run_internal(self, fds: list[IO[Any]], selector: selectors.BaseSelector) -> None:
212 """Tails pv log files and periodically logs aggregated progress."""
214 @final
215 class Sample(NamedTuple):
216 """Sliding window entry for throughput calculation."""
218 sent_bytes: int
219 timestamp_nanos: int
221 log = self._log
222 update_interval_secs, sliding_window_secs = (
223 self._progress_update_intervals if self._progress_update_intervals is not None else self._get_update_intervals()
224 )
225 update_interval_nanos: int = round(update_interval_secs * 1_000_000_000)
226 sliding_window_nanos: int = round(sliding_window_secs * 1_000_000_000)
227 etas: list[ProgressReporter.TransferStat.ETA] = []
228 sent_bytes, last_status_len = 0, 0
229 num_lines, num_readables = 0, 0
230 start_time_nanos: int = time.monotonic_ns()
231 next_update_nanos: int = start_time_nanos + update_interval_nanos
232 latest_samples: deque[Sample] = deque([Sample(0, start_time_nanos)]) # sliding window w/ recent measures
233 sleep_nanos: int = 0
234 while True:
235 empty_file_name_queue: set[str] = set()
236 empty_states: list[State] = []
237 with self._lock:
238 if self._is_stopping:
239 return
240 # progress reporter thread picks up pv log files that so far aren't being tailed
241 n = len(self._file_name_queue)
242 m = len(self._file_name_set)
243 self._file_name_set.update(self._file_name_queue) # union
244 assert len(self._file_name_set) == n + m # aka assert (previous) file_name_set.isdisjoint(file_name_queue)
245 local_file_name_queue: set[str] = self._file_name_queue
246 self._file_name_queue = empty_file_name_queue # exchange buffers
247 states: list[State] = self._states
248 self._states = empty_states # exchange buffers
249 for state in states:
250 if state is State.IS_PAUSING:
251 next_update_nanos = time.monotonic_ns() + 10 * 365 * 86400 * 1_000_000_000 # infinity
252 sleep_nanos = next_update_nanos
253 else:
254 assert state is State.IS_RESETTING
255 sent_bytes, last_status_len = 0, 0
256 num_lines, num_readables = 0, 0
257 start_time_nanos = time.monotonic_ns()
258 next_update_nanos = start_time_nanos + update_interval_nanos
259 latest_samples = deque([Sample(0, start_time_nanos)]) # sliding window w/ recent measures
260 sleep_nanos = round(update_interval_nanos / 2.5)
261 for pv_log_file in local_file_name_queue:
262 try:
263 Path(pv_log_file).touch(mode=FILE_PERMISSIONS)
264 fd = open_nofollow(pv_log_file, mode="r", newline="", encoding="utf-8")
265 except FileNotFoundError: # a third party has somehow deleted the log file or directory
266 with self._lock:
267 self._file_name_set.discard(pv_log_file) # enable re-adding the file later via enqueue_pv_log_file()
268 log.warning("ProgressReporter: pv log file disappeared before initial open, skipping: %s", pv_log_file)
269 continue # skip to the next file in the queue
270 fds.append(fd)
271 eta = self.TransferStat.ETA(timestamp_nanos=0, seq_nr=-len(fds), line_tail="")
272 selector.register(fd, selectors.EVENT_READ, data=(iter(fd), self.TransferStat(bytes_in_flight=0, eta=eta)))
273 etas.append(eta)
274 readables: list[tuple[selectors.SelectorKey, int]] = selector.select(timeout=0) # 0 indicates "don't block"
275 has_line: bool = False
276 curr_time_nanos: int = time.monotonic_ns()
277 for selector_key, _ in readables: # for each file that's ready for non-blocking read
278 num_readables += 1
279 iter_fd, transfer_stat = selector_key.data
280 for line in iter_fd: # aka iter(fd)
281 sent_bytes += self._update_transfer_stat(line, transfer_stat, curr_time_nanos)
282 num_lines += 1
283 has_line = True
284 if curr_time_nanos >= next_update_nanos:
285 elapsed_nanos: int = curr_time_nanos - start_time_nanos
286 msg0, msg3 = self._format_sent_bytes(sent_bytes, elapsed_nanos) # throughput etc since replication starttime
287 msg1: str = self._format_duration(elapsed_nanos) # duration since replication start time
288 oldest: Sample = latest_samples[0] # throughput etc, over sliding window
289 _, msg2 = self._format_sent_bytes(sent_bytes - oldest.sent_bytes, curr_time_nanos - oldest.timestamp_nanos)
290 msg4: str = max(etas).line_tail if len(etas) > 0 else "" # progress bar, ETAs
291 timestamp: str = datetime.now().isoformat(sep=" ", timespec="seconds") # 2024-09-03 12:26:15
292 status_line: str = f"{timestamp} [I] zfs sent {msg0} {msg1} {msg2} {msg3} {msg4}".rstrip()
293 status_line = status_line.ljust(last_status_len) # "overwrite" trailing chars of previous status with spaces
295 # The Unix console skips back to the beginning of the console line when it sees this \r control char:
296 log.log(LOG_STDOUT, "%s", status_line, extra={"terminator": "\r"})
298 # log.log(log_trace, "\nnum_lines: %s, num_readables: %s", num_lines, num_readables)
299 last_status_len = len(status_line.rstrip())
300 next_update_nanos += update_interval_nanos
301 latest_samples.append(Sample(sent_bytes, curr_time_nanos))
302 if elapsed_nanos >= sliding_window_nanos:
303 latest_samples.popleft() # slide the sliding window containing recent measurements
304 if not self._first_status_line_emitted.is_set():
305 self._first_status_line_emitted.set()
306 elif not has_line:
307 # Avoid burning CPU busily spinning on I/O readiness as fds are almost always ready for non-blocking read
308 # even if no new pv log line has been written. Yet retain ability to wake up immediately on reporter.stop().
309 if self._sleeper.sleep(min(sleep_nanos, next_update_nanos - curr_time_nanos)):
310 self._sleeper.reset() # sleep was interrupted; ensure we can sleep normally again
311 if self._inject_error:
312 raise ValueError("Injected ProgressReporter error") # for testing only
314 def _update_transfer_stat(self, line: str, s: TransferStat, curr_time_nanos: int) -> int:
315 """Update ``s`` from one pv status line and return bytes delta."""
316 num_bytes, s.eta.timestamp_nanos, s.eta.line_tail = self._parse_pv_line(line, curr_time_nanos)
317 bytes_in_flight: int = s.bytes_in_flight
318 s.bytes_in_flight = num_bytes if line.endswith("\r") else 0 # intermediate vs. final status update of each transfer
319 return num_bytes - bytes_in_flight
321 NO_RATES_REGEX = re.compile(r".*/s\s*[)\]]?\s*") # matches until end of last pv rate, e.g. "834MiB/s]" or "834MiB/s)"
322 # time remaining --eta "ETA 00:00:39" or "ETA 2+0:00:39" or "ETA 2:0:00:39", followed by trailing --fineta timestamp ETA
323 TIME_REMAINING_ETA_REGEX = re.compile(r".*?ETA\s*((\d+)[+:])?(\d\d?):(\d\d):(\d\d).*(ETA|FIN).*")
325 @staticmethod
326 def _parse_pv_line(line: str, curr_time_nanos: int) -> tuple[int, int, str]:
327 """Parses a pv status line into transferred bytes and ETA timestamp."""
328 assert isinstance(line, str)
329 if ":" in line:
330 line = line.split(":", 1)[1].strip()
331 sent_bytes, line = _pv_size_to_bytes(line)
332 line = ProgressReporter.NO_RATES_REGEX.sub("", line.lstrip(), count=1) # strip --timer, --rate, --avg-rate
333 if match := ProgressReporter.TIME_REMAINING_ETA_REGEX.fullmatch(line): # extract pv --eta duration
334 _, days, hours, minutes, secs, _ = match.groups()
335 time_remaining_secs = (86400 * int(days) if days else 0) + int(hours) * 3600 + int(minutes) * 60 + int(secs)
336 curr_time_nanos += time_remaining_secs * 1_000_000_000 # ETA timestamp = now + time remaining duration
337 return sent_bytes, curr_time_nanos, line
338 return 0, curr_time_nanos, ""
340 @staticmethod
341 def _format_sent_bytes(num_bytes: int, duration_nanos: int) -> tuple[str, str]:
342 """Returns a human-readable byte count and rate."""
343 bytes_per_sec: int = round(1_000_000_000 * num_bytes / max(1, duration_nanos))
344 return f"{human_readable_bytes(num_bytes, precision=2)}", f"[{human_readable_bytes(bytes_per_sec, precision=2)}/s]"
346 @staticmethod
347 def _format_duration(duration_nanos: int) -> str:
348 """Formats ``duration_nanos`` as HH:MM:SS string."""
349 total_seconds: int = duration_nanos // 1_000_000_000
350 hours, remainder = divmod(total_seconds, 3600)
351 minutes, seconds = divmod(remainder, 60)
352 return f"{hours}:{minutes:02d}:{seconds:02d}"
354 def _get_update_intervals(self) -> tuple[float, float]:
355 """Extracts polling intervals from ``pv_program_opts``."""
356 parser = argparse.ArgumentParser(allow_abbrev=False)
357 parser.add_argument("--interval", "-i", type=float, default=1)
358 parser.add_argument("--average-rate-window", "-m", type=float, default=30)
359 args, _ = parser.parse_known_args(args=self._pv_program_opts)
360 interval: float = min(60 * 60, max(args.interval, 0.1))
361 return interval, min(60 * 60, max(args.average_rate_window, interval))
364def _pv_size_to_bytes(
365 size: str,
366) -> tuple[int, str]: # example inputs: "800B", "4.12 KiB", "510 MiB", "510 MB", "4Gb", "2TiB"
367 """Converts pv size string to bytes and returns remaining text."""
368 if match := _PV_SIZE_TO_BYTES_REGEX.fullmatch(size):
369 number: float = float(match.group(1).replace(",", ".").replace(_ARABIC_DECIMAL_SEPARATOR, "."))
370 i: int = "KMGTPEZYRQ".index(match.group(2)) if match.group(2) else -1
371 m: int = 1024 if match.group(3) == "i" else 1000
372 b: int = 1 if match.group(4) == "B" else 8
373 line_tail: str = match.group(5)
374 if line_tail and line_tail.startswith("/s"):
375 raise ValueError("Invalid pv_size: " + size) # stems from 'pv --rate' or 'pv --average-rate'
376 size_in_bytes: int = round(number * (m ** (i + 1)) / b)
377 return size_in_bytes, line_tail
378 else:
379 return 0, "" # skip partial or bad 'pv' log file line (pv process killed while writing?)
382def count_num_bytes_transferred_by_zfs_send(basis_pv_log_file: str) -> int:
383 """Scrapes the .pv log file(s) and sums up the 'pv --bytes' column."""
385 def parse_pv_line(line: str) -> int:
386 """Extracts byte count from a single pv log line."""
387 if ":" in line:
388 col: str = line.split(":", 1)[1].strip()
389 num_bytes, _ = _pv_size_to_bytes(col)
390 return num_bytes
391 return 0
393 total_bytes: int = 0
394 files: list[str] = [basis_pv_log_file] + glob.glob(glob.escape(basis_pv_log_file + PV_FILE_THREAD_SEPARATOR) + "[0-9]*")
395 for file in files:
396 if os.path.isfile(file):
397 try:
398 with open_nofollow(file, mode="r", newline="", encoding="utf-8") as fd:
399 line: str | None = None
400 for line in fd:
401 assert line is not None
402 if line.endswith("\r"):
403 continue # skip all but the most recent status update of each transfer
404 total_bytes += parse_pv_line(line)
405 line = None
406 if line is not None:
407 total_bytes += parse_pv_line(line) # consume last line of file w/ intermediate status update, if any
408 except FileNotFoundError:
409 pass # harmless
410 return total_bytes