fix: cgroup-aware DuckDB memory_limit + allocator memory release (#9245)

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/
+cp target/release/libwindmill_duckdb_ffi_internal.* ../target/debug/

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<Option<String>, 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<String>,
+    temp_directory: Option<String>,
+}
+
+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<String, 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<PrepareQueryResult> = 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<Vec<String>>), 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<Box<RawValue>>> = vec![];
     let mut column_order = None;

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<Vec<String>>,
     occupancy_metrics: &mut OccupancyMetrics,
     parent_runnable_path: Option<String>,
+    job_dir: &str,
     run_inline: bool,
 ) -> Result<Box<RawValue>> {
     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<String> {
+    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<String> {
+    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<Item = &'a str>,
@@ -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<RawValue>, Option<Vec<String>>)> {
     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<Box<RawValue>> {
     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() {

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;