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

// This file is part of Substrate.

// Copyright (C) 2022 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.

use sc_cli::Result;
use sc_service::Configuration;

use log::info;
use serde::Serialize;
use std::{env, fs, path::PathBuf};

use super::cmd::StorageParams;
use crate::shared::{Stats, UnderscoreHelper};

static VERSION: &str = env!("CARGO_PKG_VERSION");
static TEMPLATE: &str = include_str!("./weights.hbs");

/// Data consumed by Handlebar to fill out the `weights.hbs` template.
#[derive(Serialize, Default, Debug, Clone)]
pub(crate) struct TemplateData {
	/// Name of the database used.
	db_name: String,
	/// Block number that was used.
	block_number: String,
	/// Name of the runtime. Taken from the chain spec.
	runtime_name: String,
	/// Version of the benchmarking CLI used.
	version: String,
	/// Date that the template was filled out.
	date: String,
	/// Hostname of the machine that executed the benchmarks.
	hostname: String,
	/// CPU name of the machine that executed the benchmarks.
	cpuname: String,
	/// Command line arguments that were passed to the CLI.
	args: Vec<String>,
	/// Storage params of the executed command.
	params: StorageParams,
	/// The weight for one `read`.
	read_weight: u64,
	/// The weight for one `write`.
	write_weight: u64,
	/// Stats about a `read` benchmark. Contains *time* and *value size* stats.
	/// The *value size* stats are currently not used in the template.
	read: Option<(Stats, Stats)>,
	/// Stats about a `write` benchmark. Contains *time* and *value size* stats.
	/// The *value size* stats are currently not used in the template.
	write: Option<(Stats, Stats)>,
}

impl TemplateData {
	/// Returns a new [`Self`] from the given configuration.
	pub fn new(cfg: &Configuration, params: &StorageParams) -> Self {
		TemplateData {
			db_name: format!("{}", cfg.database),
			runtime_name: cfg.chain_spec.name().into(),
			version: VERSION.into(),
			date: chrono::Utc::now().format("%Y-%m-%d (Y/M/D)").to_string(),
			hostname: params.hostinfo.hostname(),
			cpuname: params.hostinfo.cpuname(),
			args: env::args().collect::<Vec<String>>(),
			params: params.clone(),
			..Default::default()
		}
	}

	/// Sets the stats and calculates the final weights.
	pub fn set_stats(
		&mut self,
		read: Option<(Stats, Stats)>,
		write: Option<(Stats, Stats)>,
	) -> Result<()> {
		if let Some(read) = read {
			self.read_weight = self.params.weight_params.calc_weight(&read.0)?;
			self.read = Some(read);
		}
		if let Some(write) = write {
			self.write_weight = self.params.weight_params.calc_weight(&write.0)?;
			self.write = Some(write);
		}
		Ok(())
	}

	/// Sets the block id that was used.
	pub fn set_block_number(&mut self, block_number: String) {
		self.block_number = block_number
	}

	/// Fills out the `weights.hbs` or specified HBS template with its own data.
	/// Writes the result to `path` which can be a directory or file.
	pub fn write(&self, path: &Option<PathBuf>, hbs_template: &Option<PathBuf>) -> Result<()> {
		let mut handlebars = handlebars::Handlebars::new();
		// Format large integers with underscore.
		handlebars.register_helper("underscore", Box::new(UnderscoreHelper));
		// Don't HTML escape any characters.
		handlebars.register_escape_fn(|s| -> String { s.to_string() });
		// Use custom template if provided.
		let template = match hbs_template {
			Some(template) if template.is_file() => fs::read_to_string(template)?,
			Some(_) => return Err("Handlebars template is not a valid file!".into()),
			None => TEMPLATE.to_string(),
		};

		let out_path = self.build_path(path);
		let mut fd = fs::File::create(&out_path)?;
		info!("Writing weights to {:?}", fs::canonicalize(&out_path)?);

		handlebars
			.render_template_to_write(&template, &self, &mut fd)
			.map_err(|e| format!("HBS template write: {:?}", e).into())
	}

	/// Builds a path for the weight file.
	fn build_path(&self, weight_out: &Option<PathBuf>) -> PathBuf {
		let mut path = match weight_out {
			Some(p) => PathBuf::from(p),
			None => PathBuf::new(),
		};

		if path.is_dir() || path.as_os_str().is_empty() {
			path.push(format!("{}_weights", self.db_name.to_lowercase()));
			path.set_extension("rs");
		}
		path
	}
}