diff --git a/tests/test_transfer.py b/tests/test_transfer.py index c056c7a..637b734 100644 --- a/tests/test_transfer.py +++ b/tests/test_transfer.py @@ -270,3 +270,15 @@ def test_speed_stats_empty(): def test_speed_stats_single(): cur, peak, avg = transfer.speed_stats([8 * 1024 * 1024]) assert (cur, peak, avg) == (8 * 1024 * 1024, 8 * 1024 * 1024, 8 * 1024 * 1024) + +# ---------- 目录大小轮询 (rsync 输出缓冲替代方案) ---------- + +def test_dir_size_returns_bytes(tmp_path): + """du 目录返回字节数""" + (tmp_path / "a.bin").write_bytes(b"x" * 1000) + (tmp_path / "sub").mkdir() + (tmp_path / "sub" / "b.bin").write_bytes(b"y" * 500) + assert transfer.dir_size(str(tmp_path)) == 1500 + +def test_dir_size_missing_returns_zero(): + assert transfer.dir_size("/nonexistent/path/xyz") == 0 diff --git a/transfer.py b/transfer.py index 572a502..79d764c 100644 --- a/transfer.py +++ b/transfer.py @@ -55,6 +55,18 @@ def human_speed(bps: float) -> str: return f"{human_size(bps)}/s" +def dir_size(path: str) -> int: + """du 目录字节数 (轮询进度用, 绕过 rsync 输出缓冲)""" + if not os.path.isdir(path): + return 0 + try: + proc = subprocess.run(["du", "-sb", path], + capture_output=True, text=True, timeout=30) + return int(proc.stdout.split()[0]) + except (ValueError, OSError, subprocess.TimeoutExpired): + return 0 + + def speed_stats(history) -> tuple[float, float, float]: """速度统计: (当前, 峰值, 平均) 字节/秒""" if not history: @@ -547,40 +559,41 @@ class TaskRunner: # -- 执行 def run(self) -> int: - """同步执行传输, 返回退出码""" + """同步执行传输, 返回退出码 (进度由目录轮询驱动, 绕过后端输出缓冲)""" cmd = self.build_command() self.lm.open_session(f"{self.cfg.backend}_{os.path.basename(self.cfg.source.rstrip('/'))}") self.lm.write(f"cmd: {' '.join(cmd)}") self.lm.write(f"mode: {self.cfg.backend} bwlimit={self.cfg.bwlimit or 'off'} " f"resume={self.cfg.resume} checksum={self.cfg.checksum}") + # 源总大小 (进度分母) + self.total_bytes = dir_size(self.cfg.source.rstrip("/")) + self.lm.write(f"source size: {human_size(self.total_bytes)}") + # stdout 丢弃 (进度帧缓冲不可靠), stderr 单独读错误 self.proc = subprocess.Popen( - cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, - text=True, bufsize=0, + cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, + text=True, bufsize=1, ) - # rsync progress2 / rclone --progress 用 \r 刷新帧(不换行), - # 故按 \r 与 \n 均作分隔符分帧解析, 避免进度只取到末帧 - assert self.proc.stdout is not None - buf = "" - while True: - chunk = self.proc.stdout.read(8192) - if not chunk: + # 轮询线程: 每秒统计目标目录大小驱动进度 + self._stop.clear() + poll = threading.Thread(target=self._poll_loop, daemon=True) + poll.start() + assert self.proc.stderr is not None + for line in self.proc.stderr: + if self._stop.is_set(): break - buf += chunk - while True: - m = re.search(r"[\r\n]", buf) - if not m: - break - seg, buf = buf[: m.start()], buf[m.start() + 1:] - if self._stop.is_set(): - self.proc.kill() - break - self._dispatch(seg.rstrip()) - if buf.strip() and not self._stop.is_set(): - self._dispatch(buf.rstrip()) - if self._stop.is_set(): - self.proc.kill() + line = line.rstrip() + if line: + self._error(line) rc = self.proc.wait() - self.lm.write(f"exit code: {rc} (进度帧 {self.stats_frames} 次)") + self._stop.set() + poll.join(timeout=3) + # 最终统计 + final = dir_size(self.cfg.dest.rstrip("/")) + self.stats.bytes_done = final + if self.total_bytes: + self.stats.pct = min(100.0, final / self.total_bytes * 100) + self.lm.write(f"exit code: {rc} (帧 {self.stats_frames} 次, " + f"目标 {human_size(final)}/{human_size(self.total_bytes)})") removed = self.lm.prune() if removed: self.lm.write(f"pruned old logs: {len(removed)}") @@ -588,6 +601,30 @@ class TaskRunner: self.on_done(rc) return rc + def _poll_loop(self) -> None: + """每秒轮询目标目录大小, 更新进度 (独立线程)""" + dst = self.cfg.dest.rstrip("/") + last = dir_size(dst) + last_t = time.time() + poll = 0 + while not self._stop.is_set(): + time.sleep(1) + done = dir_size(dst) + now = time.time() + speed = (done - last) / max(0.1, now - last_t) + self.stats.bytes_done = done + self.stats.bytes_total = self.total_bytes + self.stats.speed_bps = max(0, speed) + if self.total_bytes: + self.stats.pct = min(100.0, done / self.total_bytes * 100) + last, last_t = done, now + self._emit_stats() + poll += 1 + # 每 5 秒一条进度快照到日志栏 + if self.on_log and poll % 5 == 0: + self.on_log(f"已传 {human_size(done)} / {human_size(self.total_bytes)} " + f"({self.stats.pct:.0f}%) {human_speed(self.stats.speed_bps)}") + def stop(self) -> None: self._stop.set() if self.proc and self.proc.poll() is None: