Skip to main content

sui_core/
overload_monitor.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::authority::AuthorityState;
5use mysten_metrics::monitored_scope;
6use std::cmp::{max, min};
7use std::hash::Hasher;
8use std::sync::Weak;
9use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
10use std::time::Duration;
11use std::time::{SystemTime, UNIX_EPOCH};
12use sui_config::node::AuthorityOverloadConfig;
13use sui_types::digests::TransactionDigest;
14use sui_types::error::SuiErrorKind;
15use sui_types::error::SuiResult;
16use sui_types::fp_bail;
17use tokio::time::sleep;
18use tracing::{debug, info};
19use twox_hash::XxHash64;
20
21#[derive(Default)]
22pub struct AuthorityOverloadInfo {
23    /// Whether the authority is overloaded.
24    pub is_overload: AtomicBool,
25
26    /// The calculated percentage of transactions to drop.
27    pub load_shedding_percentage: AtomicU32,
28}
29
30impl AuthorityOverloadInfo {
31    pub fn set_overload(&self, load_shedding_percentage: u32) {
32        self.is_overload.store(true, Ordering::Relaxed);
33        self.load_shedding_percentage
34            .store(min(load_shedding_percentage, 100), Ordering::Relaxed);
35    }
36
37    pub fn clear_overload(&self) {
38        self.is_overload.store(false, Ordering::Relaxed);
39        self.load_shedding_percentage.store(0, Ordering::Relaxed);
40    }
41}
42
43const STEADY_OVERLOAD_REDUCTION_PERCENTAGE: u32 = 10;
44const EXECUTION_RATE_RATIO_FOR_COMPARISON: f64 = 0.95;
45const ADDITIONAL_LOAD_SHEDDING: f64 = 0.02;
46
47// The update interval of the random seed used to determine whether a txn should be rejected.
48const SEED_UPDATE_DURATION_SECS: u64 = 30;
49
50// Monitors the overload signals in `authority_state` periodically, and updates its `overload_info`
51// when the signals indicates overload.
52pub async fn overload_monitor(
53    authority_state: Weak<AuthorityState>,
54    config: AuthorityOverloadConfig,
55) {
56    info!("Starting system overload monitor.");
57
58    loop {
59        let authority_exist = check_authority_overload(&authority_state, &config);
60        if !authority_exist {
61            // `authority_state` doesn't exist anymore. Quit overload monitor.
62            break;
63        }
64        sleep(config.overload_monitor_interval).await;
65    }
66
67    info!("Shut down system overload monitor.");
68}
69
70// Checks authority overload signals, and updates authority's `overload_info`.
71// Returns whether the authority state exists.
72fn check_authority_overload(
73    authority_state: &Weak<AuthorityState>,
74    config: &AuthorityOverloadConfig,
75) -> bool {
76    let _scope = monitored_scope("OverloadMonitor::check_authority_overload");
77    let authority_arc = authority_state.upgrade();
78    if authority_arc.is_none() {
79        // `authority_state` doesn't exist anymore.
80        return false;
81    }
82
83    let authority = authority_arc.unwrap();
84    let queueing_latency = authority
85        .metrics
86        .execution_queueing_latency
87        .latency()
88        .unwrap_or_default();
89    let txn_ready_rate = authority.metrics.txn_ready_rate_tracker.lock().rate();
90    let execution_rate = authority.metrics.execution_rate_tracker.lock().rate();
91
92    debug!(
93        "Check authority overload signal, queueing latency {:?}, ready rate {:?}, execution rate {:?}.",
94        queueing_latency, txn_ready_rate, execution_rate
95    );
96
97    let (is_overload, load_shedding_percentage) = check_overload_signals(
98        config,
99        authority
100            .overload_info
101            .load_shedding_percentage
102            .load(Ordering::Relaxed),
103        queueing_latency,
104        txn_ready_rate,
105        execution_rate,
106    );
107
108    if is_overload {
109        authority
110            .overload_info
111            .set_overload(load_shedding_percentage);
112    } else {
113        authority.overload_info.clear_overload();
114    }
115
116    authority
117        .metrics
118        .authority_overload_status
119        .set(is_overload as i64);
120    authority
121        .metrics
122        .authority_load_shedding_percentage
123        .set(load_shedding_percentage as i64);
124    true
125}
126
127// Calculates the percentage of transactions to drop in order to reduce execution queue.
128// Returns the integer percentage between 0 and 100.
129fn calculate_load_shedding_percentage(txn_ready_rate: f64, execution_rate: f64) -> u32 {
130    // When transaction ready rate is practically 0, we aren't adding more load to the
131    // execution driver, so no shedding.
132    // TODO: consensus handler or transaction manager can also be overloaded.
133    if txn_ready_rate < 1e-10 {
134        return 0;
135    }
136
137    // Deflate the execution rate to account for the case that execution_rate is close to
138    // txn_ready_rate.
139    if execution_rate * EXECUTION_RATE_RATIO_FOR_COMPARISON > txn_ready_rate {
140        return 0;
141    }
142
143    // In order to maintain execution queue length, we need to drop at least (1 - executionRate / readyRate).
144    // To reduce the queue length, here we add 10% more transactions to drop.
145    (((1.0 - execution_rate * EXECUTION_RATE_RATIO_FOR_COMPARISON / txn_ready_rate)
146        + ADDITIONAL_LOAD_SHEDDING)
147        .min(1.0)
148        * 100.0)
149        .round() as u32
150}
151
152// Given overload signals (`queueing_latency`, `txn_ready_rate`, `execution_rate`), return whether
153// the authority server should enter load shedding mode, and how much percentage of transactions to drop.
154// Note that the final load shedding percentage should also take the current load shedding percentage
155// into consideration. If we are already shedding 40% load, based on the current txn_ready_rate
156// and execution_rate, we need to shed 10% more, the outcome is that we need to shed
157// 40% + (1 - 40%) * 10% = 46%.
158// When txn_ready_rate is less than execution_rate, we gradually reduce load shedding percentage until
159// the queueing latency is back to normal.
160fn check_overload_signals(
161    config: &AuthorityOverloadConfig,
162    current_load_shedding_percentage: u32,
163    queueing_latency: Duration,
164    txn_ready_rate: f64,
165    execution_rate: f64,
166) -> (bool, u32) {
167    // First, we calculate based on the current `txn_ready_rate` and `execution_rate`,
168    // what's the percentage of traffic to shed from `txn_ready_rate`.
169    let additional_load_shedding_percentage =
170        if queueing_latency > config.execution_queue_latency_hard_limit {
171            let calculated_load_shedding_percentage =
172                calculate_load_shedding_percentage(txn_ready_rate, execution_rate);
173
174            if calculated_load_shedding_percentage > 0
175                || txn_ready_rate >= config.safe_transaction_ready_rate as f64
176            {
177                max(
178                    calculated_load_shedding_percentage,
179                    config.min_load_shedding_percentage_above_hard_limit,
180                )
181            } else {
182                0
183            }
184        } else if queueing_latency > config.execution_queue_latency_soft_limit {
185            calculate_load_shedding_percentage(txn_ready_rate, execution_rate)
186        } else {
187            0
188        };
189
190    // Next, we calculate the new load shedding percentage.
191    let load_shedding_percentage = if additional_load_shedding_percentage > 0 {
192        // When we need to shed more load, since the `txn_ready_rate` is already influenced
193        // by `current_load_shedding_percentage`, we need to calculate the new load shedding
194        // percentage from `current_load_shedding_percentage` and
195        // `additional_load_shedding_percentage`.
196        current_load_shedding_percentage
197            + (100 - current_load_shedding_percentage) * additional_load_shedding_percentage / 100
198    } else if txn_ready_rate > config.safe_transaction_ready_rate as f64
199        && current_load_shedding_percentage > 10
200    {
201        // We don't need to shed more load. However, the enqueue rate is still not minimal.
202        // We gradually reduce load shedding percentage (10% at a time) to gracefully accept
203        // more load.
204        current_load_shedding_percentage - STEADY_OVERLOAD_REDUCTION_PERCENTAGE
205    } else {
206        // The current transaction ready rate is considered very low. Turn off load shedding mode.
207        0
208    };
209
210    let load_shedding_percentage = min(
211        load_shedding_percentage,
212        config.max_load_shedding_percentage,
213    );
214    let overload_status = load_shedding_percentage > 0;
215    (overload_status, load_shedding_percentage)
216}
217
218// Return true if we should reject the txn with `tx_digest`.
219fn should_reject_tx(
220    load_shedding_percentage: u32,
221    tx_digest: TransactionDigest,
222    temporal_seed: u64,
223) -> bool {
224    // TODO: we also need to add a secret salt (e.g. first consensus commit in the current epoch),
225    // to prevent gaming the system.
226    let mut hasher = XxHash64::with_seed(temporal_seed);
227    hasher.write(tx_digest.inner());
228    let value = hasher.finish();
229    value % 100 < load_shedding_percentage as u64
230}
231
232// Checks if we can accept the transaction with `tx_digest`.
233pub fn overload_monitor_accept_tx(
234    load_shedding_percentage: u32,
235    tx_digest: TransactionDigest,
236) -> SuiResult {
237    // Derive a random seed from the epoch time for transaction selection. Changing the seed every
238    // `SEED_UPDATE_DURATION_SECS` interval allows rejected transaction's retry to have a chance
239    // to go through in the future.
240    // Also, using the epoch time instead of randomly generating a seed allows that all validators
241    // makes the same decision.
242    let temporal_seed = SystemTime::now()
243        .duration_since(UNIX_EPOCH)
244        .expect("Sui did not exist prior to 1970")
245        .as_secs()
246        / SEED_UPDATE_DURATION_SECS;
247
248    if should_reject_tx(load_shedding_percentage, tx_digest, temporal_seed) {
249        // TODO: using `SEED_UPDATE_DURATION_SECS` is a safe suggestion that the time based seed
250        // is definitely different by then. However, a shorter suggestion may be available.
251        fp_bail!(
252            SuiErrorKind::ValidatorOverloadedRetryAfter {
253                retry_after_secs: SEED_UPDATE_DURATION_SECS
254            }
255            .into()
256        );
257    }
258    Ok(())
259}
260
261#[cfg(test)]
262#[allow(clippy::disallowed_methods)] // allow unbounded_channel() since tests are simulating txn manager execution driver interaction.
263mod tests {
264    use super::*;
265
266    use crate::authority::test_authority_builder::TestAuthorityBuilder;
267    use rand::{
268        Rng, SeedableRng,
269        rngs::{OsRng, StdRng},
270    };
271    use std::sync::Arc;
272    use sui_macros::sim_test;
273    use tokio::sync::mpsc::UnboundedReceiver;
274    use tokio::sync::mpsc::UnboundedSender;
275    use tokio::sync::mpsc::unbounded_channel;
276    use tokio::sync::oneshot;
277    use tokio::task::JoinHandle;
278    use tokio::time::{Instant, MissedTickBehavior, interval};
279
280    #[test]
281    pub fn test_authority_overload_info() {
282        let overload_info = AuthorityOverloadInfo::default();
283        assert!(!overload_info.is_overload.load(Ordering::Relaxed));
284        assert_eq!(
285            overload_info
286                .load_shedding_percentage
287                .load(Ordering::Relaxed),
288            0
289        );
290
291        {
292            overload_info.set_overload(20);
293            assert!(overload_info.is_overload.load(Ordering::Relaxed));
294            assert_eq!(
295                overload_info
296                    .load_shedding_percentage
297                    .load(Ordering::Relaxed),
298                20
299            );
300        }
301
302        // Tests that load shedding percentage can't go beyond 100%.
303        {
304            overload_info.set_overload(110);
305            assert!(overload_info.is_overload.load(Ordering::Relaxed));
306            assert_eq!(
307                overload_info
308                    .load_shedding_percentage
309                    .load(Ordering::Relaxed),
310                100
311            );
312        }
313
314        {
315            overload_info.clear_overload();
316            assert!(!overload_info.is_overload.load(Ordering::Relaxed));
317            assert_eq!(
318                overload_info
319                    .load_shedding_percentage
320                    .load(Ordering::Relaxed),
321                0
322            );
323        }
324    }
325
326    #[test]
327    pub fn test_calculate_load_shedding_ratio() {
328        assert_eq!(calculate_load_shedding_percentage(95.0, 100.1), 0);
329        assert_eq!(calculate_load_shedding_percentage(95.0, 100.0), 2);
330        assert_eq!(calculate_load_shedding_percentage(100.0, 100.0), 7);
331        assert_eq!(calculate_load_shedding_percentage(110.0, 100.0), 16);
332        assert_eq!(calculate_load_shedding_percentage(180.0, 100.0), 49);
333        assert_eq!(calculate_load_shedding_percentage(100.0, 0.0), 100);
334        assert_eq!(calculate_load_shedding_percentage(0.0, 1.0), 0);
335    }
336
337    #[test]
338    pub fn test_check_overload_signals() {
339        let config = AuthorityOverloadConfig {
340            execution_queue_latency_hard_limit: Duration::from_secs(10),
341            execution_queue_latency_soft_limit: Duration::from_secs(1),
342            max_load_shedding_percentage: 90,
343            ..Default::default()
344        };
345
346        // When execution queueing latency is within soft limit, don't start overload protection.
347        assert_eq!(
348            check_overload_signals(&config, 0, Duration::from_millis(500), 1000.0, 10.0),
349            (false, 0)
350        );
351
352        // When execution queueing latency hits soft limit and execution rate is higher, don't
353        // start overload protection.
354        assert_eq!(
355            check_overload_signals(&config, 0, Duration::from_secs(2), 100.0, 120.0),
356            (false, 0)
357        );
358
359        // When execution queueing latency hits soft limit, but not hard limit, start overload
360        // protection.
361        assert_eq!(
362            check_overload_signals(&config, 0, Duration::from_secs(2), 100.0, 100.0),
363            (true, 7)
364        );
365
366        // When execution queueing latency hits hard limit, start more aggressive overload
367        // protection.
368        assert_eq!(
369            check_overload_signals(&config, 0, Duration::from_secs(11), 100.0, 100.0),
370            (true, 50)
371        );
372
373        // When execution queueing latency hits hard limit and calculated shedding percentage
374        // is higher than min_load_shedding_percentage_above_hard_limit.
375        assert_eq!(
376            check_overload_signals(&config, 0, Duration::from_secs(11), 240.0, 100.0),
377            (true, 62)
378        );
379
380        // When execution queueing latency hits hard limit, but transaction ready rate
381        // is within safe_transaction_ready_rate, don't start overload protection.
382        assert_eq!(
383            check_overload_signals(&config, 0, Duration::from_secs(11), 20.0, 100.0),
384            (false, 0)
385        );
386
387        // Maximum transactions shed is cap by `max_load_shedding_percentage` config.
388        assert_eq!(
389            check_overload_signals(&config, 0, Duration::from_secs(11), 100.0, 0.0),
390            (true, 90)
391        );
392
393        // When the system is already shedding 50% of load, and the current txn ready rate
394        // and execution rate require another 20%, the final shedding rate is 60%.
395        assert_eq!(
396            check_overload_signals(&config, 50, Duration::from_secs(2), 116.0, 100.0),
397            (true, 60)
398        );
399
400        // Load shedding percentage is gradually reduced when txn ready rate is lower than
401        // execution rate.
402        assert_eq!(
403            check_overload_signals(&config, 90, Duration::from_secs(2), 200.0, 300.0),
404            (true, 80)
405        );
406
407        // When queueing delay is above hard limit, we shed additional 50% every time.
408        assert_eq!(
409            check_overload_signals(&config, 50, Duration::from_secs(11), 100.0, 100.0),
410            (true, 75)
411        );
412    }
413
414    #[tokio::test(flavor = "current_thread")]
415    pub async fn test_check_authority_overload() {
416        telemetry_subscribers::init_for_testing();
417
418        let config = AuthorityOverloadConfig {
419            safe_transaction_ready_rate: 0,
420            ..Default::default()
421        };
422        let state = TestAuthorityBuilder::new()
423            .with_authority_overload_config(config.clone())
424            .build()
425            .await;
426
427        // Initialize latency reporter.
428        for _ in 0..1000 {
429            state
430                .metrics
431                .execution_queueing_latency
432                .report(Duration::from_secs(20));
433        }
434
435        // Creates a simple case to see if authority state overload_info can be updated
436        // correctly by check_authority_overload.
437        let authority = Arc::downgrade(&state);
438        assert!(check_authority_overload(&authority, &config));
439        assert!(state.overload_info.is_overload.load(Ordering::Relaxed));
440        assert_eq!(
441            state
442                .overload_info
443                .load_shedding_percentage
444                .load(Ordering::Relaxed),
445            config.min_load_shedding_percentage_above_hard_limit
446        );
447
448        // Checks that check_authority_overload should return false when the input
449        // authority state doesn't exist.
450        let authority = Arc::downgrade(&state);
451        drop(state);
452        assert!(!check_authority_overload(&authority, &config));
453    }
454
455    // Creates an AuthorityState and starts an overload monitor that monitors its metrics.
456    async fn start_overload_monitor() -> (Arc<AuthorityState>, JoinHandle<()>) {
457        let overload_config = AuthorityOverloadConfig::default();
458        let state = TestAuthorityBuilder::new()
459            .with_authority_overload_config(overload_config.clone())
460            .build()
461            .await;
462        let authority_state = Arc::downgrade(&state);
463        let monitor_handle = tokio::spawn(async move {
464            overload_monitor(authority_state, overload_config).await;
465        });
466        (state, monitor_handle)
467    }
468
469    // Starts a load generator that generates a steady workload, and also allow it to accept
470    // burst of request through `burst_rx`.
471    // Request tracking is done by the overload monitor inside `authority`.
472    fn start_load_generator(
473        steady_rate: f64,
474        tx: UnboundedSender<Instant>,
475        mut burst_rx: UnboundedReceiver<u32>,
476        authority: Arc<AuthorityState>,
477        enable_load_shedding: bool,
478        total_requests_arc: Arc<AtomicU32>,
479        dropped_requests_arc: Arc<AtomicU32>,
480    ) -> JoinHandle<()> {
481        tokio::spawn(async move {
482            let mut interval = interval(Duration::from_secs_f64(1.0 / steady_rate));
483            let mut rng = StdRng::from_rng(&mut OsRng).unwrap();
484            let mut total_requests: u32 = 0;
485            let mut total_dropped_requests: u32 = 0;
486
487            // Helper function to check whether we should send a request.
488            let mut do_send =
489                |enable_load_shedding: bool, authority: Arc<AuthorityState>| -> bool {
490                    if enable_load_shedding {
491                        let shedding_percentage = authority
492                            .overload_info
493                            .load_shedding_percentage
494                            .load(Ordering::Relaxed);
495                        !(shedding_percentage > 0 && rng.gen_range(0..100) < shedding_percentage)
496                    } else {
497                        true
498                    }
499                };
500
501            loop {
502                tokio::select! {
503                    now = interval.tick() => {
504                        total_requests += 1;
505                        if do_send(enable_load_shedding, authority.clone()) {
506                            if tx.send(now).is_err() {
507                                info!("Load generator stopping. Total requests {:?}, total dropped requests {:?}.", total_requests, total_dropped_requests);
508                                total_requests_arc.store(total_requests, Ordering::SeqCst);
509                                dropped_requests_arc.store(total_dropped_requests, Ordering::SeqCst);
510                                return;
511                            }
512                            authority.metrics.txn_ready_rate_tracker.lock().record();
513                        } else {
514                            total_dropped_requests += 1;
515                        }
516                    }
517                    Some(burst) = burst_rx.recv() => {
518                        let now = Instant::now();
519                        total_requests += burst;
520                        for _ in 0..burst {
521                            if do_send(enable_load_shedding, authority.clone()) {
522                                if tx.send(now).is_err() {
523                                    info!("Load generator stopping. Total requests {:?}, total dropped requests {:?}.", total_requests, total_dropped_requests);
524                                    total_requests_arc.store(total_requests, Ordering::SeqCst);
525                                    dropped_requests_arc.store(total_dropped_requests, Ordering::SeqCst);
526                                    return;
527                                }
528                                authority.metrics.txn_ready_rate_tracker.lock().record();
529                            } else {
530                                total_dropped_requests += 1;
531                            }
532                        }
533                    }
534                }
535            }
536        })
537    }
538
539    // Starts a request executor that can consume request based on `execution_rate`.
540    // Request tracking is done by the overload monitor inside `authority`.
541    fn start_executor(
542        execution_rate: f64,
543        mut rx: UnboundedReceiver<Instant>,
544        mut stop_rx: oneshot::Receiver<()>,
545        authority: Arc<AuthorityState>,
546    ) -> JoinHandle<()> {
547        tokio::spawn(async move {
548            let mut interval = interval(Duration::from_secs_f64(1.0 / execution_rate));
549            interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
550            loop {
551                tokio::select! {
552                    Some(start_time) = rx.recv() => {
553                        authority.metrics.execution_rate_tracker.lock().record();
554                        authority.metrics.execution_queueing_latency.report(start_time.elapsed());
555                        interval.tick().await;
556                    }
557                    _ = &mut stop_rx => {
558                        info!("Executor stopping");
559                        return;
560                    }
561                }
562            }
563        })
564    }
565
566    // Helper fundtion to periodically print the current overload info.
567    async fn sleep_and_print_stats(state: Arc<AuthorityState>, seconds: u32) {
568        for _ in 0..seconds {
569            info!(
570                "Overload: {:?}. Shedding percentage: {:?}. Queue: {:?}, Ready rate: {:?}. Exec rate: {:?}.",
571                state.overload_info.is_overload.load(Ordering::Relaxed),
572                state
573                    .overload_info
574                    .load_shedding_percentage
575                    .load(Ordering::Relaxed),
576                state.metrics.execution_queueing_latency.latency(),
577                state.metrics.txn_ready_rate_tracker.lock().rate(),
578                state.metrics.execution_rate_tracker.lock().rate(),
579            );
580            sleep(Duration::from_secs(1)).await;
581        }
582    }
583
584    // Running a workload with consistent steady `generator_rate` and `executor_rate`.
585    // It checks that the dropped requests should in between min_dropping_rate and max_dropping_rate.
586    async fn run_consistent_workload_test(
587        generator_rate: f64,
588        executor_rate: f64,
589        min_dropping_rate: f64,
590        max_dropping_rate: f64,
591    ) {
592        let (state, monitor_handle) = start_overload_monitor().await;
593
594        let (tx, rx) = unbounded_channel();
595        let (_burst_tx, burst_rx) = unbounded_channel();
596        let total_requests = Arc::new(AtomicU32::new(0));
597        let dropped_requests = Arc::new(AtomicU32::new(0));
598        let load_generator = start_load_generator(
599            generator_rate,
600            tx.clone(),
601            burst_rx,
602            state.clone(),
603            true,
604            total_requests.clone(),
605            dropped_requests.clone(),
606        );
607
608        let (stop_tx, stop_rx) = oneshot::channel();
609        let executor = start_executor(executor_rate, rx, stop_rx, state.clone());
610
611        sleep_and_print_stats(state.clone(), 300).await;
612
613        stop_tx.send(()).unwrap();
614        let _ = tokio::join!(load_generator, executor);
615
616        let dropped_ratio = dropped_requests.load(Ordering::SeqCst) as f64
617            / total_requests.load(Ordering::SeqCst) as f64;
618        assert!(min_dropping_rate <= dropped_ratio);
619        assert!(dropped_ratio <= max_dropping_rate);
620
621        monitor_handle.abort();
622        let _ = monitor_handle.await;
623    }
624
625    // Tests that when request generation rate is slower than execution rate, no requests should be dropped.
626    #[tokio::test(flavor = "current_thread", start_paused = true)]
627    pub async fn test_workload_consistent_no_overload() {
628        telemetry_subscribers::init_for_testing();
629        run_consistent_workload_test(900.0, 1000.0, 0.0, 0.0).await;
630    }
631
632    // Tests that when request generation rate is slightly above execution rate, a small portion of
633    // requests should be dropped.
634    #[tokio::test(flavor = "current_thread", start_paused = true)]
635    pub async fn test_workload_consistent_slightly_overload() {
636        telemetry_subscribers::init_for_testing();
637        // Dropping rate should be around 15%.
638        run_consistent_workload_test(1100.0, 1000.0, 0.05, 0.25).await;
639    }
640
641    // Tests that when request generation rate is much higher than execution rate, a large portion of
642    // requests should be dropped.
643    #[tokio::test(flavor = "current_thread", start_paused = true)]
644    pub async fn test_workload_consistent_overload() {
645        telemetry_subscribers::init_for_testing();
646        // Dropping rate should be around 70%.
647        run_consistent_workload_test(3000.0, 1000.0, 0.6, 0.8).await;
648    }
649
650    // Tests that when there is a very short single spike, no request should be dropped.
651    #[tokio::test(flavor = "current_thread", start_paused = true)]
652    pub async fn test_workload_single_spike() {
653        telemetry_subscribers::init_for_testing();
654        let (state, monitor_handle) = start_overload_monitor().await;
655
656        let (tx, rx) = unbounded_channel();
657        let (burst_tx, burst_rx) = unbounded_channel();
658        let total_requests = Arc::new(AtomicU32::new(0));
659        let dropped_requests = Arc::new(AtomicU32::new(0));
660        let load_generator = start_load_generator(
661            10.0,
662            tx.clone(),
663            burst_rx,
664            state.clone(),
665            true,
666            total_requests.clone(),
667            dropped_requests.clone(),
668        );
669
670        let (stop_tx, stop_rx) = oneshot::channel();
671        let executor = start_executor(1000.0, rx, stop_rx, state.clone());
672
673        sleep_and_print_stats(state.clone(), 10).await;
674        // Send out a burst of 5000 requests.
675        burst_tx.send(5000).unwrap();
676        sleep_and_print_stats(state.clone(), 20).await;
677
678        stop_tx.send(()).unwrap();
679        let _ = tokio::join!(load_generator, executor);
680
681        // No requests should be dropped.
682        assert_eq!(dropped_requests.load(Ordering::SeqCst), 0);
683
684        monitor_handle.abort();
685        let _ = monitor_handle.await;
686    }
687
688    // Tests that when there are regular spikes that keep queueing latency consistently high,
689    // overload monitor should kick in and shed load.
690    #[tokio::test(flavor = "current_thread", start_paused = true)]
691    pub async fn test_workload_consistent_short_spike() {
692        telemetry_subscribers::init_for_testing();
693        let (state, monitor_handle) = start_overload_monitor().await;
694
695        let (tx, rx) = unbounded_channel();
696        let (burst_tx, burst_rx) = unbounded_channel();
697        let total_requests = Arc::new(AtomicU32::new(0));
698        let dropped_requests = Arc::new(AtomicU32::new(0));
699        let load_generator = start_load_generator(
700            10.0,
701            tx.clone(),
702            burst_rx,
703            state.clone(),
704            true,
705            total_requests.clone(),
706            dropped_requests.clone(),
707        );
708
709        let (stop_tx, stop_rx) = oneshot::channel();
710        let executor = start_executor(1000.0, rx, stop_rx, state.clone());
711
712        sleep_and_print_stats(state.clone(), 15).await;
713        for _ in 0..16 {
714            // Regularly send out a burst of request.
715            burst_tx.send(10000).unwrap();
716            sleep_and_print_stats(state.clone(), 5).await;
717        }
718
719        stop_tx.send(()).unwrap();
720        let _ = tokio::join!(load_generator, executor);
721        let dropped_ratio = dropped_requests.load(Ordering::SeqCst) as f64
722            / total_requests.load(Ordering::SeqCst) as f64;
723
724        // We should drop about 50% of request because the burst throughput is about 2x of
725        // execution rate.
726        assert!(0.4 < dropped_ratio);
727        assert!(dropped_ratio < 0.6);
728
729        monitor_handle.abort();
730        let _ = monitor_handle.await;
731    }
732
733    // Tests that the ratio of rejected transactions created randomly matches load shedding percentage in
734    // the overload monitor.
735    #[test]
736    fn test_txn_rejection_rate() {
737        for rejection_percentage in 0..=100 {
738            let mut reject_count = 0;
739            for _ in 0..10000 {
740                let digest = TransactionDigest::random();
741                if should_reject_tx(rejection_percentage, digest, 28455473) {
742                    reject_count += 1;
743                }
744            }
745
746            debug!(
747                "Rejection percentage: {:?}, reject count: {:?}.",
748                rejection_percentage, reject_count
749            );
750            // Give it a 3% fluctuation.
751            assert!(rejection_percentage as f32 / 100.0 - 0.03 < reject_count as f32 / 10000.0);
752            assert!(reject_count as f32 / 10000.0 < rejection_percentage as f32 / 100.0 + 0.03);
753        }
754    }
755
756    // Tests that rejected transaction will have a chance to be accepted in the future.
757    #[sim_test]
758    async fn test_txn_rejection_over_time() {
759        let start_time = Instant::now();
760        let mut digest = TransactionDigest::random();
761        let mut temporal_seed = 1708108277 / SEED_UPDATE_DURATION_SECS;
762        let load_shedding_percentage = 50;
763
764        // Find a rejected transaction with 50% rejection rate.
765        while !should_reject_tx(load_shedding_percentage, digest, temporal_seed)
766            && start_time.elapsed() < Duration::from_secs(30)
767        {
768            digest = TransactionDigest::random();
769        }
770
771        // It should always be rejected using the current temporal_seed.
772        for _ in 0..100 {
773            assert!(should_reject_tx(
774                load_shedding_percentage,
775                digest,
776                temporal_seed
777            ));
778        }
779
780        // It will be accepted in the future.
781        temporal_seed += 1;
782        while should_reject_tx(load_shedding_percentage, digest, temporal_seed)
783            && start_time.elapsed() < Duration::from_secs(30)
784        {
785            temporal_seed += 1;
786        }
787
788        // Make sure that the tests can finish within 30 seconds.
789        assert!(start_time.elapsed() < Duration::from_secs(30));
790    }
791}