dvs/
utils.rs

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

dvs/
utils.rs
