| 1 | //! Cross-process admission for expensive local commands. |
| 2 | //! |
| 3 | //! Fleet and Workflow workers execute in separate Codewhale processes, so an |
| 4 | //! in-process semaphore cannot protect the host. Heavy shell commands instead |
| 5 | //! take one of a small number of filesystem-backed permits under |
| 6 | //! `CODEWHALE_HOME`. The default of two permits is deliberately conservative |
| 7 | //! for the 36 GiB laptop class from #4864. |
| 8 | |
| 9 | use std::fs::{File, OpenOptions}; |
| 10 | use std::io; |
| 11 | use std::path::{Path, PathBuf}; |
| 12 | use std::time::{Duration, Instant}; |
| 13 | |
| 14 | use anyhow::{Context, Result, anyhow}; |
| 15 | use fd_lock::{RwLock, RwLockWriteGuard}; |
| 16 | use tokio_util::sync::CancellationToken; |
| 17 | |
| 18 | pub(crate) const DEFAULT_HEAVY_COMMAND_LIMIT: usize = 2; |
| 19 | const MAX_HEAVY_COMMAND_LIMIT: usize = 16; |
| 20 | const ADMISSION_POLL_INTERVAL: Duration = Duration::from_millis(50); |
| 21 | |
| 22 | /// When the host free-RAM fraction drops to/below these thresholds the |
| 23 | /// effective heavy-command admission limit tightens so a saturated host stops |
| 24 | /// admitting new link graphs (#4864 req 7). Values are deliberately generous |
| 25 | /// because the measurement is advisory, not authoritative. |
| 26 | const CONSTRAINED_FREE_FRACTION: f64 = 0.30; |
| 27 | const CRITICAL_FREE_FRACTION: f64 = 0.15; |
| 28 | |
| 29 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 30 | pub(crate) enum CommandExpense { |
| 31 | Normal, |
| 32 | Heavy, |
| 33 | } |
| 34 | |
| 35 | /// Measured host memory pressure used to tighten heavy-command admission. |
| 36 | /// |
| 37 | /// `Unknown` means "could not be measured"; admission then fails open (uses the |
| 38 | /// configured limit) rather than risk blocking on an unmeasurable host. This |
| 39 | /// keeps the gate safe on any CI runner where the probe is unavailable. |
| 40 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 41 | pub(crate) enum MemoryPressure { |
| 42 | Unknown, |
| 43 | Nominal, |
| 44 | Constrained, |
| 45 | Critical, |
| 46 | } |
| 47 | |
| 48 | /// Pluggable memory probe so the admission policy is unit-testable without |
| 49 | /// having to drive the host into real memory pressure. |
| 50 | pub(crate) trait MemoryProbe: Send + Sync { |
| 51 | /// Free-RAM fraction in `0.0..=1.0`, or `None` when it cannot be measured. |
| 52 | fn free_fraction(&self) -> Option<f64>; |
| 53 | } |
| 54 | |
| 55 | struct HostMemoryProbe; |
| 56 | |
| 57 | impl MemoryProbe for HostMemoryProbe { |
| 58 | fn free_fraction(&self) -> Option<f64> { |
| 59 | host_memory_free_fraction() |
| 60 | } |
| 61 | } |
| 62 | |
| 63 | fn classify_memory_pressure(free_fraction: Option<f64>) -> MemoryPressure { |
| 64 | match free_fraction { |
| 65 | None => MemoryPressure::Unknown, |
| 66 | Some(fraction) if fraction <= CRITICAL_FREE_FRACTION => MemoryPressure::Critical, |
| 67 | Some(fraction) if fraction <= CONSTRAINED_FREE_FRACTION => MemoryPressure::Constrained, |
| 68 | Some(_) => MemoryPressure::Nominal, |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | /// Effective admission limit after applying host memory pressure. `Critical` |
| 73 | /// yields zero so queued heavy commands wait for the host to recover instead of |
| 74 | /// snowballing; `Constrained` halves the budget (never below one). |
| 75 | fn effective_admission_limit(configured: usize, pressure: MemoryPressure) -> usize { |
| 76 | match pressure { |
| 77 | MemoryPressure::Critical => 0, |
| 78 | MemoryPressure::Constrained => configured.div_ceil(2).max(1), |
| 79 | MemoryPressure::Nominal | MemoryPressure::Unknown => configured, |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | #[derive(Debug)] |
| 84 | struct HeavyPermitSlot { |
| 85 | _guard: RwLockWriteGuard<'static, File>, |
| 86 | // The guard borrows the lock. Keeping the boxed lock here makes that |
| 87 | // allocation outlive the guard; field drop order is guard, then lock. |
| 88 | _lock: Box<RwLock<File>>, |
| 89 | } |
| 90 | |
| 91 | /// A held cross-process heavy-command permit. |
| 92 | #[derive(Debug)] |
| 93 | pub(crate) struct HeavyCommandPermit { |
| 94 | _slot: HeavyPermitSlot, |
| 95 | queued_for: Duration, |
| 96 | limit: usize, |
| 97 | memory_pressure: MemoryPressure, |
| 98 | } |
| 99 | |
| 100 | impl HeavyCommandPermit { |
| 101 | pub(crate) fn queued_for(&self) -> Duration { |
| 102 | self.queued_for |
| 103 | } |
| 104 | |
| 105 | pub(crate) fn limit(&self) -> usize { |
| 106 | self.limit |
| 107 | } |
| 108 | |
| 109 | pub(crate) fn memory_pressure(&self) -> MemoryPressure { |
| 110 | self.memory_pressure |
| 111 | } |
| 112 | } |
| 113 | |
| 114 | pub(crate) fn infer_command_expense(command: &str) -> CommandExpense { |
| 115 | let heavy = command |
| 116 | .split(['\n', '\r', ';', '|', '&']) |
| 117 | .any(segment_is_heavy); |
| 118 | |
| 119 | if heavy { |
| 120 | CommandExpense::Heavy |
| 121 | } else { |
| 122 | CommandExpense::Normal |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | /// Commands that run another command from their arguments. The classifier |
| 127 | /// looks through them (#6560 D01-11): `sudo cargo build`, `timeout 600 |
| 128 | /// cargo test`, `xargs cargo check` and `bash -c "cargo build"` compile |
| 129 | /// exactly like the bare command and need the same permit. This is resource |
| 130 | /// admission, not a security boundary: shell policy still judges the whole |
| 131 | /// command. |
| 132 | const COMMAND_WRAPPERS: &[&str] = &[ |
| 133 | "sudo", |
| 134 | "doas", |
| 135 | "env", |
| 136 | "nice", |
| 137 | "nohup", |
| 138 | "time", |
| 139 | "timeout", |
| 140 | "gtimeout", |
| 141 | "xargs", |
| 142 | "command", |
| 143 | "exec", |
| 144 | "stdbuf", |
| 145 | "ionice", |
| 146 | "caffeinate", |
| 147 | "watch", |
| 148 | "eval", |
| 149 | "bash", |
| 150 | "sh", |
| 151 | "zsh", |
| 152 | "dash", |
| 153 | ]; |
| 154 | |
| 155 | fn segment_is_heavy(segment: &str) -> bool { |
| 156 | let tokens: Vec<String> = segment |
| 157 | .split_whitespace() |
| 158 | .map(|token| { |
| 159 | token |
| 160 | .trim_matches(['"', '\'', '(', ')', '{', '}', '`', '$']) |
| 161 | .to_string() |
| 162 | }) |
| 163 | .collect(); |
| 164 | let stem = |token: &str| { |
| 165 | Path::new(token) |
| 166 | .file_stem() |
| 167 | .and_then(|name| name.to_str()) |
| 168 | .unwrap_or_default() |
| 169 | .to_ascii_lowercase() |
| 170 | }; |
| 171 | let Some(mut index) = tokens |
| 172 | .iter() |
| 173 | .position(|token| !token.is_empty() && !token.contains('=') && token != "env") |
| 174 | else { |
| 175 | return false; |
| 176 | }; |
| 177 | if COMMAND_WRAPPERS.contains(&stem(&tokens[index]).as_str()) { |
| 178 | // Wrapper options and their values (`-u user`, `600`, `-n1`) sit |
| 179 | // between the wrapper and the wrapped command; take the first |
| 180 | // compiler that follows. |
| 181 | let Some(offset) = tokens[index + 1..] |
| 182 | .iter() |
| 183 | .position(|token| matches!(stem(token).as_str(), "cargo" | "rustc")) |
| 184 | else { |
| 185 | return false; |
| 186 | }; |
| 187 | index += 1 + offset; |
| 188 | } |
| 189 | let executable = stem(&tokens[index]); |
| 190 | if !matches!(executable.as_str(), "cargo" | "rustc") { |
| 191 | return false; |
| 192 | } |
| 193 | if executable == "rustc" { |
| 194 | return true; |
| 195 | } |
| 196 | tokens[index + 1..] |
| 197 | .iter() |
| 198 | .map(|arg| arg.trim().to_ascii_lowercase()) |
| 199 | .find(|arg| !arg.is_empty() && !arg.starts_with('-') && !arg.contains('=')) |
| 200 | .is_some_and(|subcommand| { |
| 201 | matches!( |
| 202 | subcommand.as_str(), |
| 203 | "build" | "test" | "check" | "clippy" | "rustc" |
| 204 | ) |
| 205 | }) |
| 206 | } |
| 207 | |
| 208 | pub(crate) async fn acquire_heavy_command_permit( |
| 209 | command: &str, |
| 210 | cancel: Option<&CancellationToken>, |
| 211 | ) -> Result<Option<HeavyCommandPermit>> { |
| 212 | if infer_command_expense(command) == CommandExpense::Normal { |
| 213 | return Ok(None); |
| 214 | } |
| 215 | |
| 216 | let limit = configured_heavy_command_limit(); |
| 217 | let root = admission_root(); |
| 218 | let probe = HostMemoryProbe; |
| 219 | acquire_heavy_command_permit_at(&root, limit, cancel, &probe) |
| 220 | .await |
| 221 | .map(Some) |
| 222 | } |
| 223 | |
| 224 | async fn acquire_heavy_command_permit_at( |
| 225 | root: &Path, |
| 226 | limit: usize, |
| 227 | cancel: Option<&CancellationToken>, |
| 228 | probe: &dyn MemoryProbe, |
| 229 | ) -> Result<HeavyCommandPermit> { |
| 230 | tokio::fs::create_dir_all(root) |
| 231 | .await |
| 232 | .with_context(|| format!("creating resource admission directory {}", root.display()))?; |
| 233 | let started = Instant::now(); |
| 234 | |
| 235 | loop { |
| 236 | if cancel.is_some_and(|token| token.is_cancelled()) { |
| 237 | return Err(anyhow!( |
| 238 | "heavy command canceled while queued for resource admission" |
| 239 | )); |
| 240 | } |
| 241 | // Re-measure each iteration: under memory pressure the effective limit |
| 242 | // tightens so a saturated host stops admitting new heavy link graphs |
| 243 | // (#4864 req 7). Critical pressure yields zero slots, so the command |
| 244 | // waits for recovery instead of snowballing. |
| 245 | let pressure = classify_memory_pressure(probe.free_fraction()); |
| 246 | let effective = effective_admission_limit(limit, pressure); |
| 247 | for slot in 0..effective { |
| 248 | let path = root.join(format!("heavy-{slot}.lock")); |
| 249 | match try_lock_slot(&path) { |
| 250 | Ok(Some(slot)) => { |
| 251 | return Ok(HeavyCommandPermit { |
| 252 | _slot: slot, |
| 253 | queued_for: started.elapsed(), |
| 254 | limit, |
| 255 | memory_pressure: pressure, |
| 256 | }); |
| 257 | } |
| 258 | Ok(None) => {} |
| 259 | Err(error) => { |
| 260 | return Err(error).with_context(|| { |
| 261 | format!("acquiring heavy command permit {}", path.display()) |
| 262 | }); |
| 263 | } |
| 264 | } |
| 265 | } |
| 266 | tokio::time::sleep(ADMISSION_POLL_INTERVAL).await; |
| 267 | } |
| 268 | } |
| 269 | |
| 270 | fn try_lock_slot(path: &Path) -> io::Result<Option<HeavyPermitSlot>> { |
| 271 | let file = OpenOptions::new() |
| 272 | .create(true) |
| 273 | .read(true) |
| 274 | .write(true) |
| 275 | .truncate(false) |
| 276 | .open(path)?; |
| 277 | let lock = Box::new(RwLock::new(file)); |
| 278 | let lock_ptr = Box::into_raw(lock); |
| 279 | // SAFETY: `lock_ptr` remains allocated in `HeavyPermitSlot::_lock` for the |
| 280 | // lifetime of `_guard`, and the guard is dropped before that box. |
| 281 | let guard = match unsafe { (&mut *lock_ptr).try_write() } { |
| 282 | Ok(guard) => guard, |
| 283 | Err(error) if error.kind() == io::ErrorKind::WouldBlock => { |
| 284 | // SAFETY: no guard was created, so reclaim the allocation now. |
| 285 | unsafe { drop(Box::from_raw(lock_ptr)) }; |
| 286 | return Ok(None); |
| 287 | } |
| 288 | Err(error) => { |
| 289 | // SAFETY: no guard was created, so reclaim the allocation now. |
| 290 | unsafe { drop(Box::from_raw(lock_ptr)) }; |
| 291 | return Err(error); |
| 292 | } |
| 293 | }; |
| 294 | // SAFETY: the allocation is owned exactly once by this box and is stable on |
| 295 | // the heap even if `HeavyPermitSlot` moves. |
| 296 | let lock = unsafe { Box::from_raw(lock_ptr) }; |
| 297 | // SAFETY: the boxed lock remains alive until after `_guard` is dropped. |
| 298 | let guard = unsafe { |
| 299 | std::mem::transmute::<RwLockWriteGuard<'_, File>, RwLockWriteGuard<'static, File>>(guard) |
| 300 | }; |
| 301 | Ok(Some(HeavyPermitSlot { |
| 302 | _guard: guard, |
| 303 | _lock: lock, |
| 304 | })) |
| 305 | } |
| 306 | |
| 307 | fn configured_heavy_command_limit() -> usize { |
| 308 | std::env::var("CODEWHALE_HEAVY_COMMAND_LIMIT") |
| 309 | .ok() |
| 310 | .and_then(|value| value.trim().parse::<usize>().ok()) |
| 311 | .filter(|limit| *limit > 0) |
| 312 | .unwrap_or(DEFAULT_HEAVY_COMMAND_LIMIT) |
| 313 | .min(MAX_HEAVY_COMMAND_LIMIT) |
| 314 | } |
| 315 | |
| 316 | fn admission_root() -> PathBuf { |
| 317 | if let Some(home) = codewhale_paths::codewhale_home_override().ok().flatten() { |
| 318 | return home.join("resource-admission"); |
| 319 | } |
| 320 | if let Some(home) = codewhale_paths::user_home() { |
| 321 | return home.join(".codewhale").join("resource-admission"); |
| 322 | } |
| 323 | std::env::temp_dir().join("codewhale-resource-admission") |
| 324 | } |
| 325 | |
| 326 | /// Host free-RAM fraction in `0.0..=1.0`, or `None` when it cannot be measured. |
| 327 | /// |
| 328 | /// Each implementation shells out to a standard, always-present tool so no new |
| 329 | /// crate or build-feature dependency is introduced, and every error path returns |
| 330 | /// `None` so admission fails open (never blocks on an unmeasurable host). This |
| 331 | /// keeps the gate safe on any CI runner, while protecting the macOS dogfood host |
| 332 | /// and Linux/Windows machines where the tool exists. |
| 333 | #[cfg(target_os = "linux")] |
| 334 | fn host_memory_free_fraction() -> Option<f64> { |
| 335 | let meminfo = std::fs::read_to_string("/proc/meminfo").ok()?; |
| 336 | let total = parse_meminfo_kb(&meminfo, "MemTotal:")?; |
| 337 | let available = parse_meminfo_kb(&meminfo, "MemAvailable:")?; |
| 338 | (total > 0).then(|| (available as f64 / total as f64).clamp(0.0, 1.0)) |
| 339 | } |
| 340 | |
| 341 | #[cfg(target_os = "linux")] |
| 342 | fn parse_meminfo_kb(meminfo: &str, prefix: &str) -> Option<u64> { |
| 343 | meminfo |
| 344 | .lines() |
| 345 | .find(|line| line.starts_with(prefix)) |
| 346 | .and_then(|line| line.split_whitespace().nth(1)) |
| 347 | .and_then(|value| value.parse::<u64>().ok()) |
| 348 | } |
| 349 | |
| 350 | #[cfg(target_os = "macos")] |
| 351 | fn host_memory_free_fraction() -> Option<f64> { |
| 352 | let total = run_capture("/usr/sbin/sysctl", &["-n", "hw.memsize"]) |
| 353 | .and_then(|bytes| bytes.trim().parse::<u64>().ok())?; |
| 354 | let page_size = run_capture("/usr/bin/pagesize", &[])? |
| 355 | .trim() |
| 356 | .parse::<u64>() |
| 357 | .ok()?; |
| 358 | let stats = run_capture("/usr/bin/vm_stat", &[])?; |
| 359 | let free_pages = memory_pages_from_vm_stat(&stats, &["Pages free:", "Pages inactive:"])?; |
| 360 | let free_bytes = free_pages.checked_mul(page_size)?; |
| 361 | (total > 0).then(|| (free_bytes as f64 / total as f64).clamp(0.0, 1.0)) |
| 362 | } |
| 363 | |
| 364 | #[cfg(target_os = "macos")] |
| 365 | fn memory_pages_from_vm_stat(vm_stat: &str, prefixes: &[&str]) -> Option<u64> { |
| 366 | let mut total = 0u64; |
| 367 | for prefix in prefixes { |
| 368 | let pages = vm_stat |
| 369 | .lines() |
| 370 | .find(|line| line.trim_start().starts_with(prefix)) |
| 371 | .and_then(|line| { |
| 372 | line.split('.') |
| 373 | .nth(1) |
| 374 | .and_then(|rest| rest.trim().parse::<u64>().ok()) |
| 375 | })?; |
| 376 | total = total.checked_add(pages)?; |
| 377 | } |
| 378 | Some(total) |
| 379 | } |
| 380 | |
| 381 | #[cfg(windows)] |
| 382 | fn host_memory_free_fraction() -> Option<f64> { |
| 383 | // `wmic` is deprecated but present on every supported Windows runner and |
| 384 | // avoids adding a GlobalMemoryStatusEx build dependency. Fail open on error. |
| 385 | let out = run_capture( |
| 386 | "C:\\Windows\\System32\\wbem\\wmic.exe", |
| 387 | &[ |
| 388 | "OS", |
| 389 | "get", |
| 390 | "FreePhysicalMemory,TotalVisibleMemorySize", |
| 391 | "/value", |
| 392 | ], |
| 393 | )?; |
| 394 | let free_kb = wmic_value(&out, "FreePhysicalMemory=")?; |
| 395 | let total_kb = wmic_value(&out, "TotalVisibleMemorySize=")?; |
| 396 | (total_kb > 0).then(|| (free_kb as f64 / total_kb as f64).clamp(0.0, 1.0)) |
| 397 | } |
| 398 | |
| 399 | #[cfg(windows)] |
| 400 | fn wmic_value(output: &str, key: &str) -> Option<u64> { |
| 401 | output |
| 402 | .lines() |
| 403 | .find_map(|line| line.trim().strip_prefix(key)) |
| 404 | .and_then(|value| value.trim().parse::<u64>().ok()) |
| 405 | } |
| 406 | |
| 407 | #[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))] |
| 408 | fn host_memory_free_fraction() -> Option<f64> { |
| 409 | None |
| 410 | } |
| 411 | |
| 412 | #[cfg(any(target_os = "macos", target_os = "windows"))] |
| 413 | fn run_capture(program: &str, args: &[&str]) -> Option<String> { |
| 414 | let output = std::process::Command::new(program) |
| 415 | .args(args) |
| 416 | .output() |
| 417 | .ok()?; |
| 418 | if !output.status.success() { |
| 419 | return None; |
| 420 | } |
| 421 | String::from_utf8(output.stdout).ok() |
| 422 | } |
| 423 | |
| 424 | #[cfg(test)] |
| 425 | mod tests { |
| 426 | use std::sync::Arc; |
| 427 | use std::sync::atomic::{AtomicUsize, Ordering}; |
| 428 | |
| 429 | use super::*; |
| 430 | |
| 431 | /// Deterministic probe returning a fixed fraction so admission behavior is |
| 432 | /// independent of the host running the test suite. |
| 433 | struct StaticMemoryProbe(Option<f64>); |
| 434 | impl MemoryProbe for StaticMemoryProbe { |
| 435 | fn free_fraction(&self) -> Option<f64> { |
| 436 | self.0 |
| 437 | } |
| 438 | } |
| 439 | |
| 440 | const NOMINAL_PROBE: StaticMemoryProbe = StaticMemoryProbe(Some(0.9)); |
| 441 | |
| 442 | #[test] |
| 443 | fn whitespace_codewhale_home_uses_shared_user_home_for_admission_state() { |
| 444 | let _lock = crate::test_support::lock_test_env(); |
| 445 | let tmp = tempfile::tempdir().expect("temporary root"); |
| 446 | let home = tmp.path().join("home"); |
| 447 | let userprofile = tmp.path().join("userprofile"); |
| 448 | let _home = crate::test_support::EnvVarGuard::set("HOME", &home); |
| 449 | let _userprofile = crate::test_support::EnvVarGuard::set("USERPROFILE", &userprofile); |
| 450 | let _codewhale_home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", " \t "); |
| 451 | |
| 452 | assert_eq!( |
| 453 | admission_root(), |
| 454 | home.join(".codewhale").join("resource-admission") |
| 455 | ); |
| 456 | } |
| 457 | |
| 458 | #[test] |
| 459 | fn memory_pressure_classification_and_effective_limit() { |
| 460 | // Unknown (unmeasurable) fails open to the configured limit. |
| 461 | assert_eq!(classify_memory_pressure(None), MemoryPressure::Unknown); |
| 462 | assert_eq!(effective_admission_limit(2, MemoryPressure::Unknown), 2); |
| 463 | assert_eq!(classify_memory_pressure(Some(0.9)), MemoryPressure::Nominal); |
| 464 | assert_eq!( |
| 465 | classify_memory_pressure(Some(0.30)), |
| 466 | MemoryPressure::Constrained |
| 467 | ); |
| 468 | assert_eq!( |
| 469 | classify_memory_pressure(Some(0.15)), |
| 470 | MemoryPressure::Critical |
| 471 | ); |
| 472 | assert_eq!( |
| 473 | classify_memory_pressure(Some(0.0)), |
| 474 | MemoryPressure::Critical |
| 475 | ); |
| 476 | assert_eq!(effective_admission_limit(4, MemoryPressure::Nominal), 4); |
| 477 | assert_eq!(effective_admission_limit(4, MemoryPressure::Constrained), 2); |
| 478 | assert_eq!(effective_admission_limit(1, MemoryPressure::Constrained), 1); |
| 479 | assert_eq!(effective_admission_limit(4, MemoryPressure::Critical), 0); |
| 480 | } |
| 481 | |
| 482 | #[test] |
| 483 | fn infers_only_expensive_rust_compilation_commands() { |
| 484 | for command in [ |
| 485 | "cargo test -p codewhale-tui shell::tests", |
| 486 | "env CARGO_BUILD_JOBS=2 cargo build --workspace", |
| 487 | "cargo check", |
| 488 | "cargo clippy --all-targets", |
| 489 | "/usr/bin/rustc src/main.rs", |
| 490 | "printf ok && cargo rustc -- --emit=metadata", |
| 491 | // #6560 D01-11: wrappers run the same compilation. |
| 492 | "sudo cargo build --release", |
| 493 | "sudo -u builder cargo test", |
| 494 | "timeout 600 cargo test --workspace", |
| 495 | "nice -n 10 cargo check", |
| 496 | "bash -c \"cargo build\"", |
| 497 | "sh -c 'cargo clippy --all-targets'", |
| 498 | "git ls-files | xargs cargo check", |
| 499 | "(cargo build)", |
| 500 | "env RUSTFLAGS=-Dwarnings nohup cargo build", |
| 501 | ] { |
| 502 | assert_eq!( |
| 503 | infer_command_expense(command), |
| 504 | CommandExpense::Heavy, |
| 505 | "{command}" |
| 506 | ); |
| 507 | } |
| 508 | for command in [ |
| 509 | "cargo fmt --check", |
| 510 | "cargo metadata", |
| 511 | "git status", |
| 512 | "echo cargo test", |
| 513 | "sudo cargo fmt", |
| 514 | "timeout 5 git status", |
| 515 | "bash -c \"ls\"", |
| 516 | "xargs rm", |
| 517 | ] { |
| 518 | assert_eq!( |
| 519 | infer_command_expense(command), |
| 520 | CommandExpense::Normal, |
| 521 | "{command}" |
| 522 | ); |
| 523 | } |
| 524 | } |
| 525 | |
| 526 | #[tokio::test(flavor = "multi_thread", worker_threads = 4)] |
| 527 | async fn cross_task_heavy_admission_never_exceeds_limit() { |
| 528 | let temp = tempfile::tempdir().expect("tempdir"); |
| 529 | let active = Arc::new(AtomicUsize::new(0)); |
| 530 | let peak = Arc::new(AtomicUsize::new(0)); |
| 531 | let mut tasks = Vec::new(); |
| 532 | |
| 533 | for _ in 0..12 { |
| 534 | let root = temp.path().to_path_buf(); |
| 535 | let active = Arc::clone(&active); |
| 536 | let peak = Arc::clone(&peak); |
| 537 | tasks.push(tokio::spawn(async move { |
| 538 | let _permit = acquire_heavy_command_permit_at(&root, 2, None, &NOMINAL_PROBE) |
| 539 | .await |
| 540 | .expect("permit"); |
| 541 | let current = active.fetch_add(1, Ordering::SeqCst) + 1; |
| 542 | peak.fetch_max(current, Ordering::SeqCst); |
| 543 | tokio::time::sleep(Duration::from_millis(30)).await; |
| 544 | active.fetch_sub(1, Ordering::SeqCst); |
| 545 | })); |
| 546 | } |
| 547 | for task in tasks { |
| 548 | task.await.expect("admission task"); |
| 549 | } |
| 550 | |
| 551 | assert_eq!(active.load(Ordering::SeqCst), 0); |
| 552 | assert_eq!(peak.load(Ordering::SeqCst), 2); |
| 553 | } |
| 554 | |
| 555 | #[tokio::test] |
| 556 | async fn queued_admission_observes_cancellation() { |
| 557 | let temp = tempfile::tempdir().expect("tempdir"); |
| 558 | let _held = acquire_heavy_command_permit_at(temp.path(), 1, None, &NOMINAL_PROBE) |
| 559 | .await |
| 560 | .expect("initial permit"); |
| 561 | let cancel = CancellationToken::new(); |
| 562 | let wait_cancel = cancel.clone(); |
| 563 | let root = temp.path().to_path_buf(); |
| 564 | let waiter = tokio::spawn(async move { |
| 565 | acquire_heavy_command_permit_at(&root, 1, Some(&wait_cancel), &NOMINAL_PROBE).await |
| 566 | }); |
| 567 | |
| 568 | tokio::time::sleep(Duration::from_millis(75)).await; |
| 569 | cancel.cancel(); |
| 570 | let error = tokio::time::timeout(Duration::from_secs(1), waiter) |
| 571 | .await |
| 572 | .expect("bounded cancellation") |
| 573 | .expect("waiter task") |
| 574 | .expect_err("queued command must cancel"); |
| 575 | assert!(error.to_string().contains("canceled while queued")); |
| 576 | } |
| 577 | |
| 578 | #[tokio::test] |
| 579 | async fn critical_memory_pressure_queues_without_admitting() { |
| 580 | let temp = tempfile::tempdir().expect("tempdir"); |
| 581 | // Critical pressure -> zero effective slots -> the command cannot be |
| 582 | // admitted and must observe cancellation rather than spin forever. |
| 583 | let critical = StaticMemoryProbe(Some(0.05)); |
| 584 | let cancel = CancellationToken::new(); |
| 585 | let wait_cancel = cancel.clone(); |
| 586 | let root = temp.path().to_path_buf(); |
| 587 | let waiter = tokio::spawn(async move { |
| 588 | acquire_heavy_command_permit_at(&root, 2, Some(&wait_cancel), &critical).await |
| 589 | }); |
| 590 | tokio::time::sleep(Duration::from_millis(120)).await; |
| 591 | cancel.cancel(); |
| 592 | let error = tokio::time::timeout(Duration::from_secs(2), waiter) |
| 593 | .await |
| 594 | .expect("bounded cancellation") |
| 595 | .expect("waiter task") |
| 596 | .expect_err("critical-pressure command must not be admitted"); |
| 597 | assert!(error.to_string().contains("canceled while queued")); |
| 598 | } |
| 599 | |
| 600 | #[test] |
| 601 | fn host_memory_probe_is_fail_safe() { |
| 602 | // On any supported host the probe either measures a plausible fraction |
| 603 | // or admits it cannot; it must never panic or return an out-of-range. |
| 604 | if let Some(fraction) = host_memory_free_fraction() { |
| 605 | assert!( |
| 606 | (0.0..=1.0).contains(&fraction), |
| 607 | "measured free fraction out of range: {fraction}" |
| 608 | ); |
| 609 | } |
| 610 | } |
| 611 | |
| 612 | /// Opt-in real cargo acceptance (#4864 req 8): drive an actual heavy-class |
| 613 | /// cargo invocation through the admission path end to end. `cargo check |
| 614 | /// --version` classifies as heavy (subcommand `check`) but exits instantly, |
| 615 | /// so this is fast. Gated behind an env var so CI never runs a real cargo. |
| 616 | #[cfg(unix)] |
| 617 | #[tokio::test] |
| 618 | async fn real_cargo_command_is_admitted_and_released() { |
| 619 | if std::env::var_os("CODEWHALE_RESOURCE_ADMISSION_RUST_ACCEPTANCE").is_none() { |
| 620 | tracing::info!( |
| 621 | "skipping opt-in real-rust admission acceptance; \ |
| 622 | set CODEWHALE_RESOURCE_ADMISSION_RUST_ACCEPTANCE=1 to run" |
| 623 | ); |
| 624 | return; |
| 625 | } |
| 626 | let temp = tempfile::tempdir().expect("tempdir"); |
| 627 | let permit = acquire_heavy_command_permit_at(temp.path(), 2, None, &NOMINAL_PROBE) |
| 628 | .await |
| 629 | .expect("heavy permit for real cargo"); |
| 630 | assert_eq!(permit.limit(), 2); |
| 631 | assert_eq!(permit.memory_pressure(), MemoryPressure::Nominal); |
| 632 | let cargo = std::process::Command::new("cargo") |
| 633 | .arg("check") |
| 634 | .arg("--version") |
| 635 | .output() |
| 636 | .expect("run cargo"); |
| 637 | assert!(cargo.status.success(), "cargo check --version failed"); |
| 638 | drop(permit); |
| 639 | let _again = acquire_heavy_command_permit_at(temp.path(), 2, None, &NOMINAL_PROBE) |
| 640 | .await |
| 641 | .expect("re-acquire after release"); |
| 642 | } |
| 643 | } |
| 644 |