Support reference-counting state backend. (#5769)

* Optimize pinning

* Ref counting state backend

* Style

Co-Authored-By: Wei Tang <hi@that.world>

* Update Cargo.lock

* Handle empty node

Co-authored-by: Wei Tang <hi@that.world>

substrate/client/db/src/lib.rs

@@ -109,8 +109,9 @@ pub type DbState<B> = sp_state_machine::TrieBackend<
 	Arc<dyn sp_state_machine::Storage<HashFor<B>>>, HashFor<B>
 >;
 
+const DB_HASH_LEN: usize = 32;
 /// Hash type that this backend uses for the database.
-pub type DbHash = [u8; 32];
+pub type DbHash = [u8; DB_HASH_LEN];
 
 /// A reference tracking state.
 ///
@@ -314,6 +315,13 @@ impl DatabaseSettingsSrc {
 			DatabaseSettingsSrc::Custom(_) => None,
 		}
 	}
+	/// Check if database supports internal ref counting for state data.
+	pub fn supports_ref_counting(&self) -> bool {
+		match self {
+			DatabaseSettingsSrc::ParityDb { .. } => true,
+			_ => false,
+		}
+	}
 }
 
 /// Create an instance of db-backed client.
@@ -716,13 +724,18 @@ impl<Block: BlockT> sc_client_api::backend::BlockImportOperation<Block> for Bloc
 struct StorageDb<Block: BlockT> {
 	pub db: Arc<dyn Database<DbHash>>,
 	pub state_db: StateDb<Block::Hash, Vec<u8>>,
+	prefix_keys: bool,
 }
 
 impl<Block: BlockT> sp_state_machine::Storage<HashFor<Block>> for StorageDb<Block> {
 	fn get(&self, key: &Block::Hash, prefix: Prefix) -> Result<Option<DBValue>, String> {
-		let key = prefixed_key::<HashFor<Block>>(key, prefix);
-		self.state_db.get(&key, self)
-			.map_err(|e| format!("Database backend error: {:?}", e))
+		if self.prefix_keys {
+			let key = prefixed_key::<HashFor<Block>>(key, prefix);
+			self.state_db.get(&key, self)
+		} else {
+			self.state_db.get(key.as_ref(), self)
+		}
+		.map_err(|e| format!("Database backend error: {:?}", e))
 	}
 }
 
@@ -843,11 +856,15 @@ impl<Block: BlockT> Backend<Block> {
 		let map_e = |e: sc_state_db::Error<io::Error>| sp_blockchain::Error::from(
 			format!("State database error: {:?}", e)
 		);
-		let state_db: StateDb<_, _> = StateDb::new(config.pruning.clone(), &StateMetaDb(&*db))
-			.map_err(map_e)?;
+		let state_db: StateDb<_, _> = StateDb::new(
+			config.pruning.clone(),
+			!config.source.supports_ref_counting(),
+			&StateMetaDb(&*db),
+		).map_err(map_e)?;
 		let storage_db = StorageDb {
 			db: db.clone(),
 			state_db,
+			prefix_keys: !config.source.supports_ref_counting(),
 		};
 		let offchain_storage = offchain::LocalStorage::new(db.clone());
 		let changes_tries_storage = DbChangesTrieStorage::new(
@@ -1112,17 +1129,32 @@ impl<Block: BlockT> Backend<Block> {
 				let mut bytes: u64 = 0;
 				let mut removal: u64 = 0;
 				let mut bytes_removal: u64 = 0;
-				for (key, (val, rc)) in operation.db_updates.drain() {
+				for (mut key, (val, rc)) in operation.db_updates.drain() {
+					if !self.storage.prefix_keys {
+						// Strip prefix
+						key.drain(0 .. key.len() - DB_HASH_LEN);
+					};
 					if rc > 0 {
 						ops += 1;
 						bytes += key.len() as u64 + val.len() as u64;
-
-						changeset.inserted.push((key, val.to_vec()));
+						if rc == 1 {
+							changeset.inserted.push((key, val.to_vec()));
+						} else {
+							changeset.inserted.push((key.clone(), val.to_vec()));
+							for _ in 0 .. rc - 1 {
+								changeset.inserted.push((key.clone(), Default::default()));
+							}
+						}
 					} else if rc < 0 {
 						removal += 1;
 						bytes_removal += key.len() as u64;
-
-						changeset.deleted.push(key);
+						if rc == -1 {
+							changeset.deleted.push(key);
+						} else {
+							for _ in 0 .. -rc {
+								changeset.deleted.push(key.clone());
+							}
+						}
 					}
 				}
 				self.state_usage.tally_writes_nodes(ops, bytes);

substrate/client/db/src/parity_db.rs

@@ -17,6 +17,8 @@
 /// A `Database` adapter for parity-db.
 
 use sp_database::{Database, Change, Transaction, ColumnId};
+use crate::utils::NUM_COLUMNS;
+use crate::columns;
 
 struct DbAdapter(parity_db::Db);
 
@@ -30,8 +32,13 @@ fn handle_err<T>(result: parity_db::Result<T>) -> T {
 }
 
 /// Wrap RocksDb database into a trait object that implements `sp_database::Database`
-pub fn open<H: Clone>(path: &std::path::Path, num_columns: u32) -> parity_db::Result<std::sync::Arc<dyn Database<H>>> {
-	let db = parity_db::Db::with_columns(path, num_columns as u8)?;
+pub fn open<H: Clone>(path: &std::path::Path) -> parity_db::Result<std::sync::Arc<dyn Database<H>>> {
+	let mut config = parity_db::Options::with_columns(path, NUM_COLUMNS as u8);
+	let mut state_col = &mut config.columns[columns::STATE as usize];
+	state_col.ref_counted = true;
+	state_col.preimage = true;
+	state_col.uniform = true;
+	let db = parity_db::Db::open(&config)?;
 	Ok(std::sync::Arc::new(DbAdapter(db)))
 }
 

substrate/client/db/src/utils.rs

@@ -254,7 +254,7 @@ pub fn open_database<Block: BlockT>(
 		},
 		#[cfg(feature = "parity-db")]
 		DatabaseSettingsSrc::ParityDb { path } => {
-			crate::parity_db::open(&path, NUM_COLUMNS)
+			crate::parity_db::open(&path)
 				.map_err(|e| sp_blockchain::Error::Backend(format!("{:?}", e)))?
 		},
 		DatabaseSettingsSrc::Custom(db) => db.clone(),