From a61db26698a056784544689dc078eeeafd29f806 Mon Sep 17 00:00:00 2001 From: Alex Metelli Date: Thu, 2 Jul 2026 04:18:18 +0800 Subject: [PATCH 1/3] fix: avoid global SQL stack guard mutation --- Cargo.lock | 1 + Cargo.toml | 1 + datafusion/sql/Cargo.toml | 3 +- datafusion/sql/src/query.rs | 7 +- datafusion/sql/src/set_expr.rs | 111 ++++++++++++++-------------- datafusion/sql/src/stack.rs | 59 +++++---------- datafusion/sql/src/unparser/expr.rs | 36 ++++----- 7 files changed, 97 insertions(+), 121 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fa0f5fe50ece2..d4512288bb658 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2624,6 +2624,7 @@ dependencies = [ "regex", "rstest", "sqlparser", + "stacker", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index df9cdfed40e00..769e5bd710949 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -195,6 +195,7 @@ rstest = "0.26.1" serde_json = "1" sha2 = "^0.11.0" sqlparser = { version = "0.62.0", default-features = false, features = ["std", "visitor"] } +stacker = "0.1.24" strum = "0.28.0" strum_macros = "0.28.0" tempfile = "3" diff --git a/datafusion/sql/Cargo.toml b/datafusion/sql/Cargo.toml index cc299ce507099..318f8934639aa 100644 --- a/datafusion/sql/Cargo.toml +++ b/datafusion/sql/Cargo.toml @@ -44,7 +44,7 @@ name = "datafusion_sql" default = ["unicode_expressions", "unparser"] unicode_expressions = [] unparser = [] -recursive_protection = ["dep:recursive"] +recursive_protection = ["dep:recursive", "dep:stacker"] # Note the sql planner should not depend directly on the datafusion-function packages # so that it can be used in a standalone manner with other function implementations. @@ -62,6 +62,7 @@ log = { workspace = true } recursive = { workspace = true, optional = true } regex = { workspace = true } sqlparser = { workspace = true } +stacker = { workspace = true, optional = true } [dev-dependencies] ctor = { workspace = true } diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index 76124cbc7eb59..e2b9e4d2d5305 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -19,7 +19,6 @@ use std::sync::Arc; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use crate::stack::StackGuard; use datafusion_common::{Constraints, DFSchema, Result, not_impl_err}; use datafusion_expr::expr::{Sort, WildcardOptions}; @@ -81,11 +80,9 @@ impl SqlToRel<'_, S> { // The functions called from `set_expr_to_plan()` need more than 128KB // stack in debug builds as investigated in: // https://github.com/apache/datafusion/pull/13310#discussion_r1836813902 - let plan = { - // scope for dropping _guard - let _guard = StackGuard::new(256 * 1024); + let plan = crate::stack::maybe_grow(|| { self.set_expr_to_plan(other, planner_context) - }?; + })?; let oby_exprs = to_order_by_exprs(order_by)?; let order_by_rex = self.order_by_to_sort_expr( oby_exprs, diff --git a/datafusion/sql/src/set_expr.rs b/datafusion/sql/src/set_expr.rs index dc8e4f14d1ee8..51b11f3087095 100644 --- a/datafusion/sql/src/set_expr.rs +++ b/datafusion/sql/src/set_expr.rs @@ -25,70 +25,71 @@ use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; use sqlparser::ast::{SetExpr, SetOperator, SetQuantifier, Spanned}; impl SqlToRel<'_, S> { - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] pub(super) fn set_expr_to_plan( &self, set_expr: SetExpr, planner_context: &mut PlannerContext, ) -> Result { - let set_expr_span = Span::try_from_sqlparser_span(set_expr.span()); - match set_expr { - SetExpr::Select(s) => self.select_to_plan(*s, None, planner_context), - SetExpr::Values(v) => self.sql_values_to_plan(v, planner_context), - SetExpr::SetOperation { - op, - left, - right, - set_quantifier, - } => { - let left_span = Span::try_from_sqlparser_span(left.span()); - let right_span = Span::try_from_sqlparser_span(right.span()); - let left_plan = self.set_expr_to_plan(*left, planner_context); - // Store the left plan's schema so that the right side can - // alias duplicate expressions to match. Skip for BY NAME - // operations since those match columns by name, not position. - if let Ok(plan) = &left_plan - && plan.schema().fields().len() > 1 - && !matches!( - set_quantifier, - SetQuantifier::ByName - | SetQuantifier::AllByName - | SetQuantifier::DistinctByName - ) - { - planner_context - .set_set_expr_left_schema(Some(Arc::clone(plan.schema()))); - } - let right_plan = self.set_expr_to_plan(*right, planner_context); - planner_context.set_set_expr_left_schema(None); - let (left_plan, right_plan) = match (left_plan, right_plan) { - (Ok(left_plan), Ok(right_plan)) => (left_plan, right_plan), - (Err(left_err), Err(right_err)) => { - return Err(DataFusionError::Collection(vec![ - left_err, right_err, - ])); + crate::stack::maybe_grow(|| { + let set_expr_span = Span::try_from_sqlparser_span(set_expr.span()); + match set_expr { + SetExpr::Select(s) => self.select_to_plan(*s, None, planner_context), + SetExpr::Values(v) => self.sql_values_to_plan(v, planner_context), + SetExpr::SetOperation { + op, + left, + right, + set_quantifier, + } => { + let left_span = Span::try_from_sqlparser_span(left.span()); + let right_span = Span::try_from_sqlparser_span(right.span()); + let left_plan = self.set_expr_to_plan(*left, planner_context); + // Store the left plan's schema so that the right side can + // alias duplicate expressions to match. Skip for BY NAME + // operations since those match columns by name, not position. + if let Ok(plan) = &left_plan + && plan.schema().fields().len() > 1 + && !matches!( + set_quantifier, + SetQuantifier::ByName + | SetQuantifier::AllByName + | SetQuantifier::DistinctByName + ) + { + planner_context + .set_set_expr_left_schema(Some(Arc::clone(plan.schema()))); } - (Err(err), _) | (_, Err(err)) => { - return Err(err); + let right_plan = self.set_expr_to_plan(*right, planner_context); + planner_context.set_set_expr_left_schema(None); + let (left_plan, right_plan) = match (left_plan, right_plan) { + (Ok(left_plan), Ok(right_plan)) => (left_plan, right_plan), + (Err(left_err), Err(right_err)) => { + return Err(DataFusionError::Collection(vec![ + left_err, right_err, + ])); + } + (Err(err), _) | (_, Err(err)) => { + return Err(err); + } + }; + if !(set_quantifier == SetQuantifier::ByName + || set_quantifier == SetQuantifier::AllByName) + { + self.validate_set_expr_num_of_columns( + op, + left_span, + right_span, + &left_plan, + &right_plan, + set_expr_span, + )?; } - }; - if !(set_quantifier == SetQuantifier::ByName - || set_quantifier == SetQuantifier::AllByName) - { - self.validate_set_expr_num_of_columns( - op, - left_span, - right_span, - &left_plan, - &right_plan, - set_expr_span, - )?; + self.set_operation_to_plan(op, left_plan, right_plan, set_quantifier) } - self.set_operation_to_plan(op, left_plan, right_plan, set_quantifier) + SetExpr::Query(q) => self.query_to_plan(*q, planner_context), + _ => not_impl_err!("Query {set_expr} not implemented yet"), } - SetExpr::Query(q) => self.query_to_plan(*q, planner_context), - _ => not_impl_err!("Query {set_expr} not implemented yet"), - } + }) } pub(super) fn is_union_all(set_quantifier: SetQuantifier) -> Result { diff --git a/datafusion/sql/src/stack.rs b/datafusion/sql/src/stack.rs index b7d5eebdd7188..d27c62e8f7ef0 100644 --- a/datafusion/sql/src/stack.rs +++ b/datafusion/sql/src/stack.rs @@ -15,49 +15,28 @@ // specific language governing permissions and limitations // under the License. -pub use inner::StackGuard; +/// The local red zone used by SQL recursive entry points. +/// +/// Some SQL planner and unparser recursion paths need more than `recursive`'s +/// default 128 KiB red zone in debug builds. Keep this value local to each +/// stack-growth checkpoint rather than mutating `recursive`'s process-global +/// minimum stack size. +pub(crate) const SQL_RECURSION_RED_ZONE: usize = 256 * 1024; -/// A guard that sets the minimum stack size for the current thread to `min_stack_size` bytes. +/// Runs `callback` on a stack with enough space for SQL recursive entry points. #[cfg(feature = "recursive_protection")] -mod inner { - /// Sets the stack size to `min_stack_size` bytes on call to `new()` and - /// resets to the previous value when this structure is dropped. - pub struct StackGuard { - previous_stack_size: usize, - } - - impl StackGuard { - /// Sets the stack size to `min_stack_size` bytes on call to `new()` and - /// resets to the previous value when this structure is dropped. - pub fn new(min_stack_size: usize) -> Self { - let previous_stack_size = recursive::get_minimum_stack_size(); - recursive::set_minimum_stack_size(min_stack_size); - Self { - previous_stack_size, - } - } - } - - impl Drop for StackGuard { - fn drop(&mut self) { - recursive::set_minimum_stack_size(self.previous_stack_size); - } - } +#[inline] +pub(crate) fn maybe_grow(callback: impl FnOnce() -> R) -> R { + stacker::maybe_grow( + SQL_RECURSION_RED_ZONE, + recursive::get_stack_allocation_size(), + callback, + ) } -/// A stub implementation of the stack guard when the recursive protection -/// feature is not enabled +/// Runs `callback` without stack growth when recursive protection is disabled. #[cfg(not(feature = "recursive_protection"))] -mod inner { - /// A stub implementation of the stack guard when the recursive protection - /// feature is not enabled that does nothing - pub struct StackGuard; - - impl StackGuard { - /// A stub implementation of the stack guard when the recursive protection - /// feature is not enabled - pub fn new(_min_stack_size: usize) -> Self { - Self - } - } +#[inline] +pub(crate) fn maybe_grow(callback: impl FnOnce() -> R) -> R { + callback() } diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index f5690ee797598..46886d06339f7 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -31,7 +31,6 @@ use std::vec; use super::Unparser; use super::dialect::{DistinctFromStyle, IntervalStyle}; -use crate::stack::StackGuard; use arrow::array::{ ArrayRef, Date32Array, Date64Array, PrimitiveArray, types::{ @@ -100,26 +99,25 @@ impl Unparser<'_> { // default `recursive` red zone, so without raising the minimum stack // size the stack-growing trampoline engages too late and the OS stack // overflows on deeply nested expressions (issue #23056). The size - // mirrors the planner's `StackGuard` usage in `query.rs`. - let _guard = StackGuard::new(256 * 1024); - self.expr_to_sql_with_nesting(expr) + // mirrors the planner's stack-growth usage in `query.rs`. + crate::stack::maybe_grow(|| self.expr_to_sql_with_nesting(expr)) } /// Recursive entry point shared by the public [`Self::expr_to_sql`] and the /// internal recursion sites (scalar-function arguments, arrays, maps, and /// dialect scalar-function overrides). /// - /// This carries the `recursive` annotation so every nesting level becomes a - /// stack-growth checkpoint. Internal recursion must call this rather than - /// the public [`Self::expr_to_sql`]: the public entry point is not - /// annotated and would re-install the [`StackGuard`] on every level. - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] + /// This is a stack-growth checkpoint. Internal recursion must call this + /// rather than the public [`Self::expr_to_sql`]: the public entry point + /// would re-enter the public stack-growth boundary on every level. pub(crate) fn expr_to_sql_with_nesting(&self, expr: &Expr) -> Result { - let mut root_expr = self.expr_to_sql_inner(expr)?; - if self.pretty { - root_expr = self.remove_unnecessary_nesting(root_expr, LOWEST, LOWEST); - } - Ok(root_expr) + crate::stack::maybe_grow(|| { + let mut root_expr = self.expr_to_sql_inner(expr)?; + if self.pretty { + root_expr = self.remove_unnecessary_nesting(root_expr, LOWEST, LOWEST); + } + Ok(root_expr) + }) } fn distinct_from_to_sql( @@ -155,9 +153,8 @@ impl Unparser<'_> { } } - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn expr_to_sql_inner(&self, expr: &Expr) -> Result { - match expr { + crate::stack::maybe_grow(|| match expr { Expr::InList(InList { expr, list, @@ -658,7 +655,7 @@ impl Unparser<'_> { Expr::LambdaVariable(l) => Ok(ast::Expr::Identifier( self.new_ident_quoted_if_needs(l.name.clone()), )), - } + }) } pub fn scalar_function_to_sql( @@ -1016,14 +1013,13 @@ impl Unparser<'_> { /// /// Also note that when fetching the precedence of a nested expression, we ignore other nested /// expressions, so precedence of expr `(a * (b + c))` equals `*` and not `+`. - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn remove_unnecessary_nesting( &self, expr: ast::Expr, left_op: &BinaryOperator, right_op: &BinaryOperator, ) -> ast::Expr { - match expr { + crate::stack::maybe_grow(|| match expr { ast::Expr::Nested(nested) => { let surrounding_precedence = self .sql_op_precedence(left_op) @@ -1076,7 +1072,7 @@ impl Unparser<'_> { self.remove_unnecessary_nesting(*expr, left_op, IS), )), _ => expr, - } + }) } fn inner_precedence(&self, expr: &ast::Expr) -> u8 { From b0259aef78043847947145155984d514e85b77a3 Mon Sep 17 00:00:00 2001 From: Alex Metelli Date: Thu, 2 Jul 2026 04:55:50 +0800 Subject: [PATCH 2/3] test: cover SQL stack growth race --- datafusion/sql/src/stack.rs | 14 +++++++ datafusion/sql/src/unparser/expr.rs | 59 +++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/datafusion/sql/src/stack.rs b/datafusion/sql/src/stack.rs index d27c62e8f7ef0..60c7d9b241dae 100644 --- a/datafusion/sql/src/stack.rs +++ b/datafusion/sql/src/stack.rs @@ -40,3 +40,17 @@ pub(crate) fn maybe_grow(callback: impl FnOnce() -> R) -> R { pub(crate) fn maybe_grow(callback: impl FnOnce() -> R) -> R { callback() } + +#[cfg(all(test, feature = "recursive_protection"))] +mod tests { + use super::*; + + #[test] + fn maybe_grow_does_not_mutate_recursive_minimum_stack_size() { + let before = recursive::get_minimum_stack_size(); + let observed = maybe_grow(recursive::get_minimum_stack_size); + + assert_eq!(observed, before); + assert_eq!(recursive::get_minimum_stack_size(), before); + } +} diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 46886d06339f7..14ac65b32ed6a 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -3397,6 +3397,65 @@ mod tests { handle.join().expect("unparsing thread should not panic"); } + /// Regression test for https://github.com/apache/datafusion/issues/23246 + /// + /// A concurrent `StackGuard` drop used to restore `recursive`'s + /// process-global red zone to 128 KiB while another thread was still + /// relying on DataFusion's 256 KiB SQL red zone. Keep the global value + /// churning at the default while deep unparsing runs to ensure the unparser + /// uses local stack-growth boundaries instead. + #[cfg(feature = "recursive_protection")] + #[test] + fn test_deeply_nested_expr_with_concurrent_recursive_minimum_churn() { + const DEPTH: usize = 2_000; + const DEFAULT_RECURSIVE_RED_ZONE: usize = 128 * 1024; + + let previous_minimum = recursive::get_minimum_stack_size(); + recursive::set_minimum_stack_size(DEFAULT_RECURSIVE_RED_ZONE); + + let running = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)); + let barrier = std::sync::Arc::new(std::sync::Barrier::new(2)); + + let churner = std::thread::spawn({ + let running = std::sync::Arc::clone(&running); + let barrier = std::sync::Arc::clone(&barrier); + move || { + barrier.wait(); + while running.load(std::sync::atomic::Ordering::Relaxed) { + recursive::set_minimum_stack_size(DEFAULT_RECURSIVE_RED_ZONE); + std::thread::yield_now(); + } + } + }); + + let worker = std::thread::Builder::new() + .stack_size(2 * 1024 * 1024) + .spawn({ + let barrier = std::sync::Arc::clone(&barrier); + move || { + barrier.wait(); + + let mut nested_fn: Expr = col("c"); + for _ in 0..DEPTH { + nested_fn = array_has(nested_fn, lit("x")); + } + + let pg = PostgreSqlDialect {}; + Unparser::new(&pg) + .expr_to_sql(&nested_fn) + .expect("deeply nested scalar function should unparse"); + } + }) + .unwrap(); + + let worker_result = worker.join(); + running.store(false, std::sync::atomic::Ordering::Relaxed); + churner.join().expect("churner thread should not panic"); + recursive::set_minimum_stack_size(previous_minimum); + + worker_result.expect("unparsing thread should not panic"); + } + #[test] fn test_window_func_support_window_frame() -> Result<()> { let default_dialect: Arc = From db44a880d320280dd98340943ae414fc7ea92993 Mon Sep 17 00:00:00 2001 From: Alex Metelli Date: Sat, 4 Jul 2026 03:43:38 +0800 Subject: [PATCH 3/3] test: prove unparser stack guard isolation --- datafusion/sql/src/stack.rs | 1 + datafusion/sql/src/unparser/expr.rs | 71 ++++++++++------------------- 2 files changed, 26 insertions(+), 46 deletions(-) diff --git a/datafusion/sql/src/stack.rs b/datafusion/sql/src/stack.rs index 60c7d9b241dae..ed3bf1553ebfc 100644 --- a/datafusion/sql/src/stack.rs +++ b/datafusion/sql/src/stack.rs @@ -21,6 +21,7 @@ /// default 128 KiB red zone in debug builds. Keep this value local to each /// stack-growth checkpoint rather than mutating `recursive`'s process-global /// minimum stack size. +#[cfg(feature = "recursive_protection")] pub(crate) const SQL_RECURSION_RED_ZONE: usize = 256 * 1024; /// Runs `callback` on a stack with enough space for SQL recursive entry points. diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 14ac65b32ed6a..33457a1515645 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -3397,63 +3397,42 @@ mod tests { handle.join().expect("unparsing thread should not panic"); } - /// Regression test for https://github.com/apache/datafusion/issues/23246 - /// - /// A concurrent `StackGuard` drop used to restore `recursive`'s - /// process-global red zone to 128 KiB while another thread was still - /// relying on DataFusion's 256 KiB SQL red zone. Keep the global value - /// churning at the default while deep unparsing runs to ensure the unparser - /// uses local stack-growth boundaries instead. #[cfg(feature = "recursive_protection")] #[test] - fn test_deeply_nested_expr_with_concurrent_recursive_minimum_churn() { - const DEPTH: usize = 2_000; + fn test_expr_to_sql_does_not_mutate_recursive_minimum_stack_size() -> Result<()> { const DEFAULT_RECURSIVE_RED_ZONE: usize = 128 * 1024; let previous_minimum = recursive::get_minimum_stack_size(); recursive::set_minimum_stack_size(DEFAULT_RECURSIVE_RED_ZONE); - let running = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)); - let barrier = std::sync::Arc::new(std::sync::Barrier::new(2)); - - let churner = std::thread::spawn({ - let running = std::sync::Arc::clone(&running); - let barrier = std::sync::Arc::clone(&barrier); - move || { - barrier.wait(); - while running.load(std::sync::atomic::Ordering::Relaxed) { - recursive::set_minimum_stack_size(DEFAULT_RECURSIVE_RED_ZONE); - std::thread::yield_now(); - } - } - }); - - let worker = std::thread::Builder::new() - .stack_size(2 * 1024 * 1024) - .spawn({ - let barrier = std::sync::Arc::clone(&barrier); - move || { - barrier.wait(); - - let mut nested_fn: Expr = col("c"); - for _ in 0..DEPTH { - nested_fn = array_has(nested_fn, lit("x")); - } - - let pg = PostgreSqlDialect {}; - Unparser::new(&pg) - .expr_to_sql(&nested_fn) - .expect("deeply nested scalar function should unparse"); + let observed_minimum = Arc::new(std::sync::atomic::AtomicUsize::new(usize::MAX)); + let dialect = DuckDBDialect::new().with_custom_scalar_overrides(vec![( + "dummy_udf", + Box::new({ + let observed_minimum = Arc::clone(&observed_minimum); + move |unparser: &Unparser, args: &[Expr]| { + observed_minimum.store( + recursive::get_minimum_stack_size(), + std::sync::atomic::Ordering::Relaxed, + ); + unparser.scalar_function_to_sql("dummy_udf", args).map(Some) } - }) - .unwrap(); + }) as ScalarFnToSqlHandler, + )]); + let expr = ScalarUDF::new_from_impl(DummyUDF::new()).call(vec![col("a")]); - let worker_result = worker.join(); - running.store(false, std::sync::atomic::Ordering::Relaxed); - churner.join().expect("churner thread should not panic"); + let result = Unparser::new(&dialect).expr_to_sql(&expr); + let final_minimum = recursive::get_minimum_stack_size(); recursive::set_minimum_stack_size(previous_minimum); - worker_result.expect("unparsing thread should not panic"); + result?; + assert_eq!( + observed_minimum.load(std::sync::atomic::Ordering::Relaxed), + DEFAULT_RECURSIVE_RED_ZONE + ); + assert_eq!(final_minimum, DEFAULT_RECURSIVE_RED_ZONE); + + Ok(()) } #[test]