diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml index 161e173611..645f154e05 100644 --- a/backend/windmill-worker/Cargo.toml +++ b/backend/windmill-worker/Cargo.toml @@ -114,7 +114,8 @@ hmac.workspace = true pem = { workspace = true, optional = true } rsa = { workspace = true, optional = true } urlencoding.workspace = true -nix.workspace = true +# `fs` adds flock(2) for the cross-process Python install lock (shared cache mounts) +nix = { workspace = true, features = ["fs"] } bytes.workspace = true reqwest.workspace = true reqwest-middleware.workspace = true diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs index 6c8506f0d2..9ca79c6852 100644 --- a/backend/windmill-worker/src/python_executor.rs +++ b/backend/windmill-worker/src/python_executor.rs @@ -62,6 +62,13 @@ lazy_static::lazy_static! { static ref PY_CONCURRENT_DOWNLOADS: usize = var("PY_CONCURRENT_DOWNLOADS").ok().map(|flag| flag.parse().unwrap_or(20)).unwrap_or(20); + // uv's HTTP request timeout (seconds). spawn_uv_install uses env_clear(), so a + // UV_HTTP_TIMEOUT set on the worker is dropped unless forwarded explicitly. + // Only forwarded when set; otherwise uv keeps its own default. Lets operators + // raise it for slow/contended private registries ("operation timed out"). + static ref UV_HTTP_TIMEOUT: Option = + var("UV_HTTP_TIMEOUT").ok().filter(|v| !v.is_empty()); + static ref NON_ALPHANUM_CHAR: Regex = regex::Regex::new(r"[^0-9A-Za-z=.-]").unwrap(); @@ -2059,6 +2066,24 @@ lazy_static::lazy_static! { /// in-memory marker and pays only the original single stat. static ref VERIFIED_VENVS: tokio::sync::Mutex> = tokio::sync::Mutex::new(HashSet::new()); + + // In-process locks serializing concurrent installs into the same shared + // `venv_p` cache dir; `uv --reinstall` removes a package's .dist-info/RECORD + // before rewriting it, so a sibling install/verify racing it corrupts the + // dir. Keyed by venv_p so distinct deps still install in parallel. + static ref PY_INSTALL_LOCKS: tokio::sync::Mutex>>> = + tokio::sync::Mutex::new(std::collections::HashMap::new()); +} + +/// Returns the in-process install lock for a given target cache dir, creating it +/// on first use. Idle entries (only the map holds a reference) are pruned each +/// call so the map stays bounded by the number of in-flight installs. +async fn get_venv_install_lock(venv_p: &str) -> Arc> { + let mut map = PY_INSTALL_LOCKS.lock().await; + map.retain(|_, v| Arc::strong_count(v) > 1); + map.entry(venv_p.to_string()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() } /// Spawn process of uv install @@ -2104,6 +2129,9 @@ async fn spawn_uv_install( if *NATIVE_CERT { vars.push(("UV_NATIVE_TLS", "true")); } + if let Some(timeout) = UV_HTTP_TIMEOUT.as_ref() { + vars.push(("UV_HTTP_TIMEOUT", timeout.as_str())); + } let _owner; if let Some(py_path) = py_path.as_ref() { @@ -2207,6 +2235,9 @@ async fn spawn_uv_install( let mut envs = vec![("PATH", PATH_ENV.as_str())]; envs.push(("HOME", HOME_ENV.as_str())); envs.push(("UV_INDEX_STRATEGY", uv_index_strategy)); + if let Some(timeout) = UV_HTTP_TIMEOUT.as_ref() { + envs.push(("UV_HTTP_TIMEOUT", timeout.as_str())); + } if let Some(mirror) = uv_python_install_mirror.as_ref() { envs.push(("UV_PYTHON_INSTALL_MIRROR", mirror)); } @@ -2780,6 +2811,96 @@ pub async fn handle_python_reqs( ); let start = std::time::Instant::now(); + + // Lock the shared target dir (see PY_INSTALL_LOCKS). In-process lock + // first; only one task per dir then contends the cross-process file + // lock below. Both guards drop on every return path. + let venv_lock = get_venv_install_lock(&venv_p).await; + let _venv_guard = tokio::select! { + _ = kill_rx.recv() => { + pids.lock().await.get_mut(i).and_then(|e| e.take()); + return Err(Error::from(anyhow::anyhow!( + "install of {venv_p} canceled while waiting for venv lock" + ))); + } + guard = venv_lock.lock_owned() => guard, + }; + + // Cross-process advisory lock. Best-effort: if the filesystem doesn't + // support flock we log and proceed — verify_wheel_record + job retry + // still guard correctness, just without the dedup. + #[cfg(unix)] + let _venv_file_lock: Option = { + use std::os::unix::io::AsRawFd; + let lock_path = format!("{venv_p}.lock"); + if let Some(parent) = std::path::Path::new(&lock_path).parent() { + let _ = std::fs::create_dir_all(parent); + } + match std::fs::OpenOptions::new().create(true).write(true).open(&lock_path) { + Ok(f) => { + // Bounded wait: a holder that crashes releases the lock (the + // kernel drops it on fd close), but a live-but-stuck holder + // (e.g. uv wedged on a hung mount) would otherwise block us + // forever. After the cap, proceed degraded rather than hang — + // verify_wheel_record + retry still guard correctness. + const MAX_WAIT: std::time::Duration = std::time::Duration::from_secs(300); + let waited_since = std::time::Instant::now(); + loop { + match nix::fcntl::flock(f.as_raw_fd(), nix::fcntl::FlockArg::LockExclusiveNonblock) { + Ok(()) => break Some(f), + // EWOULDBLOCK == EAGAIN on Linux: another holder has the lock. + Err(nix::errno::Errno::EWOULDBLOCK) => { + if waited_since.elapsed() >= MAX_WAIT { + tracing::warn!( + workspace_id = %w_id, + "venv install lock {lock_path} still held after {}s, proceeding without cross-process install lock", + MAX_WAIT.as_secs() + ); + break Some(f); + } + tokio::select! { + _ = kill_rx.recv() => { + pids.lock().await.get_mut(i).and_then(|e| e.take()); + return Err(Error::from(anyhow::anyhow!( + "install of {venv_p} canceled while waiting for venv file lock" + ))); + } + _ = tokio::time::sleep(std::time::Duration::from_millis(200)) => {} + } + } + Err(e) => { + tracing::warn!( + workspace_id = %w_id, + "could not flock {lock_path}, proceeding without cross-process install lock: {e}" + ); + break Some(f); + } + } + } + } + Err(e) => { + tracing::warn!( + workspace_id = %w_id, + "could not open install lock file {lock_path}, proceeding without cross-process install lock: {e}" + ); + None + } + } + }; + + // Double-checked: another job (this process or another sharing the + // mount) may have installed this exact dep while we waited on the + // locks. Reuse it instead of reinstalling. + if metadata(format!("{venv_p}/.valid.windmill")).await.is_ok() { + print_success( + false, false, &job_id, &w_id, &req, req_tl, counter_arc, + total_to_install, start, &conn, + ) + .await; + pids.lock().await.get_mut(i).and_then(|e| e.take()); + return Ok(()); + } + #[cfg(all(feature = "enterprise", feature = "parquet"))] if is_not_pro { if let Some(os) = windmill_object_store::get_object_store().await { @@ -3533,11 +3654,7 @@ mod tests { } let dist_info = root.join("pkg-1.0.0.dist-info"); std::fs::create_dir_all(&dist_info).unwrap(); - std::fs::write( - dist_info.join("RECORD"), - record_entries.join("\n") + "\n", - ) - .unwrap(); + std::fs::write(dist_info.join("RECORD"), record_entries.join("\n") + "\n").unwrap(); } #[tokio::test] @@ -3605,4 +3722,99 @@ mod tests { .await .is_ok()); } + + // Regression tests for the concurrent-install guard. Two jobs installing the + // same uncached dep into the shared `venv_p` used to race uv's `--reinstall`, + // corrupting the on-disk wheel and failing with "Env installation did not + // succeed". The guard serializes those installs. + + #[tokio::test] + async fn test_venv_install_lock_serializes_same_path() { + use std::sync::atomic::{AtomicUsize, Ordering}; + // Same target path => one shared lock => no two tasks install at once. + let active = Arc::new(AtomicUsize::new(0)); + let max_seen = Arc::new(AtomicUsize::new(0)); + let mut handles = vec![]; + for _ in 0..8 { + let active = active.clone(); + let max_seen = max_seen.clone(); + handles.push(tokio::spawn(async move { + let lock = get_venv_install_lock("/cache/py/3.11/samedep==1.0").await; + let _g = lock.lock_owned().await; + let cur = active.fetch_add(1, Ordering::SeqCst) + 1; + max_seen.fetch_max(cur, Ordering::SeqCst); + // Yield so any concurrency would be observed by another task. + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + active.fetch_sub(1, Ordering::SeqCst); + })); + } + for h in handles { + h.await.unwrap(); + } + assert_eq!( + max_seen.load(Ordering::SeqCst), + 1, + "installs into the same target dir must be serialized" + ); + } + + #[tokio::test] + async fn test_venv_install_lock_distinct_paths_are_independent() { + // Different target paths get different locks and never block each other. + let a = get_venv_install_lock("/cache/py/3.11/depA==1.0").await; + let b = get_venv_install_lock("/cache/py/3.11/depB==1.0").await; + let _ga = a.lock_owned().await; + // Holding depA's lock must not prevent acquiring depB's. + assert!( + b.try_lock().is_ok(), + "distinct deps must install in parallel" + ); + // Same path returns the same underlying lock. + let a2 = get_venv_install_lock("/cache/py/3.11/depA==1.0").await; + assert!( + a2.try_lock().is_err(), + "same target dir must map to the same lock" + ); + } + + #[cfg(unix)] + #[tokio::test] + async fn test_venv_file_lock_excludes_across_descriptions() { + // The cross-process layer: flock on a sibling `.lock` excludes a second + // independent open file description (i.e. another worker process) while + // held, and frees it on close. Mirrors the loop in handle_python_reqs. + use nix::fcntl::{flock, FlockArg}; + use std::os::unix::io::AsRawFd; + + let dir = std::env::temp_dir().join("wm_venv_lock_test"); + std::fs::create_dir_all(&dir).unwrap(); + let lock_path = dir.join("dep==1.0.lock"); + + let f1 = std::fs::OpenOptions::new() + .create(true) + .write(true) + .open(&lock_path) + .unwrap(); + flock(f1.as_raw_fd(), FlockArg::LockExclusiveNonblock).unwrap(); + + // A second descriptor (stand-in for another process) cannot take it. + let f2 = std::fs::OpenOptions::new() + .create(true) + .write(true) + .open(&lock_path) + .unwrap(); + assert_eq!( + flock(f2.as_raw_fd(), FlockArg::LockExclusiveNonblock), + Err(nix::errno::Errno::EWOULDBLOCK), + "a second holder must be blocked while the lock is held" + ); + + // Releasing the first lets the second acquire it. + drop(f1); + flock(f2.as_raw_fd(), FlockArg::LockExclusiveNonblock) + .expect("lock must be acquirable once the holder releases it"); + + drop(f2); + let _ = std::fs::remove_file(&lock_path); + } }