Skip to main content

mollusk_svm_bencher/
lib.rs

1//! The Mollusk Compute Unit Bencher can be used to benchmark the compute unit
2//! usage of Solana programs. It provides a simple API for developers to write
3//! benchmarks for their programs, or compare multiple implementations of their
4//! programs in a matrix, which can be checked while making changes to the
5//! program.
6//!
7//! A markdown file is generated, which captures all of the compute unit
8//! benchmarks. In the case of single program if a benchmark has a previous
9//! value, the delta is also recorded. This can be useful for developers to
10//! check the implications of changes to the program on compute unit usage.
11//!
12//! ```rust,ignore
13//! use {
14//!     mollusk_svm_bencher::MolluskComputeUnitBencher,
15//!     mollusk_svm::Mollusk,
16//!     /* ... */
17//! };
18//!
19//! // Optionally disable logging.
20//! solana_logger::setup_with("");
21//!
22//! /* Instruction & accounts setup ... */
23//!
24//! let mollusk = Mollusk::new(&program_id, "my_program");
25//!
26//! MolluskComputeUnitBencher::new(mollusk)
27//!     .bench(("bench0", &instruction0, &accounts0))
28//!     .bench(("bench1", &instruction1, &accounts1))
29//!     .bench(("bench2", &instruction2, &accounts2))
30//!     .bench(("bench3", &instruction3, &accounts3))
31//!     .must_pass(true)
32//!     .out_dir("../target/benches")
33//!     .execute();
34//! ```
35//!
36//! The `must_pass` argument can be provided to trigger a panic if any defined
37//! benchmark tests do not pass. `out_dir` specifies the directory where the
38//! markdown file will be written.
39//!
40//! Developers can invoke this benchmark test with `cargo bench`. They may need
41//! to add a bench to the project's `Cargo.toml`.
42//!
43//! ```toml
44//! [[bench]]
45//! name = "compute_units"
46//! harness = false
47//! ```
48//!
49//! The markdown file will contain entries according to the defined benchmarks.
50//!
51//! ```markdown
52//! | Name   | CUs   | Delta  |
53//! |--------|-------|--------|
54//! | bench0 | 450   | --     |
55//! | bench1 | 579   | -129   |
56//! | bench2 | 1,204 | +754   |
57//! | bench3 | 2,811 | +2,361 |
58//! ```
59//! ### Matrix Benchmarking
60//!
61//! If you want to compare multiple program implementations (e.g., comparing
62//! an optimized version against a baseline), use
63//! `MolluskComputeUnitMatrixBencher`. This generates a table where each program
64//! is a column.
65//!
66//! ```rust,ignore
67//! use {
68//!     mollusk_svm_bencher::MolluskComputeUnitMatrixBencher,
69//!     mollusk_svm::Mollusk,
70//!     /* ... */
71//! };
72//!
73//! /* Instruction & accounts setup ... */
74//!
75//! let mollusk = Mollusk::new(&program_id, "program_v1");
76//!
77//! MolluskComputeUnitMatrixBencher::new(mollusk)
78//!     .programs(&["program_v1", "program_v2", "program_v3"])
79//!     .bench(("bench0", &instruction0, &accounts0))
80//!     .bench(("bench1", &instruction1, &accounts1))
81//!     .must_pass(true)
82//!     .out_dir("../target/benches")
83//!     .execute();
84//! ```
85//! The matrix markdown file will contain entries comparing all provided
86//! programs.
87//!
88//! ```markdown
89//! | Name     | CU (`program_v1`) | CU (`program_v2`) | CU (`program_v3`) |
90//! |----------|-------------------|-------------------|-------------------|
91//! | `bench0` | 1,400             | 1,390             | 1,385             |
92//! | `bench1` | 2,100             | 2,050             | 2,045             |
93//! ```
94
95pub mod result;
96
97use {
98    chrono::Utc,
99    mollusk_svm::{result::ProgramResult, Mollusk},
100    result::{
101        mx_write_results, write_results, MolluskComputeUnitBenchResult,
102        MolluskComputeUnitMatrixBenchResult,
103    },
104    solana_account::Account,
105    solana_instruction::Instruction,
106    solana_pubkey::Pubkey,
107    std::{path::PathBuf, process::Command},
108};
109
110/// A bench is a tuple of a name, an instruction, and a list of accounts.
111pub type Bench<'a> = (&'a str, &'a Instruction, &'a [(Pubkey, Account)]);
112
113/// Mollusk's compute unit bencher.
114///
115/// Allows developers to bench test compute unit usage on their programs.
116pub struct MolluskComputeUnitBencher<'a> {
117    benches: Vec<Bench<'a>>,
118    mollusk: Mollusk,
119    must_pass: bool,
120    out_dir: PathBuf,
121}
122
123impl<'a> MolluskComputeUnitBencher<'a> {
124    /// Create a new bencher, to which benches and configurations can be added.
125    pub fn new(mollusk: Mollusk) -> Self {
126        let mut out_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
127        out_dir.push("benches");
128        Self {
129            benches: Vec::new(),
130            mollusk,
131            must_pass: false,
132            out_dir,
133        }
134    }
135
136    /// Add a bench to the bencher.
137    pub fn bench(mut self, bench: Bench<'a>) -> Self {
138        self.benches.push(bench);
139        self
140    }
141
142    /// Set whether the bencher should panic if a program execution fails.
143    pub const fn must_pass(mut self, must_pass: bool) -> Self {
144        self.must_pass = must_pass;
145        self
146    }
147
148    /// Set the output directory for the results.
149    pub fn out_dir(mut self, out_dir: &str) -> Self {
150        self.out_dir = PathBuf::from(out_dir);
151        self
152    }
153
154    /// Execute the benches.
155    pub fn execute(&mut self) {
156        let table_header = Utc::now().to_string();
157        let solana_version = get_solana_version();
158        let bench_results = std::mem::take(&mut self.benches)
159            .into_iter()
160            .map(|(name, instruction, accounts)| {
161                let result = self.mollusk.process_instruction(instruction, accounts);
162                match result.program_result {
163                    ProgramResult::Success => (),
164                    _ => {
165                        if self.must_pass {
166                            panic!(
167                                "Program execution failed, but `must_pass` was set. Error: {:?}",
168                                result.program_result
169                            );
170                        }
171                    }
172                }
173                MolluskComputeUnitBenchResult::new(name, result)
174            })
175            .collect::<Vec<_>>();
176        write_results(&self.out_dir, &table_header, &solana_version, bench_results);
177    }
178}
179
180/// Mollusk's matrix compute unit bencher.
181///
182/// Allows developers to bench test compute unit usage on multiple
183/// implementations of their programs.
184pub struct MolluskComputeUnitMatrixBencher<'a> {
185    mollusk: &'a mut Mollusk,
186    program_names: Vec<&'a str>,
187    benches: Vec<Bench<'a>>,
188    must_pass: bool,
189    out_dir: PathBuf,
190}
191
192impl<'a> MolluskComputeUnitMatrixBencher<'a> {
193    /// Create a new matrix bencher, to which benches and configurations can be
194    /// added.
195    pub fn new(mollusk: &'a mut Mollusk) -> Self {
196        let mut out_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
197        out_dir.push("benches");
198        Self {
199            mollusk,
200            program_names: Vec::new(),
201            benches: Vec::new(),
202            must_pass: false,
203            out_dir,
204        }
205    }
206
207    /// Add the program names to be benched.
208    pub fn programs(mut self, names: &[&'a str]) -> Self {
209        self.program_names = names.to_vec();
210        self
211    }
212
213    /// Add a bench to the bencher.
214    pub fn bench(mut self, bench: Bench<'a>) -> Self {
215        self.benches.push(bench);
216        self
217    }
218
219    /// Set whether the bencher should panic if a program execution fails.
220    pub fn must_pass(mut self, must_pass: bool) -> Self {
221        self.must_pass = must_pass;
222        self
223    }
224
225    /// Set the output directory for the results.
226    pub fn out_dir(mut self, out_dir: &str) -> Self {
227        self.out_dir = PathBuf::from(out_dir);
228        self
229    }
230
231    /// Execute the benches.
232    pub fn execute(&mut self) {
233        let table_header = Utc::now().to_string();
234        let solana_version = get_solana_version();
235
236        let mut bench_results: Vec<MolluskComputeUnitMatrixBenchResult> = Vec::new();
237        for program_name in &self.program_names {
238            // Extract the program ID from the first instruction.
239            if let Some((_, first_instruction, _)) = self.benches.first() {
240                self.mollusk
241                    .add_program(&first_instruction.program_id, program_name);
242            }
243
244            let mut ix_results = MolluskComputeUnitMatrixBenchResult::new(program_name);
245
246            for (ix_name, instruction, accounts) in &self.benches {
247                let result = self.mollusk.process_instruction(instruction, accounts);
248                match result.program_result {
249                    ProgramResult::Success => (),
250                    _ => {
251                        if self.must_pass {
252                            panic!(
253                                "Program execution failed, but `must_pass` was set. Error: {:?}",
254                                result.program_result
255                            );
256                        }
257                    }
258                }
259                ix_results.add_result(ix_name, result);
260            }
261            bench_results.push(ix_results);
262        }
263
264        mx_write_results(
265            &self.out_dir,
266            &table_header,
267            &solana_version,
268            &bench_results,
269        );
270    }
271}
272
273pub fn get_solana_version() -> String {
274    match Command::new("solana").arg("--version").output() {
275        Ok(output) if output.status.success() => {
276            String::from_utf8_lossy(&output.stdout).trim().to_string()
277        }
278        _ => "Unknown".to_string(),
279    }
280}