gadget_sdk/benchmark/
mod.rs

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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
pub use tokio;

/// The runtime trait that all runtimes must implement.
pub trait Runtime {
    /// Runs the given future to completion on the runtime.
    fn block_on<F>(&self, future: F) -> F::Output
    where
        F: std::future::Future;
}

/// The [`tokio`](https://crates.io/crates/tokio) runtime.
///
/// This will execute the benchmark using the `tokio` runtime.
#[derive(Debug, Clone, Copy)]
pub struct TokioRuntime;

impl Runtime for TokioRuntime {
    fn block_on<F>(&self, future: F) -> F::Output
    where
        F: std::future::Future,
    {
        let rt = tokio::runtime::Handle::current();
        rt.block_on(future)
    }
}

/// A benchmarking harness.
#[derive(Debug)]
#[allow(dead_code)]
pub struct Bencher<R> {
    /// The runtime to use for running benchmarks.
    runtime: R,
    /// The time at which the benchmark started.
    started_at: std::time::Instant,
    /// The max number of cores for this benchmark.
    cores: usize,
}

/// The results of a benchmark.
///
/// This implements [`Display`] to provide a human-readable summary of the benchmark.
#[derive(Debug, Clone)]
pub struct BenchmarkSummary {
    /// The name of the benchmark.
    pub name: String,
    /// The job identifier.
    pub job_id: u8,
    /// The duration of the benchmark.
    pub elapsed: std::time::Duration,
    /// The number of cores the benchmark was run with.
    pub cores: usize,
    /// The amount of memory used by the benchmark (in bytes).
    pub ram_usage: u64,
}

impl<R: Runtime> Bencher<R> {
    /// Create a new benchmark harness.
    ///
    /// # Examples
    ///
    /// ```
    /// use gadget_sdk::benchmark::{Bencher, TokioRuntime};
    ///
    /// const THREADS: usize = 4;
    ///
    /// let bencher = Bencher::new(THREADS, TokioRuntime);
    /// ```
    pub fn new(threads: usize, runtime: R) -> Self {
        Self {
            runtime,
            started_at: std::time::Instant::now(),
            cores: threads,
        }
    }

    /// Runs the given future on the [`Runtime`].
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use gadget_sdk::benchmark::{Bencher, TokioRuntime};
    ///
    /// const THREADS: usize = 4;
    ///
    /// let bencher = Bencher::new(THREADS, TokioRuntime);
    /// bencher.block_on(async {
    ///     // Do some work...
    /// });
    /// ```
    pub fn block_on<F>(&self, future: F) -> F::Output
    where
        F: std::future::Future,
    {
        self.runtime.block_on(future)
    }

    /// Ends the benchmark and returns a summary.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use gadget_sdk::benchmark::{Bencher, TokioRuntime};
    /// const THREADS: usize = 4;
    ///
    /// let bencher = Bencher::new(THREADS, TokioRuntime);
    /// bencher.block_on(async {
    ///     // Do some work...
    /// });
    ///
    /// let summary = bencher.stop("my_benchmark", 0);
    /// println!("{}", summary);
    /// ```
    #[cfg(feature = "std")] // TODO: Benchmark execution time for WASM?
    pub fn stop<N: ToString>(&self, name: N, job_id: u8) -> BenchmarkSummary {
        let pid = sysinfo::get_current_pid().expect("Failed to get current process ID");
        let s = sysinfo::System::new_all();
        let process = s
            .process(pid)
            .expect("Failed to get current process from the system");
        let ram_usage = process.memory();
        BenchmarkSummary {
            name: name.to_string(),
            job_id,
            elapsed: self.started_at.elapsed(),
            cores: self.cores,
            ram_usage,
        }
    }
}

impl std::fmt::Display for BenchmarkSummary {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        const KB: f32 = 1024.00;
        const MB: f32 = 1024.00 * KB;
        const GB: f32 = 1024.00 * MB;
        let ram_usage = self.ram_usage as f32;
        let (ram_usage, unit) = if ram_usage < KB {
            (ram_usage, "B")
        } else if ram_usage < MB {
            (ram_usage / KB, "KB")
        } else if ram_usage < GB {
            (ram_usage / MB, "MB")
        } else {
            (ram_usage / GB, "GB")
        };

        write!(
            f,
            "Benchmark: {}\nJob ID: {}\nElapsed: {:?}\nvCPU: {}\nRAM Usage: {ram_usage:.2} {unit}\n",
            self.name, self.job_id, self.elapsed, self.cores,
        )
    }
}