Skip to content
Merged
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
48 changes: 30 additions & 18 deletions crates/dbsp/src/operator/dynamic/balance/balancer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -778,26 +778,38 @@ impl BalancerInner {
// );
let cluster_index = *self.stream_to_cluster.get(&node_id).unwrap();

let (_exchange_sender, _persistent_id, hints) = self
.integrals
.get_mut(&node_id)
.ok_or(BalancerError::NotRegisteredWithBalancer(node_id))?;

let mut hints = hints.clone();

match hint {
BalancerHint::Policy(policy) => {
if let Some(policy) = policy {
hints.policy_hint = Some(policy);
self.solve_cluster(cluster_index, false)?;
} else {
hints.policy_hint = None;
}
// Apply the hint to the stored hints, remembering the previous value so we
// can roll back if the hint turns out to be unenforceable. The hint must be
// stored *before* validating below, because `solve_cluster` reads the hint
// from `self.integrals`; validating against the pre-hint state would accept
// hints that make the cluster unsatisfiable and crash at the first step.
let old_hints = {
let (_exchange_sender, _persistent_id, hints) = self
.integrals
.get_mut(&node_id)
.ok_or(BalancerError::NotRegisteredWithBalancer(node_id))?;

let old_hints = hints.clone();
match hint {
BalancerHint::Policy(policy) => hints.policy_hint = policy,
BalancerHint::Size(size) => hints.size_hint = size,
BalancerHint::Skew(skew) => hints.skew_hint = skew,
}
BalancerHint::Size(size) => hints.size_hint = size,
BalancerHint::Skew(skew) => hints.skew_hint = skew,
old_hints
};

if let BalancerHint::Policy(Some(policy)) = hint
&& self.solve_cluster(cluster_index, false).is_err()
{
self.integrals.get_mut(&node_id).unwrap().2 = old_hints;
return Err(BalancerError::InvalidPolicyHint(
policy,
format!(
"'{policy}' policy is incompatible with the partitioning constraints of the streams it joins with"
),
));
}
self.integrals.get_mut(&node_id).unwrap().2 = hints.clone();

self.clusters[cluster_index].refresh_requested = true;

Ok(())
Expand Down
45 changes: 45 additions & 0 deletions crates/dbsp/src/operator/dynamic/balance/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1147,6 +1147,51 @@ fn test_accumulate_trace_with_balancer_policy_returns_to_trace_policy() {
assert_eq!(output_trace, expected_output_trace);
}

/// A `Broadcast` policy hint on the outer (left) input of a left join cannot be
/// honored: the left side of a left join must be Shard or Balance. The hint must
/// be rejected when it is installed, rather than accepted and then crashing the
/// circuit with `NoSolution` on the first step.
///
/// Regression test: `set_hint` used to validate the cluster *before* persisting
/// the hint, so it checked a still-satisfiable (hint-free) instance, accepted the
/// hint, and the first step then solved the now-unsatisfiable instance and
/// panicked in `solve_all_clusters().unwrap()`. This is the crash produced by a
/// SQL `/*+ broadcast(pt) */` hint where `pt` is the outer input of a left join.
#[test]
fn test_reject_broadcast_hint_on_left_join_outer_input() {
init_test_logger();

let workers = 4;
let (
mut circuit,
(left_input_handle, right_input_handle, left_input_node_id, _right_input_node_id, _output),
) = Runtime::init_circuit(
CircuitConfig::from(workers)
.with_splitter_chunk_size_records(2)
.with_balancer_min_absolute_improvement_threshold(0)
.with_mode(Mode::Persistent),
left_join_with_balancer_test_circuit,
)
.unwrap();

// Broadcasting the outer input of a left join is infeasible, so the hint must
// be rejected here instead of being accepted and crashing at the first step.
let result = circuit.set_balancer_hint(
&left_input_node_id,
BalancerHint::Policy(Some(PartitioningPolicy::Broadcast)),
);
assert!(
result.is_err(),
"broadcast hint on the outer input of a left join must be rejected, got {result:?}"
);

// With the hint rejected and rolled back, the circuit runs normally instead of
// crashing with `NoSolution` on the first step.
left_input_handle.push(1, (1, 1));
right_input_handle.push(1, (1, 1));
circuit.transaction().unwrap();
}

/// Join a large left collection with a small right collection. Both are skewed.
/// The balancer should balance the left collection using Policy::Balance and the
/// right collection using Policy::Broadcast.
Expand Down