use std::cell::RefCell;
use std::path::PathBuf;
use std::time::{Duration, Instant};
pub const HELD_ENV: &str = "AMONT_HOST_SLOT";
pub const DIR_ENV: &str = "AMONT_SLOT_DIR";
pub const HEAVY: &[&str] = &[
"pre-commit-clippy",
"pre-commit-go-vet",
"pre-commit-pyright",
"pre-push-run-tests-js",
"pre-push-cargo-test",
"pre-push-go-test",
"pre-push-pytest",
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Weight {
Light,
Heavy,
}
pub fn weight_of(name: &str) -> Weight {
if HEAVY.contains(&name) {
Weight::Heavy
} else {
Weight::Light
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Reason {
Disabled,
Nested,
DirUnsafe(String),
Unsupported,
TimedOut(Duration),
Io(String),
}
pub enum Acquired {
Slot(Slot),
Unqueued(Reason),
}
pub struct Slot {
_file: std::fs::File,
}
pub fn configured(settings: &crate::config::Settings) -> u64 {
let cores = std::thread::available_parallelism().map_or(1, |n| n.get() as i64);
crate::config::host_integer_or(settings, "amont.hostSlots", (cores / 4).max(1), 0..=64) as u64
}
pub fn dir() -> PathBuf {
if let Some(d) = std::env::var_os(DIR_ENV).filter(|d| !d.is_empty()) {
return PathBuf::from(d);
}
#[cfg(target_os = "linux")]
if let Some(d) = std::env::var_os("XDG_RUNTIME_DIR").filter(|d| !d.is_empty()) {
return PathBuf::from(d).join("amont-slots");
}
PathBuf::from(format!("/tmp/amont-slots-{}", platform::uid()))
}
pub fn safe_dir(dir: &std::path::Path) -> Result<(), String> {
platform::make_private(dir)?;
platform::check_private(dir)
}
pub fn acquire(n: u64, budget: Duration) -> Acquired {
if n == 0 {
return Acquired::Unqueued(Reason::Disabled);
}
if !platform::SUPPORTED {
return Acquired::Unqueued(Reason::Unsupported);
}
let dir = dir();
if let Err(why) = safe_dir(&dir) {
return Acquired::Unqueued(Reason::DirUnsafe(why));
}
let started = Instant::now();
loop {
for i in 0..n {
let path = dir.join(format!("slot-{i}"));
match platform::try_lock(&path) {
Ok(Some(file)) => return Acquired::Slot(Slot { _file: file }),
Ok(None) => {}
Err(e) => {
return Acquired::Unqueued(Reason::Io(format!("{}: {e}", path.display())))
}
}
}
if started.elapsed() >= budget {
return Acquired::Unqueued(Reason::TimedOut(started.elapsed()));
}
std::thread::sleep(Duration::from_millis(200));
}
}
#[derive(Default)]
struct Current {
name: String,
heavy: bool,
settled: bool,
held: Option<Slot>,
queued: Duration,
}
thread_local! {
static CURRENT: RefCell<Current> = RefCell::new(Current::default());
}
pub struct CheckGuard;
impl CheckGuard {
pub fn queued(&self) -> Duration {
CURRENT.with(|c| c.borrow().queued)
}
}
impl Drop for CheckGuard {
fn drop(&mut self) {
CURRENT.with(|c| *c.borrow_mut() = Current::default());
}
}
pub fn enter_check(name: &str) -> CheckGuard {
CURRENT.with(|c| {
*c.borrow_mut() = Current {
name: name.to_string(),
heavy: weight_of(name) == Weight::Heavy,
..Current::default()
}
});
CheckGuard
}
pub fn current_check() -> Option<String> {
CURRENT.with(|c| {
let c = c.borrow();
(!c.name.is_empty()).then(|| crate::short_name(&c.name).to_string())
})
}
fn nested() -> bool {
std::env::var_os(HELD_ENV).is_some_and(|v| !v.is_empty())
}
pub fn before_spawn(settings: &crate::config::Settings) -> bool {
let pending = CURRENT.with(|c| {
let c = c.borrow();
c.heavy && !c.settled
});
if pending {
let n = if nested() { 0 } else { configured(settings) };
let outcome = if nested() {
Acquired::Unqueued(Reason::Nested)
} else if n == 0 {
Acquired::Unqueued(Reason::Disabled)
} else {
let budget = match crate::hooks::common::check_timeout(settings) {
0 => 3600,
s => s,
};
let sink = crate::live::current_sink();
if let Some((stage, idx)) = &sink {
stage.queue(*idx, Some((Instant::now(), n)));
}
let started = Instant::now();
let got = acquire(n, Duration::from_secs(budget));
if let Some((stage, idx)) = &sink {
stage.queue(*idx, None);
}
let waited = started.elapsed();
CURRENT.with(|c| c.borrow_mut().queued += waited);
if matches!(got, Acquired::Slot(_)) && waited >= Duration::from_secs(1) {
crate::hooks::common::say(&format!(
" (waited {} for a host slot, {} {n})",
crate::hooks::common::human_secs(waited.as_secs()),
crate::ui::highlight("amont.hostSlots")
));
}
got
};
let held = match outcome {
Acquired::Slot(slot) => Some(slot),
Acquired::Unqueued(reason) => {
note(&reason, n);
None
}
};
CURRENT.with(|c| {
let mut c = c.borrow_mut();
c.held = held;
c.settled = true;
});
}
CURRENT.with(|c| c.borrow().held.is_some())
}
fn note(reason: &Reason, n: u64) {
match reason {
Reason::DirUnsafe(why) => crate::hooks::common::warn(&format!(
"host slots are off for this check: {why}. Remove it, or set {} to a directory \
only you own",
crate::ui::highlight(DIR_ENV)
)),
Reason::TimedOut(waited) => crate::hooks::common::warn(&format!(
"no host slot freed in {} ({} heavy checks were running on this machine, \
{} {n}); running it anyway",
crate::hooks::common::human_secs(waited.as_secs()),
n,
crate::ui::highlight("amont.hostSlots")
)),
Reason::Io(why) => {
crate::hooks::common::warn(&format!("host slots are off for this check: {why}"))
}
Reason::Disabled | Reason::Nested | Reason::Unsupported => {}
}
}
#[cfg(unix)]
mod platform {
use std::os::unix::fs::{DirBuilderExt, MetadataExt, PermissionsExt};
use std::os::unix::io::AsRawFd;
use std::path::Path;
pub const SUPPORTED: bool = true;
extern "C" {
#[link_name = "flock"]
fn libc_flock(fd: i32, op: i32) -> i32;
#[link_name = "getuid"]
fn libc_getuid() -> u32;
}
const LOCK_EX: i32 = 2;
const LOCK_NB: i32 = 4;
pub fn uid() -> u32 {
unsafe { libc_getuid() }
}
pub fn make_private(dir: &Path) -> Result<(), String> {
if dir.symlink_metadata().is_ok() {
return Ok(());
}
std::fs::DirBuilder::new()
.recursive(true)
.mode(0o700)
.create(dir)
.map_err(|e| format!("cannot create {}: {e}", dir.display()))
}
pub fn check_private(dir: &Path) -> Result<(), String> {
let meta = dir
.symlink_metadata()
.map_err(|e| format!("cannot read {}: {e}", dir.display()))?;
if meta.file_type().is_symlink() {
return Err(format!("{} is a symbolic link", dir.display()));
}
if !meta.is_dir() {
return Err(format!("{} is not a directory", dir.display()));
}
if meta.uid() != uid() {
return Err(format!("{} belongs to another user", dir.display()));
}
let mode = meta.permissions().mode() & 0o777;
if mode != 0o700 {
return Err(format!("{} has mode {mode:o}, not 700", dir.display()));
}
Ok(())
}
pub fn try_lock(path: &Path) -> std::io::Result<Option<std::fs::File>> {
let file = std::fs::OpenOptions::new()
.create(true)
.truncate(false)
.write(true)
.open(path)?;
if unsafe { libc_flock(file.as_raw_fd(), LOCK_EX | LOCK_NB) } == 0 {
return Ok(Some(file));
}
let e = std::io::Error::last_os_error();
if e.kind() == std::io::ErrorKind::WouldBlock {
Ok(None)
} else {
Err(e)
}
}
}
#[cfg(not(unix))]
mod platform {
use std::path::Path;
pub const SUPPORTED: bool = false;
pub fn uid() -> u32 {
0
}
pub fn make_private(_dir: &Path) -> Result<(), String> {
Ok(())
}
pub fn check_private(_dir: &Path) -> Result<(), String> {
Ok(())
}
pub fn try_lock(_path: &Path) -> std::io::Result<Option<std::fs::File>> {
Ok(None)
}
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
fn scratch(tag: &str) -> PathBuf {
let d = std::env::temp_dir().join(format!("amont-slots-test-{}-{tag}", std::process::id()));
let _ = std::fs::remove_dir_all(&d);
d
}
#[test]
fn only_a_private_directory_is_accepted() {
use std::os::unix::fs::PermissionsExt;
let d = scratch("private");
assert_eq!(safe_dir(&d), Ok(()));
std::fs::set_permissions(&d, std::fs::Permissions::from_mode(0o755)).unwrap();
assert!(safe_dir(&d).unwrap_err().contains("mode 755"));
let link = scratch("link");
std::os::unix::fs::symlink(&d, &link).unwrap();
assert!(safe_dir(&link).unwrap_err().contains("symbolic link"));
let _ = std::fs::remove_file(&link);
let _ = std::fs::remove_dir_all(&d);
}
#[test]
fn slots_are_exclusive_and_released_on_drop() {
let d = scratch("exclusive");
safe_dir(&d).unwrap();
let a = platform::try_lock(&d.join("slot-0"))
.expect("no error")
.expect("first");
assert!(platform::try_lock(&d.join("slot-0"))
.expect("busy is not an error")
.is_none());
drop(a);
let freed = (0..100).any(|_| {
matches!(platform::try_lock(&d.join("slot-0")), Ok(Some(_))) || {
std::thread::sleep(Duration::from_millis(20));
false
}
});
assert!(freed, "a dropped slot was not released within 2 s");
let _ = std::fs::remove_dir_all(&d);
}
#[test]
fn heavy_is_the_list() {
assert_eq!(weight_of("pre-commit-clippy"), Weight::Heavy);
assert_eq!(weight_of("pre-commit-cargo-fmt"), Weight::Light);
}
#[test]
fn every_heavy_name_is_a_builtin() {
for name in HEAVY {
assert!(
crate::registry::CHECKS.iter().any(|c| c.name == *name),
"{name} is not a built-in check"
);
}
}
#[test]
fn cargo_pins_the_host_knobs_for_tests() {
assert_eq!(std::env::var(HELD_ENV).as_deref(), Ok("held"));
assert_eq!(std::env::var("AMONT_IDLE_LOAD_SCALE").as_deref(), Ok("1"));
}
#[test]
fn a_light_check_never_takes_a_slot() {
let settings = crate::config::Settings::default();
let _g = enter_check("pre-commit-cargo-fmt");
assert!(!before_spawn(&settings));
}
#[test]
fn a_nested_heavy_check_settles_without_a_slot() {
let settings = crate::config::Settings::default();
let _g = enter_check("pre-commit-clippy");
assert_eq!(current_check().as_deref(), Some("clippy"));
assert!(!before_spawn(&settings));
assert!(CURRENT.with(|c| c.borrow().settled));
}
#[test]
fn an_unopenable_slot_is_an_error_not_busy() {
let d = scratch("missing").join("absent-subdir").join("slot-0");
assert!(platform::try_lock(&d).is_err());
}
}