1use std::sync::Arc;
5
6use consensus_config::{AuthorityIndex, Stake};
7use mysten_common::ZipDebugEqIteratorExt;
8use parking_lot::RwLock;
9use tracing::{debug, info};
10
11use crate::{
12 context::Context, dag_state::DagState, leader_scoring::ReputationScores,
13 round_tracker::QuorumRound,
14};
15
16#[derive(Debug, Clone, Copy, Eq, PartialEq)]
17pub(crate) enum AncestorState {
18 Include,
19 Exclude(u64),
21}
22
23#[derive(Clone)]
24struct AncestorInfo {
25 state: AncestorState,
26 lock_until_round: u32,
29}
30
31impl AncestorInfo {
32 fn new() -> Self {
33 Self {
34 state: AncestorState::Include,
35 lock_until_round: 0,
36 }
37 }
38
39 fn is_locked(&self, current_clock_round: u32) -> bool {
40 self.lock_until_round >= current_clock_round
41 }
42
43 fn set_lock(&mut self, lock_until_round: u32) {
44 self.lock_until_round = lock_until_round;
45 }
46}
47
48#[derive(Debug)]
49struct StateTransition {
50 authority_id: AuthorityIndex,
51 score: u64,
53 stake: u64,
55 high_quorum_round: u32,
58}
59
60pub(crate) struct AncestorStateManager {
61 context: Arc<Context>,
62 dag_state: Arc<RwLock<DagState>>,
63 state_map: Vec<AncestorInfo>,
64 excluded_nodes_stake_threshold: u64,
65 total_excluded_stake: Stake,
68 pub(crate) propagation_scores: ReputationScores,
71}
72
73impl AncestorStateManager {
74 #[cfg(not(test))]
77 const STATE_LOCK_CLOCK_ROUNDS: u32 = 450;
78 #[cfg(test)]
79 const STATE_LOCK_CLOCK_ROUNDS: u32 = 5;
80
81 const SCORE_EXCLUSION_THRESHOLD_PERCENTAGE: u64 = 20;
83
84 pub(crate) fn new(context: Arc<Context>, dag_state: Arc<RwLock<DagState>>) -> Self {
85 let state_map = vec![AncestorInfo::new(); context.committee.size()];
86
87 let excluded_nodes_stake_threshold_percentage =
90 2 * context.protocol_config.bad_nodes_stake_threshold() / 3;
91
92 let excluded_nodes_stake_threshold = (excluded_nodes_stake_threshold_percentage
93 * context.committee.total_stake())
94 / 100 as Stake;
95
96 Self {
97 context,
98 dag_state,
99 state_map,
100 excluded_nodes_stake_threshold,
101 total_excluded_stake: 0,
103 propagation_scores: ReputationScores::default(),
104 }
105 }
106
107 pub(crate) fn set_propagation_scores(&mut self, scores: ReputationScores) {
108 self.propagation_scores = scores;
109 }
110
111 pub(crate) fn get_ancestor_states(&self) -> Vec<AncestorState> {
112 self.state_map.iter().map(|info| info.state).collect()
113 }
114
115 pub(crate) fn update_all_ancestors_state(&mut self, accepted_quorum_rounds: &[QuorumRound]) {
117 let network_high_quorum_round =
121 self.calculate_network_high_quorum_round(accepted_quorum_rounds);
122
123 let current_clock_round = self.dag_state.read().threshold_clock_round();
124 let low_score_threshold = (self.propagation_scores.highest_score()
125 * Self::SCORE_EXCLUSION_THRESHOLD_PERCENTAGE)
126 / 100;
127
128 debug!(
129 "Updating all ancestor state at round {current_clock_round} using network high quorum round of {network_high_quorum_round}, low score threshold of {low_score_threshold}, and exclude stake threshold of {}",
130 self.excluded_nodes_stake_threshold
131 );
132
133 let mut exclude_to_include = Vec::new();
137 let mut include_to_exclude = Vec::new();
138
139 for (idx, score) in self
143 .propagation_scores
144 .scores_per_authority
145 .iter()
146 .enumerate()
147 {
148 let authority_id = self
149 .context
150 .committee
151 .to_authority_index(idx)
152 .expect("Index should be valid");
153 let ancestor_info = &self.state_map[idx];
154 let (_low, authority_high_quorum_round) = accepted_quorum_rounds[idx];
155 let stake = self.context.committee.authority(authority_id).stake;
156
157 if ancestor_info.is_locked(current_clock_round) {
159 continue;
160 }
161
162 match ancestor_info.state {
163 AncestorState::Include => {
164 if *score <= low_score_threshold {
165 include_to_exclude.push(StateTransition {
166 authority_id,
167 score: *score,
168 stake,
169 high_quorum_round: authority_high_quorum_round,
170 });
171 }
172 }
173 AncestorState::Exclude(_) => {
174 if *score > low_score_threshold
175 || authority_high_quorum_round >= network_high_quorum_round
176 {
177 exclude_to_include.push(StateTransition {
178 authority_id,
179 score: *score,
180 stake,
181 high_quorum_round: authority_high_quorum_round,
182 });
183 }
184 }
185 }
186 }
187
188 for transition in exclude_to_include {
191 self.apply_state_change(transition, AncestorState::Include, current_clock_round);
192 }
193
194 include_to_exclude.sort_by_key(|t| t.score);
197
198 for transition in include_to_exclude {
202 if self.total_excluded_stake + transition.stake <= self.excluded_nodes_stake_threshold {
206 let new_state = AncestorState::Exclude(transition.score);
207 self.apply_state_change(transition, new_state, current_clock_round);
208 } else {
209 info!(
210 "Authority {} would have moved to {:?} state with score {} & quorum_round {} but we would have exceeded total excluded stake threshold. current_excluded_stake {} + authority_stake {} > exclude_stake_threshold {}",
211 transition.authority_id,
212 AncestorState::Exclude(transition.score),
213 transition.score,
214 transition.high_quorum_round,
215 self.total_excluded_stake,
216 transition.stake,
217 self.excluded_nodes_stake_threshold
218 );
219 }
220 }
221 }
222
223 fn apply_state_change(
224 &mut self,
225 transition: StateTransition,
226 new_state: AncestorState,
227 current_clock_round: u32,
228 ) {
229 let block_hostname = &self
230 .context
231 .committee
232 .authority(transition.authority_id)
233 .hostname;
234 let ancestor_info = &mut self.state_map[transition.authority_id.value()];
235
236 match (ancestor_info.state, new_state) {
237 (AncestorState::Exclude(_), AncestorState::Include) => {
238 self.total_excluded_stake = self.total_excluded_stake
239 .checked_sub(transition.stake)
240 .expect("total_excluded_stake underflow - trying to subtract more stake than we're tracking as excluded");
241 }
242 (AncestorState::Include, AncestorState::Exclude(_)) => {
243 self.total_excluded_stake += transition.stake;
244 }
245 _ => {
246 panic!("Calls to this function should only be made for state transition.")
247 }
248 }
249
250 ancestor_info.state = new_state;
251 let lock_until_round = current_clock_round + Self::STATE_LOCK_CLOCK_ROUNDS;
252 ancestor_info.set_lock(lock_until_round);
253
254 info!(
255 "Authority {} moved to {new_state:?} state with score {} & quorum_round {} and locked until round {lock_until_round}. Total excluded stake: {}",
256 transition.authority_id,
257 transition.score,
258 transition.high_quorum_round,
259 self.total_excluded_stake
260 );
261
262 self.context
263 .metrics
264 .node_metrics
265 .ancestor_state_change_by_authority
266 .with_label_values(&[
267 block_hostname.as_str(),
268 match new_state {
269 AncestorState::Include => "include",
270 AncestorState::Exclude(_) => "exclude",
271 },
272 ])
273 .inc();
274 }
275
276 fn calculate_network_high_quorum_round(&self, accepted_quorum_rounds: &[QuorumRound]) -> u32 {
284 let committee = &self.context.committee;
285
286 let mut high_quorum_rounds_with_stake = accepted_quorum_rounds
287 .iter()
288 .zip_debug_eq(committee.authorities())
289 .map(|((_low, high), (_, authority))| (*high, authority.stake))
290 .collect::<Vec<_>>();
291 high_quorum_rounds_with_stake.sort();
292
293 let mut total_stake = 0;
294 let mut network_high_quorum_round = 0;
295
296 for (round, stake) in high_quorum_rounds_with_stake.iter() {
297 total_stake += stake;
298 if total_stake >= self.context.committee.quorum_threshold() {
299 network_high_quorum_round = *round;
300 break;
301 }
302 }
303
304 network_high_quorum_round
305 }
306}
307
308#[cfg(test)]
309mod test {
310 use super::*;
311 use crate::{
312 leader_scoring::ReputationScores, storage::mem_store::MemStore,
313 test_dag_builder::DagBuilder,
314 };
315
316 #[tokio::test]
317 async fn test_calculate_network_high_accepted_quorum_round() {
318 telemetry_subscribers::init_for_testing();
319
320 let (context, _key_pairs) = Context::new_for_test(4);
321 let context = Arc::new(context);
322 let store = Arc::new(MemStore::new());
323 let dag_state = Arc::new(RwLock::new(DagState::new(context.clone(), store.clone())));
324
325 let scores = ReputationScores::new((1..=300).into(), vec![1, 2, 4, 3]);
326 let mut ancestor_state_manager =
327 AncestorStateManager::new(context.clone(), dag_state.clone());
328 ancestor_state_manager.set_propagation_scores(scores);
329
330 let network_high_quorum_round =
333 ancestor_state_manager.calculate_network_high_quorum_round(&[(0, 0); 4]);
334 assert_eq!(network_high_quorum_round, 0);
335
336 let accepted_quorum_rounds = vec![(50, 229), (175, 229), (179, 229), (179, 300)];
337
338 let network_high_quorum_round =
339 ancestor_state_manager.calculate_network_high_quorum_round(&accepted_quorum_rounds);
340 assert_eq!(network_high_quorum_round, 229);
341 }
342
343 #[tokio::test]
350 async fn test_update_all_ancestor_state_using_accepted_rounds() {
351 telemetry_subscribers::init_for_testing();
352 let (mut context, _key_pairs) = Context::new_for_test(5);
353 context
354 .protocol_config
355 .set_bad_nodes_stake_threshold_for_testing(33);
356 let context = Arc::new(context);
357 let store = Arc::new(MemStore::new());
358 let dag_state = Arc::new(RwLock::new(DagState::new(context.clone(), store.clone())));
359 let mut dag_builder = DagBuilder::new(context.clone());
360
361 let scores = ReputationScores::new((1..=300).into(), vec![1, 2, 4, 3, 4]);
362 let mut ancestor_state_manager = AncestorStateManager::new(context, dag_state.clone());
363 ancestor_state_manager.set_propagation_scores(scores);
364
365 let accepted_quorum_rounds =
366 vec![(225, 229), (225, 229), (229, 300), (229, 300), (229, 300)];
367 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
368
369 let state_map = ancestor_state_manager.get_ancestor_states();
372 for state in state_map.iter() {
373 assert_eq!(*state, AncestorState::Include);
374 }
375
376 let scores = ReputationScores::new((1..=300).into(), vec![10, 9, 100, 100, 100]);
377 ancestor_state_manager.set_propagation_scores(scores);
378 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
379
380 let state_map = ancestor_state_manager.get_ancestor_states();
386 for (authority, state) in state_map.iter().enumerate() {
387 if authority == 1 {
388 assert_eq!(*state, AncestorState::Exclude(9));
389 } else {
390 assert_eq!(*state, AncestorState::Include);
391 }
392 }
393
394 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
395
396 let state_map = ancestor_state_manager.get_ancestor_states();
399 for (authority, state) in state_map.iter().enumerate() {
400 if authority == 1 {
401 assert_eq!(*state, AncestorState::Exclude(9));
402 } else {
403 assert_eq!(*state, AncestorState::Include);
404 }
405 }
406
407 dag_builder.layers(1..=6).build();
410 let blocks = dag_builder.blocks.values().cloned().collect::<Vec<_>>();
411 dag_state.write().accept_blocks(blocks);
412
413 let accepted_quorum_rounds =
414 vec![(225, 229), (229, 300), (229, 300), (229, 300), (229, 300)];
415 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
416
417 let state_map = ancestor_state_manager.get_ancestor_states();
421 for (authority, state) in state_map.iter().enumerate() {
422 if authority == 0 {
423 assert_eq!(*state, AncestorState::Exclude(10));
424 } else {
425 assert_eq!(*state, AncestorState::Include);
426 }
427 }
428
429 let accepted_quorum_rounds =
430 vec![(229, 300), (229, 300), (229, 300), (229, 300), (229, 300)];
431 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
432
433 let state_map = ancestor_state_manager.get_ancestor_states();
437
438 for (authority, state) in state_map.iter().enumerate() {
439 if authority == 0 {
440 assert_eq!(*state, AncestorState::Exclude(10));
441 } else {
442 assert_eq!(*state, AncestorState::Include);
443 }
444 }
445
446 dag_builder.layers(7..=12).build();
448 let blocks = dag_builder.blocks.values().cloned().collect::<Vec<_>>();
449 dag_state.write().accept_blocks(blocks);
450
451 let scores = ReputationScores::new((1..=300).into(), vec![10, 100, 100, 100, 100]);
452 ancestor_state_manager.set_propagation_scores(scores);
453 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
454
455 let state_map = ancestor_state_manager.get_ancestor_states();
458 for state in state_map.iter() {
459 assert_eq!(*state, AncestorState::Include);
460 }
461 }
462}