dvs/
utils.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use anyhow::{Result, bail};
4
5const DEFAULT_THREADS_PER_CPU: usize = 4;
6/// Maximum thread count threshold when `DVS_NUM_THREADS` is unset
7const DEFAULT_MAX_THREADS: usize = 16;
8/// Maximum thread count threshold if `DVS_NUM_THREADS` is set
9const ENV_MAX_THREADS: usize = 32;
10
11/// Global thread count override. 0 = unset (use environment variable or default).
12static NUM_THREADS: AtomicUsize = AtomicUsize::new(0);
13
14/// Set the number of threads for DVS parallel operations.
15/// Pass 0 to clear (revert to environment variable / automatic detection).
16pub fn set_num_threads(n: usize) {
17    NUM_THREADS.store(n, Ordering::Relaxed);
18}
19
20/// Returns the configured thread count, or `None` if unset.
21pub(crate) fn get_num_threads() -> Option<usize> {
22    match NUM_THREADS.load(Ordering::Relaxed) {
23        0 => None,
24        n => Some(n),
25    }
26}
27
28/// Which tier of the priority chain produced the thread count.
29#[derive(Debug, Clone, Copy)]
30enum ThreadSource {
31    Override,
32    Environment,
33    Default,
34}
35
36impl std::fmt::Display for ThreadSource {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        f.write_str(match self {
39            Self::Override => "override",
40            Self::Environment => "environment",
41            Self::Default => "default",
42        })
43    }
44}
45
46const KB: u64 = 1_024;
47const MB: u64 = 1_024 * 1_024;
48const GB: u64 = 1_024 * 1_024 * 1_024;
49const TB: u64 = 1_024 * 1_024 * 1_024 * 1_024;
50
51/// Formats a byte count into a human-readable string (e.g. "10.5 MB").
52/// Uses base-1024 divisors to match `ls -h` / `du -h` output.
53pub fn format_size(bytes: u64) -> String {
54    if bytes >= TB {
55        format!("{:.1} TB", bytes as f64 / TB as f64)
56    } else if bytes >= GB {
57        format!("{:.1} GB", bytes as f64 / GB as f64)
58    } else if bytes >= MB {
59        format!("{:.1} MB", bytes as f64 / MB as f64)
60    } else if bytes >= KB {
61        format!("{:.1} KB", bytes as f64 / KB as f64)
62    } else {
63        format!("{bytes} B")
64    }
65}
66
67/// Takes a human formatted byte size and convert it to the number of bytes
68pub fn parse_size(size: &str) -> Result<u64> {
69    let s = size.trim();
70
71    if s == "0" {
72        return Ok(0);
73    }
74
75    let num_end = s
76        .find(|c: char| !c.is_ascii_digit() && c != '.')
77        .unwrap_or(s.len());
78    let (num_str, unit_str) = s.split_at(num_end);
79    let num_str = num_str.trim();
80    let unit_str = unit_str.trim().to_ascii_uppercase();
81
82    if num_str.is_empty() {
83        bail!("Invalid size string: {s:?} (no number found)");
84    }
85
86    let num: f64 = num_str
87        .parse()
88        .map_err(|_| anyhow::anyhow!("Invalid number in size string: {num_str:?}"))?;
89
90    let multiplier = match unit_str.as_str() {
91        "" | "B" => 1u64,
92        "K" | "KB" | "KIB" => KB,
93        "M" | "MB" | "MIB" => MB,
94        "G" | "GB" | "GIB" => GB,
95        "T" | "TB" | "TIB" => TB,
96        _ => bail!("Unknown size unit: {unit_str:?}"),
97    };
98
99    Ok((num * multiplier as f64) as u64)
100}
101
102/// Creates a rayon thread pool with a thread count resolved from a 3-tier
103/// priority chain, each clamped to `work_items`:
104///
105/// 1. **Override** — `set_num_threads(n)` (capped at `ENV_MAX_THREADS`)
106/// 2. **Environment** — `DVS_NUM_THREADS` environment variable, must parse to
107///    `usize > 0` (capped at `ENV_MAX_THREADS`)
108/// 3. **Default** — `available_cpus * DEFAULT_THREADS_PER_CPU`
109///    (capped at `DEFAULT_MAX_THREADS`)
110pub fn get_threadpool(work_items: usize) -> Result<rayon::ThreadPool> {
111    debug_assert_ne!(
112        work_items, 0,
113        "the thread pool should not be instantiated when there are no work items to process"
114    );
115    let work_limit = work_items.max(1);
116    let available_cpus = std::thread::available_parallelism()
117        .map(|n| n.get())
118        .unwrap_or(1)
119        .max(1);
120
121    let override_threads = get_num_threads();
122    let environment_threads = std::env::var("DVS_NUM_THREADS")
123        .ok()
124        .and_then(|v| v.parse::<usize>().ok())
125        .filter(|&n| n > 0);
126
127    let num_threads = match (override_threads, environment_threads) {
128        (Some(n), _) | (None, Some(n)) => n.min(ENV_MAX_THREADS).min(work_limit),
129        (None, None) => available_cpus
130            .saturating_mul(DEFAULT_THREADS_PER_CPU)
131            .min(DEFAULT_MAX_THREADS)
132            .min(work_limit),
133    };
134
135    let source = if override_threads.is_some() {
136        ThreadSource::Override
137    } else if environment_threads.is_some() {
138        ThreadSource::Environment
139    } else {
140        ThreadSource::Default
141    };
142
143    log::debug!("thread pool: {num_threads} threads (source: {source}, work_items={work_items})",);
144
145    let pool = rayon::ThreadPoolBuilder::new()
146        .num_threads(num_threads)
147        .build()?;
148    Ok(pool)
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    #[test]
155    fn parse_size_basic() {
156        assert_eq!(parse_size("0").unwrap(), 0);
157        assert_eq!(parse_size("500MB").unwrap(), 500 * MB);
158        assert_eq!(parse_size("500 MB").unwrap(), 500 * MB);
159        assert_eq!(parse_size("1GB").unwrap(), GB);
160        assert_eq!(parse_size("1 gb").unwrap(), GB);
161        assert_eq!(parse_size("2TB").unwrap(), 2 * TB);
162        assert_eq!(parse_size("100KB").unwrap(), 100 * KB);
163        assert_eq!(parse_size("1024B").unwrap(), 1024);
164        assert_eq!(parse_size("1.5GB").unwrap(), (1.5 * GB as f64) as u64);
165        assert!(parse_size("").is_err());
166        assert!(parse_size("MB").is_err());
167        assert!(parse_size("500XB").is_err());
168    }
169}
dvs/utils.rs

dvs/
utils.rs
