everything_ipc/folder/
size.rs

1/*!
2Folder size batch lookup and cache.
3
4- Batch lookup of all folder sizes in the parent folder.
5- Thread-local cache with [`get_folder_size()`].
6- For drivers (like `C:\`), it uses Windows API directly.
7
8## References
9- [`IbDOpusExt/ViewerPlugin/DOpusExt.cpp`](https://github.com/Chaoses-Ib/IbDOpusExt/blob/421397f1f73d49b1351ec6cebdf35a74dddb9019/ViewerPlugin/DOpusExt.cpp#L40-L113)
10*/
11
12use std::{
13    cell::UnsafeCell,
14    io,
15    path::{Path, PathBuf},
16    sync::Arc,
17    time::Duration,
18};
19
20use bon::{bon, builder};
21use rapidhash::{HashMapExt, RapidHashMap as HashMap};
22use thiserror::Error;
23use tracing::{debug, info, warn};
24use widestring::U16CString;
25use windows::{Win32::Storage::FileSystem::GetDiskFreeSpaceExW, core::PCWSTR};
26
27use crate::{
28    search,
29    wm::{self, EverythingClient, RequestFlags, SearchFlags},
30};
31
32#[derive(Error, Debug)]
33pub enum Error {
34    #[error("path is relative")]
35    RelativePath,
36
37    #[error("folder not found")]
38    NotFound,
39
40    #[error(transparent)]
41    Io(#[from] io::Error),
42
43    #[error(transparent)]
44    Ipc(#[from] wm::IpcError),
45}
46
47// #[cfg(any(doc, not(feature = "drop-join-thread")))]
48thread_local! {
49    static CLIENT: UnsafeCell<FolderSizeClient> = const { UnsafeCell::new(FolderSizeClient::new()) };
50}
51
52/// Get the size of a folder.
53///
54/// Uses a thread-local [`FolderSizeClient`] for caching.
55///
56/// See [`folder::size`](super::size) for details.
57///
58/// ## Arguments
59/// - `path`: An absolute path to a folder
60/// - `timeout`: Optional timeout for IPC queries.
61///   If `None`, uses the default timeout.
62/// - `parent_max_size`: Optional mutable reference to receive the maximum size of folders in the parent folder.
63///
64///   You may also want to set `eager_get_links`.
65///
66/// - `eager_get_links`: If `true`, for folders in the parent folder with size 0,
67///   eagerly resolve symlinks/junctions and query the resolved path's size.
68///
69/// ## Returns
70/// - `Ok(u64)`: The size in bytes
71/// - `Err(Error)`: If the path is invalid
72///
73/// ## Note
74/// If `drop-join-thread` feature is enabled, you need to call
75/// [`wm::EverythingClient::shared_quit_join_thread()`]
76/// before process exit / DLL unload to avoid deadlock,
77/// or use [`FolderSizeClient`] instead.
78///
79// #[cfg(any(doc, not(feature = "drop-join-thread")))]
80#[builder]
81pub fn get_folder_size(
82    #[builder(start_fn)] path: &Path,
83    timeout: Option<Duration>,
84    parent_max_size: Option<&mut u64>,
85    #[builder(default)] eager_get_links: bool,
86) -> Result<u64, Error> {
87    CLIENT.with(|cell| {
88        let client = unsafe { &mut *cell.get() };
89        client
90            .get_folder_size(path)
91            .maybe_timeout(timeout)
92            .maybe_parent_max_size(parent_max_size)
93            .eager_get_links(eager_get_links)
94            .call()
95    })
96}
97
98/// Folder size client with parent directory cache.
99#[derive(Default)]
100pub struct FolderSizeClient {
101    everything: Option<Arc<EverythingClient>>,
102    last_parent: PathBuf,
103    /// `HashMap::new()` is not const.
104    result_map: Option<HashMap<String, u64>>,
105}
106
107#[bon]
108impl FolderSizeClient {
109    pub const fn new() -> Self {
110        Self {
111            everything: None,
112            last_parent: PathBuf::new(),
113            result_map: None,
114        }
115    }
116
117    /*
118    /// Get or create the Everything client
119    fn everything(&mut self) -> Result<&EverythingClient, wm::IpcError> {
120        // TODO: get_or_try_insert_with()
121        let everything = match self.everything.as_ref() {
122            Some(everything) => everything,
123            None => self.everything.insert(EverythingClient::shared()?),
124        };
125        Ok(everything)
126    }
127    */
128
129    /// Get the size of a folder.
130    ///
131    /// See [`folder::size`](super::size) for details.
132    ///
133    /// ## Arguments
134    /// - `path`: An absolute path to a folder
135    /// - `timeout`: Optional timeout for IPC queries.
136    ///   If `None`, uses the default timeout.
137    /// - `parent_max_size`: Optional mutable reference to receive the maximum size of folders in the parent folder.
138    ///
139    ///   You may also want to set `eager_get_links`.
140    ///
141    /// - `eager_get_links`: If `true`, for folders in the parent folder with size 0,
142    ///   eagerly resolve symlinks/junctions and query the resolved path's size.
143    ///
144    /// ## Returns
145    /// - `Ok(u64)`: The size in bytes
146    /// - `Err(Error)`: If the path is invalid
147    #[builder]
148    pub fn get_folder_size(
149        &mut self,
150        #[builder(start_fn)] path: &Path,
151        timeout: Option<Duration>,
152        parent_max_size: Option<&mut u64>,
153        #[builder(default)] eager_get_links: bool,
154    ) -> Result<u64, Error> {
155        debug_assert_eq!(search::normalize_path_ev(path), path);
156
157        // Get the parent directory
158        let parent = match path.parent() {
159            Some(p) if p.as_os_str().is_empty() => return Err(Error::RelativePath),
160            Some(p) => p,
161            None => {
162                // Handle 3-character paths (e.g., "C:\")
163                if path.as_os_str().len() == 3 {
164                    let path_u16 = U16CString::from_os_str(path).unwrap();
165                    let mut size = 0u64;
166                    if unsafe {
167                        GetDiskFreeSpaceExW(PCWSTR(path_u16.as_ptr()), None, Some(&mut size), None)
168                    }
169                    .is_ok()
170                    {
171                        return Ok(size);
172                    }
173                }
174                return Err(Error::RelativePath);
175            }
176        };
177
178        // Get or create the Everything client
179        // TODO: get_or_try_insert_with()
180        let everything = match self.everything.as_ref() {
181            Some(everything) => everything,
182            None => self.everything.insert(EverythingClient::shared()?),
183        };
184
185        // Check if we need to query for a new parent
186        let needs_query = self.last_parent != parent;
187
188        if needs_query {
189            // Clear and rebuild cache for new parent
190            self.last_parent = parent.to_path_buf();
191            self.result_map = None;
192
193            // Query Everything for files in the folder
194            let search_query = format!(r#"folder:infolder:"{}""#, parent.display());
195            let query_list = everything
196                .query_wait(&search_query)
197                .search_flags(SearchFlags::empty())
198                .request_flags(RequestFlags::FileName | RequestFlags::Size)
199                .maybe_timeout(timeout)
200                .call()
201                .inspect_err(|e| warn!(%e, ?parent, "query failed"))?;
202            info!(len = query_list.len(), "query");
203
204            // Build result map from query results
205            let mut result_map = HashMap::with_capacity(query_list.len());
206            for item in query_list.iter() {
207                if let (Some(filename), Some(mut file_size)) = (
208                    item.get_str(RequestFlags::FileName),
209                    item.get_size(RequestFlags::Size),
210                ) {
211                    let filename_str = filename.to_string_lossy();
212
213                    // For folders with size 0, check if it's a symlink/junction when eager_get_links is true
214                    if eager_get_links && file_size == 0 {
215                        let path = parent.join(&filename_str);
216                        match search::canonicalize_path_ev(&path) {
217                            // Resolve path is different, query for the resolved path size
218                            Ok(realpath) if realpath != path => {
219                                debug!(dir = filename_str, ?realpath);
220                                match everything
221                                    .get_folder_size(&realpath)
222                                    .maybe_timeout(timeout)
223                                    .call()
224                                {
225                                    Ok(size) => {
226                                        file_size = size;
227                                    }
228                                    e => warn!(?e, ?realpath, "query realpath failed"),
229                                }
230                            }
231                            Ok(_) => (),
232                            Err(e) => warn!(%e, ?path, "realpath failed"),
233                        }
234                    }
235
236                    result_map.insert(filename_str, file_size);
237                }
238            }
239            self.result_map = Some(result_map);
240        }
241
242        // Look up the file in the result map
243        // Or PathFindFileNameW()
244        let filename = path
245            .file_name()
246            .and_then(|f| f.to_str())
247            .ok_or(Error::RelativePath)?
248            .to_string();
249
250        match self.result_map.as_ref().and_then(|m| {
251            if let Some(max_size) = parent_max_size {
252                *max_size = m.values().max().copied().unwrap_or_default();
253            }
254
255            m.get(&filename).copied()
256        }) {
257            // Empty folder
258            Some(0) if eager_get_links => {
259                return Ok(0);
260            }
261            // If size is 0, try with realpath
262            Some(0) if !eager_get_links => {
263                let realpath = search::canonicalize_path_ev(path)?;
264                if realpath != path {
265                    debug!(?realpath);
266                    // TODO: pipe?
267                    let size = everything
268                        .get_folder_size(&realpath)
269                        .maybe_timeout(timeout)
270                        .call()?;
271
272                    // Cache realpath size
273                    // We got Some(0)
274                    self.result_map.as_mut().unwrap().insert(filename, size);
275
276                    return Ok(size);
277                }
278                // Empty folder
279                return Ok(0);
280            }
281            Some(size) => return Ok(size),
282            None => {
283                debug!(filename, map = ?self.result_map);
284            }
285        }
286
287        // TODO: May be new folder
288        Err(Error::NotFound)
289    }
290}
291
292#[cfg(not(feature = "drop-join-thread"))]
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    #[test_log::test]
298    #[test_log(default_log_filter = "trace")]
299    fn get_folder_size_root() {
300        let r = get_folder_size(Path::new(r"C:\")).call();
301        dbg!(&r);
302        // Just verify it returns something without panicking
303        assert!(r.unwrap() > 0);
304    }
305
306    #[test_log::test]
307    #[test_log(default_log_filter = "trace")]
308    fn get_folder_size_ev() {
309        let r = get_folder_size(Path::new(r"C:\Windows")).call();
310        dbg!(&r);
311        assert!(r.unwrap() > 0);
312
313        let r = get_folder_size(Path::new(r"C:\Users")).call();
314        dbg!(&r);
315        assert!(r.unwrap() > 0);
316    }
317
318    #[test_log::test]
319    #[test_log(default_log_filter = "trace")]
320    fn get_folder_size_ev_max() {
321        let mut max_size: u64 = 0;
322        let r = get_folder_size(Path::new(r"C:\Windows"))
323            .parent_max_size(&mut max_size)
324            .call();
325        dbg!(&r, max_size);
326        assert!(r.unwrap() > 0);
327
328        let mut max_size2: u64 = 0;
329        let r = get_folder_size(Path::new(r"C:\Users"))
330            .parent_max_size(&mut max_size2)
331            .call();
332        dbg!(&r, max_size2);
333        assert!(r.unwrap() > 0);
334
335        assert_eq!(max_size, max_size2);
336    }
337
338    #[test_log::test]
339    #[test_log(default_log_filter = "trace")]
340    fn get_folder_size_ev_realpath() {
341        // Test realpath resolution: "C:\Documents and Settings" -> "C:\Users"
342        let r = get_folder_size(Path::new(r"C:\Documents and Settings"))
343            .call()
344            .unwrap();
345        dbg!(&r);
346        assert!(r > 0);
347        let r1 = get_folder_size(Path::new(r"C:\Documents and Settings"))
348            .call()
349            .unwrap();
350        dbg!(&r1);
351        assert_eq!(r, r1);
352
353        let r2 = get_folder_size(Path::new(r"C:\Users")).call().unwrap();
354        dbg!(&r2);
355        assert!(r2 > 0);
356        assert_eq!(r, r2);
357    }
358
359    #[test_log::test]
360    #[test_log(default_log_filter = "trace")]
361    fn get_folder_size_ev_realpath_eager() {
362        // Test realpath resolution: "C:\Documents and Settings" -> "C:\Users"
363        let mut max_size: u64 = 0;
364        let r = get_folder_size(Path::new(r"C:\Documents and Settings"))
365            .parent_max_size(&mut max_size)
366            .eager_get_links(true)
367            .call()
368            .unwrap();
369        info!(r, max_size);
370        assert!(r > 0);
371        let r1 = get_folder_size(Path::new(r"C:\Documents and Settings"))
372            .parent_max_size(&mut max_size)
373            .eager_get_links(true)
374            .call()
375            .unwrap();
376        info!(r1, max_size);
377        assert_eq!(r, r1);
378
379        let r2 = get_folder_size(Path::new(r"C:\Users"))
380            .parent_max_size(&mut max_size)
381            .eager_get_links(true)
382            .call()
383            .unwrap();
384        info!(r2, max_size);
385        assert!(r2 > 0);
386        assert_eq!(r, r2);
387    }
388}
everything_ipc/folder/size.rs

everything_ipc/folder/
size.rs
