diff --git a/backend/Cargo.lock b/backend/Cargo.lock index f912c01a79..7943dee150 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -17796,6 +17796,7 @@ dependencies = [ "windmill-runtime-nativets", "windmill-types", "windmill-worker-volumes", + "windows 0.61.3", "x509-parser 0.16.0", "yaml-rust", ] diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index 7a200fb44b..d4479e06d1 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -157,6 +157,8 @@ lazy_static::lazy_static! { pub static ref NATIVE_MODE: bool = std::env::var("NATIVE_MODE").ok().is_some_and(|x| x == "1" || x == "true"); + pub static ref LIMIT_WINDOWS_TO_1CU: bool = std::env::var("LIMIT_WINDOWS_TO_1CU").ok().is_some_and(|x| x == "1" || x == "true"); + pub static ref CGROUP_V2_PATH_RE: Regex = Regex::new(r#"(?m)^0::(/.*)$"#).unwrap(); pub static ref CGROUP_V2_CPU_RE: Regex = Regex::new(r#"(?m)^(\d+) \S+$"#).unwrap(); pub static ref CGROUP_V1_INACTIVE_FILE_RE: Regex = Regex::new(r#"(?m)^total_inactive_file (\d+)$"#).unwrap(); @@ -1052,6 +1054,9 @@ pub fn get_vcpus() -> Option { #[cfg(windows)] pub fn get_vcpus() -> Option { + if *LIMIT_WINDOWS_TO_1CU { + return Some(100000); // 1 vCPU + } let mut sys = System::new(); sys.refresh_cpu_all(); (sys.cpus().len() * 100000).try_into().ok() @@ -1101,6 +1106,9 @@ pub fn get_memory() -> Option { #[cfg(windows)] pub fn get_memory() -> Option { + if *LIMIT_WINDOWS_TO_1CU { + return Some(2 * 1024 * 1024 * 1024); // 2 GB + } let mut sys = System::new(); sys.refresh_memory(); Some(sys.total_memory() as i64) diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml index 4e2357b14e..62cb9de4bb 100644 --- a/backend/windmill-worker/Cargo.toml +++ b/backend/windmill-worker/Cargo.toml @@ -146,6 +146,9 @@ hyper-tls = { workspace = true, optional = true } hyper-util = { workspace = true, optional = true } rcgen = { workspace = true, optional = true } +[target.'cfg(windows)'.dependencies] +windows = { version = "0.61", features = ["Win32_System_JobObjects", "Win32_System_Threading"] } + [dev-dependencies] tempfile.workspace = true x509-parser.workspace = true diff --git a/backend/windmill-worker/src/common.rs b/backend/windmill-worker/src/common.rs index d303771cc8..80c3673254 100644 --- a/backend/windmill-worker/src/common.rs +++ b/backend/windmill-worker/src/common.rs @@ -702,6 +702,114 @@ lazy_static! { static ref DISABLE_PROCESS_GROUP: bool = std::env::var("DISABLE_PROCESS_GROUP").is_ok(); } +/// 2 GB memory limit in bytes for LIMIT_WINDOWS_TO_1CU +#[cfg(windows)] +const MEMORY_LIMIT_1CU: usize = 2 * 1024 * 1024 * 1024; + +/// Wrapper that holds a Windows Job Object handle alongside the child process. +/// The job object enforces memory limits and is closed when the child is dropped. +#[cfg(windows)] +struct MemoryLimitedChild { + inner: Box, + _job_handle: Win32JobHandle, +} + +/// RAII wrapper for a raw Win32 HANDLE that closes it on drop. +#[cfg(windows)] +struct Win32JobHandle(windows::Win32::Foundation::HANDLE); + +// SAFETY: Win32 HANDLEs are plain pointer-sized values with no thread affinity; +// the kernel ref-counts the underlying object, so sending/sharing the handle is safe. +#[cfg(windows)] +unsafe impl Send for Win32JobHandle {} +#[cfg(windows)] +unsafe impl Sync for Win32JobHandle {} + +#[cfg(windows)] +impl Drop for Win32JobHandle { + fn drop(&mut self) { + let _ = unsafe { windows::Win32::Foundation::CloseHandle(self.0) }; + } +} + +#[cfg(windows)] +impl process_wrap::tokio::TokioChildWrapper for MemoryLimitedChild { + fn inner(&self) -> &tokio::process::Child { + self.inner.inner() + } + fn inner_mut(&mut self) -> &mut tokio::process::Child { + self.inner.inner_mut() + } + fn into_inner(self: Box) -> tokio::process::Child { + self.inner.into_inner() + } + fn start_kill(&mut self) -> std::io::Result<()> { + self.inner.start_kill() + } + fn wait( + &mut self, + ) -> Box> + Send + '_> + { + self.inner.wait() + } + fn try_wait(&mut self) -> std::io::Result> { + self.inner.try_wait() + } +} + +/// Create a Windows Job Object with a memory limit and assign the process to it. +#[cfg(windows)] +fn apply_job_memory_limit(pid: u32, memory_limit: usize) -> Result { + use windows::Win32::Foundation::HANDLE; + use windows::Win32::System::JobObjects::*; + use windows::Win32::System::Threading::{OpenProcess, PROCESS_SET_QUOTA, PROCESS_TERMINATE}; + + unsafe { + let job = CreateJobObjectW(None, None).map_err(|e| { + std::io::Error::new(std::io::ErrorKind::Other, format!("CreateJobObjectW: {e}")) + })?; + + let mut info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); + info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_JOB_MEMORY; + info.JobMemoryLimit = memory_limit; + + SetInformationJobObject( + job, + JobObjectExtendedLimitInformation, + &info as *const _ as _, + std::mem::size_of::() as u32, + ) + .map_err(|e| { + let _ = windows::Win32::Foundation::CloseHandle(job); + std::io::Error::new( + std::io::ErrorKind::Other, + format!("SetInformationJobObject: {e}"), + ) + })?; + + let process_handle = OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, false, pid) + .map_err(|e| { + let _ = windows::Win32::Foundation::CloseHandle(job); + std::io::Error::new( + std::io::ErrorKind::Other, + format!("OpenProcess({pid}): {e}"), + ) + })?; + + let assign_result = AssignProcessToJobObject(job, process_handle); + let _ = windows::Win32::Foundation::CloseHandle(process_handle); + assign_result.map_err(|e| { + let _ = windows::Win32::Foundation::CloseHandle(job); + std::io::Error::new( + std::io::ErrorKind::Other, + format!("AssignProcessToJobObject: {e}"), + ) + })?; + + Ok(Win32JobHandle(job)) + } +} + pub fn build_command_with_isolation(program: &str, args: &[&str]) -> Command { use tokio::process::Command; @@ -767,9 +875,31 @@ pub async fn start_child_process( } } - return cmd + let child: Box = cmd .spawn() - .map_err(|err| tentatively_improve_error(err.into(), executable)); + .map_err(|err| tentatively_improve_error(err.into(), executable))?; + + #[cfg(windows)] + if *windmill_common::worker::LIMIT_WINDOWS_TO_1CU { + if let Some(pid) = child.inner().id() { + match apply_job_memory_limit(pid, MEMORY_LIMIT_1CU) { + Ok(job_handle) => { + tracing::info!( + "Applied 2GB memory limit (LIMIT_WINDOWS_TO_1CU) to child process {pid}" + ); + return Ok(Box::new(MemoryLimitedChild { + inner: child, + _job_handle: job_handle, + })); + } + Err(e) => { + tracing::warn!("Failed to apply memory limit to child process {pid}: {e}"); + } + } + } + } + + Ok(child) } pub async fn resolve_job_timeout(