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

#![forbid(unsafe_code)]

use std::{
    fmt::Display,
    ops::{Add, AddAssign},
    sync::atomic::{AtomicBool, AtomicU64, Ordering},
    time::Duration,
};

use serde::Serialize;

mod env;
mod metrics;
mod result;
mod runner;
pub use result::ShumaiResult;
pub use runner::run;
pub use shumai_config_impl::{config, ShumaiConfig};

pub mod __dep {
    pub use colored;
    pub use once_cell;
    pub use regex;
    pub use serde;
    pub use serde_json;
    pub use toml;
}

/// The context send to MultiBench::run()
pub struct Context<'a, C: BenchConfig> {
    pub thread_id: usize,
    pub thread_cnt: usize,
    pub config: &'a C,
    ready_thread: &'a AtomicU64,
    running: &'a AtomicBool,
}

impl<C: BenchConfig> Context<'_, C> {
    /// A barrier to ensure all threads start at exactly the same time,
    /// every run() should call context.wait_for_start() right after initialization or it will block forever.
    pub fn wait_for_start(&self) {
        self.ready_thread.fetch_add(1, Ordering::Relaxed);
        while !self.is_running() {
            std::hint::spin_loop();
        }
    }

    /// Main thread will let each bencher know whether to stop running
    pub fn is_running(&self) -> bool {
        self.running.load(Ordering::Relaxed)
    }
}

pub trait BenchResult:
    serde::Serialize + Default + AddAssign + Add<Output = Self> + Clone + Send + Sync + Display
{
    fn short_value(&self) -> usize;

    #[must_use]
    fn normalize_time(self, dur: &Duration) -> Self;
}

impl BenchResult for usize {
    fn short_value(&self) -> usize {
        *self
    }

    fn normalize_time(self, dur: &Duration) -> usize {
        ((self as f64) / dur.as_secs_f64()) as usize
    }
}

pub trait BenchConfig: Clone + Serialize + Send + Sync {
    fn name(&self) -> &String;
    fn thread(&self) -> &[usize];
    fn bench_sec(&self) -> usize;
}

/// The call chain of a MultiThreadBench:
/// load() -> run() [thread t1] -> run() [thread t2] -> ... -> cleanup()
pub trait ShumaiBench: Send + Sync {
    type Result: BenchResult;
    type Config: BenchConfig;

    /// The benchmark should init their code, load the necessary data and warmup the resources
    /// Note that the `load` will only be called once, no matter what `sample_size` is.
    fn load(&mut self) -> Option<serde_json::Value>;

    /// Run concurrent benchmark
    /// Inside this function should call context.wait_for_start() to notify the main thread;
    /// it also blocks current thread until every thread is ready (i.e. issued context.wait_for_start())
    fn run(&self, context: Context<Self::Config>) -> Self::Result;

    fn on_iteration_finished(&mut self, _cur_iter: usize) {}

    fn on_thread_finished(&mut self, _cur_thread: usize) {}

    /// clean up resources, if necessary
    fn cleanup(&mut self) -> Option<serde_json::Value>;
}