use std::sync::Arc;
use tokio::sync::Semaphore;
const MAX_PARALLEL_ENV: &str = "ZCCACHE_MAX_PARALLEL_COMPILES";
pub(super) fn resolve_pool(is_ci: bool) -> Option<Arc<Semaphore>> {
resolve_pool_with_env(is_ci, |name| std::env::var(name).ok(), num_cpus_default())
}
pub(super) fn resolve_pool_with_env<F>(
is_ci: bool,
env_lookup: F,
num_cpus: usize,
) -> Option<Arc<Semaphore>>
where
F: Fn(&str) -> Option<String>,
{
let cap = match env_lookup(MAX_PARALLEL_ENV) {
Some(raw) => match parse_override(&raw) {
CapOverride::Unlimited => return None,
CapOverride::Explicit(n) => n,
CapOverride::Invalid => {
tracing::warn!(
env = MAX_PARALLEL_ENV,
value = raw,
"invalid {MAX_PARALLEL_ENV}; falling back to default"
);
default_cap(is_ci, num_cpus)
}
},
None => default_cap(is_ci, num_cpus),
};
Some(Arc::new(Semaphore::new(cap)))
}
#[derive(Debug, PartialEq, Eq)]
enum CapOverride {
Unlimited,
Explicit(usize),
Invalid,
}
fn parse_override(raw: &str) -> CapOverride {
let trimmed = raw.trim();
if trimmed.eq_ignore_ascii_case("unlimited") || trimmed == "0" {
return CapOverride::Unlimited;
}
match trimmed.parse::<usize>() {
Ok(n) => CapOverride::Explicit(n),
Err(_) => CapOverride::Invalid,
}
}
fn default_cap(is_ci: bool, num_cpus: usize) -> usize {
if is_ci {
num_cpus.max(1)
} else {
num_cpus.saturating_sub(1).max(1)
}
}
fn num_cpus_default() -> usize {
std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1)
}
#[cfg(test)]
mod tests {
use super::*;
fn no_env(_: &str) -> Option<String> {
None
}
fn fixed_env(val: &'static str) -> impl Fn(&str) -> Option<String> {
move |name: &str| {
if name == MAX_PARALLEL_ENV {
Some(val.to_string())
} else {
None
}
}
}
#[test]
fn unset_env_interactive_defaults_to_num_cpus_minus_one() {
let pool = resolve_pool_with_env(false, no_env, 16);
let sem = pool.expect("pool should exist when not explicitly disabled");
assert_eq!(
sem.available_permits(),
15,
"interactive cap = num_cpus - 1"
);
}
#[test]
fn unset_env_interactive_floors_at_one_for_single_core() {
let pool = resolve_pool_with_env(false, no_env, 1);
let sem = pool.expect("pool should exist");
assert_eq!(sem.available_permits(), 1, "single-core box floors to 1");
}
#[test]
fn unset_env_ci_defaults_to_num_cpus() {
let pool = resolve_pool_with_env(true, no_env, 8);
let sem = pool.expect("pool should exist");
assert_eq!(sem.available_permits(), 8, "CI cap = num_cpus");
}
#[test]
fn explicit_override_wins() {
let pool = resolve_pool_with_env(false, fixed_env("4"), 64);
let sem = pool.expect("pool should exist");
assert_eq!(sem.available_permits(), 4);
}
#[test]
fn explicit_override_wins_on_ci_too() {
let pool = resolve_pool_with_env(true, fixed_env("2"), 64);
let sem = pool.expect("pool should exist");
assert_eq!(sem.available_permits(), 2);
}
#[test]
fn zero_disables_cap() {
let pool = resolve_pool_with_env(false, fixed_env("0"), 16);
assert!(pool.is_none(), "0 must opt out of the cap");
}
#[test]
fn unlimited_keyword_disables_cap() {
let pool = resolve_pool_with_env(false, fixed_env("unlimited"), 16);
assert!(pool.is_none());
let pool = resolve_pool_with_env(true, fixed_env("UNLIMITED"), 16);
assert!(pool.is_none(), "case-insensitive");
}
#[test]
fn invalid_value_falls_back_to_default() {
let pool = resolve_pool_with_env(false, fixed_env("abc"), 8);
let sem = pool.expect("should still create a pool with the default");
assert_eq!(sem.available_permits(), 7);
}
#[test]
fn parse_override_classifies() {
assert_eq!(parse_override("0"), CapOverride::Unlimited);
assert_eq!(parse_override("unlimited"), CapOverride::Unlimited);
assert_eq!(parse_override(" Unlimited "), CapOverride::Unlimited);
assert_eq!(parse_override("12"), CapOverride::Explicit(12));
assert_eq!(parse_override("not-a-number"), CapOverride::Invalid);
}
#[test]
fn default_cap_arithmetic() {
assert_eq!(default_cap(true, 16), 16);
assert_eq!(default_cap(false, 16), 15);
assert_eq!(default_cap(false, 1), 1, "floor");
assert_eq!(default_cap(true, 1), 1);
assert_eq!(default_cap(false, 0), 1, "saturate");
}
}