1use std::{sync::Arc, time::Duration};
20
21use consensus_types::block::Round;
22use futures::stream::{FuturesUnordered, StreamExt as _};
23use mysten_common::sync::notify_once::NotifyOnce;
24use mysten_metrics::monitored_scope;
25use parking_lot::RwLock;
26use tokio::{task::JoinHandle, time::MissedTickBehavior};
27
28use crate::{
29 BlockAPI as _, context::Context, core_thread::CoreThreadDispatcher, dag_state::DagState,
30 network::ValidatorNetworkClient, round_tracker::RoundTracker, task::join_and_propagate_panic,
31};
32
33pub(crate) struct RoundProberHandle {
35 prober_task: JoinHandle<()>,
36 shutdown_notify: Arc<NotifyOnce>,
37}
38
39impl RoundProberHandle {
40 pub(crate) async fn stop(self) {
41 let _ = self.shutdown_notify.notify();
42 join_and_propagate_panic(self.prober_task).await;
44 }
45}
46
47pub(crate) struct RoundProber<C: ValidatorNetworkClient> {
48 context: Arc<Context>,
49 core_thread_dispatcher: Arc<dyn CoreThreadDispatcher>,
50 round_tracker: Arc<RwLock<RoundTracker>>,
51 dag_state: Arc<RwLock<DagState>>,
52 network_client: Arc<C>,
53 shutdown_notify: Arc<NotifyOnce>,
54}
55
56impl<C: ValidatorNetworkClient> RoundProber<C> {
57 pub(crate) fn new(
58 context: Arc<Context>,
59 core_thread_dispatcher: Arc<dyn CoreThreadDispatcher>,
60 round_tracker: Arc<RwLock<RoundTracker>>,
61 dag_state: Arc<RwLock<DagState>>,
62 network_client: Arc<C>,
63 ) -> Self {
64 Self {
65 context,
66 core_thread_dispatcher,
67 round_tracker,
68 dag_state,
69 network_client,
70 shutdown_notify: Arc::new(NotifyOnce::new()),
71 }
72 }
73
74 pub(crate) fn start(self) -> RoundProberHandle {
75 let shutdown_notify = self.shutdown_notify.clone();
76 let loop_shutdown_notify = shutdown_notify.clone();
77 let prober_task = tokio::spawn(async move {
78 let mut interval = tokio::time::interval(Duration::from_millis(
82 self.context.parameters.round_prober_interval_ms,
83 ));
84 interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
85 loop {
86 tokio::select! {
87 _ = interval.tick() => {
88 self.probe().await;
89 }
90 _ = loop_shutdown_notify.wait() => {
91 break;
92 }
93 }
94 }
95 });
96 RoundProberHandle {
97 prober_task,
98 shutdown_notify,
99 }
100 }
101
102 pub(crate) async fn probe(&self) -> Round {
105 let _scope = monitored_scope("RoundProber");
106
107 let node_metrics = &self.context.metrics.node_metrics;
108 let request_timeout =
109 Duration::from_millis(self.context.parameters.round_prober_request_timeout_ms);
110 let own_index = self.context.own_index;
111 let mut requests = FuturesUnordered::new();
112
113 for (peer, _) in self.context.committee.authorities() {
114 if peer == own_index {
115 continue;
116 }
117 let network_client = self.network_client.clone();
118 requests.push(async move {
119 let result = tokio::time::timeout(
120 request_timeout,
121 network_client.get_latest_rounds(peer, request_timeout),
122 )
123 .await;
124 (peer, result)
125 });
126 }
127
128 let mut highest_received_rounds =
129 vec![vec![0; self.context.committee.size()]; self.context.committee.size()];
130 let mut highest_accepted_rounds =
131 vec![vec![0; self.context.committee.size()]; self.context.committee.size()];
132
133 let blocks = self
134 .dag_state
135 .read()
136 .get_last_cached_block_per_authority(Round::MAX);
137 let local_highest_accepted_rounds = blocks
138 .into_iter()
139 .map(|(block, _)| block.round())
140 .collect::<Vec<_>>();
141 let last_proposed_round = local_highest_accepted_rounds[own_index];
142
143 highest_received_rounds[own_index] =
146 self.round_tracker.read().local_highest_received_rounds();
147 highest_accepted_rounds[own_index] = local_highest_accepted_rounds;
148 highest_received_rounds[own_index][own_index] = last_proposed_round;
149 highest_accepted_rounds[own_index][own_index] = last_proposed_round;
150
151 loop {
152 tokio::select! {
153 result = requests.next() => {
154 let Some((peer, result)) = result else { break };
155 let peer_name = &self.context.committee.authority(peer).hostname;
156 match result {
157 Ok(Ok((received, accepted))) => {
158 if received.len() == self.context.committee.size()
159 {
160 highest_received_rounds[peer] = received;
161 } else {
162 node_metrics.round_prober_request_errors.with_label_values(&["invalid_received_rounds"]).inc();
163 tracing::warn!("Received invalid number of received rounds from peer {}", peer_name);
164 }
165
166 if accepted.len() == self.context.committee.size() {
167 highest_accepted_rounds[peer] = accepted;
168 } else {
169 node_metrics.round_prober_request_errors.with_label_values(&["invalid_accepted_rounds"]).inc();
170 tracing::warn!("Received invalid number of accepted rounds from peer {}", peer_name);
171 }
172 },
173 Ok(Err(err)) => {
184 node_metrics.round_prober_request_errors.with_label_values(&["failed_fetch"]).inc();
185 tracing::debug!("Failed to get latest rounds from peer {}: {:?}", peer_name, err);
186 },
187 Err(_) => {
188 node_metrics.round_prober_request_errors.with_label_values(&["timeout"]).inc();
189 tracing::debug!("Timeout while getting latest rounds from peer {}", peer_name);
190 },
191 }
192 }
193 _ = self.shutdown_notify.wait() => break,
194 }
195 }
196
197 self.round_tracker
198 .write()
199 .update_from_probe(highest_accepted_rounds, highest_received_rounds);
200 let propagation_delay = self
201 .round_tracker
202 .read()
203 .calculate_propagation_delay(last_proposed_round);
204
205 let _ = self
206 .core_thread_dispatcher
207 .set_propagation_delay(propagation_delay);
208
209 propagation_delay
210 }
211}
212
213#[cfg(test)]
214mod test {
215 use std::{collections::BTreeSet, sync::Arc, time::Duration};
216
217 use async_trait::async_trait;
218 use bytes::Bytes;
219 use consensus_config::AuthorityIndex;
220 use consensus_types::block::{BlockRef, Round};
221 use parking_lot::RwLock;
222
223 use crate::{
224 TestBlock, VerifiedBlock,
225 commit::{CertifiedCommits, CommitRange},
226 context::Context,
227 core_thread::{CoreError, CoreThreadDispatcher},
228 dag_state::DagState,
229 error::{ConsensusError, ConsensusResult},
230 network::{BlockStream, ValidatorNetworkClient},
231 round_prober::RoundProber,
232 round_tracker::RoundTracker,
233 storage::mem_store::MemStore,
234 };
235
236 struct FakeThreadDispatcher {}
237
238 impl FakeThreadDispatcher {
239 fn new() -> Self {
240 Self {}
241 }
242 }
243
244 #[async_trait]
245 impl CoreThreadDispatcher for FakeThreadDispatcher {
246 async fn add_blocks(
247 &self,
248 _blocks: Vec<VerifiedBlock>,
249 ) -> Result<BTreeSet<BlockRef>, CoreError> {
250 unimplemented!()
251 }
252
253 async fn check_block_refs(
254 &self,
255 _block_refs: Vec<BlockRef>,
256 ) -> Result<BTreeSet<BlockRef>, CoreError> {
257 unimplemented!()
258 }
259
260 async fn add_certified_commits(
261 &self,
262 _commits: CertifiedCommits,
263 ) -> Result<BTreeSet<BlockRef>, CoreError> {
264 unimplemented!()
265 }
266
267 async fn new_block(&self, _round: Round, _force: bool) -> Result<(), CoreError> {
268 unimplemented!()
269 }
270
271 async fn get_missing_blocks(&self) -> Result<BTreeSet<BlockRef>, CoreError> {
272 unimplemented!()
273 }
274
275 fn set_propagation_delay(&self, _propagation_delay: Round) -> Result<(), CoreError> {
276 Ok(())
277 }
278
279 fn set_last_known_proposed_round(&self, _round: Round) -> Result<(), CoreError> {
280 unimplemented!()
281 }
282 }
283
284 struct FakeNetworkClient {
285 highest_received_rounds: Vec<Vec<Round>>,
286 highest_accepted_rounds: Vec<Vec<Round>>,
287 }
288
289 impl FakeNetworkClient {
290 fn new(
291 highest_received_rounds: Vec<Vec<Round>>,
292 highest_accepted_rounds: Vec<Vec<Round>>,
293 ) -> Self {
294 Self {
295 highest_received_rounds,
296 highest_accepted_rounds,
297 }
298 }
299 }
300
301 #[async_trait]
302 impl ValidatorNetworkClient for FakeNetworkClient {
303 async fn send_block(
304 &self,
305 _peer: AuthorityIndex,
306 _serialized_block: &VerifiedBlock,
307 _timeout: Duration,
308 ) -> ConsensusResult<()> {
309 unimplemented!("Unimplemented")
310 }
311
312 async fn subscribe_blocks(
313 &self,
314 _peer: AuthorityIndex,
315 _last_received: Round,
316 _timeout: Duration,
317 ) -> ConsensusResult<BlockStream> {
318 unimplemented!("Unimplemented")
319 }
320
321 async fn fetch_blocks(
322 &self,
323 _peer: AuthorityIndex,
324 _block_refs: Vec<BlockRef>,
325 _fetch_after_rounds: Vec<Round>,
326 _fetch_missing_ancestors: bool,
327 _timeout: Duration,
328 ) -> ConsensusResult<Vec<Bytes>> {
329 unimplemented!("Unimplemented")
330 }
331
332 async fn fetch_commits(
333 &self,
334 _peer: AuthorityIndex,
335 _commit_range: CommitRange,
336 _timeout: Duration,
337 ) -> ConsensusResult<(Vec<Bytes>, Vec<Bytes>)> {
338 unimplemented!("Unimplemented")
339 }
340
341 async fn fetch_latest_blocks(
342 &self,
343 _peer: AuthorityIndex,
344 _authorities: Vec<AuthorityIndex>,
345 _timeout: Duration,
346 ) -> ConsensusResult<Vec<Bytes>> {
347 unimplemented!("Unimplemented")
348 }
349
350 async fn get_latest_rounds(
351 &self,
352 peer: AuthorityIndex,
353 _timeout: Duration,
354 ) -> ConsensusResult<(Vec<Round>, Vec<Round>)> {
355 let received_rounds = self.highest_received_rounds[peer].clone();
356 let accepted_rounds = self.highest_accepted_rounds[peer].clone();
357 if received_rounds.is_empty() && accepted_rounds.is_empty() {
358 Err(ConsensusError::NetworkRequestTimeout("test".to_string()))
359 } else {
360 Ok((received_rounds, accepted_rounds))
361 }
362 }
363 }
364
365 #[tokio::test]
366 async fn test_round_prober() {
367 telemetry_subscribers::init_for_testing();
368 const NUM_AUTHORITIES: usize = 7;
369 let context = Arc::new(Context::new_for_test(NUM_AUTHORITIES).0);
370 let core_thread_dispatcher = Arc::new(FakeThreadDispatcher::new());
371 let store = Arc::new(MemStore::new());
372 let dag_state = Arc::new(RwLock::new(DagState::new(context.clone(), store)));
373 let network_client = Arc::new(FakeNetworkClient::new(
375 vec![
376 vec![],
377 vec![109, 121, 131, 0, 151, 161, 171],
378 vec![101, 0, 103, 104, 105, 166, 107],
379 vec![],
380 vec![100, 102, 133, 0, 155, 106, 177],
381 vec![105, 115, 103, 0, 125, 126, 127],
382 vec![10, 20, 30, 40, 50, 60],
383 ], vec![
385 vec![],
386 vec![0, 121, 131, 0, 151, 161, 171],
387 vec![1, 0, 103, 104, 105, 166, 107],
388 vec![],
389 vec![0, 102, 133, 0, 155, 106, 177],
390 vec![1, 115, 103, 0, 125, 126, 127],
391 vec![1, 20, 30, 40, 50, 60],
392 ], ));
394
395 let round_tracker = Arc::new(RwLock::new(RoundTracker::new(
397 context.clone(),
398 vec![110, 120, 130, 140, 150, 160, 170],
399 )));
400 let prober = RoundProber::new(
401 context.clone(),
402 core_thread_dispatcher.clone(),
403 round_tracker.clone(),
404 dag_state.clone(),
405 network_client.clone(),
406 );
407
408 let blocks = (0..NUM_AUTHORITIES)
410 .map(|authority| {
411 let round = 110 + (authority as u32 * 10);
412 VerifiedBlock::new_for_test(TestBlock::new(round, authority as u32).build())
413 })
414 .collect::<Vec<_>>();
415
416 dag_state.write().accept_blocks(blocks);
417
418 let propagation_delay = prober.probe().await;
429
430 assert_eq!(propagation_delay, 10);
431 }
432}