Skip to content

Commit 4d51e14

Browse files
sokrasokra
committed
use double locking to balance edges
1 parent e9d3828 commit 4d51e14

File tree

13 files changed

+306
-593
lines changed

13 files changed

+306
-593
lines changed

‎turbopack/crates/turbo-tasks-memory/src/aggregation/balance_edge.rs

Lines changed: 151 additions & 169 deletions
Original file line numberDiff line numberDiff line change
@@ -2,41 +2,29 @@ use std::cmp::Ordering;
22

33
use super::{
44
balance_queue::BalanceQueue,
5-
followers::{
6-
add_follower_count, remove_follower_count, remove_positive_follower_count,
7-
RemovePositveFollowerCountResult,
8-
},
9-
in_progress::is_in_progress,
5+
in_progress::{is_in_progress, start_in_progress_all, start_in_progress_count},
106
increase::IncreaseReason,
117
increase_aggregation_number_internal,
12-
uppers::{
13-
add_upper_count, remove_positive_upper_count, remove_upper_count,
14-
RemovePositiveUpperCountResult,
15-
},
16-
AggregationContext, AggregationNode,
8+
notify_lost_follower::notify_lost_follower,
9+
notify_new_follower::notify_new_follower,
10+
util::{get_aggregated_add_change, get_aggregated_remove_change, get_followers_or_children},
11+
AggregationContext, AggregationNode, PreparedInternalOperation, PreparedOperation, StackVec,
1712
};
1813

19-
// Migrated followers to uppers or uppers to followers depending on the
14+
// Migrate followers to uppers or uppers to followers depending on the
2015
// aggregation numbers of the nodes involved in the edge. Might increase targets
2116
// aggregation number if they are equal.
2217
pub(super) fn balance_edge<C: AggregationContext>(
2318
ctx: &C,
2419
balance_queue: &mut BalanceQueue<C::NodeRef>,
2520
upper_id: &C::NodeRef,
26-
mut upper_aggregation_number: u32,
2721
target_id: &C::NodeRef,
28-
mut target_aggregation_number: u32,
29-
) -> (u32, u32) {
30-
// too many uppers on target
31-
let mut extra_uppers = 0;
32-
// too many followers on upper
33-
let mut extra_followers = 0;
34-
// The last info about uppers
35-
let mut uppers_count: Option<isize> = None;
36-
// The last info about followers
37-
let mut followers_count = None;
38-
22+
) {
3923
loop {
24+
let (mut upper, mut target) = ctx.node_pair(upper_id, target_id);
25+
let upper_aggregation_number = upper.aggregation_number();
26+
let target_aggregation_number = target.aggregation_number();
27+
4028
let root = upper_aggregation_number == u32::MAX || target_aggregation_number == u32::MAX;
4129
let order = if root {
4230
Ordering::Greater
@@ -45,164 +33,158 @@ pub(super) fn balance_edge<C: AggregationContext>(
4533
};
4634
match order {
4735
Ordering::Equal => {
48-
// we probably want to increase the aggregation number of target
49-
let upper = ctx.node(upper_id);
50-
upper_aggregation_number = upper.aggregation_number();
5136
drop(upper);
52-
if upper_aggregation_number != u32::MAX
53-
&& upper_aggregation_number == target_aggregation_number
54-
{
55-
let target = ctx.node(target_id);
56-
target_aggregation_number = target.aggregation_number();
57-
if upper_aggregation_number == target_aggregation_number {
58-
// increase target aggregation number
59-
increase_aggregation_number_internal(
60-
ctx,
61-
balance_queue,
62-
target,
63-
target_id,
64-
target_aggregation_number + 1,
65-
target_aggregation_number + 1,
66-
IncreaseReason::EqualAggregationNumberOnBalance,
67-
);
68-
}
69-
}
37+
// increase target aggregation number
38+
increase_aggregation_number_internal(
39+
ctx,
40+
balance_queue,
41+
target,
42+
target_id,
43+
target_aggregation_number + 1,
44+
target_aggregation_number + 1,
45+
IncreaseReason::EqualAggregationNumberOnBalance,
46+
);
7047
}
7148
Ordering::Less => {
72-
// target should probably be a follower of upper
73-
if uppers_count.map_or(false, |count| count <= 0) {
74-
// We already removed all uppers, maybe too many
49+
if is_in_progress(ctx, upper_id) {
50+
drop(target);
51+
let AggregationNode::Aggegating(aggregating) = &mut *upper else {
52+
unreachable!();
53+
};
54+
aggregating
55+
.enqueued_balancing
56+
.push((upper_id.clone(), target_id.clone()));
57+
drop(upper);
58+
// Somebody else will balance this edge
7559
break;
76-
} else if extra_followers == 0 {
77-
let upper = ctx.node(upper_id);
78-
upper_aggregation_number = upper.aggregation_number();
79-
if upper_aggregation_number < target_aggregation_number {
80-
// target should be a follower of upper
81-
// add some extra followers
82-
let count = uppers_count.unwrap_or(1) as usize;
83-
extra_followers += count;
84-
followers_count = Some(add_follower_count(
85-
ctx,
86-
balance_queue,
87-
upper,
88-
upper_id,
89-
target_id,
90-
count,
91-
true,
92-
));
93-
}
60+
}
61+
62+
// target should be a follower of upper
63+
let count = target
64+
.uppers_mut()
65+
.remove_all_positive_clonable_count(upper_id);
66+
if count == 0 {
67+
break;
68+
}
69+
let added = upper
70+
.followers_mut()
71+
.unwrap()
72+
.add_clonable_count(target_id, count);
73+
74+
// target removed as upper
75+
let remove_change = get_aggregated_remove_change(ctx, &target);
76+
let followers = get_followers_or_children(ctx, &target);
77+
78+
let upper_uppers = if added {
79+
// target added as follower
80+
let uppers = upper.uppers().iter().cloned().collect::<StackVec<_>>();
81+
start_in_progress_all(ctx, &uppers);
82+
uppers
9483
} else {
95-
// we already have extra followers, remove some uppers to balance
96-
let count = extra_followers + extra_uppers;
97-
let target = ctx.node(target_id);
98-
if is_in_progress(ctx, upper_id) {
99-
drop(target);
100-
let mut upper = ctx.node(upper_id);
101-
if is_in_progress(ctx, upper_id) {
102-
let AggregationNode::Aggegating(aggregating) = &mut *upper else {
103-
unreachable!();
104-
};
105-
aggregating.enqueued_balancing.push((
106-
upper_id.clone(),
107-
upper_aggregation_number,
108-
target_id.clone(),
109-
target_aggregation_number,
110-
));
111-
drop(upper);
112-
// Somebody else will balance this edge
113-
return (upper_aggregation_number, target_aggregation_number);
114-
}
115-
} else {
116-
let RemovePositiveUpperCountResult {
117-
removed_count,
118-
remaining_count,
119-
} = remove_positive_upper_count(
120-
ctx,
121-
balance_queue,
122-
target,
123-
upper_id,
124-
count,
125-
);
126-
decrease_numbers(removed_count, &mut extra_uppers, &mut extra_followers);
127-
uppers_count = Some(remaining_count);
128-
}
84+
Default::default()
85+
};
86+
87+
drop(target);
88+
89+
// target removed as upper
90+
let remove_prepared =
91+
remove_change.and_then(|remove_change| upper.apply_change(ctx, remove_change));
92+
start_in_progress_count(ctx, upper_id, followers.len() as u32);
93+
let prepared = followers
94+
.into_iter()
95+
.map(|child_id| {
96+
upper.notify_lost_follower(ctx, balance_queue, upper_id, &child_id)
97+
})
98+
.collect::<StackVec<_>>();
99+
drop(upper);
100+
101+
// target added as follower
102+
for upper_id in upper_uppers {
103+
notify_new_follower(
104+
ctx,
105+
balance_queue,
106+
ctx.node(&upper_id),
107+
&upper_id,
108+
target_id,
109+
false,
110+
);
129111
}
112+
113+
// target removed as upper
114+
remove_prepared.apply(ctx);
115+
prepared.apply(ctx, balance_queue);
116+
117+
break;
130118
}
131119
Ordering::Greater => {
132-
// target should probably be an inner node of upper
133-
if followers_count.map_or(false, |count| count <= 0) {
134-
// We already removed all followers, maybe too many
120+
if is_in_progress(ctx, upper_id) {
121+
let AggregationNode::Aggegating(aggregating) = &mut *upper else {
122+
unreachable!();
123+
};
124+
aggregating
125+
.enqueued_balancing
126+
.push((upper_id.clone(), target_id.clone()));
127+
drop(upper);
128+
// Somebody else will balance this edge
129+
break;
130+
}
131+
132+
// target should be a inner node of upper
133+
let count = upper
134+
.followers_mut()
135+
.unwrap()
136+
.remove_all_positive_clonable_count(target_id);
137+
if count == 0 {
135138
break;
136-
} else if extra_uppers == 0 {
137-
let target = ctx.node(target_id);
138-
target_aggregation_number = target.aggregation_number();
139-
if root || target_aggregation_number < upper_aggregation_number {
140-
// target should be a inner node of upper
141-
if is_in_progress(ctx, upper_id) {
142-
drop(target);
143-
let mut upper = ctx.node(upper_id);
144-
if is_in_progress(ctx, upper_id) {
145-
let AggregationNode::Aggegating(aggregating) = &mut *upper else {
146-
unreachable!();
147-
};
148-
aggregating.enqueued_balancing.push((
149-
upper_id.clone(),
150-
upper_aggregation_number,
151-
target_id.clone(),
152-
target_aggregation_number,
153-
));
154-
drop(upper);
155-
// Somebody else will balance this edge
156-
return (upper_aggregation_number, target_aggregation_number);
157-
}
158-
} else {
159-
// add some extra uppers
160-
let count = followers_count.unwrap_or(1) as usize;
161-
extra_uppers += count;
162-
uppers_count = Some(
163-
add_upper_count(
164-
ctx,
165-
balance_queue,
166-
target,
167-
target_id,
168-
upper_id,
169-
count,
170-
true,
171-
)
172-
.new_count,
173-
);
174-
}
175-
}
139+
}
140+
let added = target.uppers_mut().add_clonable_count(upper_id, count);
141+
142+
// target removed as follower
143+
let uppers = upper.uppers().iter().cloned().collect::<StackVec<_>>();
144+
start_in_progress_all(ctx, &uppers);
145+
146+
let (add_change, followers) = if added {
147+
// target added as upper
148+
let add_change = get_aggregated_add_change(ctx, &target);
149+
let followers = get_followers_or_children(ctx, &target);
150+
start_in_progress_count(ctx, upper_id, followers.len() as u32);
151+
(add_change, followers)
176152
} else {
177-
// we already have extra uppers, try to remove some followers to balance
178-
let count = extra_followers + extra_uppers;
179-
let upper = ctx.node(upper_id);
180-
let RemovePositveFollowerCountResult {
181-
removed_count,
182-
remaining_count,
183-
} = remove_positive_follower_count(ctx, balance_queue, upper, target_id, count);
184-
decrease_numbers(removed_count, &mut extra_followers, &mut extra_uppers);
185-
followers_count = Some(remaining_count);
153+
(None, Default::default())
154+
};
155+
156+
drop(target);
157+
158+
// target added as upper
159+
let add_prepared =
160+
add_change.and_then(|add_change| upper.apply_change(ctx, add_change));
161+
let prepared = followers
162+
.into_iter()
163+
.filter_map(|child_id| {
164+
upper.notify_new_follower(ctx, balance_queue, upper_id, &child_id, false)
165+
})
166+
.collect::<StackVec<_>>();
167+
168+
drop(upper);
169+
170+
add_prepared.apply(ctx);
171+
for prepared in prepared {
172+
prepared.apply(ctx, balance_queue);
186173
}
174+
175+
// target removed as follower
176+
for upper_id in uppers {
177+
notify_lost_follower(
178+
ctx,
179+
balance_queue,
180+
ctx.node(&upper_id),
181+
&upper_id,
182+
target_id,
183+
);
184+
}
185+
186+
break;
187187
}
188188
}
189189
}
190-
if extra_followers > 0 {
191-
let upper = ctx.node(upper_id);
192-
remove_follower_count(ctx, balance_queue, upper, target_id, extra_followers);
193-
}
194-
if extra_uppers > 0 {
195-
let target = ctx.node(target_id);
196-
remove_upper_count(ctx, balance_queue, target, upper_id, extra_uppers);
197-
}
198-
(upper_aggregation_number, target_aggregation_number)
199-
}
200-
201-
fn decrease_numbers(amount: usize, a: &mut usize, b: &mut usize) {
202-
if *a >= amount {
203-
*a -= amount;
204-
} else {
205-
*b -= amount - *a;
206-
*a = 0;
207-
}
208190
}

0 commit comments

Comments
 (0)