hotshot_example_types/
testable_delay.rs

1use std::{collections::HashMap, time::Duration};
2
3use async_trait::async_trait;
4use rand::Rng;
5use tokio::time::sleep;
6
7#[derive(Eq, Hash, PartialEq, Debug, Clone)]
8/// What type of delay we want to apply to
9pub enum DelayOptions {
10    None,
11    Random,
12    Fixed,
13}
14
15#[derive(Eq, Hash, PartialEq, Debug, Clone)]
16/// Current implementations that are supported for testing async delays
17pub enum SupportedTraitTypesForAsyncDelay {
18    Storage,
19    ValidatedState,
20    BlockHeader,
21}
22
23#[derive(Eq, Hash, PartialEq, Debug, Clone)]
24/// Config for each supported type
25pub struct DelaySettings {
26    // Option to tell the async function what to do
27    pub delay_option: DelayOptions,
28    // Rng min time
29    pub min_time_in_milliseconds: u64,
30    // Rng max time
31    pub max_time_in_milliseconds: u64,
32    // Fixed time for fixed delay option
33    pub fixed_time_in_milliseconds: u64,
34}
35
36impl Default for DelaySettings {
37    fn default() -> Self {
38        DelaySettings {
39            delay_option: DelayOptions::None,
40            min_time_in_milliseconds: 0,
41            max_time_in_milliseconds: 0,
42            fixed_time_in_milliseconds: 0,
43        }
44    }
45}
46
47#[derive(Eq, PartialEq, Debug, Clone, Default)]
48/// Settings for each type
49pub struct DelayConfig {
50    config: HashMap<SupportedTraitTypesForAsyncDelay, DelaySettings>,
51}
52
53impl DelayConfig {
54    pub fn new(config: HashMap<SupportedTraitTypesForAsyncDelay, DelaySettings>) -> Self {
55        DelayConfig { config }
56    }
57
58    pub fn add_settings_for_all_types(&mut self, settings: DelaySettings) {
59        let iterator = SupportedTraitTypesForAsyncDelayIterator::new();
60
61        for supported_type in iterator {
62            self.config.insert(supported_type, settings.clone());
63        }
64    }
65
66    pub fn add_setting(
67        &mut self,
68        supported_type: SupportedTraitTypesForAsyncDelay,
69        settings: &DelaySettings,
70    ) {
71        self.config.insert(supported_type, settings.clone());
72    }
73
74    pub fn get_setting(
75        &self,
76        supported_type: &SupportedTraitTypesForAsyncDelay,
77    ) -> Option<&DelaySettings> {
78        self.config.get(supported_type)
79    }
80}
81
82#[async_trait]
83/// Implement this method to add some delay to async call
84pub trait TestableDelay {
85    /// Add a delay from settings
86    async fn handle_async_delay(settings: &DelaySettings) {
87        match settings.delay_option {
88            DelayOptions::None => {},
89            DelayOptions::Fixed => {
90                sleep(Duration::from_millis(settings.fixed_time_in_milliseconds)).await;
91            },
92            DelayOptions::Random => {
93                let sleep_in_millis = rand::thread_rng().gen_range(
94                    settings.min_time_in_milliseconds..=settings.max_time_in_milliseconds,
95                );
96                sleep(Duration::from_millis(sleep_in_millis)).await;
97            },
98        }
99    }
100
101    /// Look for settings in the config and run it
102    async fn run_delay_settings_from_config(delay_config: &DelayConfig);
103}
104
105/// Iterator to iterate over enum
106struct SupportedTraitTypesForAsyncDelayIterator {
107    index: usize,
108}
109
110impl SupportedTraitTypesForAsyncDelayIterator {
111    fn new() -> Self {
112        SupportedTraitTypesForAsyncDelayIterator { index: 0 }
113    }
114}
115
116impl Iterator for SupportedTraitTypesForAsyncDelayIterator {
117    type Item = SupportedTraitTypesForAsyncDelay;
118
119    fn next(&mut self) -> Option<Self::Item> {
120        let supported_type = match self.index {
121            0 => Some(SupportedTraitTypesForAsyncDelay::Storage),
122            1 => Some(SupportedTraitTypesForAsyncDelay::ValidatedState),
123            2 => Some(SupportedTraitTypesForAsyncDelay::BlockHeader),
124            _ => {
125                assert_eq!(
126                    self.index, 3,
127                    "Need to ensure that newly added or removed \
128                     `SupportedTraitTypesForAsyncDelay` enum is handled in iterator"
129                );
130                return None;
131            },
132        };
133        self.index += 1;
134        supported_type
135    }
136}