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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
// This file is part of Substrate.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Interfaces, types and utils for benchmarking a FRAME runtime.
use alloc::vec::Vec;
use codec::{Decode, Encode};
use frame_support::{dispatch::DispatchErrorWithPostInfo, pallet_prelude::*, traits::StorageInfo};
use scale_info::TypeInfo;
#[cfg(feature = "std")]
use serde::{Deserialize, Serialize};
use sp_io::hashing::blake2_256;
use sp_runtime::{traits::TrailingZeroInput, DispatchError};
use sp_storage::TrackedStorageKey;

/// An alphabet of possible parameters to use for benchmarking.
#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
#[derive(Encode, Decode, Clone, Copy, PartialEq, Debug, TypeInfo)]
#[allow(missing_docs)]
#[allow(non_camel_case_types)]
pub enum BenchmarkParameter {
	a,
	b,
	c,
	d,
	e,
	f,
	g,
	h,
	i,
	j,
	k,
	l,
	m,
	n,
	o,
	p,
	q,
	r,
	s,
	t,
	u,
	v,
	w,
	x,
	y,
	z,
}

#[cfg(feature = "std")]
impl std::fmt::Display for BenchmarkParameter {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "{:?}", self)
	}
}

/// The results of a single of benchmark.
#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
#[derive(Encode, Decode, Clone, PartialEq, Debug, TypeInfo)]
pub struct BenchmarkBatch {
	/// The pallet containing this benchmark.
	#[cfg_attr(feature = "std", serde(with = "serde_as_str"))]
	pub pallet: Vec<u8>,
	/// The instance of this pallet being benchmarked.
	#[cfg_attr(feature = "std", serde(with = "serde_as_str"))]
	pub instance: Vec<u8>,
	/// The extrinsic (or benchmark name) of this benchmark.
	#[cfg_attr(feature = "std", serde(with = "serde_as_str"))]
	pub benchmark: Vec<u8>,
	/// The results from this benchmark.
	pub results: Vec<BenchmarkResult>,
}

// TODO: could probably make API cleaner here.
/// The results of a single of benchmark, where time and db results are separated.
#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
#[derive(Encode, Decode, Clone, PartialEq, Debug)]
pub struct BenchmarkBatchSplitResults {
	/// The pallet containing this benchmark.
	#[cfg_attr(feature = "std", serde(with = "serde_as_str"))]
	pub pallet: Vec<u8>,
	/// The instance of this pallet being benchmarked.
	#[cfg_attr(feature = "std", serde(with = "serde_as_str"))]
	pub instance: Vec<u8>,
	/// The extrinsic (or benchmark name) of this benchmark.
	#[cfg_attr(feature = "std", serde(with = "serde_as_str"))]
	pub benchmark: Vec<u8>,
	/// The extrinsic timing results from this benchmark.
	pub time_results: Vec<BenchmarkResult>,
	/// The db tracking results from this benchmark.
	pub db_results: Vec<BenchmarkResult>,
}

/// Result from running benchmarks on a FRAME pallet.
/// Contains duration of the function call in nanoseconds along with the benchmark parameters
/// used for that benchmark result.
#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
#[derive(Encode, Decode, Default, Clone, PartialEq, Debug, TypeInfo)]
pub struct BenchmarkResult {
	pub components: Vec<(BenchmarkParameter, u32)>,
	pub extrinsic_time: u128,
	pub storage_root_time: u128,
	pub reads: u32,
	pub repeat_reads: u32,
	pub writes: u32,
	pub repeat_writes: u32,
	pub proof_size: u32,
	#[cfg_attr(feature = "std", serde(skip))]
	pub keys: Vec<(Vec<u8>, u32, u32, bool)>,
}

impl BenchmarkResult {
	pub fn from_weight(w: Weight) -> Self {
		Self { extrinsic_time: (w.ref_time() / 1_000) as u128, ..Default::default() }
	}
}

/// Helper module to make serde serialize `Vec<u8>` as strings.
#[cfg(feature = "std")]
mod serde_as_str {
	pub fn serialize<S>(value: &Vec<u8>, serializer: S) -> Result<S::Ok, S::Error>
	where
		S: serde::Serializer,
	{
		let s = std::str::from_utf8(value).map_err(serde::ser::Error::custom)?;
		serializer.collect_str(s)
	}

	pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
	where
		D: serde::de::Deserializer<'de>,
	{
		let s: &str = serde::de::Deserialize::deserialize(deserializer)?;
		Ok(s.into())
	}
}

/// Possible errors returned from the benchmarking pipeline.
#[derive(Clone, PartialEq, Debug)]
pub enum BenchmarkError {
	/// The benchmarking pipeline should stop and return the inner string.
	Stop(&'static str),
	/// The benchmarking pipeline is allowed to fail here, and we should use the
	/// included weight instead.
	Override(BenchmarkResult),
	/// The benchmarking pipeline is allowed to fail here, and we should simply
	/// skip processing these results.
	Skip,
	/// No weight can be determined; set the weight of this call to zero.
	///
	/// You can also use `Override` instead, but this is easier to use since `Override` expects the
	/// correct components to be present.
	Weightless,
}

impl From<BenchmarkError> for &'static str {
	fn from(e: BenchmarkError) -> Self {
		match e {
			BenchmarkError::Stop(s) => s,
			BenchmarkError::Override(_) => "benchmark override",
			BenchmarkError::Skip => "benchmark skip",
			BenchmarkError::Weightless => "benchmark weightless",
		}
	}
}

impl From<&'static str> for BenchmarkError {
	fn from(s: &'static str) -> Self {
		Self::Stop(s)
	}
}

impl From<DispatchErrorWithPostInfo> for BenchmarkError {
	fn from(e: DispatchErrorWithPostInfo) -> Self {
		Self::Stop(e.into())
	}
}

impl From<DispatchError> for BenchmarkError {
	fn from(e: DispatchError) -> Self {
		Self::Stop(e.into())
	}
}

/// Configuration used to setup and run runtime benchmarks.
#[derive(Encode, Decode, Default, Clone, PartialEq, Debug, TypeInfo)]
pub struct BenchmarkConfig {
	/// The encoded name of the pallet to benchmark.
	pub pallet: Vec<u8>,
	/// The encoded name of the benchmark/extrinsic to run.
	pub benchmark: Vec<u8>,
	/// The selected component values to use when running the benchmark.
	pub selected_components: Vec<(BenchmarkParameter, u32)>,
	/// Enable an extra benchmark iteration which runs the verification logic for a benchmark.
	pub verify: bool,
	/// Number of times to repeat benchmark within the Wasm environment. (versus in the client)
	pub internal_repeats: u32,
}

/// A list of benchmarks available for a particular pallet and instance.
///
/// All `Vec<u8>` must be valid utf8 strings.
#[derive(Encode, Decode, Default, Clone, PartialEq, Debug, TypeInfo)]
pub struct BenchmarkList {
	pub pallet: Vec<u8>,
	pub instance: Vec<u8>,
	pub benchmarks: Vec<BenchmarkMetadata>,
}

#[derive(Encode, Decode, Default, Clone, PartialEq, Debug, TypeInfo)]
pub struct BenchmarkMetadata {
	pub name: Vec<u8>,
	pub components: Vec<(BenchmarkParameter, u32, u32)>,
	pub pov_modes: Vec<(Vec<u8>, Vec<u8>)>,
}

sp_api::decl_runtime_apis! {
	/// Runtime api for benchmarking a FRAME runtime.
	pub trait Benchmark {
		/// Get the benchmark metadata available for this runtime.
		///
		/// Parameters
		/// - `extra`: Also list benchmarks marked "extra" which would otherwise not be
		///            needed for weight calculation.
		fn benchmark_metadata(extra: bool) -> (Vec<BenchmarkList>, Vec<StorageInfo>);

		/// Dispatch the given benchmark.
		fn dispatch_benchmark(config: BenchmarkConfig) -> Result<Vec<BenchmarkBatch>, sp_runtime::RuntimeString>;
	}
}

/// Interface that provides functions for benchmarking the runtime.
#[sp_runtime_interface::runtime_interface]
pub trait Benchmarking {
	/// Get the number of nanoseconds passed since the UNIX epoch
	///
	/// WARNING! This is a non-deterministic call. Do not use this within
	/// consensus critical logic.
	fn current_time() -> u128 {
		std::time::SystemTime::now()
			.duration_since(std::time::SystemTime::UNIX_EPOCH)
			.expect("Unix time doesn't go backwards; qed")
			.as_nanos()
	}

	/// Reset the trie database to the genesis state.
	fn wipe_db(&mut self) {
		self.wipe()
	}

	/// Commit pending storage changes to the trie database and clear the database cache.
	fn commit_db(&mut self) {
		self.commit()
	}

	/// Get the read/write count.
	fn read_write_count(&self) -> (u32, u32, u32, u32) {
		self.read_write_count()
	}

	/// Reset the read/write count.
	fn reset_read_write_count(&mut self) {
		self.reset_read_write_count()
	}

	/// Get the DB whitelist.
	fn get_whitelist(&self) -> Vec<TrackedStorageKey> {
		self.get_whitelist()
	}

	/// Set the DB whitelist.
	fn set_whitelist(&mut self, new: Vec<TrackedStorageKey>) {
		self.set_whitelist(new)
	}

	// Add a new item to the DB whitelist.
	fn add_to_whitelist(&mut self, add: TrackedStorageKey) {
		let mut whitelist = self.get_whitelist();
		match whitelist.iter_mut().find(|x| x.key == add.key) {
			// If we already have this key in the whitelist, update to be the most constrained
			// value.
			Some(item) => {
				item.reads += add.reads;
				item.writes += add.writes;
				item.whitelisted = item.whitelisted || add.whitelisted;
			},
			// If the key does not exist, add it.
			None => {
				whitelist.push(add);
			},
		}
		self.set_whitelist(whitelist);
	}

	// Remove an item from the DB whitelist.
	fn remove_from_whitelist(&mut self, remove: Vec<u8>) {
		let mut whitelist = self.get_whitelist();
		whitelist.retain(|x| x.key != remove);
		self.set_whitelist(whitelist);
	}

	fn get_read_and_written_keys(&self) -> Vec<(Vec<u8>, u32, u32, bool)> {
		self.get_read_and_written_keys()
	}

	/// Get current estimated proof size.
	fn proof_size(&self) -> Option<u32> {
		self.proof_size()
	}
}

/// The pallet benchmarking trait.
pub trait Benchmarking {
	/// Get the benchmarks available for this pallet. Generally there is one benchmark per
	/// extrinsic, so these are sometimes just called "extrinsics".
	///
	/// Parameters
	/// - `extra`: Also return benchmarks marked "extra" which would otherwise not be needed for
	///   weight calculation.
	fn benchmarks(extra: bool) -> Vec<BenchmarkMetadata>;

	/// Run the benchmarks for this pallet.
	fn run_benchmark(
		name: &[u8],
		selected_components: &[(BenchmarkParameter, u32)],
		whitelist: &[TrackedStorageKey],
		verify: bool,
		internal_repeats: u32,
	) -> Result<Vec<BenchmarkResult>, BenchmarkError>;
}

/// The recording trait used to mark the start and end of a benchmark.
pub trait Recording {
	/// Start the benchmark.
	fn start(&mut self) {}

	// Stop the benchmark.
	fn stop(&mut self) {}
}

/// A no-op recording, used for unit test.
struct NoopRecording;
impl Recording for NoopRecording {}

/// A no-op recording, used for tests that should setup some state before running the benchmark.
struct TestRecording<'a> {
	on_before_start: Option<&'a dyn Fn()>,
}

impl<'a> TestRecording<'a> {
	fn new(on_before_start: &'a dyn Fn()) -> Self {
		Self { on_before_start: Some(on_before_start) }
	}
}

impl<'a> Recording for TestRecording<'a> {
	fn start(&mut self) {
		(self.on_before_start.take().expect("start called more than once"))();
	}
}

/// Records the time and proof size of a single benchmark iteration.
pub struct BenchmarkRecording<'a> {
	on_before_start: Option<&'a dyn Fn()>,
	start_extrinsic: Option<u128>,
	finish_extrinsic: Option<u128>,
	start_pov: Option<u32>,
	end_pov: Option<u32>,
}

impl<'a> BenchmarkRecording<'a> {
	pub fn new(on_before_start: &'a dyn Fn()) -> Self {
		Self {
			on_before_start: Some(on_before_start),
			start_extrinsic: None,
			finish_extrinsic: None,
			start_pov: None,
			end_pov: None,
		}
	}
}

impl<'a> Recording for BenchmarkRecording<'a> {
	fn start(&mut self) {
		(self.on_before_start.take().expect("start called more than once"))();
		self.start_pov = crate::benchmarking::proof_size();
		self.start_extrinsic = Some(crate::benchmarking::current_time());
	}

	fn stop(&mut self) {
		self.finish_extrinsic = Some(crate::benchmarking::current_time());
		self.end_pov = crate::benchmarking::proof_size();
	}
}

impl<'a> BenchmarkRecording<'a> {
	pub fn start_pov(&self) -> Option<u32> {
		self.start_pov
	}

	pub fn end_pov(&self) -> Option<u32> {
		self.end_pov
	}

	pub fn diff_pov(&self) -> Option<u32> {
		self.start_pov.zip(self.end_pov).map(|(start, end)| end.saturating_sub(start))
	}

	pub fn elapsed_extrinsic(&self) -> Option<u128> {
		self.start_extrinsic
			.zip(self.finish_extrinsic)
			.map(|(start, end)| end.saturating_sub(start))
	}
}

/// The required setup for creating a benchmark.
///
/// Instance generic parameter is optional and can be used in order to capture unused generics for
/// instantiable pallets.
pub trait BenchmarkingSetup<T, I = ()> {
	/// Return the components and their ranges which should be tested in this benchmark.
	fn components(&self) -> Vec<(BenchmarkParameter, u32, u32)>;

	/// Set up the storage, and prepare a closure to run the benchmark.
	fn instance(
		&self,
		recording: &mut impl Recording,
		components: &[(BenchmarkParameter, u32)],
		verify: bool,
	) -> Result<(), BenchmarkError>;

	/// Same as `instance` but passing a closure to run before the benchmark starts.
	fn test_instance(
		&self,
		components: &[(BenchmarkParameter, u32)],
		on_before_start: &dyn Fn(),
	) -> Result<(), BenchmarkError> {
		return self.instance(&mut TestRecording::new(on_before_start), components, true);
	}

	/// Same as `instance` but passing a no-op recording for unit tests.
	fn unit_test_instance(
		&self,
		components: &[(BenchmarkParameter, u32)],
	) -> Result<(), BenchmarkError> {
		return self.instance(&mut NoopRecording {}, components, true);
	}
}

/// Grab an account, seeded by a name and index.
pub fn account<AccountId: Decode>(name: &'static str, index: u32, seed: u32) -> AccountId {
	let entropy = (name, index, seed).using_encoded(blake2_256);
	Decode::decode(&mut TrailingZeroInput::new(entropy.as_ref()))
		.expect("infinite length input; no invalid inputs for type; qed")
}

/// This caller account is automatically whitelisted for DB reads/writes by the benchmarking macro.
pub fn whitelisted_caller<AccountId: Decode>() -> AccountId {
	account::<AccountId>("whitelisted_caller", 0, 0)
}

#[macro_export]
macro_rules! whitelist_account {
	($acc:ident) => {
		frame_benchmarking::benchmarking::add_to_whitelist(
			frame_system::Account::<T>::hashed_key_for(&$acc).into(),
		);
	};
}