shumai/
lib.rs

1#![forbid(unsafe_code)]
2
3use std::{
4    fmt::Display,
5    ops::{Add, AddAssign},
6    sync::atomic::{AtomicBool, AtomicU64, Ordering},
7    time::Duration,
8};
9
10use serde::Serialize;
11
12mod env;
13mod metrics;
14mod result;
15mod runner;
16pub use result::ShumaiResult;
17pub use runner::run;
18pub use shumai_config_impl::{config, ShumaiConfig};
19
20pub mod __dep {
21    pub use colored;
22    pub use regex;
23    pub use serde;
24    pub use serde_json;
25    pub use toml;
26}
27
28/// The context send to MultiBench::run()
29pub struct Context<'a, C: BenchConfig> {
30    running: &'a AtomicBool,
31    ready_thread: &'a AtomicU64,
32    pub thread_id: usize,
33    pub thread_cnt: usize,
34    pub config: &'a C,
35}
36
37impl<'a, C: BenchConfig> Context<'a, C> {
38    /// A barrier to ensure all threads start at exactly the same time,
39    /// every run() should call context.wait_for_start() right after initialization or it will block forever.
40    pub fn wait_for_start(&self) {
41        self.ready_thread.fetch_add(1, Ordering::Relaxed);
42        while !self.is_running() {
43            std::hint::spin_loop();
44        }
45    }
46
47    /// Main thread will let each bencher know whether to stop running
48    pub fn is_running(&self) -> bool {
49        self.running.load(Ordering::Relaxed)
50    }
51
52    pub(crate) fn new(
53        thread_id: usize,
54        thread_cnt: usize,
55        config: &'a C,
56        ready_thread: &'a AtomicU64,
57        running: &'a AtomicBool,
58    ) -> Self {
59        Context {
60            running,
61            ready_thread,
62            thread_id,
63            thread_cnt,
64            config,
65        }
66    }
67}
68
69pub trait BenchResult:
70    serde::Serialize + Default + AddAssign + Add<Output = Self> + Clone + Send + Sync + Display
71{
72    fn short_value(&self) -> usize;
73
74    #[must_use]
75    fn normalize_time(self, dur: &Duration) -> Self;
76}
77
78impl BenchResult for usize {
79    fn short_value(&self) -> usize {
80        *self
81    }
82
83    fn normalize_time(self, dur: &Duration) -> usize {
84        ((self as f64) / dur.as_secs_f64()) as usize
85    }
86}
87
88pub trait BenchConfig: Clone + Serialize + Send + Sync {
89    fn name(&self) -> &String;
90    fn thread(&self) -> &[usize];
91    fn bench_sec(&self) -> usize;
92}
93
94/// The call chain of a MultiThreadBench:
95/// load() -> run() [thread t1] -> run() [thread t2] -> ... -> cleanup()
96pub trait ShumaiBench: Send + Sync {
97    type Result: BenchResult;
98    type Config: BenchConfig;
99
100    /// The benchmark should init their code, load the necessary data and warmup the resources
101    /// Note that the `load` will only be called once, no matter what `sample_size` is.
102    fn load(&mut self) -> Option<serde_json::Value>;
103
104    /// Run concurrent benchmark
105    /// Inside this function should call context.wait_for_start() to notify the main thread;
106    /// it also blocks current thread until every thread is ready (i.e. issued context.wait_for_start())
107    fn run(&self, context: Context<Self::Config>) -> Self::Result;
108
109    fn on_iteration_finished(&mut self, _cur_iter: usize) -> Option<serde_json::Value> {
110        None
111    }
112
113    fn on_thread_finished(&mut self, _cur_thread: usize) -> Option<serde_json::Value> {
114        None
115    }
116
117    /// clean up resources, if necessary
118    fn cleanup(&mut self) -> Option<serde_json::Value>;
119}
shumai/lib.rs

shumai/
lib.rs
