From 6a9e7da8a3a382e52cc31ec93ebf3f7c2788f900 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=89=E5=92=B2=E9=9B=85=20=C2=B7=20Misaki=20Masa?= Date: Mon, 28 Apr 2025 19:51:12 +0800 Subject: [PATCH] perf: double directory size calculation speed (#2683) --- yazi-config/src/keymap/key.rs | 2 +- yazi-fs/src/calculator.rs | 96 +++++++++++++++++++++++++++ yazi-fs/src/fns.rs | 26 +------- yazi-fs/src/lib.rs | 2 +- yazi-scheduler/src/file/file.rs | 8 +-- yazi-scheduler/src/prework/prework.rs | 4 +- 6 files changed, 105 insertions(+), 33 deletions(-) create mode 100644 yazi-fs/src/calculator.rs diff --git a/yazi-config/src/keymap/key.rs b/yazi-config/src/keymap/key.rs index 5aeb9872..62e2e8a9 100644 --- a/yazi-config/src/keymap/key.rs +++ b/yazi-config/src/keymap/key.rs @@ -198,6 +198,6 @@ impl Display for Key { _ => "Unknown", }; - write!(f, "{}>", code) + write!(f, "{code}>") } } diff --git a/yazi-fs/src/calculator.rs b/yazi-fs/src/calculator.rs new file mode 100644 index 00000000..91ecb8ab --- /dev/null +++ b/yazi-fs/src/calculator.rs @@ -0,0 +1,96 @@ +use std::{collections::VecDeque, fs::ReadDir, future::poll_fn, mem, path::{Path, PathBuf}, pin::Pin, task::{Poll, ready}, time::{Duration, Instant}}; + +use tokio::task::JoinHandle; +use yazi_shared::Either; + +type Task = Either; + +pub enum SizeCalculator { + Idle((VecDeque, Option)), + Pending(JoinHandle<(VecDeque, Option)>), +} + +impl SizeCalculator { + pub async fn new(path: impl AsRef) -> std::io::Result { + let p = path.as_ref().to_owned(); + tokio::task::spawn_blocking(|| { + let mut buf = VecDeque::from([Either::Right(std::fs::read_dir(p)?)]); + let size = Self::next_chunk(&mut buf); + Ok(Self::Idle((buf, size))) + }) + .await? + } + + pub async fn total(path: impl AsRef) -> std::io::Result { + let mut it = Self::new(path).await?; + let mut total = 0; + while let Some(n) = it.next().await? { + total += n; + } + Ok(total) + } + + pub async fn next(&mut self) -> std::io::Result> { + poll_fn(|cx| { + loop { + match self { + Self::Idle((buf, size)) => { + if let Some(s) = size.take() { + return Poll::Ready(Ok(Some(s))); + } else if buf.is_empty() { + return Poll::Ready(Ok(None)); + } + + let mut buf = mem::take(buf); + *self = Self::Pending(tokio::task::spawn_blocking(move || { + let size = Self::next_chunk(&mut buf); + (buf, size) + })); + } + Self::Pending(handle) => { + *self = Self::Idle(ready!(Pin::new(handle).poll(cx))?); + } + } + } + }) + .await + } + + fn next_chunk(buf: &mut VecDeque>) -> Option { + let (mut i, mut size, now) = (0, 0, Instant::now()); + macro_rules! pop_and_continue { + () => {{ + buf.pop_front(); + if buf.is_empty() { + return Some(size); + } + continue; + }}; + } + + while i < 5000 && now.elapsed() < Duration::from_millis(50) { + i += 1; + let front = buf.front_mut()?; + + if let Either::Left(p) = front { + *front = match std::fs::read_dir(p) { + Ok(it) => Either::Right(it), + Err(_) => pop_and_continue!(), + }; + } + + let Some(next) = front.right_mut()?.next() else { + pop_and_continue!(); + }; + + let Ok(ent) = next else { continue }; + let Ok(ft) = ent.file_type() else { continue }; + if ft.is_dir() { + buf.push_back(Either::Left(ent.path())); + } else { + size += ent.metadata().map(|m| m.len()).unwrap_or(0); + } + } + Some(size) + } +} diff --git a/yazi-fs/src/fns.rs b/yazi-fs/src/fns.rs index cf73a577..91b59b4e 100644 --- a/yazi-fs/src/fns.rs +++ b/yazi-fs/src/fns.rs @@ -1,4 +1,4 @@ -use std::{borrow::Cow, collections::{HashMap, HashSet, VecDeque}, ffi::{OsStr, OsString}, path::{Path, PathBuf}}; +use std::{borrow::Cow, collections::{HashMap, HashSet}, ffi::{OsStr, OsString}, path::{Path, PathBuf}}; use anyhow::{Result, bail}; use tokio::{fs, io::{self, AsyncWriteExt}, select, sync::{mpsc, oneshot}, time}; @@ -181,30 +181,6 @@ pub async fn realname_unchecked<'a>( } } -pub async fn calculate_size(path: &Path) -> u64 { - let mut total = 0; - let mut stack = VecDeque::from([path.to_path_buf()]); - while let Some(path) = stack.pop_front() { - let Ok(meta) = fs::symlink_metadata(&path).await else { continue }; - if !meta.is_dir() { - total += meta.len(); - continue; - } - - let Ok(mut it) = fs::read_dir(path).await else { continue }; - while let Ok(Some(entry)) = it.next_entry().await { - let Ok(meta) = entry.metadata().await else { continue }; - - if meta.is_dir() { - stack.push_back(entry.path()); - } else { - total += meta.len(); - } - } - } - total -} - pub fn copy_with_progress( from: &Path, to: &Path, diff --git a/yazi-fs/src/lib.rs b/yazi-fs/src/lib.rs index 22488e75..7df8a8ae 100644 --- a/yazi-fs/src/lib.rs +++ b/yazi-fs/src/lib.rs @@ -2,7 +2,7 @@ yazi_macro::mod_pub!(cha mounts); -yazi_macro::mod_flat!(cwd file files filter fns op path sorter sorting stage step xdg); +yazi_macro::mod_flat!(calculator cwd file files filter fns op path sorter sorting stage step xdg); pub fn init() { CWD.init(<_>::default()); diff --git a/yazi-scheduler/src/file/file.rs b/yazi-scheduler/src/file/file.rs index 61595663..358998c9 100644 --- a/yazi-scheduler/src/file/file.rs +++ b/yazi-scheduler/src/file/file.rs @@ -4,7 +4,7 @@ use anyhow::{Result, anyhow}; use tokio::{fs::{self, DirEntry}, io::{self, ErrorKind::{AlreadyExists, NotFound}}, sync::mpsc}; use tracing::warn; use yazi_config::YAZI; -use yazi_fs::{calculate_size, cha::Cha, copy_with_progress, maybe_exists, ok_or_not_found, path_relative_to, skip_path}; +use yazi_fs::{SizeCalculator, cha::Cha, copy_with_progress, maybe_exists, ok_or_not_found, path_relative_to, skip_path}; use yazi_shared::url::Url; use super::{FileOp, FileOpDelete, FileOpHardlink, FileOpLink, FileOpPaste, FileOpTrash}; @@ -49,7 +49,7 @@ impl File { && matches!(e.raw_os_error(), Some(1) | Some(93)) => { task.retry += 1; - self.log(task.id, format!("Paste task retry: {:?}", task))?; + self.log(task.id, format!("Paste task retry: {task:?}"))?; self.queue(FileOp::Paste(task), LOW).await?; return Ok(()); } @@ -122,7 +122,7 @@ impl File { FileOp::Delete(task) => { if let Err(e) = fs::remove_file(&task.target).await { if e.kind() != NotFound && maybe_exists(&task.target).await { - self.fail(task.id, format!("Delete task failed: {:?}, {e}", task))?; + self.fail(task.id, format!("Delete task failed: {task:?}, {e}"))?; Err(e)? } } @@ -318,7 +318,7 @@ impl File { pub async fn trash(&self, mut task: FileOpTrash) -> Result<()> { let id = task.id; - task.length = calculate_size(&task.target).await; + task.length = SizeCalculator::total(&task.target).await?; self.prog.send(TaskProg::New(id, task.length))?; self.queue(FileOp::Trash(task), LOW).await?; diff --git a/yazi-scheduler/src/prework/prework.rs b/yazi-scheduler/src/prework/prework.rs index 7584c690..9a978371 100644 --- a/yazi-scheduler/src/prework/prework.rs +++ b/yazi-scheduler/src/prework/prework.rs @@ -6,7 +6,7 @@ use parking_lot::{Mutex, RwLock}; use tokio::sync::mpsc; use tracing::error; use yazi_config::Priority; -use yazi_fs::{FilesOp, calculate_size}; +use yazi_fs::{FilesOp, SizeCalculator}; use yazi_plugin::isolate; use yazi_shared::{event::CmdCow, url::Url}; @@ -73,7 +73,7 @@ impl Prework { self.prog.send(TaskProg::Adv(task.id, 1, 0))?; } PreworkOp::Size(task) => { - let length = calculate_size(&task.target).await; + let length = SizeCalculator::total(&task.target).await?; task.throttle.done((task.target, length), |buf| { { let mut loading = self.size_loading.write();