Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
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
119 changes: 73 additions & 46 deletions core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,21 +36,22 @@ pub struct ModelAnalyzeRule {
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
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<Mutex<HashSet<String>>>,
}

/// 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<Mutex<HashSet<String>>>,
/// 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<HashSet<String>>;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Now that the stack is per-invocation and never crosses a thread boundary, the Mutex is permanently uncontended — and, more importantly, the type now says the opposite of what this PR just established. Mutex reads as "shared across threads", which is precisely the property being removed.

RefCell would encode the new invariant in the type system: it is !Sync, so any future attempt to move this back onto ModelAnalyzeRule as a field, or to share it across threads, becomes a compile error rather than a silently reintroduced race. That seems worth having right before step 3 removes call_lock and real concurrent traffic starts arriving here.

Secondary benefit: the scoped block in build_model_plan_node that drops the borrow before recursing is load-bearing. If someone later widens it across the recursive call, parking_lot::Mutex (non-reentrant) deadlocks, whereas RefCell panics with a location. The latter is far easier to diagnose.

I prototyped it to make sure this is not just theory — it is a 4-line change:

type ModelStack = RefCell<HashSet<String>>;
//  cycle_stack.lock()      -> cycle_stack.borrow_mut()
//  self.stack.lock()       -> self.stack.borrow_mut()
//  use parking_lot::Mutex  -> use std::cell::RefCell

Result: clippy --all-targets --all-features -- -D warnings clean (exit 0, zero warnings, forced fresh analysis), cargo test --lib 148/148, concurrency regression still passes. No Send/Sync obstacles.

(Unrelated and pre-existing, just noting it while we are here: ModelStack / ModelStackGuard are named "stack" but the underlying type is a HashSet, which has no ordering.)


/// 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);
}
Expand All @@ -64,16 +65,16 @@ impl Debug for ModelAnalyzeRule {

impl AnalyzerRule for ModelAnalyzeRule {
fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) -> Result<LogicalPlan> {
// 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| {
Expand Down Expand Up @@ -103,7 +104,6 @@ impl ModelAnalyzeRule {
analyzed_wren_mdl,
session_state,
properties,
building_models: Arc::new(Mutex::new(HashSet::new())),
}
}

Expand Down Expand Up @@ -349,10 +349,16 @@ impl ModelAnalyzeRule {
plan: LogicalPlan,
scope_manager: &mut ScopeManager,
current_scope_id: ScopeId,
cycle_stack: &ModelStack,
) -> Result<Transformed<LogicalPlan>> {
plan.transform_up(&mut |plan| -> Result<Transformed<LogicalPlan>> {
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| {
Expand All @@ -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(
Expand All @@ -383,6 +390,7 @@ impl ModelAnalyzeRule {
plan: LogicalPlan,
scope_manager: &mut ScopeManager,
current_scope_id: ScopeId,
cycle_stack: &ModelStack,
) -> Result<Transformed<LogicalPlan>> {
match plan {
LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) => {
Expand All @@ -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 =
Expand All @@ -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
}
Expand All @@ -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
}
Expand All @@ -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).
Expand All @@ -490,10 +496,11 @@ impl ModelAnalyzeRule {
model: Arc<wren_core_base::mdl::Model>,
required_fields: Vec<Expr>,
original_table_scan: Option<LogicalPlan>,
cycle_stack: &ModelStack,
) -> Result<ModelPlanNode> {
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 `{}`",
Expand All @@ -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,
};

Expand All @@ -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)
}
Expand All @@ -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<Expr> {
fn analyze_rlac_subqueries(
&self,
expr: Expr,
cycle_stack: &ModelStack,
) -> Result<Expr> {
expr.transform_down(|expr| -> Result<Transformed<Expr>> {
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,
Expand All @@ -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 {
Expand All @@ -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),
Expand All @@ -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<LogicalPlan> {
fn analyze_subquery_plan(
&self,
plan: LogicalPlan,
cycle_stack: &ModelStack,
) -> Result<LogicalPlan> {
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| {
Expand All @@ -597,20 +619,19 @@ impl ModelAnalyzeRule {

fn analyze_table_scan(
&self,
analyzed_wren_mdl: Arc<AnalyzedWrenMDL>,
session_state_ref: SessionStateRef,
table_scan: TableScan,
alias: Option<TableReference>,
scope_manager: &mut ScopeManager,
current_scope_id: ScopeId,
cycle_stack: &ModelStack,
) -> Result<Transformed<LogicalPlan>> {
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<Expr> = if let Some(used_columns) =
scope_manager.try_get_required_columns(current_scope_id, &table_ref)
Expand All @@ -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),
Expand All @@ -660,6 +682,7 @@ impl ModelAnalyzeRule {
scope_manager: &mut ScopeManager,
current_scope_id: ScopeId,
alias: TableReference,
cycle_stack: &ModelStack,
) -> Result<Transformed<LogicalPlan>> {
let SubqueryAlias { input, .. } = subquery_alias;
if let LogicalPlan::Extension(Extension { node }) =
Expand Down Expand Up @@ -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),
});
Expand Down
Loading
Loading