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 last_suppressed: Vec<AuthorityIndex>,
72 pub(crate) propagation_scores: ReputationScores,
75}
76
77impl AncestorStateManager {
78 #[cfg(not(test))]
81 const STATE_LOCK_CLOCK_ROUNDS: u32 = 450;
82 #[cfg(test)]
83 const STATE_LOCK_CLOCK_ROUNDS: u32 = 5;
84
85 const SCORE_EXCLUSION_THRESHOLD_PERCENTAGE: u64 = 20;
87
88 pub(crate) fn new(context: Arc<Context>, dag_state: Arc<RwLock<DagState>>) -> Self {
89 let state_map = vec![AncestorInfo::new(); context.committee.size()];
90
91 let excluded_nodes_stake_threshold_percentage =
94 2 * context.protocol_config.bad_nodes_stake_threshold() / 3;
95
96 let excluded_nodes_stake_threshold = (excluded_nodes_stake_threshold_percentage
97 * context.committee.total_stake())
98 / 100 as Stake;
99
100 Self {
101 context,
102 dag_state,
103 state_map,
104 excluded_nodes_stake_threshold,
105 total_excluded_stake: 0,
107 last_suppressed: Vec::new(),
108 propagation_scores: ReputationScores::default(),
109 }
110 }
111
112 pub(crate) fn set_propagation_scores(&mut self, scores: ReputationScores) {
113 self.propagation_scores = scores;
114 }
115
116 pub(crate) fn get_ancestor_states(&self) -> Vec<AncestorState> {
117 self.state_map.iter().map(|info| info.state).collect()
118 }
119
120 pub(crate) fn update_all_ancestors_state(&mut self, accepted_quorum_rounds: &[QuorumRound]) {
122 let network_high_quorum_round =
126 self.calculate_network_high_quorum_round(accepted_quorum_rounds);
127
128 let current_clock_round = self.dag_state.read().threshold_clock_round();
129 let low_score_threshold = (self.propagation_scores.highest_score()
130 * Self::SCORE_EXCLUSION_THRESHOLD_PERCENTAGE)
131 / 100;
132
133 debug!(
134 "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 {}",
135 self.excluded_nodes_stake_threshold
136 );
137
138 let mut exclude_to_include = Vec::new();
142 let mut include_to_exclude = Vec::new();
143
144 for (idx, score) in self
148 .propagation_scores
149 .scores_per_authority
150 .iter()
151 .enumerate()
152 {
153 let authority_id = self
154 .context
155 .committee
156 .to_authority_index(idx)
157 .expect("Index should be valid");
158 let ancestor_info = &self.state_map[idx];
159 let (_low, authority_high_quorum_round) = accepted_quorum_rounds[idx];
160 let stake = self.context.committee.authority(authority_id).stake;
161
162 if ancestor_info.is_locked(current_clock_round) {
164 continue;
165 }
166
167 match ancestor_info.state {
168 AncestorState::Include => {
169 if *score <= low_score_threshold {
170 include_to_exclude.push(StateTransition {
171 authority_id,
172 score: *score,
173 stake,
174 high_quorum_round: authority_high_quorum_round,
175 });
176 }
177 }
178 AncestorState::Exclude(_) => {
179 if *score > low_score_threshold
180 || authority_high_quorum_round >= network_high_quorum_round
181 {
182 exclude_to_include.push(StateTransition {
183 authority_id,
184 score: *score,
185 stake,
186 high_quorum_round: authority_high_quorum_round,
187 });
188 }
189 }
190 }
191 }
192
193 for transition in exclude_to_include {
196 self.apply_state_change(transition, AncestorState::Include, current_clock_round);
197 }
198
199 include_to_exclude.sort_by_key(|t| t.score);
202
203 let mut suppressed = Vec::new();
207 for transition in include_to_exclude {
208 if self.total_excluded_stake + transition.stake <= self.excluded_nodes_stake_threshold {
212 let new_state = AncestorState::Exclude(transition.score);
213 self.apply_state_change(transition, new_state, current_clock_round);
214 } else {
215 suppressed.push(transition);
216 }
217 }
218
219 self.context
220 .metrics
221 .node_metrics
222 .ancestor_exclusion_suppressed_by_stake_threshold
223 .set(suppressed.len() as i64);
224
225 let mut suppressed_ids = suppressed
227 .iter()
228 .map(|t| t.authority_id)
229 .collect::<Vec<_>>();
230 suppressed_ids.sort_unstable();
231
232 if suppressed_ids != self.last_suppressed {
233 info!(
234 "Could not exclude {} authorities at round {current_clock_round} with current_excluded_stake {} and exclude_stake_threshold {}: {:?}",
235 suppressed.len(),
236 self.total_excluded_stake,
237 self.excluded_nodes_stake_threshold,
238 suppressed
239 .iter()
240 .map(|t| (t.authority_id, t.score))
241 .collect::<Vec<_>>(),
242 );
243 self.last_suppressed = suppressed_ids;
244 }
245 }
246
247 fn apply_state_change(
248 &mut self,
249 transition: StateTransition,
250 new_state: AncestorState,
251 current_clock_round: u32,
252 ) {
253 let block_hostname = &self
254 .context
255 .committee
256 .authority(transition.authority_id)
257 .hostname;
258 let ancestor_info = &mut self.state_map[transition.authority_id.value()];
259
260 match (ancestor_info.state, new_state) {
261 (AncestorState::Exclude(_), AncestorState::Include) => {
262 self.total_excluded_stake = self.total_excluded_stake
263 .checked_sub(transition.stake)
264 .expect("total_excluded_stake underflow - trying to subtract more stake than we're tracking as excluded");
265 }
266 (AncestorState::Include, AncestorState::Exclude(_)) => {
267 self.total_excluded_stake += transition.stake;
268 }
269 _ => {
270 panic!("Calls to this function should only be made for state transition.")
271 }
272 }
273
274 ancestor_info.state = new_state;
275 let lock_until_round = current_clock_round + Self::STATE_LOCK_CLOCK_ROUNDS;
276 ancestor_info.set_lock(lock_until_round);
277
278 info!(
279 "Authority {} moved to {new_state:?} state with score {} & quorum_round {} and locked until round {lock_until_round}. Total excluded stake: {}",
280 transition.authority_id,
281 transition.score,
282 transition.high_quorum_round,
283 self.total_excluded_stake
284 );
285
286 self.context
287 .metrics
288 .node_metrics
289 .ancestor_state_change_by_authority
290 .with_label_values(&[
291 block_hostname.as_str(),
292 match new_state {
293 AncestorState::Include => "include",
294 AncestorState::Exclude(_) => "exclude",
295 },
296 ])
297 .inc();
298 }
299
300 fn calculate_network_high_quorum_round(&self, accepted_quorum_rounds: &[QuorumRound]) -> u32 {
308 let committee = &self.context.committee;
309
310 let mut high_quorum_rounds_with_stake = accepted_quorum_rounds
311 .iter()
312 .zip_debug_eq(committee.authorities())
313 .map(|((_low, high), (_, authority))| (*high, authority.stake))
314 .collect::<Vec<_>>();
315 high_quorum_rounds_with_stake.sort();
316
317 let mut total_stake = 0;
318 let mut network_high_quorum_round = 0;
319
320 for (round, stake) in high_quorum_rounds_with_stake.iter() {
321 total_stake += stake;
322 if total_stake >= self.context.committee.quorum_threshold() {
323 network_high_quorum_round = *round;
324 break;
325 }
326 }
327
328 network_high_quorum_round
329 }
330}
331
332#[cfg(test)]
333mod test {
334 use super::*;
335 use crate::{
336 leader_scoring::ReputationScores, storage::mem_store::MemStore,
337 test_dag_builder::DagBuilder,
338 };
339
340 #[tokio::test]
341 async fn test_calculate_network_high_accepted_quorum_round() {
342 telemetry_subscribers::init_for_testing();
343
344 let (context, _key_pairs) = Context::new_for_test(4);
345 let context = Arc::new(context);
346 let store = Arc::new(MemStore::new());
347 let dag_state = Arc::new(RwLock::new(DagState::new(context.clone(), store.clone())));
348
349 let scores = ReputationScores::new((1..=300).into(), vec![1, 2, 4, 3]);
350 let mut ancestor_state_manager =
351 AncestorStateManager::new(context.clone(), dag_state.clone());
352 ancestor_state_manager.set_propagation_scores(scores);
353
354 let network_high_quorum_round =
357 ancestor_state_manager.calculate_network_high_quorum_round(&[(0, 0); 4]);
358 assert_eq!(network_high_quorum_round, 0);
359
360 let accepted_quorum_rounds = vec![(50, 229), (175, 229), (179, 229), (179, 300)];
361
362 let network_high_quorum_round =
363 ancestor_state_manager.calculate_network_high_quorum_round(&accepted_quorum_rounds);
364 assert_eq!(network_high_quorum_round, 229);
365 }
366
367 #[tokio::test]
374 async fn test_update_all_ancestor_state_using_accepted_rounds() {
375 telemetry_subscribers::init_for_testing();
376 let (mut context, _key_pairs) = Context::new_for_test(5);
377 context
378 .protocol_config
379 .set_bad_nodes_stake_threshold_for_testing(33);
380 let context = Arc::new(context);
381 let store = Arc::new(MemStore::new());
382 let dag_state = Arc::new(RwLock::new(DagState::new(context.clone(), store.clone())));
383 let mut dag_builder = DagBuilder::new(context.clone());
384
385 let scores = ReputationScores::new((1..=300).into(), vec![1, 2, 4, 3, 4]);
386 let mut ancestor_state_manager =
387 AncestorStateManager::new(context.clone(), dag_state.clone());
388 ancestor_state_manager.set_propagation_scores(scores);
389
390 let accepted_quorum_rounds =
391 vec![(225, 229), (225, 229), (229, 300), (229, 300), (229, 300)];
392 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
393
394 let state_map = ancestor_state_manager.get_ancestor_states();
397 for state in state_map.iter() {
398 assert_eq!(*state, AncestorState::Include);
399 }
400
401 let scores = ReputationScores::new((1..=300).into(), vec![10, 9, 100, 100, 100]);
402 ancestor_state_manager.set_propagation_scores(scores);
403 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
404
405 let state_map = ancestor_state_manager.get_ancestor_states();
411 for (authority, state) in state_map.iter().enumerate() {
412 if authority == 1 {
413 assert_eq!(*state, AncestorState::Exclude(9));
414 } else {
415 assert_eq!(*state, AncestorState::Include);
416 }
417 }
418
419 assert_eq!(
421 context
422 .metrics
423 .node_metrics
424 .ancestor_exclusion_suppressed_by_stake_threshold
425 .get(),
426 1
427 );
428
429 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
430
431 let state_map = ancestor_state_manager.get_ancestor_states();
434 for (authority, state) in state_map.iter().enumerate() {
435 if authority == 1 {
436 assert_eq!(*state, AncestorState::Exclude(9));
437 } else {
438 assert_eq!(*state, AncestorState::Include);
439 }
440 }
441
442 assert_eq!(
445 context
446 .metrics
447 .node_metrics
448 .ancestor_exclusion_suppressed_by_stake_threshold
449 .get(),
450 1
451 );
452
453 dag_builder.layers(1..=6).build();
456 let blocks = dag_builder.blocks.values().cloned().collect::<Vec<_>>();
457 dag_state.write().accept_blocks(blocks);
458
459 let accepted_quorum_rounds =
460 vec![(225, 229), (229, 300), (229, 300), (229, 300), (229, 300)];
461 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
462
463 let state_map = ancestor_state_manager.get_ancestor_states();
467 for (authority, state) in state_map.iter().enumerate() {
468 if authority == 0 {
469 assert_eq!(*state, AncestorState::Exclude(10));
470 } else {
471 assert_eq!(*state, AncestorState::Include);
472 }
473 }
474
475 assert_eq!(
478 context
479 .metrics
480 .node_metrics
481 .ancestor_exclusion_suppressed_by_stake_threshold
482 .get(),
483 0
484 );
485
486 let accepted_quorum_rounds =
487 vec![(229, 300), (229, 300), (229, 300), (229, 300), (229, 300)];
488 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
489
490 let state_map = ancestor_state_manager.get_ancestor_states();
494
495 for (authority, state) in state_map.iter().enumerate() {
496 if authority == 0 {
497 assert_eq!(*state, AncestorState::Exclude(10));
498 } else {
499 assert_eq!(*state, AncestorState::Include);
500 }
501 }
502
503 dag_builder.layers(7..=12).build();
505 let blocks = dag_builder.blocks.values().cloned().collect::<Vec<_>>();
506 dag_state.write().accept_blocks(blocks);
507
508 let scores = ReputationScores::new((1..=300).into(), vec![10, 100, 100, 100, 100]);
509 ancestor_state_manager.set_propagation_scores(scores);
510 ancestor_state_manager.update_all_ancestors_state(&accepted_quorum_rounds);
511
512 let state_map = ancestor_state_manager.get_ancestor_states();
515 for state in state_map.iter() {
516 assert_eq!(*state, AncestorState::Include);
517 }
518 }
519}