diff --git a/backend/windmill-duckdb-ffi-internal/build_dev.sh b/backend/windmill-duckdb-ffi-internal/build_dev.sh index 1fb13057a2..d4ae2b7583 100755 --- a/backend/windmill-duckdb-ffi-internal/build_dev.sh +++ b/backend/windmill-duckdb-ffi-internal/build_dev.sh @@ -1,3 +1,3 @@ CARGO_NET_GIT_FETCH_WITH_CLI=true cargo build --release -p windmill_duckdb_ffi_internal mkdir -p ../target/debug/ -cp target/release/libwindmill_duckdb_ffi_internal.* ../target/debug/ \ No newline at end of file +cp target/release/libwindmill_duckdb_ffi_internal.* ../target/debug/ diff --git a/backend/windmill-duckdb-ffi-internal/src/lib.rs b/backend/windmill-duckdb-ffi-internal/src/lib.rs index 2701319d6e..594b0f6df4 100644 --- a/backend/windmill-duckdb-ffi-internal/src/lib.rs +++ b/backend/windmill-duckdb-ffi-internal/src/lib.rs @@ -11,6 +11,22 @@ use rust_decimal::{prelude::FromPrimitive, Decimal}; use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; +// Worker passes "" for "no override" — saves an extra C string nullability dance. +// Returns an owned String so the value outlives the raw pointer's lifetime. +fn ptr_to_opt_str(ptr: *const c_char) -> Result, String> { + if ptr.is_null() { + return Ok(None); + } + let s = unsafe { CStr::from_ptr(ptr) } + .to_str() + .map_err(|e| format!("Invalid string in duckdb ffi: {}", e))?; + Ok(if s.is_empty() { + None + } else { + Some(s.to_owned()) + }) +} + #[derive(Deserialize, Clone, Debug, PartialEq, Default)] pub struct Arg { pub name: String, @@ -34,7 +50,7 @@ pub extern "C" fn get_version() -> c_uint { // Increment when making breaking changes to the FFI interface. // The windmill worker will check that the version matches or else refuse to call // the FFI functions to avoid undefined behavior. - return 1; + return 2; } #[unsafe(no_mangle)] @@ -45,10 +61,16 @@ pub extern "C" fn run_duckdb_ffi( token: *const c_char, base_internal_url: *const c_char, w_id: *const c_char, + memory_limit: *const c_char, + temp_directory: *const c_char, column_order_ptr: *mut *mut c_char, collect_last_only: bool, collect_first_row_only: bool, ) -> *mut c_char { + let resource_limits = match (ptr_to_opt_str(memory_limit), ptr_to_opt_str(temp_directory)) { + (Ok(m), Ok(t)) => Ok(ResourceLimits { memory_limit: m, temp_directory: t }), + (Err(e), _) | (_, Err(e)) => Err(e), + }; let (r, column_order) = match convert_args( query_block_list, query_block_list_count, @@ -57,8 +79,9 @@ pub extern "C" fn run_duckdb_ffi( base_internal_url, w_id, ) + .and_then(|args| resource_limits.map(|r| (args, r))) .and_then( - |(query_block_list, job_args, token, base_internal_url, w_id)| { + |((query_block_list, job_args, token, base_internal_url, w_id), limits)| { run_duckdb_internal( query_block_list, query_block_list_count, @@ -66,6 +89,7 @@ pub extern "C" fn run_duckdb_ffi( token, base_internal_url, w_id, + limits, collect_last_only, collect_first_row_only, ) @@ -150,7 +174,13 @@ pub extern "C" fn prepare_duckdb_ffi( token: *const c_char, base_internal_url: *const c_char, w_id: *const c_char, + memory_limit: *const c_char, + temp_directory: *const c_char, ) -> *mut c_char { + let resource_limits = match (ptr_to_opt_str(memory_limit), ptr_to_opt_str(temp_directory)) { + (Ok(m), Ok(t)) => Ok(ResourceLimits { memory_limit: m, temp_directory: t }), + (Err(e), _) | (_, Err(e)) => Err(e), + }; let r = match convert_prepare_args( query_block_list, query_block_list_count, @@ -158,9 +188,12 @@ pub extern "C" fn prepare_duckdb_ffi( base_internal_url, w_id, ) - .and_then(|(query_block_list, token, base_internal_url, w_id)| { - prepare_duckdb_internal(query_block_list, token, base_internal_url, w_id) - }) { + .and_then(|args| resource_limits.map(|r| (args, r))) + .and_then( + |((query_block_list, token, base_internal_url, w_id), limits)| { + prepare_duckdb_internal(query_block_list, token, base_internal_url, w_id, limits) + }, + ) { Ok(result) => result, Err(err) => { let err = serde_json::to_string(&err) @@ -175,12 +208,58 @@ pub extern "C" fn prepare_duckdb_ffi( }) } +#[derive(Clone, Default)] +struct ResourceLimits { + memory_limit: Option, + temp_directory: Option, +} + +fn sql_single_quote(s: &str) -> String { + s.replace('\'', "''") +} + +// Bounds memory so DuckDB spills to disk before blowing the cgroup cap and +// getting the worker SIGKILLed. Spill goes to the job dir (when set) so it is +// cleaned up with the job, otherwise DuckDB's default temp_directory is kept. +fn configure_duckdb_resource_limits( + conn: &duckdb::Connection, + limits: &ResourceLimits, +) -> Result<(), String> { + let mut config_sql = String::new(); + // jemalloc-specific setting bundled with the Linux DuckDB build. macOS and + // Windows builds may not accept it; gated to avoid breaking those workers. + if cfg!(target_os = "linux") { + config_sql.push_str("SET allocator_background_threads=true;\n"); + } + if let Some(mem) = limits.memory_limit.as_deref() { + config_sql.push_str(&format!("SET memory_limit='{}';\n", sql_single_quote(mem))); + } + if let Some(tmp) = limits.temp_directory.as_deref() { + config_sql.push_str(&format!( + "SET temp_directory='{}';\n", + sql_single_quote(tmp) + )); + } + if config_sql.is_empty() { + return Ok(()); + } + conn.execute_batch(&config_sql).map_err(|e| { + format!( + "Error configuring DuckDB resource limits: {}", + e.to_string() + ) + }) +} + fn setup_duckdb_connection( conn: &duckdb::Connection, token: &str, base_internal_url: &str, w_id: &str, + limits: &ResourceLimits, ) -> Result<(), String> { + configure_duckdb_resource_limits(conn, limits)?; + let (s3_access_key, s3_secret_key) = token.rsplit_once('.').unwrap_or(("", token)); let (s3_endpoint_ssl, s3_endpoint) = base_internal_url .split_once("://") @@ -249,10 +328,11 @@ fn prepare_duckdb_internal( token: &str, base_internal_url: &str, w_id: &str, + limits: ResourceLimits, ) -> Result { let conn = duckdb::Connection::open_in_memory().map_err(|e| e.to_string())?; - setup_duckdb_connection(&conn, token, base_internal_url, w_id)?; + setup_duckdb_connection(&conn, token, base_internal_url, w_id, &limits)?; let mut results: Vec = vec![]; @@ -379,12 +459,13 @@ fn run_duckdb_internal<'a>( token: &str, base_internal_url: &str, w_id: &str, + limits: ResourceLimits, collect_last_only: bool, collect_first_row_only: bool, ) -> Result<(String, Option>), String> { let conn = duckdb::Connection::open_in_memory().map_err(|e| e.to_string())?; - setup_duckdb_connection(&conn, token, base_internal_url, w_id)?; + setup_duckdb_connection(&conn, token, base_internal_url, w_id, &limits)?; let mut results: Vec>> = vec![]; let mut column_order = None; diff --git a/backend/windmill-worker/src/duckdb_executor.rs b/backend/windmill-worker/src/duckdb_executor.rs index d8c40f6c4d..4a019be9ab 100644 --- a/backend/windmill-worker/src/duckdb_executor.rs +++ b/backend/windmill-worker/src/duckdb_executor.rs @@ -11,7 +11,7 @@ use serde_json::{json, Value}; use uuid::Uuid; use windmill_common::error::{to_anyhow, Error, Result}; use windmill_common::utils::sanitize_string_from_password; -use windmill_common::worker::{Connection, SqlResultCollectionStrategy}; +use windmill_common::worker::{get_memory, Connection, SqlResultCollectionStrategy}; use windmill_common::workspaces::{ get_datatable_resource_from_db_unchecked, get_ducklake_from_db_unchecked, DucklakeCatalogResourceType, @@ -44,6 +44,7 @@ pub async fn do_duckdb( #[allow(unused_variables)] column_order_ref: &mut Option>, occupancy_metrics: &mut OccupancyMetrics, parent_runnable_path: Option, + job_dir: &str, run_inline: bool, ) -> Result> { let annotations = windmill_common::worker::SqlAnnotations::parse(query); @@ -160,6 +161,7 @@ pub async fn do_duckdb( let base_internal_url = client.base_internal_url.clone(); let w_id = job.workspace_id.clone(); + let job_dir = job_dir.to_string(); if annotations.prepare { let result = tokio::task::spawn_blocking(move || { @@ -168,6 +170,7 @@ pub async fn do_duckdb( &token, &base_internal_url, &w_id, + &job_dir, ) }) .await @@ -185,6 +188,7 @@ pub async fn do_duckdb( &token, &base_internal_url, &w_id, + &job_dir, collection_strategy, ) }) @@ -259,6 +263,8 @@ struct DuckDbFfiLib { token: *const c_char, base_internal_url: *const c_char, w_id: *const c_char, + memory_limit: *const c_char, + temp_directory: *const c_char, column_order_ptr: *mut *mut c_char, collect_last_only: bool, collect_first_row_only: bool, @@ -273,6 +279,8 @@ struct DuckDbFfiLib { token: *const c_char, base_internal_url: *const c_char, w_id: *const c_char, + memory_limit: *const c_char, + temp_directory: *const c_char, ) -> *mut c_char, >, >, @@ -319,7 +327,7 @@ impl DuckDbFfiLib { // Version mismatch should only be possible on Windows agent workers // We check for it because FFI interface mismatch will cause undefined behavior / crashes unsafe { - let expected_version: c_uint = 1; + let expected_version: c_uint = 2; let get_version: Symbol<'static, unsafe extern "C" fn() -> c_uint> = lib.get(b"get_version") .map_err(|e| return Error::ExecutionErr(format!("Could not find get_version in the duckdb ffi library. If you are not using docker, consider manually upgrading windmill_duckdb_ffi_lib. {}", e.to_string())))?; @@ -345,6 +353,35 @@ impl DuckDbFfiLib { } } +// 20% headroom for Rust runtime + DuckDB's untracked allocations. Mirrors +// DuckDB's own default ratio, but applied to the worker's cgroup budget +// instead of host RAM. +const DUCKDB_MEMORY_FRACTION: f64 = 0.8; +// Treat cgroup values above 1 PiB as "unlimited" (kernels report page-aligned +// huge numbers when uncapped). get_memory() falls back to host RAM in that +// case, which is exactly what we want to leave to DuckDB's own default. +const CGROUP_UNLIMITED_THRESHOLD: i64 = 1024 * 1024 * 1024 * 1024 * 1024; + +// `DUCKDB_MEMORY_LIMIT` env override, else fraction of the worker's cgroup +// memory (as reported by windmill-common), else None (keep DuckDB's default). +fn resolve_duckdb_memory_limit() -> Option { + if let Ok(v) = env::var("DUCKDB_MEMORY_LIMIT") { + let v = v.trim(); + if !v.is_empty() { + return Some(v.to_string()); + } + } + cgroup_bytes_to_duckdb_memory_limit(get_memory()?) +} + +fn cgroup_bytes_to_duckdb_memory_limit(bytes: i64) -> Option { + if bytes <= 0 || bytes >= CGROUP_UNLIMITED_THRESHOLD { + return None; + } + let mib = ((bytes as f64 * DUCKDB_MEMORY_FRACTION) as i64) / (1024 * 1024); + Some(format!("{}MiB", mib.max(64))) +} + // Read backend/windmill-duckdb-ffi-internal/README_DEV.md for details about why we use FFI fn run_duckdb_ffi_safe<'a>( query_block_list: impl Iterator, @@ -353,6 +390,7 @@ fn run_duckdb_ffi_safe<'a>( token: &str, base_internal_url: &str, w_id: &str, + job_dir: &str, collection_strategy: SqlResultCollectionStrategy, ) -> Result<(Box, Option>)> { let query_block_list = query_block_list @@ -372,6 +410,9 @@ fn run_duckdb_ffi_safe<'a>( let token = CString::new(token).map_err(to_anyhow)?; let base_internal_url = CString::new(base_internal_url).map_err(to_anyhow)?; let w_id = CString::new(w_id).map_err(to_anyhow)?; + let memory_limit = + CString::new(resolve_duckdb_memory_limit().unwrap_or_default()).map_err(to_anyhow)?; + let temp_directory = CString::new(job_dir).map_err(to_anyhow)?; let run_duckdb_ffi = &DuckDbFfiLib::get_singleton()?.run_duckdb_ffi; let free_cstr = &DuckDbFfiLib::get_singleton()?.free_cstr; @@ -384,6 +425,8 @@ fn run_duckdb_ffi_safe<'a>( token.as_ptr(), base_internal_url.as_ptr(), w_id.as_ptr(), + memory_limit.as_ptr(), + temp_directory.as_ptr(), &mut column_order, collection_strategy.collect_last_statement_only(query_block_list_count), collection_strategy.collect_first_row_only(), @@ -424,6 +467,7 @@ fn prepare_duckdb_ffi_safe<'a>( token: &str, base_internal_url: &str, w_id: &str, + job_dir: &str, ) -> Result> { let query_block_list = query_block_list .map(|s| { @@ -440,6 +484,9 @@ fn prepare_duckdb_ffi_safe<'a>( let token = CString::new(token).map_err(to_anyhow)?; let base_internal_url = CString::new(base_internal_url).map_err(to_anyhow)?; let w_id = CString::new(w_id).map_err(to_anyhow)?; + let memory_limit = + CString::new(resolve_duckdb_memory_limit().unwrap_or_default()).map_err(to_anyhow)?; + let temp_directory = CString::new(job_dir).map_err(to_anyhow)?; let lib = DuckDbFfiLib::get_singleton()?; let prepare_fn = lib.prepare_duckdb_ffi.as_ref().ok_or_else(|| { @@ -456,6 +503,8 @@ fn prepare_duckdb_ffi_safe<'a>( token.as_ptr(), base_internal_url.as_ptr(), w_id.as_ptr(), + memory_limit.as_ptr(), + temp_directory.as_ptr(), ); let str = CStr::from_ptr(ptr).to_string_lossy().to_string(); free_cstr(ptr); @@ -783,6 +832,44 @@ pub struct Arg { mod tests { use super::*; + #[test] + fn cgroup_bytes_unlimited_or_invalid_returns_none() { + assert_eq!(cgroup_bytes_to_duckdb_memory_limit(0), None); + assert_eq!(cgroup_bytes_to_duckdb_memory_limit(-1), None); + // 1 PiB sentinel: cgroup v1 reports ~i64::MAX when uncapped. + assert_eq!( + cgroup_bytes_to_duckdb_memory_limit(CGROUP_UNLIMITED_THRESHOLD), + None + ); + } + + #[test] + fn cgroup_bytes_real_values_take_80_percent() { + // 1 GiB -> 80% -> 819 MiB (floored to MiB) + assert_eq!( + cgroup_bytes_to_duckdb_memory_limit(1024 * 1024 * 1024), + Some("819MiB".to_string()) + ); + // 4 GiB -> 3276 MiB + assert_eq!( + cgroup_bytes_to_duckdb_memory_limit(4 * 1024 * 1024 * 1024), + Some("3276MiB".to_string()) + ); + } + + #[test] + fn cgroup_bytes_tiny_values_floored_to_64mib() { + // Tiny cgroup must not produce a 0/unusable limit. + assert_eq!( + cgroup_bytes_to_duckdb_memory_limit(1024 * 1024), + Some("64MiB".to_string()) + ); + assert_eq!( + cgroup_bytes_to_duckdb_memory_limit(1), + Some("64MiB".to_string()) + ); + } + // Tests for parse_attach_db_resource function #[test] fn test_parse_attach_db_resource_postgres_res_prefix() { diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 7798e73c98..ff5a06a022 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -4700,6 +4700,7 @@ pub async fn run_language_executor( column_order, occupancy_metrics, parent_runnable_path, + job_dir, run_inline, )) .await;