Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix recursive-protection feature flag #13887

Merged
merged 4 commits into from
Dec 24, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion datafusion/common/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ name = "datafusion_common"
path = "src/lib.rs"

[features]
default = ["recursive-protection"]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

having the feature enabled by default in subcrates meant I think it basically was not possible to disable when using datafusion core crate

avro = ["apache-avro"]
backtrace = []
pyarrow = ["pyo3", "arrow/pyarrow", "parquet"]
Expand Down
7 changes: 7 additions & 0 deletions datafusion/core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,13 @@ pyarrow = ["datafusion-common/pyarrow", "parquet"]
regex_expressions = [
"datafusion-functions/regex_expressions",
]
recusive-protection = [
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This adds the recursive-protection feature to the datafusion (core) crate which then activates the feature in the subcrates

"datafusion-common/recursive-protection",
"datafusion-expr/recursive-protection",
"datafusion-optimizer/recursive-protection",
"datafusion-physical-optimizer/recursive-protection",
"datafusion-sql/recursive-protection"
]
serde = ["arrow-schema/serde"]
string_expressions = ["datafusion-functions/string_expressions"]
unicode_expressions = [
Expand Down
1 change: 0 additions & 1 deletion datafusion/expr/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ name = "datafusion_expr"
path = "src/lib.rs"

[features]
default = ["recursive-protection"]
recursive-protection = ["dep:recursive"]

[dependencies]
Expand Down
1 change: 0 additions & 1 deletion datafusion/optimizer/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ name = "datafusion_optimizer"
path = "src/lib.rs"

[features]
default = ["recursive-protection"]
recursive-protection = ["dep:recursive"]

[dependencies]
Expand Down
1 change: 0 additions & 1 deletion datafusion/physical-optimizer/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@ rust-version = { workspace = true }
workspace = true

[features]
default = ["recursive-protection"]
recursive-protection = ["dep:recursive"]

[dependencies]
Expand Down
2 changes: 1 addition & 1 deletion datafusion/sql/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ name = "datafusion_sql"
path = "src/lib.rs"

[features]
default = ["unicode_expressions", "unparser", "recursive-protection"]
default = ["unicode_expressions", "unparser"]
unicode_expressions = []
unparser = []
recursive-protection = ["dep:recursive"]
Expand Down
1 change: 1 addition & 0 deletions datafusion/sql/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ mod query;
mod relation;
mod select;
mod set_expr;
mod stack;
mod statement;
#[cfg(feature = "unparser")]
pub mod unparser;
Expand Down
10 changes: 6 additions & 4 deletions datafusion/sql/src/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ use std::sync::Arc;

use crate::planner::{ContextProvider, PlannerContext, SqlToRel};

use crate::stack::StackGuard;
use datafusion_common::{not_impl_err, Constraints, DFSchema, Result};
use datafusion_expr::expr::Sort;
use datafusion_expr::{
Expand Down Expand Up @@ -62,10 +63,11 @@ impl<S: ContextProvider> 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 min_stack_size = recursive::get_minimum_stack_size();
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

here is pretty good evidence that the feature flag is not working as expected, as this code requires recursive to compile but it wasn't failing the build -- I pulled the logic its own structure

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Late LGTM. StackGuard is a very clever fix.

recursive::set_minimum_stack_size(256 * 1024);
let plan = self.set_expr_to_plan(other, planner_context)?;
recursive::set_minimum_stack_size(min_stack_size);
let plan = {
// scope for dropping _guard
let _guard = StackGuard::new(256 * 1024);
self.set_expr_to_plan(other, planner_context)
}?;
let oby_exprs = to_order_by_exprs(query.order_by)?;
let order_by_rex = self.order_by_to_sort_expr(
oby_exprs,
Expand Down
63 changes: 63 additions & 0 deletions datafusion/sql/src/stack.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

pub use inner::StackGuard;

/// A guard that sets the minimum stack size for the current thread to `min_stack_size` bytes.
#[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);
}
}
}

/// A stub implementation of the stack guard when the recursive protection
/// feature is not enabled
#[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
}
}
}
Loading