From 6aca526463f543edb71c5bb664a87b9f23d28880 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89lo=C3=AFs?= <c@elo.tf>
Date: Tue, 23 Apr 2024 13:16:03 +0200
Subject: [PATCH] add cli param keys_limit

---
 .../utils/frame/benchmarking-cli/src/storage/cmd.rs   | 11 ++++++++++-
 .../utils/frame/benchmarking-cli/src/storage/read.rs  |  8 ++++++--
 .../utils/frame/benchmarking-cli/src/storage/write.rs |  8 ++++++--
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs b/substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs
index 307c9207fdaf..aa224c853afc 100644
--- a/substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs
+++ b/substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs
@@ -116,6 +116,11 @@ pub struct StorageParams {
 	/// Include child trees in benchmark.
 	#[arg(long)]
 	pub include_child_trees: bool,
+
+	/// Maximum number of keys to read
+	/// (All keys if not define)
+	#[arg(long)]
+	pub keys_limit: Option<usize>,
 }
 
 impl StorageCmd {
@@ -191,7 +196,11 @@ impl StorageCmd {
 		BA: ClientBackend<B>,
 	{
 		let hash = client.usage_info().chain.best_hash;
-		let mut keys: Vec<_> = client.storage_keys(hash, None, None)?.collect();
+		let mut keys: Vec<_> = if let Some(keys_limit) = self.params.keys_limit {
+			client.storage_keys(hash, None, None)?.take(keys_limit).collect()
+		} else {
+			client.storage_keys(hash, None, None)?.collect()
+		};
 		let (mut rng, _) = new_rng(None);
 		keys.shuffle(&mut rng);
 
diff --git a/substrate/utils/frame/benchmarking-cli/src/storage/read.rs b/substrate/utils/frame/benchmarking-cli/src/storage/read.rs
index fe72364269d5..ba2d6e4b5b19 100644
--- a/substrate/utils/frame/benchmarking-cli/src/storage/read.rs
+++ b/substrate/utils/frame/benchmarking-cli/src/storage/read.rs
@@ -40,8 +40,12 @@ impl StorageCmd {
 		let best_hash = client.usage_info().chain.best_hash;
 
 		info!("Preparing keys from block {}", best_hash);
-		// Load all keys and randomly shuffle them.
-		let mut keys: Vec<_> = client.storage_keys(best_hash, None, None)?.collect();
+		// Load keys and randomly shuffle them.
+		let mut keys: Vec<_> = if let Some(keys_limit) = self.params.keys_limit {
+			client.storage_keys(best_hash, None, None)?.take(keys_limit).collect()
+		} else {
+			client.storage_keys(best_hash, None, None)?.collect()
+		};
 		let (mut rng, _) = new_rng(None);
 		keys.shuffle(&mut rng);
 
diff --git a/substrate/utils/frame/benchmarking-cli/src/storage/write.rs b/substrate/utils/frame/benchmarking-cli/src/storage/write.rs
index 65941497bda4..cc024ef1f6db 100644
--- a/substrate/utils/frame/benchmarking-cli/src/storage/write.rs
+++ b/substrate/utils/frame/benchmarking-cli/src/storage/write.rs
@@ -60,8 +60,12 @@ impl StorageCmd {
 		let trie = DbStateBuilder::<Block>::new(storage.clone(), original_root).build();
 
 		info!("Preparing keys from block {}", best_hash);
-		// Load all KV pairs and randomly shuffle them.
-		let mut kvs: Vec<_> = trie.pairs(Default::default())?.collect();
+		// Load KV pairs and randomly shuffle them.
+		let mut kvs: Vec<_> = if let Some(keys_limit) = self.params.keys_limit {
+			trie.pairs(Default::default())?.take(keys_limit).collect()
+		} else {
+			trie.pairs(Default::default())?.collect()
+		};
 		let (mut rng, _) = new_rng(None);
 		kvs.shuffle(&mut rng);
 		info!("Writing {} keys", kvs.len());