balter 0.1.1

A load/stress testing framework.
Documentation 
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

use serde::{Deserialize, Serialize};
use serde_with::{serde_as, DurationSeconds};
use std::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};

mod goal_tps;
mod saturate;

use goal_tps::run_goal_tps;
use saturate::run_saturate;

const DEFAULT_SATURATE_ERROR_RATE: f64 = 0.03;

// TODO: We should _not_ need to use a Boxed future! Every single function call for any load
// testing is boxed which *sucks*. Unfortunately I haven't figured out how to appease the Type
// system.
pub type BoxedFut = Pin<Box<dyn Future<Output = ()> + Send>>;

// TODO: Have a separate builder
#[serde_as]
#[derive(Clone, Serialize, Deserialize, Debug)]
pub(crate) struct ScenarioConfig {
    pub name: String,
    #[serde_as(as = "DurationSeconds")]
    pub duration: Duration,
    pub kind: ScenarioKind,
}

impl ScenarioConfig {
    fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
            duration: Default::default(),
            kind: Default::default(),
        }
    }
}

#[derive(Default, Clone, Copy, Serialize, Deserialize, Debug)]
pub(crate) enum ScenarioKind {
    #[default]
    Once,
    Tps {
        goal_tps: u32,
    },
    Saturate {
        error_rate: f64,
    },
}

#[pin_project::pin_project]
pub struct Scenario {
    fut: fn() -> BoxedFut,
    runner_fut: Option<Pin<Box<dyn Future<Output = ()> + Send>>>,
    config: ScenarioConfig,
}

impl Scenario {
    pub fn new(name: &str, fut: fn() -> BoxedFut) -> Self {
        Self {
            fut,
            runner_fut: None,
            config: ScenarioConfig::new(name),
        }
    }

    pub fn saturate(mut self) -> Self {
        self.config.kind = ScenarioKind::Saturate {
            error_rate: DEFAULT_SATURATE_ERROR_RATE,
        };
        self
    }

    pub fn saturate_error_rate(mut self, error_rate: f64) -> Self {
        self.config.kind = ScenarioKind::Saturate { error_rate };
        self
    }

    pub fn tps(mut self, tps: u32) -> Self {
        self.config.kind = ScenarioKind::Tps { goal_tps: tps };
        self
    }

    pub fn duration(mut self, duration: Duration) -> Self {
        self.config.duration = duration;
        self
    }

    pub(crate) fn set_config(mut self, config: ScenarioConfig) -> Self {
        self.config = config;
        self
    }
}

impl Future for Scenario {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // TODO: Surely there is a cleaner way to do this...
        if self.runner_fut.is_none() {
            let fut = self.fut;
            let config = self.config.clone();
            // TODO: There must be a way to run this future without boxing it. I'm feel like I'm
            // missing something really simple here.
            self.runner_fut = Some(Box::pin(async move { run_scenario(fut, config).await }));
        }

        if let Some(runner) = &mut self.runner_fut {
            runner.as_mut().poll(cx)
        } else {
            unreachable!()
        }
    }
}

async fn run_scenario(scenario: fn() -> BoxedFut, config: ScenarioConfig) {
    match config.kind {
        ScenarioKind::Once => scenario().await,
        ScenarioKind::Tps { goal_tps } => run_goal_tps(scenario, config, goal_tps).await,
        ScenarioKind::Saturate { error_rate } => run_saturate(scenario, config, error_rate).await,
    }
}