From 8c0b8d5f0ba222462a4287e9608d99024bc78dce Mon Sep 17 00:00:00 2001 From: Mithun Chicklore Yogendra Date: Sun, 9 Aug 2026 10:44:29 +0530 Subject: [PATCH 1/2] fix(proto): prevent logical plan serialization stack overflow (#24124) - Closes #23823. `LogicalPlanNode::try_from_logical_plan` recursively serializes logical plans to protobuf. In debug builds, its large `match` compiled to a 202,304 B stack frame, so ten nested `SubqueryAlias` nodes over an `EmptyRelation` overflowed a 2 MiB thread stack. Each arm was isolated independently against the original dispatcher: | Independent frame effect | Arms | Examples | | --- | ---: | --- | | >=10 KiB reduction | 4 | TableScan: 34,608 B; Join: 11,392 B; Dml / RecursiveQuery: 10,320 B | | 5-8 KiB reduction | 17 | Projection, Filter, Aggregate, Repartition, Unnest, Copy | | 2-4 KiB reduction | 7 | Values, EmptyRelation, Union, Extension | | <=64 B effect | 9 | Several DDL/statement arms; Subquery adds 16 B | These effects are not additive: each change alters the compiler's layout of the same `match` frame. Isolating the 15 arms with the largest independent reductions still left a 58,560 B frame; isolating all 37 arms reduced it to 1,680 B. - Isolate every `try_from_logical_plan` match arm behind a debug-only non-inlined helper, preventing arm-local temporaries from inflating the recursive dispatcher frame. - Add the opt-in `datafusion-proto/recursive_protection` feature using the existing DataFusion recursion pattern. - Add child-process stack-safety regressions for the original 2 MiB-stack reproducer and feature-gated stack growth. With `recursive_protection`, the dispatcher frame measures 1,648 B. The helper is only forced out of line in debug builds. Yes. - `cargo fmt --all --check` - `cargo check -p datafusion-proto --all-features` - `cargo test -p datafusion-proto --lib` - Stack-safety regression: 100 nested aliases on a 2 MiB stack. - Feature-gated stack-growth regression: 2,000 nested aliases with `recursive_protection`. `datafusion-proto` gains an opt-in `recursive_protection` feature. Existing protobuf wire format, conversion behavior, and default features are unchanged. --- Cargo.lock | 1 + datafusion/proto/Cargo.toml | 2 + datafusion/proto/src/logical_plan/mod.rs | 30 +++++- datafusion/proto/tests/cases/mod.rs | 1 + datafusion/proto/tests/cases/stack_safety.rs | 99 ++++++++++++++++++++ 5 files changed, 130 insertions(+), 3 deletions(-) create mode 100644 datafusion/proto/tests/cases/stack_safety.rs diff --git a/Cargo.lock b/Cargo.lock index 6914453b3da2c..18a7cf7863d34 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2531,6 +2531,7 @@ dependencies = [ "pretty_assertions", "prost", "rand 0.9.2", + "recursive", "serde", "serde_json", "tokio", diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index a2a0a4546188d..5ff4a8bd5dab8 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -37,6 +37,7 @@ name = "datafusion_proto" [features] default = ["parquet"] +recursive_protection = ["dep:recursive"] json = ["pbjson", "serde", "serde_json", "datafusion-proto-common/json"] parquet = ["datafusion-datasource-parquet", "datafusion-common/parquet", "datafusion/parquet"] avro = ["datafusion-datasource-avro", "datafusion-common/avro"] @@ -69,6 +70,7 @@ pbjson = { workspace = true, optional = true } prost = { workspace = true } rand = { workspace = true } serde = { version = "1.0", optional = true } +recursive = { workspace = true, optional = true } serde_json = { workspace = true, optional = true } [dev-dependencies] diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index cd19cb7bcf61e..be5db34de0207 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -104,7 +104,30 @@ pub trait AsLogicalPlan: Debug + Send + Sync + Clone { Self: Sized; } -pub trait LogicalExtensionCodec: Debug + Send + Sync { +// In debug builds, keep each serializer arm's local temporaries out of the +// recursive dispatcher frame. Without this call boundary, they inflate the +// frame of every recursive invocation. +#[cfg_attr(debug_assertions, inline(never))] +fn serialize_logical_plan_arm(serializer: F) -> Result +where + F: FnOnce() -> Result, +{ + serializer() +} + +macro_rules! dispatch_logical_plan { + ($plan:expr, { $($pattern:pat => $body:expr $(,)?)+ }) => { + match $plan { + $( + $pattern => serialize_logical_plan_arm(|| -> Result { + $body + }), + )+ + } + }; +} + +pub trait LogicalExtensionCodec: Debug + Send + Sync + std::any::Any { fn try_decode( &self, buf: &[u8], @@ -1087,6 +1110,7 @@ impl AsLogicalPlan for LogicalPlanNode { } } + #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn try_from_logical_plan( plan: &LogicalPlan, extension_codec: &dyn LogicalExtensionCodec, @@ -1094,7 +1118,7 @@ impl AsLogicalPlan for LogicalPlanNode { where Self: Sized, { - match plan { + dispatch_logical_plan!(plan, { LogicalPlan::Values(Values { values, .. }) => { let n_cols = if values.is_empty() { 0 @@ -1909,6 +1933,6 @@ impl AsLogicalPlan for LogicalPlanNode { ))), }) } - } + }) } } diff --git a/datafusion/proto/tests/cases/mod.rs b/datafusion/proto/tests/cases/mod.rs index aec6c1de30309..18885f90ef409 100644 --- a/datafusion/proto/tests/cases/mod.rs +++ b/datafusion/proto/tests/cases/mod.rs @@ -34,6 +34,7 @@ use std::sync::Arc; mod roundtrip_logical_plan; mod roundtrip_physical_plan; mod serialize; +mod stack_safety; #[derive(Debug, PartialEq, Eq, Hash)] struct MyRegexUdf { diff --git a/datafusion/proto/tests/cases/stack_safety.rs b/datafusion/proto/tests/cases/stack_safety.rs new file mode 100644 index 0000000000000..5caf4119a7186 --- /dev/null +++ b/datafusion/proto/tests/cases/stack_safety.rs @@ -0,0 +1,99 @@ +// 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. + +use std::process::Command; +use std::sync::Arc; + +use datafusion_common::DFSchema; +use datafusion_expr::logical_plan::{EmptyRelation, LogicalPlan, LogicalPlanBuilder}; +use datafusion_proto::bytes::logical_plan_to_bytes; + +const CHILD_ENV: &str = "DATAFUSION_PROTO_ISSUE_23823_CHILD"; +const ALIAS_DEPTH_ENV: &str = "DATAFUSION_PROTO_ISSUE_23823_ALIAS_DEPTH"; +const TWO_MIB_TEST_NAME: &str = + "cases::stack_safety::logical_plan_serialization_fits_a_two_mib_stack"; +#[cfg(feature = "recursive_protection")] +const GROWABLE_STACK_TEST_NAME: &str = + "cases::stack_safety::deeply_nested_logical_plan_serialization_uses_a_growable_stack"; + +fn deeply_aliased_plan(alias_depth: usize) -> LogicalPlan { + let mut plan = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(DFSchema::empty()), + }); + + for level in 0..alias_depth { + plan = LogicalPlanBuilder::from(plan) + .alias(format!("level_{level}")) + .unwrap() + .build() + .unwrap(); + } + + plan +} + +fn serialize_on_two_mib_stack(alias_depth: usize) { + let plan = deeply_aliased_plan(alias_depth); + std::thread::Builder::new() + .name("two-megabyte-stack".into()) + .stack_size(2 * 1024 * 1024) + .spawn(move || logical_plan_to_bytes(&plan).unwrap()) + .unwrap() + .join() + .unwrap(); +} + +fn run_in_child(test_name: &str, alias_depth: usize) { + if std::env::var_os(CHILD_ENV).is_some() { + let alias_depth = std::env::var(ALIAS_DEPTH_ENV).unwrap().parse().unwrap(); + serialize_on_two_mib_stack(alias_depth); + return; + } + + // A native stack overflow aborts the process. Re-run this exact test in a + // child process so a regression produces a normal test failure. + let output = Command::new(std::env::current_exe().unwrap()) + .args(["--exact", test_name, "--nocapture"]) + .env(CHILD_ENV, "1") + .env(ALIAS_DEPTH_ENV, alias_depth.to_string()) + .output() + .unwrap(); + + assert!( + output.status.success(), + "child process failed with status {}\nstdout:\n{}\nstderr:\n{}", + output.status, + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr), + ); +} + +#[test] +fn logical_plan_serialization_fits_a_two_mib_stack() { + // Ten aliases reproduce #23823. Use 100 to provide a safety margin while + // verifying the dispatcher reduction without runtime stack growth. + run_in_child(TWO_MIB_TEST_NAME, 100); +} + +#[cfg(feature = "recursive_protection")] +#[test] +fn deeply_nested_logical_plan_serialization_uses_a_growable_stack() { + // This depth exceeds the 2 MiB thread stack without recursive protection, + // exercising the `recursive` stack-growth checkpoint. + run_in_child(GROWABLE_STACK_TEST_NAME, 2_000); +} From 6f9e9211a3f9da24dd3489528db608cf0cd94478 Mon Sep 17 00:00:00 2001 From: Albert Skalt Date: Mon, 28 Sep 2026 20:38:55 +0700 Subject: [PATCH 2/2] fix(proto): decode stack overflow protection Apply the same approach as applied for the encoding part. --- datafusion/proto/src/logical_plan/mod.rs | 15 ++++---- datafusion/proto/tests/cases/stack_safety.rs | 36 +++++++++++++++----- 2 files changed, 36 insertions(+), 15 deletions(-) diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index be5db34de0207..e8819def06381 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -104,22 +104,22 @@ pub trait AsLogicalPlan: Debug + Send + Sync + Clone { Self: Sized; } -// In debug builds, keep each serializer arm's local temporaries out of the +// In debug builds, keep each [de]serializer arm's local temporaries out of the // recursive dispatcher frame. Without this call boundary, they inflate the // frame of every recursive invocation. #[cfg_attr(debug_assertions, inline(never))] -fn serialize_logical_plan_arm(serializer: F) -> Result +fn serde_logical_plan_arm(f: F) -> Result where - F: FnOnce() -> Result, + F: FnOnce() -> Result, { - serializer() + f() } macro_rules! dispatch_logical_plan { ($plan:expr, { $($pattern:pat => $body:expr $(,)?)+ }) => { match $plan { $( - $pattern => serialize_logical_plan_arm(|| -> Result { + $pattern => serde_logical_plan_arm(|| { $body }), )+ @@ -406,6 +406,7 @@ impl AsLogicalPlan for LogicalPlanNode { .map_err(|e| internal_datafusion_err!("failed to encode logical plan: {e:?}")) } + #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn try_into_logical_plan( &self, ctx: &TaskContext, @@ -416,7 +417,7 @@ impl AsLogicalPlan for LogicalPlanNode { "logical_plan::from_proto() Unsupported logical plan '{self:?}'" )) })?; - match plan { + dispatch_logical_plan!(plan, { LogicalPlanType::Values(values) => { let n_cols = values.n_cols as usize; let values: Vec> = if values.values_list.is_empty() { @@ -1107,7 +1108,7 @@ impl AsLogicalPlan for LogicalPlanNode { Arc::new(into_logical_plan!(dml_node.input, ctx, extension_codec)?), ))) } - } + }) } #[cfg_attr(feature = "recursive_protection", recursive::recursive)] diff --git a/datafusion/proto/tests/cases/stack_safety.rs b/datafusion/proto/tests/cases/stack_safety.rs index 5caf4119a7186..95e43549ee82d 100644 --- a/datafusion/proto/tests/cases/stack_safety.rs +++ b/datafusion/proto/tests/cases/stack_safety.rs @@ -19,16 +19,19 @@ use std::process::Command; use std::sync::Arc; use datafusion_common::DFSchema; +use datafusion_execution::TaskContext; use datafusion_expr::logical_plan::{EmptyRelation, LogicalPlan, LogicalPlanBuilder}; -use datafusion_proto::bytes::logical_plan_to_bytes; +use datafusion_proto::bytes::{logical_plan_from_bytes, logical_plan_to_bytes}; const CHILD_ENV: &str = "DATAFUSION_PROTO_ISSUE_23823_CHILD"; const ALIAS_DEPTH_ENV: &str = "DATAFUSION_PROTO_ISSUE_23823_ALIAS_DEPTH"; const TWO_MIB_TEST_NAME: &str = - "cases::stack_safety::logical_plan_serialization_fits_a_two_mib_stack"; + "cases::stack_safety::logical_plan_serde_fits_a_two_mib_stack"; #[cfg(feature = "recursive_protection")] const GROWABLE_STACK_TEST_NAME: &str = - "cases::stack_safety::deeply_nested_logical_plan_serialization_uses_a_growable_stack"; + "cases::stack_safety::deeply_nested_logical_plan_serde_uses_a_growable_stack"; + +const RECURSION_LIMIT_EXPECTED_DEPTH: usize = 20; fn deeply_aliased_plan(alias_depth: usize) -> LogicalPlan { let mut plan = LogicalPlan::EmptyRelation(EmptyRelation { @@ -47,12 +50,26 @@ fn deeply_aliased_plan(alias_depth: usize) -> LogicalPlan { plan } -fn serialize_on_two_mib_stack(alias_depth: usize) { +fn serde_on_two_mib_stack(alias_depth: usize) { let plan = deeply_aliased_plan(alias_depth); std::thread::Builder::new() .name("two-megabyte-stack".into()) .stack_size(2 * 1024 * 1024) - .spawn(move || logical_plan_to_bytes(&plan).unwrap()) + .spawn(move || { + let bytes = logical_plan_to_bytes(&plan).unwrap(); + match logical_plan_from_bytes(&bytes, &TaskContext::default()) { + Ok(_) => {} + Err(err) + if err.to_string().contains("recursion limit") + && alias_depth > RECURSION_LIMIT_EXPECTED_DEPTH => + { + // Ignore the error as it is expected for this depth. + } + res => { + res.unwrap(); + } + } + }) .unwrap() .join() .unwrap(); @@ -61,7 +78,7 @@ fn serialize_on_two_mib_stack(alias_depth: usize) { fn run_in_child(test_name: &str, alias_depth: usize) { if std::env::var_os(CHILD_ENV).is_some() { let alias_depth = std::env::var(ALIAS_DEPTH_ENV).unwrap().parse().unwrap(); - serialize_on_two_mib_stack(alias_depth); + serde_on_two_mib_stack(alias_depth); return; } @@ -84,15 +101,18 @@ fn run_in_child(test_name: &str, alias_depth: usize) { } #[test] -fn logical_plan_serialization_fits_a_two_mib_stack() { +fn logical_plan_serde_fits_a_two_mib_stack() { // Ten aliases reproduce #23823. Use 100 to provide a safety margin while // verifying the dispatcher reduction without runtime stack growth. run_in_child(TWO_MIB_TEST_NAME, 100); + // Check also for a depth that does not trigger a recursion limit error + // on deserialization. + run_in_child(TWO_MIB_TEST_NAME, RECURSION_LIMIT_EXPECTED_DEPTH); } #[cfg(feature = "recursive_protection")] #[test] -fn deeply_nested_logical_plan_serialization_uses_a_growable_stack() { +fn deeply_nested_logical_plan_serde_uses_a_growable_stack() { // This depth exceeds the 2 MiB thread stack without recursive protection, // exercising the `recursive` stack-growth checkpoint. run_in_child(GROWABLE_STACK_TEST_NAME, 2_000);