[Features] Reintroduce sub-du

utils/frame/sub-du/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "sub-du"
+version = "0.1.0"
+authors = ["kianenigma <kian.peymani@gmail.com>", "Parity Technologies <admin@parity.io>"]
+edition = "2021"
+
+[dependencies]
+codec = { package = "parity-scale-codec", version = "3.0.0" }
+async-std = { version = "1.12.0", features = ["attributes"] }
+jsonrpsee = { version = "0.15.1", features = ["jsonrpsee-types", "jsonrpsee-http-client", "jsonrpsee-ws-client"] }
+separator = "0.4.1"
+ansi_term = "0.12.1"
+env_logger = "0.9"
+log = "0.4.17"
+structopt = "0.3"
+
+sub-storage = { path = "../sub-storage", features = ["helpers"] }
+frame-metadata = { version = "15.0.0", default-features = false, features = ["v14"] }
+
+[features]
+default = []
+remote-test-kusama = []
+remote-test-polkadot = []
+
+[dev-dependencies]
+assert_cmd = "2"

utils/frame/sub-du/README.md

@@ -0,0 +1 @@
+# sub-du

utils/frame/sub-du/src/main.rs

@@ -0,0 +1,254 @@
+use ansi_term::{Colour::*, Style};
+use frame_metadata::{RuntimeMetadata, RuntimeMetadataPrefixed, StorageEntryType};
+use separator::Separatable;
+use structopt::StructOpt;
+use sub_storage::{get_head, get_metadata, unwrap_decoded, Hash, StorageKey};
+
+const KB: usize = 1024;
+const MB: usize = KB * KB;
+const GB: usize = MB * MB;
+
+pub const LOG_TARGET: &'static str = "sub-du";
+
+fn get_prefix(indent: usize) -> &'static str {
+	match indent {
+		1 => "├─┬",
+		2 => "│ │─┬",
+		3 => "│ │ │─",
+		_ => panic!("Unexpected indent."),
+	}
+}
+
+struct Size(usize);
+
+impl std::fmt::Display for Size {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		if self.0 <= KB {
+			write!(f, "{: <4}B", self.0)?;
+		} else if self.0 <= MB {
+			write!(f, "{: <4}K", self.0 / KB)?;
+		} else if self.0 <= GB {
+			write!(f, "{: <4}M", self.0 / MB)?;
+		}
+
+		Ok(())
+	}
+}
+
+#[derive(Debug, Clone, Default)]
+struct Module {
+	pub name: String,
+	pub size: usize,
+	pub items: Vec<Storage>,
+}
+
+impl std::fmt::Display for Module {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		let mod_style = Style::new().bold().italic().fg(Green);
+		write!(
+			f,
+			"{} {} {}\n",
+			mod_style.paint(format!("{}", Size(self.size))),
+			get_prefix(2),
+			mod_style.paint(self.name.clone())
+		)?;
+		for s in self.items.iter() {
+			write!(f, "{} {} {}\n", Size(s.size), get_prefix(3), s)?;
+		}
+		Ok(())
+	}
+}
+
+impl Module {
+	fn new(name: String) -> Self {
+		Self { name, ..Default::default() }
+	}
+}
+
+#[derive(Debug, Copy, Clone)]
+pub enum StorageItem {
+	Value(usize),
+	Map(usize, usize),
+}
+
+impl std::fmt::Display for StorageItem {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		match self {
+			Self::Value(bytes) => write!(f, "Value({} bytes)", bytes.separated_string()),
+			Self::Map(bytes, count) => {
+				write!(f, "Map({} bytes, {} keys)", bytes.separated_string(), count)
+			}
+		}
+	}
+}
+
+impl Default for StorageItem {
+	fn default() -> Self {
+		Self::Value(0)
+	}
+}
+
+#[derive(Debug, Clone, Default)]
+struct Storage {
+	pub name: String,
+	pub size: usize,
+	pub item: StorageItem,
+}
+
+impl std::fmt::Display for Storage {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		let item_style = Style::new().italic();
+		write!(f, "{} => {}", item_style.paint(self.name.clone()), self.item)
+	}
+}
+
+impl Storage {
+	fn new(name: String, item: StorageItem) -> Self {
+		let size = match item {
+			StorageItem::Value(s) => s,
+			StorageItem::Map(s, _) => s,
+		};
+		Self { name, item, size }
+	}
+}
+
+#[derive(Debug, StructOpt)]
+#[structopt(
+	name = "sub-du",
+	about = "a du-like tool that prints the map of storage usage of a substrate chain"
+)]
+struct Opt {
+	/// The block number at which the scrap should happen. Use only the hex value, no need for a
+	/// `0x` prefix.
+	#[structopt(long)]
+	at: Option<Hash>,
+
+	/// The node to connect to.
+	#[structopt(long, default_value = "ws://localhost:9944")]
+	uri: String,
+
+	/// If true, intermediate values will be printed.
+	#[structopt(long, short)]
+	progress: bool,
+
+	/// Weather to scrape all pairs or just the size of them.
+	///
+	/// If enabled, the command might take longer but then the number of keys in each map is also
+	/// scraped.
+	///
+	/// # Warning
+	///
+	/// This uses an unsafe RPC call and can only be used if the target node allows it.
+	#[structopt(long, short)]
+	scrape_pairs: bool,
+}
+
+#[async_std::main]
+async fn main() -> () {
+	env_logger::Builder::from_default_env().format_module_path(false).format_level(true).init();
+
+	let opt = Opt::from_args();
+
+	// connect to a node.
+	use jsonrpsee_ws_client::{WsClient, WsConfig};
+	let client = WsClient::new(&opt.uri, WsConfig::default()).await.unwrap();
+
+	let mut modules: Vec<Module> = vec![];
+
+	// potentially replace head with the given hash
+	let head = get_head(&client).await;
+	let at = opt.at.unwrap_or(head);
+	let runtime = sub_storage::get_runtime_version(&client, at).await;
+
+	println!("Scraping at block {:?} of {}({})", at, runtime.spec_name, runtime.spec_version,);
+
+	let raw_metadata = get_metadata(&client, at).await.0;
+	let prefixed_metadata = <RuntimeMetadataPrefixed as codec::Decode>::decode(&mut &*raw_metadata)
+		.expect("Runtime Metadata failed to decode");
+	let metadata = prefixed_metadata.1;
+
+	if let RuntimeMetadata::V12(inner) = metadata {
+		let decode_modules = unwrap_decoded(inner.modules);
+		for module in decode_modules.into_iter() {
+			let name = unwrap_decoded(module.name);
+
+			// skip, if this module has no storage items.
+			if module.storage.is_none() {
+				log::warn!(
+					target: LOG_TARGET,
+					"Module with name {:?} seem to have no storage items.",
+					name
+				);
+				continue;
+			}
+
+			let storage = unwrap_decoded(module.storage.unwrap());
+			let prefix = unwrap_decoded(storage.prefix);
+			let entries = unwrap_decoded(storage.entries);
+			let mut module_info = Module::new(name.clone());
+
+			for storage_entry in entries.into_iter() {
+				let storage_name = unwrap_decoded(storage_entry.name);
+				let ty = storage_entry.ty;
+				let key_prefix =
+					sub_storage::module_prefix_raw(prefix.as_bytes(), storage_name.as_bytes());
+
+				let (pairs, size) = if opt.scrape_pairs {
+					// this should be slower but gives more detail.
+					let pairs =
+						sub_storage::get_pairs(StorageKey(key_prefix.clone()), &client, at).await;
+					let pairs = pairs
+						.into_iter()
+						.map(|(k, v)| (k.0, v.0))
+						.collect::<Vec<(Vec<u8>, Vec<u8>)>>();
+					let size = pairs.iter().fold(0, |acc, x| acc + x.1.len());
+					(pairs, size)
+				} else {
+					// This should be faster
+					let size = sub_storage::get_storage_size(StorageKey(key_prefix), &client, at)
+						.await
+						.unwrap_or_default() as usize;
+					let pairs: Vec<_> = vec![];
+					(pairs, size)
+				};
+
+				log::debug!(
+					target: LOG_TARGET,
+					"{:?}::{:?} => count: {}, size: {} bytes",
+					name,
+					storage_name,
+					pairs.len(),
+					size
+				);
+
+				module_info.size += size;
+				let item = match ty {
+					StorageEntryType::Plain(_) => StorageItem::Value(size),
+					StorageEntryType::Map { .. } | StorageEntryType::DoubleMap { .. } => {
+						StorageItem::Map(size, pairs.len())
+					}
+				};
+				module_info.items.push(Storage::new(storage_name, item));
+			}
+			module_info.items.sort_by_key(|x| x.size);
+			module_info.items.reverse();
+			println!("Scraped module {}. Total size {}.", module_info.name, module_info.size,);
+			if opt.progress {
+				print!("{}", module_info);
+			}
+			modules.push(module_info);
+		}
+
+		println!("Scraping results done. Final sorted tree:");
+		modules.sort_by_key(|m| m.size);
+		modules.reverse();
+
+		let total: usize = modules.iter().map(|m| m.size).sum();
+		println!("{} {} {}", Size(total), get_prefix(1), runtime.spec_name,);
+		modules.into_iter().for_each(|m| {
+			print!("{}", m);
+		});
+	} else {
+		panic!("Unsupported Metadata version");
+	}
+}

utils/frame/sub-du/tests/cli.rs

@@ -0,0 +1,24 @@
+use assert_cmd::Command;
+
+#[cfg(feature = "remote-test-kusama")]
+const TEST_URI: &'static str = "wss://kusama-rpc.polkadot.io/";
+#[cfg(feature = "remote-test-polkadot")]
+const TEST_URI: &'static str = "wss://rpc.polkadot.io/";
+#[cfg(not(any(feature = "remote-test-kusama", feature = "remote-test-polkadot")))]
+const TEST_URI: &'static str = "ws://localhost:9944";
+
+#[test]
+fn sub_du_starts_to_scrape_normal() {
+	let mut cmd = Command::cargo_bin("sub-du").unwrap();
+	let stdout = cmd
+		.args(&["--uri", TEST_URI, "-p"])
+		.timeout(std::time::Duration::from_secs(10))
+		.output()
+		.unwrap()
+		.stdout;
+
+	#[cfg(feature = "remote-test-kusama")]
+	assert!(String::from_utf8_lossy(&stdout).contains("of kusama("));
+	#[cfg(feature = "remote-test-polkadot")]
+	assert!(String::from_utf8_lossy(&stdout).contains("of polkadot("));
+}

utils/frame/sub-storage/Cargo.toml

@@ -0,0 +1,45 @@
+[package]
+name = "sub-storage"
+version = "0.1.0"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2021"
+build = "build.rs"
+
+[dependencies]
+codec = { package = "parity-scale-codec", version = "3.0.0", default-features = false, features = ["derive"] }
+serde = { version = "1.0.145", features = ["derive"] }
+
+jsonrpsee = { version = "0.15.1", features = ["jsonrpsee-types", "jsonrpsee-http-client", "jsonrpsee-ws-client"] }
+
+sp-core = { version = "6.0.0", path = "../../../primitives/core" }
+sp-version = { version = "5.0.0", default-features = false, path = "../../../primitives/version" }
+frame-support = { version = "4.0.0-dev", default-features = false, path = "../../../frame/support" }
+frame-metadata = { version = "15.0.0", default-features = false, features = ["v14"] }
+
+# Optional for helpers only.
+frame-system = { version = "4.0.0-dev", path = "../../../frame/system", optional = true }
+pallet-identity = { version = "4.0.0-dev", default-features = false, path = "../../../frame/identity", optional = true }
+pallet-balances = { version = "4.0.0-dev", default-features = false, path = "../../../frame/balances", optional = true }
+ansi_term = { version = "0.12.1", optional = true }
+
+[dev-dependencies]
+async-std = { version = "1.12.0" }
+tokio = { version = "1", features = ["full"] }
+hex-literal = "0.3.4"
+hex = "0.4.3"
+frame-system = { version = "4.0.0-dev", path = "../../../frame/system"}
+sp-runtime = { version = "6.0.0", path = "../../../primitives/runtime" }
+pallet-balances = { version = "4.0.0-dev", path = "../../../frame/balances" }
+pallet-staking = { version = "4.0.0-dev", path = "../../../frame/staking" }
+pallet-proxy = { version = "4.0.0-dev", path = "../../../frame/proxy" }
+
+[features]
+remote-test-kusama = []
+remote-test-polkadot = []
+default = []
+helpers = [
+	"frame-system",
+	"pallet-identity",
+	"pallet-balances",
+	"ansi_term",
+]

utils/frame/sub-storage/README.md

@@ -0,0 +1,21 @@
+# sub-storage
+
+## Sub-Storage.
+
+A thing wrapper around substrate's RPC calls that work with storage. This module is an
+equivalent ot the polkadojs-api `api.query` and `api.const`, but written in Rust.
+
+This crate is heavily dependent on the `jsonspsee` crate and uses it internally to connect to
+nodes.
+
+
+The base functions of this crate make no assumption about the runtime. Some runtime-dependent
+functions are provided under the `helpers` module.
+
+### Unsafe RPC calls.
+
+The most useful features provided by this crate are often marked as unsafe by the substrate
+nodes. Namely, [`get_pairs`] and [`enumerate_map`] can only be used against nodes that such
+external RPCs.
+
+THIS IS A TEST.

utils/frame/sub-storage/build.rs

@@ -0,0 +1,11 @@
+#[cfg(feature = "build-docs")]
+use std::process::Command;
+
+#[cfg(not(feature = "build-docs"))]
+fn main() {}
+
+#[cfg(feature = "build-docs")]
+fn main() {
+	println!("cargo:rerun-if-changed=src/lib.rs");
+	let _ = Command::new("cargo").arg("readme").arg(">").arg("README.md").output().unwrap();
+}