diff --git a/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs b/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs index 790df3995a..399cdf9d0b 100644 --- a/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs +++ b/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs @@ -36,21 +36,22 @@ pub struct ModelAnalyzeRule { analyzed_wren_mdl: Arc, session_state: SessionStateRef, properties: SessionPropertiesRef, - /// Stack of model names currently being resolved through RLAC. Shared across - /// recursive `analyze_*` calls (including those triggered by subqueries inside an - /// RLAC condition) so we can detect cycles like A's RLAC referencing B whose RLAC - /// references A. - building_models: Arc>>, } -/// RAII guard that removes a model name from the `building_models` stack on drop, -/// regardless of how the surrounding function exits. -struct ModelStackGuard { - stack: Arc>>, +/// Cycle-detection stack for RLAC resolution (A's RLAC referencing B whose +/// RLAC references A). Allocated per `analyze` invocation and passed down +/// the recursive calls: the rule instance is shared by every — possibly +/// concurrent — query on its session context, so this must not live on `self`. +type ModelStack = Mutex>; + +/// RAII guard that removes a model name from the cycle-detection stack on +/// drop, regardless of how the surrounding function exits. +struct ModelStackGuard<'a> { + stack: &'a ModelStack, name: String, } -impl Drop for ModelStackGuard { +impl Drop for ModelStackGuard<'_> { fn drop(&mut self) { self.stack.lock().remove(&self.name); } @@ -64,16 +65,16 @@ impl Debug for ModelAnalyzeRule { impl AnalyzerRule for ModelAnalyzeRule { fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) -> Result { - // Each top-level invocation starts with a clean cycle-detection stack so the - // rule instance can be reused across queries. - self.building_models.lock().clear(); + // One stack per invocation — see [`ModelStack`] for why it must not + // be shared across queries. + let cycle_stack = ModelStack::default(); let mut scope_manager = ScopeManager::new(); let root_scope_id = scope_manager.create_root_scope(); self.analyze_scope(plan, &mut scope_manager, root_scope_id)? .map_data(|plan| { - self.analyze_model(plan, &mut scope_manager, root_scope_id) + self.analyze_model(plan, &mut scope_manager, root_scope_id, &cycle_stack) .data() })? .map_data(|plan| { @@ -103,7 +104,6 @@ impl ModelAnalyzeRule { analyzed_wren_mdl, session_state, properties, - building_models: Arc::new(Mutex::new(HashSet::new())), } } @@ -349,10 +349,16 @@ impl ModelAnalyzeRule { plan: LogicalPlan, scope_manager: &mut ScopeManager, current_scope_id: ScopeId, + cycle_stack: &ModelStack, ) -> Result> { plan.transform_up(&mut |plan| -> Result> { let plan = self - .analyze_model_internal(plan, scope_manager, current_scope_id)? + .analyze_model_internal( + plan, + scope_manager, + current_scope_id, + cycle_stack, + )? .data; // If the plan contains subquery, we should analyze the subquery recursively plan.map_subqueries(|plan| { @@ -366,6 +372,7 @@ impl ModelAnalyzeRule { Arc::unwrap_or_clone(Arc::clone(&subquery.subquery)), scope_manager, child_scope_id, + cycle_stack, )? .data; return Ok(Transformed::yes(LogicalPlan::Subquery( @@ -383,6 +390,7 @@ impl ModelAnalyzeRule { plan: LogicalPlan, scope_manager: &mut ScopeManager, current_scope_id: ScopeId, + cycle_stack: &ModelStack, ) -> Result> { match plan { LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) => { @@ -397,16 +405,16 @@ impl ModelAnalyzeRule { scope_manager, current_scope_id, alias, + cycle_stack, ), LogicalPlan::TableScan(table_scan) => { let model_plan = self .analyze_table_scan( - Arc::clone(&self.analyzed_wren_mdl), - Arc::clone(&self.session_state), table_scan, Some(alias.clone()), scope_manager, current_scope_id, + cycle_stack, )? .data; let subquery = @@ -419,23 +427,21 @@ impl ModelAnalyzeRule { } } LogicalPlan::TableScan(table_scan) => self.analyze_table_scan( - Arc::clone(&self.analyzed_wren_mdl), - Arc::clone(&self.session_state), table_scan, None, scope_manager, current_scope_id, + cycle_stack, ), LogicalPlan::Join(join) => { let left = match Arc::unwrap_or_clone(join.left) { LogicalPlan::TableScan(table_scan) => { self.analyze_table_scan( - Arc::clone(&self.analyzed_wren_mdl), - Arc::clone(&self.session_state), table_scan, None, scope_manager, current_scope_id, + cycle_stack, )? .data } @@ -445,12 +451,11 @@ impl ModelAnalyzeRule { let right = match Arc::unwrap_or_clone(join.right) { LogicalPlan::TableScan(table_scan) => { self.analyze_table_scan( - Arc::clone(&self.analyzed_wren_mdl), - Arc::clone(&self.session_state), table_scan, None, scope_manager, current_scope_id, + cycle_stack, )? .data } @@ -475,9 +480,10 @@ impl ModelAnalyzeRule { /// Construct a [`ModelPlanNode`] with cycle-aware RLAC handling. /// /// Steps: - /// 1. Push the model's name onto the shared `building_models` stack; error out if - /// it is already present (which means an RLAC condition transitively re-entered - /// this model). An RAII guard removes the name when this function exits. + /// 1. Push the model's name onto the invocation's cycle-detection stack; error out + /// if it is already present (which means an RLAC condition transitively + /// re-entered this model). An RAII guard removes the name when this function + /// exits. /// 2. Build the [`ModelPlanNode`] (the builder parses each matching RLAC condition /// via [`crate::logical_plan::analyze::access_control::RlacContextProvider`], so /// table references inside subqueries are resolved against MDL models). @@ -490,10 +496,11 @@ impl ModelAnalyzeRule { model: Arc, required_fields: Vec, original_table_scan: Option, + cycle_stack: &ModelStack, ) -> Result { let model_name = model.name().to_string(); { - let mut stack = self.building_models.lock(); + let mut stack = cycle_stack.lock(); if stack.contains(&model_name) { return plan_err!( "Detected a cycle in row level access control conditions for model `{}`", @@ -503,7 +510,7 @@ impl ModelAnalyzeRule { stack.insert(model_name.clone()); } let _guard = ModelStackGuard { - stack: Arc::clone(&self.building_models), + stack: cycle_stack, name: model_name, }; @@ -517,7 +524,8 @@ impl ModelAnalyzeRule { )?; if let Some(filter) = plan_node.rlac_filter.take() { - plan_node.rlac_filter = Some(self.analyze_rlac_subqueries(filter)?); + plan_node.rlac_filter = + Some(self.analyze_rlac_subqueries(filter, cycle_stack)?); } Ok(plan_node) } @@ -526,12 +534,18 @@ impl ModelAnalyzeRule { /// (`ScalarSubquery`, `InSubquery`, `Exists`). Each inner plan is processed with a /// fresh `ScopeManager`/scope id — RLAC subqueries are introduced after the outer /// scope analysis runs, so they don't have entries in the outer `scope_manager`. - fn analyze_rlac_subqueries(&self, expr: Expr) -> Result { + fn analyze_rlac_subqueries( + &self, + expr: Expr, + cycle_stack: &ModelStack, + ) -> Result { expr.transform_down(|expr| -> Result> { match expr { Expr::ScalarSubquery(sq) => { - let plan = - self.analyze_subquery_plan(Arc::unwrap_or_clone(sq.subquery))?; + let plan = self.analyze_subquery_plan( + Arc::unwrap_or_clone(sq.subquery), + cycle_stack, + )?; Ok(Transformed::yes(Expr::ScalarSubquery(Subquery { subquery: Arc::new(plan), outer_ref_columns: sq.outer_ref_columns, @@ -543,8 +557,10 @@ impl ModelAnalyzeRule { subquery, negated, }) => { - let plan = self - .analyze_subquery_plan(Arc::unwrap_or_clone(subquery.subquery))?; + let plan = self.analyze_subquery_plan( + Arc::unwrap_or_clone(subquery.subquery), + cycle_stack, + )?; Ok(Transformed::yes(Expr::InSubquery(InSubquery { expr, subquery: Subquery { @@ -556,8 +572,10 @@ impl ModelAnalyzeRule { }))) } Expr::Exists(Exists { subquery, negated }) => { - let plan = self - .analyze_subquery_plan(Arc::unwrap_or_clone(subquery.subquery))?; + let plan = self.analyze_subquery_plan( + Arc::unwrap_or_clone(subquery.subquery), + cycle_stack, + )?; Ok(Transformed::yes(Expr::Exists(Exists { subquery: Subquery { subquery: Arc::new(plan), @@ -577,12 +595,16 @@ impl ModelAnalyzeRule { /// on an inner subquery plan with a fresh `ScopeManager`. Used to make `TableScan`s /// introduced by RLAC condition parsing go through the same transformation as /// regular query plans. - fn analyze_subquery_plan(&self, plan: LogicalPlan) -> Result { + fn analyze_subquery_plan( + &self, + plan: LogicalPlan, + cycle_stack: &ModelStack, + ) -> Result { let mut scope_manager = ScopeManager::new(); let root_scope_id = scope_manager.create_root_scope(); self.analyze_scope(plan, &mut scope_manager, root_scope_id)? .map_data(|p| { - self.analyze_model(p, &mut scope_manager, root_scope_id) + self.analyze_model(p, &mut scope_manager, root_scope_id, cycle_stack) .data() })? .map_data(|p| { @@ -597,20 +619,19 @@ impl ModelAnalyzeRule { fn analyze_table_scan( &self, - analyzed_wren_mdl: Arc, - session_state_ref: SessionStateRef, table_scan: TableScan, alias: Option, scope_manager: &mut ScopeManager, current_scope_id: ScopeId, + cycle_stack: &ModelStack, ) -> Result> { if belong_to_mdl( - &analyzed_wren_mdl.wren_mdl(), + &self.analyzed_wren_mdl.wren_mdl(), table_scan.table_name.clone(), - Arc::clone(&session_state_ref), + self.session_state(), ) { let table_name = table_scan.table_name.table(); - if let Some(model) = analyzed_wren_mdl.wren_mdl.get_model(table_name) { + if let Some(model) = self.analyzed_wren_mdl.wren_mdl.get_model(table_name) { let table_ref = alias.unwrap_or(table_scan.table_name.clone()); let field: Vec = if let Some(used_columns) = scope_manager.try_get_required_columns(current_scope_id, &table_ref) @@ -634,6 +655,7 @@ impl ModelAnalyzeRule { Arc::clone(&model), field, Some(LogicalPlan::TableScan(table_scan.clone())), + cycle_stack, )?; let model_plan = LogicalPlan::Extension(Extension { node: Arc::new(model_plan_node), @@ -660,6 +682,7 @@ impl ModelAnalyzeRule { scope_manager: &mut ScopeManager, current_scope_id: ScopeId, alias: TableReference, + cycle_stack: &ModelStack, ) -> Result> { let SubqueryAlias { input, .. } = subquery_alias; if let LogicalPlan::Extension(Extension { node }) = @@ -688,8 +711,12 @@ impl ModelAnalyzeRule { }; vec![] }; - let model_plan_node = - self.build_model_plan_node(Arc::clone(&model), field, None)?; + let model_plan_node = self.build_model_plan_node( + Arc::clone(&model), + field, + None, + cycle_stack, + )?; let model_plan = LogicalPlan::Extension(Extension { node: Arc::new(model_plan_node), }); diff --git a/core/wren-core/core/src/mdl/mod.rs b/core/wren-core/core/src/mdl/mod.rs index c1c07e501b..d2850bd3da 100644 --- a/core/wren-core/core/src/mdl/mod.rs +++ b/core/wren-core/core/src/mdl/mod.rs @@ -5418,4 +5418,127 @@ mod test { Ok(()) } } + + /// A derived context keeps one `ModelAnalyzeRule` instance for its whole + /// lifetime, so concurrent `optimize()` calls share it — the RLAC + /// cycle-detection stack must therefore be per-invocation state. + mod analyzer_concurrency { + use super::*; + use std::sync::Barrier; + use std::thread; + + /// Acyclic chain: `customer`'s RLAC subquery selects from `allowed`, + /// so planning holds the cycle stack across the nested analysis — a + /// wide enough window for concurrent invocations to collide. + fn rlac_chain_manifest() -> Manifest { + ManifestBuilder::new() + .catalog("wren") + .schema("test") + .model( + ModelBuilder::new("customer") + .table_reference("customer_remote") + .column(ColumnBuilder::new("c_custkey", "int").build()) + .add_row_level_access_control( + "by_allowed", + vec![SessionProperty::new_required("session_user")], + "c_custkey IN (SELECT allowed_id FROM allowed WHERE allowed_user = @session_user)", + ) + .build(), + ) + .model( + ModelBuilder::new("allowed") + .table_reference("allowed_remote") + .column(ColumnBuilder::new("allowed_id", "int").build()) + .column(ColumnBuilder::new("allowed_user", "string").build()) + .build(), + ) + .build() + } + + /// Valid (acyclic) queries planned concurrently on one shared + /// derived context must never report a spurious RLAC cycle. + /// + /// Derive once and share — mirrors wren-core-py's long-lived + /// exec_ctx. Drive plans through `SessionState::optimize`: the + /// Analyzer only runs there, and `transform_sql_with_ctx` builds + /// fresh rule instances per call — either shortcut would make the + /// race unobservable. + #[test] + fn concurrent_valid_rlac_queries_never_report_spurious_cycle() -> Result<()> { + const N_THREADS: usize = 8; + const ITERATIONS: usize = 50; + const SQL: &str = "SELECT c_custkey FROM wren.test.customer"; + + let base_ctx = create_wren_ctx(None, None); + let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze( + rlac_chain_manifest(), + Arc::new(HashMap::default()), + Mode::LocalRuntime, + )?); + let headers = Arc::new(build_headers(&[( + "session_user".to_string(), + Some("'alice'".to_string()), + )])); + + let setup_rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let derived_ctx = Arc::new(setup_rt.block_on(apply_wren_on_ctx( + &base_ctx, + Arc::clone(&analyzed_mdl), + headers, + Mode::LocalRuntime, + ))?); + + // Serial sanity check: the fixture must plan cleanly without + // concurrency, so any failure below is attributable to it. + setup_rt.block_on(async { + let state = derived_ctx.state(); + let plan = state.create_logical_plan(SQL).await?; + state.optimize(&plan).map(|_| ()) + })?; + + let barrier = Arc::new(Barrier::new(N_THREADS)); + let handles: Vec<_> = (0..N_THREADS) + .map(|tid| { + let ctx = Arc::clone(&derived_ctx); + let barrier = Arc::clone(&barrier); + thread::Builder::new() + .stack_size(8 * 1024 * 1024) + .spawn(move || -> Result<()> { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + barrier.wait(); + rt.block_on(async move { + for iter in 0..ITERATIONS { + let state = ctx.state(); + let plan = state.create_logical_plan(SQL).await?; + let optimized = state.optimize(&plan); + assert!( + optimized.is_ok(), + "thread {tid} iter {iter}: valid \ + acyclic RLAC query failed under \ + same-context concurrency: {}", + optimized + .err() + .map(|e| e.to_string()) + .unwrap_or_default() + ); + } + Ok(()) + }) + }) + .expect("failed to spawn analyzer stress thread") + }) + .collect(); + + for handle in handles { + handle.join().expect("analyzer stress thread panicked")?; + } + Ok(()) + } + } }