1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use std::{collections::HashMap, time::Duration};

use async_trait::async_trait;
use rand::Rng;
use tokio::time::sleep;

#[derive(Eq, Hash, PartialEq, Debug, Clone)]
/// What type of delay we want to apply to
pub enum DelayOptions {
    None,
    Random,
    Fixed,
}

#[derive(Eq, Hash, PartialEq, Debug, Clone)]
/// Current implementations that are supported for testing async delays
pub enum SupportedTraitTypesForAsyncDelay {
    Storage,
    ValidatedState,
    BlockHeader,
}

#[derive(Eq, Hash, PartialEq, Debug, Clone)]
/// Config for each supported type
pub struct DelaySettings {
    // Option to tell the async function what to do
    pub delay_option: DelayOptions,
    // Rng min time
    pub min_time_in_milliseconds: u64,
    // Rng max time
    pub max_time_in_milliseconds: u64,
    // Fixed time for fixed delay option
    pub fixed_time_in_milliseconds: u64,
}

impl Default for DelaySettings {
    fn default() -> Self {
        DelaySettings {
            delay_option: DelayOptions::None,
            min_time_in_milliseconds: 0,
            max_time_in_milliseconds: 0,
            fixed_time_in_milliseconds: 0,
        }
    }
}

#[derive(Eq, PartialEq, Debug, Clone, Default)]
/// Settings for each type
pub struct DelayConfig {
    config: HashMap<SupportedTraitTypesForAsyncDelay, DelaySettings>,
}

impl DelayConfig {
    pub fn new(config: HashMap<SupportedTraitTypesForAsyncDelay, DelaySettings>) -> Self {
        DelayConfig { config }
    }

    pub fn add_settings_for_all_types(&mut self, settings: DelaySettings) {
        let iterator = SupportedTraitTypesForAsyncDelayIterator::new();

        for supported_type in iterator {
            self.config.insert(supported_type, settings.clone());
        }
    }

    pub fn add_setting(
        &mut self,
        supported_type: SupportedTraitTypesForAsyncDelay,
        settings: &DelaySettings,
    ) {
        self.config.insert(supported_type, settings.clone());
    }

    pub fn get_setting(
        &self,
        supported_type: &SupportedTraitTypesForAsyncDelay,
    ) -> Option<&DelaySettings> {
        self.config.get(supported_type)
    }
}

#[async_trait]
/// Implement this method to add some delay to async call
pub trait TestableDelay {
    /// Add a delay from settings
    async fn handle_async_delay(settings: &DelaySettings) {
        match settings.delay_option {
            DelayOptions::None => {}
            DelayOptions::Fixed => {
                sleep(Duration::from_millis(settings.fixed_time_in_milliseconds)).await;
            }
            DelayOptions::Random => {
                let sleep_in_millis = rand::thread_rng().gen_range(
                    settings.min_time_in_milliseconds..=settings.max_time_in_milliseconds,
                );
                sleep(Duration::from_millis(sleep_in_millis)).await;
            }
        }
    }

    /// Look for settings in the config and run it
    async fn run_delay_settings_from_config(delay_config: &DelayConfig);
}

/// Iterator to iterate over enum
struct SupportedTraitTypesForAsyncDelayIterator {
    index: usize,
}

impl SupportedTraitTypesForAsyncDelayIterator {
    fn new() -> Self {
        SupportedTraitTypesForAsyncDelayIterator { index: 0 }
    }
}

impl Iterator for SupportedTraitTypesForAsyncDelayIterator {
    type Item = SupportedTraitTypesForAsyncDelay;

    fn next(&mut self) -> Option<Self::Item> {
        let supported_type = match self.index {
            0 => Some(SupportedTraitTypesForAsyncDelay::Storage),
            1 => Some(SupportedTraitTypesForAsyncDelay::ValidatedState),
            2 => Some(SupportedTraitTypesForAsyncDelay::BlockHeader),
            _ => {
                assert_eq!(self.index, 3, "Need to ensure that newly added or removed `SupportedTraitTypesForAsyncDelay` enum is handled in iterator");
                return None;
            }
        };
        self.index += 1;
        supported_type
    }
}