Add trie cache warmup

cumulus/test/service/src/lib.rs

@@ -876,6 +876,7 @@ pub fn node_config(
 		keystore: KeystoreConfig::InMemory,
 		database: DatabaseSource::RocksDb { path: root.join("db"), cache_size: 128 },
 		trie_cache_maximum_size: Some(64 * 1024 * 1024),
+		warm_up_trie_cache: None,
 		state_pruning: Some(PruningMode::ArchiveAll),
 		blocks_pruning: BlocksPruning::KeepAll,
 		chain_spec: spec,

polkadot/node/test/service/src/lib.rs

@@ -205,6 +205,7 @@ pub fn node_config(
 		keystore: KeystoreConfig::InMemory,
 		database: DatabaseSource::RocksDb { path: root.join("db"), cache_size: 128 },
 		trie_cache_maximum_size: Some(64 * 1024 * 1024),
+		warm_up_trie_cache: None,
 		state_pruning: Default::default(),
 		blocks_pruning: BlocksPruning::KeepFinalized,
 		chain_spec: Box::new(spec),

prdoc/pr_7556.prdoc

@@ -0,0 +1,11 @@
+title: 'Add trie cache warmup'
+doc:
+- audience: Node Dev
+  description: "Warm up the Trie cache based on a CLI flag to enhance the performance of smart contracts on AssetHub by reducing storage access time."
+crates:
+- name: sc-cli
+  bump: major
+- name: sc-service
+  bump: major
+- name: sc-client-db
+  bump: minor

substrate/bin/node/cli/benches/block_production.rs

@@ -72,6 +72,7 @@ fn new_node(tokio_handle: Handle) -> node_cli::service::NewFullBase {
 		keystore: KeystoreConfig::InMemory,
 		database: DatabaseSource::RocksDb { path: root.join("db"), cache_size: 128 },
 		trie_cache_maximum_size: Some(64 * 1024 * 1024),
+		warm_up_trie_cache: None,
 		state_pruning: Some(PruningMode::ArchiveAll),
 		blocks_pruning: BlocksPruning::KeepAll,
 		chain_spec: spec,

substrate/bin/node/cli/benches/transaction_pool.rs

@@ -63,6 +63,7 @@ fn new_node(tokio_handle: Handle) -> node_cli::service::NewFullBase {
 		keystore: KeystoreConfig::InMemory,
 		database: DatabaseSource::RocksDb { path: root.join("db"), cache_size: 128 },
 		trie_cache_maximum_size: Some(64 * 1024 * 1024),
+		warm_up_trie_cache: None,
 		state_pruning: Some(PruningMode::ArchiveAll),
 		blocks_pruning: BlocksPruning::KeepAll,
 		chain_spec: spec,

substrate/client/cli/src/config.rs

@@ -252,6 +252,16 @@ pub trait CliConfiguration<DCV: DefaultConfigurationValues = ()>: Sized {
 		Ok(self.import_params().map(|x| x.trie_cache_maximum_size()).unwrap_or_default())
 	}
 
+	/// Get if we should warm up the trie cache.
+	///
+	/// By default this is retrieved from `ImportParams` if it is available. Otherwise its `None`.
+	fn warm_up_trie_cache(&self) -> Result<Option<sc_service::config::TrieCacheWarmUpStrategy>> {
+		Ok(self
+			.import_params()
+			.map(|x| x.warm_up_trie_cache().map(|x| x.into()))
+			.unwrap_or_default())
+	}
+
 	/// Get the state pruning mode.
 	///
 	/// By default this is retrieved from `PruningMode` if it is available. Otherwise its
@@ -528,6 +538,7 @@ pub trait CliConfiguration<DCV: DefaultConfigurationValues = ()>: Sized {
 			database: self.database_config(&config_dir, database_cache_size, database)?,
 			data_path: config_dir,
 			trie_cache_maximum_size: self.trie_cache_maximum_size()?,
+			warm_up_trie_cache: self.warm_up_trie_cache()?,
 			state_pruning: self.state_pruning()?,
 			blocks_pruning: self.blocks_pruning()?,
 			executor: ExecutorConfiguration {

substrate/client/cli/src/params/import_params.rs

@@ -23,7 +23,7 @@ use crate::{
 	},
 	params::{DatabaseParams, PruningParams},
 };
-use clap::Args;
+use clap::{Args, ValueEnum};
 use std::path::PathBuf;
 
 /// Parameters for block import.
@@ -80,6 +80,38 @@ pub struct ImportParams {
 	/// Providing `0` will disable the cache.
 	#[arg(long, value_name = "Bytes", default_value_t = 1024 * 1024 * 1024)]
 	pub trie_cache_size: usize,
+
+	/// Warm up the trie cache.
+	///
+	/// No warmup if flag is not present. Using flag without value chooses non-blocking warmup.
+	#[arg(long, value_name = "STRATEGY", value_enum, num_args = 0..=1, default_missing_value = "non-blocking")]
+	pub warm_up_trie_cache: Option<TrieCacheWarmUpStrategy>,
+}
+
+/// Warmup strategy for the trie cache.
+#[derive(Debug, Clone, Copy, ValueEnum)]
+pub enum TrieCacheWarmUpStrategy {
+	/// Warm up the cache in a non-blocking way.
+	#[clap(name = "non-blocking")]
+	NonBlocking,
+	/// Warm up the cache in a blocking way (not recommended for production use).
+	///
+	/// When enabled, the trie cache warm-up will block the node startup until complete.
+	/// This is not recommended for production use as it can significantly delay node startup.
+	/// Only enable this option for testing or debugging purposes.
+	#[clap(name = "blocking")]
+	Blocking,
+}
+
+impl From<TrieCacheWarmUpStrategy> for sc_service::config::TrieCacheWarmUpStrategy {
+	fn from(strategy: TrieCacheWarmUpStrategy) -> Self {
+		match strategy {
+			TrieCacheWarmUpStrategy::NonBlocking =>
+				sc_service::config::TrieCacheWarmUpStrategy::NonBlocking,
+			TrieCacheWarmUpStrategy::Blocking =>
+				sc_service::config::TrieCacheWarmUpStrategy::Blocking,
+		}
+	}
 }
 
 impl ImportParams {
@@ -92,6 +124,11 @@ impl ImportParams {
 		}
 	}
 
+	/// Specify if we should warm up the trie cache.
+	pub fn warm_up_trie_cache(&self) -> Option<TrieCacheWarmUpStrategy> {
+		self.warm_up_trie_cache
+	}
+
 	/// Get the WASM execution method from the parameters
 	pub fn wasm_method(&self) -> sc_service::config::WasmExecutionMethod {
 		self.execution_strategies.check_usage_and_print_deprecation_warning();

substrate/client/cli/src/runner.rs

@@ -252,6 +252,7 @@ mod tests {
 				keystore: sc_service::config::KeystoreConfig::InMemory,
 				database: sc_client_db::DatabaseSource::ParityDb { path: root.clone() },
 				trie_cache_maximum_size: None,
+				warm_up_trie_cache: None,
 				state_pruning: None,
 				blocks_pruning: sc_client_db::BlocksPruning::KeepAll,
 				chain_spec: Box::new(

substrate/client/db/Cargo.toml

@@ -43,6 +43,7 @@ sp-database = { workspace = true, default-features = true }
 sp-runtime = { workspace = true, default-features = true }
 sp-state-machine = { workspace = true, default-features = true }
 sp-trie = { workspace = true, default-features = true }
+sysinfo = { workspace = true }
 
 [dev-dependencies]
 array-bytes = { workspace = true, default-features = true }

substrate/client/db/src/lib.rs

@@ -1239,6 +1239,22 @@ impl<Block: BlockT> Backend<Block> {
 
 		let offchain_storage = offchain::LocalStorage::new(db.clone());
 
+		let shared_trie_cache = config.trie_cache_maximum_size.map(|maximum_size| {
+			let system_memory = sysinfo::System::new_all();
+			let used_memory = system_memory.used_memory();
+			let total_memory = system_memory.total_memory();
+
+			debug!("Initializing shared trie cache with size {} bytes, {}% of total memory", maximum_size, (maximum_size as f64 / total_memory as f64 * 100.0));
+			if maximum_size as u64 > total_memory - used_memory {
+				warn!(
+					"Not enough memory to initialize shared trie cache. Cache size: {} bytes. System memory: used {} bytes, total {} bytes",
+					maximum_size, used_memory, total_memory,
+				);
+			}
+
+			SharedTrieCache::new(sp_trie::cache::CacheSize::new(maximum_size), config.metrics_registry.as_ref())
+		});
+
 		let backend = Backend {
 			storage: Arc::new(storage_db),
 			offchain_storage,
@@ -1250,12 +1266,7 @@ impl<Block: BlockT> Backend<Block> {
 			state_usage: Arc::new(StateUsageStats::new()),
 			blocks_pruning: config.blocks_pruning,
 			genesis_state: RwLock::new(None),
-			shared_trie_cache: config.trie_cache_maximum_size.map(|maximum_size| {
-				SharedTrieCache::new(
-					sp_trie::cache::CacheSize::new(maximum_size),
-					config.metrics_registry.as_ref(),
-				)
-			}),
+			shared_trie_cache,
 		};
 
 		// Older DB versions have no last state key. Check if the state is available and set it.

substrate/client/service/src/builder.rs

@@ -27,12 +27,13 @@ use crate::{
 };
 use futures::{select, FutureExt, StreamExt};
 use jsonrpsee::RpcModule;
-use log::info;
+use log::{debug, error, info};
 use prometheus_endpoint::Registry;
 use sc_chain_spec::{get_extension, ChainSpec};
 use sc_client_api::{
 	execution_extensions::ExecutionExtensions, proof_provider::ProofProvider, BadBlocks,
-	BlockBackend, BlockchainEvents, ExecutorProvider, ForkBlocks, StorageProvider, UsageProvider,
+	BlockBackend, BlockchainEvents, ExecutorProvider, ForkBlocks, KeysIter, StorageProvider,
+	TrieCacheContext, UsageProvider,
 };
 use sc_client_db::{Backend, BlocksPruning, DatabaseSettings, PruningMode};
 use sc_consensus::import_queue::{ImportQueue, ImportQueueService};
@@ -90,6 +91,7 @@ use sp_consensus::block_validation::{
 use sp_core::traits::{CodeExecutor, SpawnNamed};
 use sp_keystore::KeystorePtr;
 use sp_runtime::traits::{Block as BlockT, BlockIdTo, NumberFor, Zero};
+use sp_storage::{ChildInfo, ChildType, PrefixedStorageKey};
 use std::{
 	str::FromStr,
 	sync::Arc,
@@ -263,12 +265,89 @@ where
 			},
 		)?;
 
+		if let Some(warm_up_strategy) = config.warm_up_trie_cache {
+			let storage_root = client.usage_info().chain.best_hash;
+			let backend_clone = backend.clone();
+
+			if warm_up_strategy.is_blocking() {
+				// We use the blocking strategy for testing purposes.
+				// So better to error out if it fails.
+				warm_up_trie_cache(backend_clone, storage_root)?;
+			} else {
+				task_manager.spawn_handle().spawn_blocking(
+					"warm-up-trie-cache",
+					None,
+					async move {
+						if let Err(e) = warm_up_trie_cache(backend_clone, storage_root) {
+							error!("Failed to warm up trie cache: {e}");
+						}
+					},
+				);
+			}
+		}
+
 		client
 	};
 
 	Ok((client, backend, keystore_container, task_manager))
 }
 
+fn child_info(key: Vec<u8>) -> Option<ChildInfo> {
+	let prefixed_key = PrefixedStorageKey::new(key);
+	ChildType::from_prefixed_key(&prefixed_key).and_then(|(child_type, storage_key)| {
+		(child_type == ChildType::ParentKeyId).then(|| ChildInfo::new_default(storage_key))
+	})
+}
+
+fn warm_up_trie_cache<TBl: BlockT>(
+	backend: Arc<TFullBackend<TBl>>,
+	storage_root: TBl::Hash,
+) -> Result<(), Error> {
+	use sc_client_api::backend::Backend;
+	use sp_state_machine::Backend as StateBackend;
+
+	let untrusted_state = || backend.state_at(storage_root, TrieCacheContext::Untrusted);
+	let trusted_state = || backend.state_at(storage_root, TrieCacheContext::Trusted);
+
+	debug!("Populating trie cache started",);
+	let start_time = std::time::Instant::now();
+	let mut keys_count = 0;
+	let mut child_keys_count = 0;
+	for key in KeysIter::<_, TBl>::new(untrusted_state()?, None, None)? {
+		if keys_count != 0 && keys_count % 100_000 == 0 {
+			debug!("{} keys and {} child keys have been warmed", keys_count, child_keys_count);
+		}
+		match child_info(key.0.clone()) {
+			Some(info) => {
+				for child_key in
+					KeysIter::<_, TBl>::new_child(untrusted_state()?, info.clone(), None, None)?
+				{
+					if trusted_state()?
+						.child_storage(&info, &child_key.0)
+						.unwrap_or_default()
+						.is_none()
+					{
+						debug!("Child storage value unexpectedly empty: {child_key:?}");
+					}
+					child_keys_count += 1;
+				}
+			},
+			None => {
+				if trusted_state()?.storage(&key.0).unwrap_or_default().is_none() {
+					debug!("Storage value unexpectedly empty: {key:?}");
+				}
+				keys_count += 1;
+			},
+		}
+	}
+	debug!(
+		"Trie cache populated with {keys_count} keys and {child_keys_count} child keys in {} s",
+		start_time.elapsed().as_secs_f32()
+	);
+
+	Ok(())
+}
+
 /// Creates a [`NativeElseWasmExecutor`](sc_executor::NativeElseWasmExecutor) according to
 /// [`Configuration`].
 #[deprecated(note = "Please switch to `new_wasm_executor`. Will be removed at end of 2024.")]

substrate/client/service/src/config.rs

@@ -70,6 +70,8 @@ pub struct Configuration {
 	///
 	/// If `None` is given the cache is disabled.
 	pub trie_cache_maximum_size: Option<usize>,
+	/// Force the trie cache to be in memory.
+	pub warm_up_trie_cache: Option<TrieCacheWarmUpStrategy>,
 	/// State pruning settings.
 	pub state_pruning: Option<PruningMode>,
 	/// Number of blocks to keep in the db.
@@ -115,6 +117,22 @@ pub struct Configuration {
 	pub base_path: BasePath,
 }
 
+/// Warmup strategy for the trie cache.
+#[derive(Debug, Clone, Copy)]
+pub enum TrieCacheWarmUpStrategy {
+	/// Warm up the cache in a non-blocking way.
+	NonBlocking,
+	/// Warm up the cache in a blocking way.
+	Blocking,
+}
+
+impl TrieCacheWarmUpStrategy {
+	/// Returns true if the warmup strategy is blocking.
+	pub(crate) fn is_blocking(&self) -> bool {
+		matches!(self, Self::Blocking)
+	}
+}
+
 /// Type for tasks spawned by the executor.
 #[derive(PartialEq)]
 pub enum TaskType {

substrate/client/service/test/src/lib.rs

@@ -235,6 +235,7 @@ fn node_config<E: ChainSpecExtension + Clone + 'static + Send + Sync>(
 		keystore: KeystoreConfig::Path { path: root.join("key"), password: None },
 		database: DatabaseSource::RocksDb { path: root.join("db"), cache_size: 128 },
 		trie_cache_maximum_size: Some(16 * 1024 * 1024),
+		warm_up_trie_cache: None,
 		state_pruning: Default::default(),
 		blocks_pruning: BlocksPruning::KeepFinalized,
 		chain_spec: Box::new((*spec).clone()),