Storage chains: indexing, renewals and reference counting (#8265)

* Transaction indexing

* Tests and fixes

* Fixed a comment

* Style

* Build

* Style

* Apply suggestions from code review

Co-authored-by: cheme <emericchevalier.pro@gmail.com>

* Code review suggestions

* Add missing impl

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* impl JoinInput

* Don't store empty slices

* JoinInput operates on slices

Co-authored-by: cheme <emericchevalier.pro@gmail.com>
Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/primitives/database/src/kvdb.rs

@@ -33,18 +33,73 @@ fn handle_err<T>(result: std::io::Result<T>) -> T {
 }
 
 /// Wrap RocksDb database into a trait object that implements `sp_database::Database`
-pub fn as_database<D: KeyValueDB + 'static, H: Clone>(db: D) -> std::sync::Arc<dyn Database<H>> {
+pub fn as_database<D, H>(db: D) -> std::sync::Arc<dyn Database<H>>
+	where D: KeyValueDB + 'static, H: Clone + AsRef<[u8]>
+{
 	std::sync::Arc::new(DbAdapter(db))
 }
 
-impl<D: KeyValueDB, H: Clone> Database<H> for DbAdapter<D> {
+impl <D: KeyValueDB> DbAdapter<D> {
+	// Returns counter key and counter value if it exists.
+	fn read_counter(&self, col: ColumnId, key: &[u8]) -> error::Result<(Vec<u8>, Option<u32>)> {
+		// Add a key suffix for the counter
+		let mut counter_key = key.to_vec();
+		counter_key.push(0);
+		Ok(match self.0.get(col, &counter_key).map_err(|e| error::DatabaseError(Box::new(e)))? {
+			Some(data) => {
+				let mut counter_data = [0; 4];
+				if data.len() != 4 {
+					return Err(error::DatabaseError(Box::new(
+								std::io::Error::new(std::io::ErrorKind::Other,
+									format!("Unexpected counter len {}", data.len())))
+					))
+				}
+				counter_data.copy_from_slice(&data);
+				let counter = u32::from_le_bytes(counter_data);
+				(counter_key, Some(counter))
+			},
+			None => (counter_key, None)
+		})
+	}
+}
+
+impl<D: KeyValueDB, H: Clone + AsRef<[u8]>> Database<H> for DbAdapter<D> {
 	fn commit(&self, transaction: Transaction<H>) -> error::Result<()> {
 		let mut tx = DBTransaction::new();
 		for change in transaction.0.into_iter() {
 			match change {
 				Change::Set(col, key, value) => tx.put_vec(col, &key, value),
 				Change::Remove(col, key) => tx.delete(col, &key),
-				_ => unimplemented!(),
+				Change::Store(col, key, value) => {
+					match self.read_counter(col, key.as_ref())? {
+						(counter_key, Some(mut counter)) => {
+							counter += 1;
+							tx.put(col, &counter_key, &counter.to_le_bytes());
+						},
+						(counter_key, None) => {
+							let d = 1u32.to_le_bytes();
+							tx.put(col, &counter_key, &d);
+							tx.put_vec(col, key.as_ref(), value);
+						},
+					}
+				}
+				Change::Reference(col, key) => {
+					if let (counter_key, Some(mut counter)) = self.read_counter(col, key.as_ref())? {
+						counter += 1;
+						tx.put(col, &counter_key, &counter.to_le_bytes());
+					}
+				}
+				Change::Release(col, key) => {
+					if let (counter_key, Some(mut counter)) = self.read_counter(col, key.as_ref())? {
+						counter -= 1;
+						if counter == 0 {
+							tx.delete(col, &counter_key);
+							tx.delete(col, key.as_ref());
+						} else {
+							tx.put(col, &counter_key, &counter.to_le_bytes());
+						}
+					}
+				}
 			}
 		}
 		self.0.write(tx).map_err(|e| error::DatabaseError(Box::new(e)))
@@ -54,7 +109,7 @@ impl<D: KeyValueDB, H: Clone> Database<H> for DbAdapter<D> {
 		handle_err(self.0.get(col, key))
 	}
 
-	fn lookup(&self, _hash: &H) -> Option<Vec<u8>> {
-		unimplemented!();
+	fn contains(&self, col: ColumnId, key: &[u8]) -> bool {
+		handle_err(self.0.has_key(col, key))
 	}
 }

substrate/primitives/database/src/lib.rs

@@ -32,16 +32,9 @@ pub type ColumnId = u32;
 pub enum Change<H> {
 	Set(ColumnId, Vec<u8>, Vec<u8>),
 	Remove(ColumnId, Vec<u8>),
-	Store(H, Vec<u8>),
-	Release(H),
-}
-
-/// An alteration to the database that references the data.
-pub enum ChangeRef<'a, H> {
-	Set(ColumnId, &'a [u8], &'a [u8]),
-	Remove(ColumnId, &'a [u8]),
-	Store(H, &'a [u8]),
-	Release(H),
+	Store(ColumnId, H, Vec<u8>),
+	Reference(ColumnId, H),
+	Release(ColumnId, H),
 }
 
 /// A series of changes to the database that can be committed atomically. They do not take effect
@@ -67,49 +60,27 @@ impl<H> Transaction<H> {
 		self.0.push(Change::Remove(col, key.to_vec()))
 	}
 	/// Store the `preimage` of `hash` into the database, so that it may be looked up later with
-	/// `Database::lookup`. This may be called multiple times, but `Database::lookup` but subsequent
+	/// `Database::get`. This may be called multiple times, but subsequent
 	/// calls will ignore `preimage` and simply increase the number of references on `hash`.
-	pub fn store(&mut self, hash: H, preimage: &[u8]) {
-		self.0.push(Change::Store(hash, preimage.to_vec()))
+	pub fn store(&mut self, col: ColumnId, hash: H, preimage: Vec<u8>) {
+		self.0.push(Change::Store(col, hash, preimage))
+	}
+	/// Increase the number of references for `hash` in the database.
+	pub fn reference(&mut self, col: ColumnId, hash: H) {
+		self.0.push(Change::Reference(col, hash))
 	}
 	/// Release the preimage of `hash` from the database. An equal number of these to the number of
-	/// corresponding `store`s must have been given before it is legal for `Database::lookup` to
+	/// corresponding `store`s must have been given before it is legal for `Database::get` to
 	/// be unable to provide the preimage.
-	pub fn release(&mut self, hash: H) {
-		self.0.push(Change::Release(hash))
+	pub fn release(&mut self, col: ColumnId, hash: H) {
+		self.0.push(Change::Release(col, hash))
 	}
 }
 
-pub trait Database<H: Clone>: Send + Sync {
+pub trait Database<H: Clone + AsRef<[u8]>>: Send + Sync {
 	/// Commit the `transaction` to the database atomically. Any further calls to `get` or `lookup`
 	/// will reflect the new state.
-	fn commit(&self, transaction: Transaction<H>) -> error::Result<()> {
-		for change in transaction.0.into_iter() {
-			match change {
-				Change::Set(col, key, value) => self.set(col, &key, &value),
-				Change::Remove(col, key) => self.remove(col, &key),
-				Change::Store(hash, preimage) => self.store(&hash, &preimage),
-				Change::Release(hash) => self.release(&hash),
-			}?;
-		}
-
-		Ok(())
-	}
-
-	/// Commit the `transaction` to the database atomically. Any further calls to `get` or `lookup`
-	/// will reflect the new state.
-	fn commit_ref<'a>(&self, transaction: &mut dyn Iterator<Item=ChangeRef<'a, H>>) -> error::Result<()> {
-		let mut tx = Transaction::new();
-		for change in transaction {
-			match change {
-				ChangeRef::Set(col, key, value) => tx.set(col, key, value),
-				ChangeRef::Remove(col, key) => tx.remove(col, key),
-				ChangeRef::Store(hash, preimage) => tx.store(hash, preimage),
-				ChangeRef::Release(hash) => tx.release(hash),
-			}
-		}
-		self.commit(tx)
-	}
+	fn commit(&self, transaction: Transaction<H>) -> error::Result<()>;
 
 	/// Retrieve the value previously stored against `key` or `None` if
 	/// `key` is not currently in the database.
@@ -120,6 +91,11 @@ pub trait Database<H: Clone>: Send + Sync {
 		self.get(col, key).is_some()
 	}
 
+	/// Check value size in the database possibly without retrieving it.
+	fn value_size(&self, col: ColumnId, key: &[u8]) -> Option<usize> {
+		self.get(col, key).map(|v| v.len())
+	}
+
 	/// Call `f` with the value previously stored against `key`.
 	///
 	/// This may be faster than `get` since it doesn't allocate.
@@ -127,50 +103,6 @@ pub trait Database<H: Clone>: Send + Sync {
 	fn with_get(&self, col: ColumnId, key: &[u8], f: &mut dyn FnMut(&[u8])) {
 		self.get(col, key).map(|v| f(&v));
 	}
-
-	/// Set the value of `key` in `col` to `value`, replacing anything that is there currently.
-	fn set(&self, col: ColumnId, key: &[u8], value: &[u8]) -> error::Result<()> {
-		let mut t = Transaction::new();
-		t.set(col, key, value);
-		self.commit(t)
-	}
-	/// Remove the value of `key` in `col`.
-	fn remove(&self, col: ColumnId, key: &[u8]) -> error::Result<()> {
-		let mut t = Transaction::new();
-		t.remove(col, key);
-		self.commit(t)
-	}
-
-	/// Retrieve the first preimage previously `store`d for `hash` or `None` if no preimage is
-	/// currently stored.
-	fn lookup(&self, hash: &H) -> Option<Vec<u8>>;
-
-	/// Call `f` with the preimage stored for `hash` and return the result, or `None` if no preimage
-	/// is currently stored.
-	///
-	/// This may be faster than `lookup` since it doesn't allocate.
-	/// Use `with_lookup` helper function if you need `f` to return a value from `f`
-	fn with_lookup(&self, hash: &H, f: &mut dyn FnMut(&[u8])) {
-		self.lookup(hash).map(|v| f(&v));
-	}
-
-	/// Store the `preimage` of `hash` into the database, so that it may be looked up later with
-	/// `Database::lookup`. This may be called multiple times, but `Database::lookup` but subsequent
-	/// calls will ignore `preimage` and simply increase the number of references on `hash`.
-	fn store(&self, hash: &H, preimage: &[u8]) -> error::Result<()> {
-		let mut t = Transaction::new();
-		t.store(hash.clone(), preimage);
-		self.commit(t)
-	}
-
-	/// Release the preimage of `hash` from the database. An equal number of these to the number of
-	/// corresponding `store`s must have been given before it is legal for `Database::lookup` to
-	/// be unable to provide the preimage.
-	fn release(&self, hash: &H) -> error::Result<()> {
-		let mut t = Transaction::new();
-		t.release(hash.clone());
-		self.commit(t)
-	}
 }
 
 impl<H> std::fmt::Debug for dyn Database<H> {
@@ -183,20 +115,13 @@ impl<H> std::fmt::Debug for dyn Database<H> {
 /// `key` is not currently in the database.
 ///
 /// This may be faster than `get` since it doesn't allocate.
-pub fn with_get<R, H: Clone>(db: &dyn Database<H>, col: ColumnId, key: &[u8], mut f: impl FnMut(&[u8]) -> R) -> Option<R> {
+pub fn with_get<R, H: Clone + AsRef<[u8]>>(
+	db: &dyn Database<H>,
+	col: ColumnId,
+	key: &[u8], mut f: impl FnMut(&[u8]) -> R
+) -> Option<R> {
 	let mut result: Option<R> = None;
 	let mut adapter = |k: &_| { result = Some(f(k)); };
 	db.with_get(col, key, &mut adapter);
 	result
 }
-
-/// Call `f` with the preimage stored for `hash` and return the result, or `None` if no preimage
-/// is currently stored.
-///
-/// This may be faster than `lookup` since it doesn't allocate.
-pub fn with_lookup<R, H: Clone>(db: &dyn Database<H>, hash: &H, mut f: impl FnMut(&[u8]) -> R) -> Option<R> {
-	let mut result: Option<R> = None;
-	let mut adapter = |k: &_| { result = Some(f(k)); };
-	db.with_lookup(hash, &mut adapter);
-	result
-}

substrate/primitives/database/src/mem.rs

@@ -17,26 +17,41 @@
 
 //! In-memory implementation of `Database`
 
-use std::collections::HashMap;
+use std::collections::{HashMap, hash_map::Entry};
 use crate::{Database, Change, ColumnId, Transaction, error};
 use parking_lot::RwLock;
 
 #[derive(Default)]
 /// This implements `Database` as an in-memory hash map. `commit` is not atomic.
-pub struct MemDb<H: Clone + Send + Sync + Eq + PartialEq + Default + std::hash::Hash>
-	(RwLock<(HashMap<ColumnId, HashMap<Vec<u8>, Vec<u8>>>, HashMap<H, Vec<u8>>)>);
+pub struct MemDb(RwLock<HashMap<ColumnId, HashMap<Vec<u8>, (u32, Vec<u8>)>>>);
 
-impl<H> Database<H> for MemDb<H>
-	where H: Clone + Send + Sync + Eq + PartialEq + Default + std::hash::Hash
+impl<H> Database<H> for MemDb
+	where H: Clone + AsRef<[u8]>
 {
 	fn commit(&self, transaction: Transaction<H>) -> error::Result<()> {
 		let mut s = self.0.write();
 		for change in transaction.0.into_iter() {
 			match change {
-				Change::Set(col, key, value) => { s.0.entry(col).or_default().insert(key, value); },
-				Change::Remove(col, key) => { s.0.entry(col).or_default().remove(&key); },
-				Change::Store(hash, preimage) => { s.1.insert(hash, preimage); },
-				Change::Release(hash) => { s.1.remove(&hash); },
+				Change::Set(col, key, value) => { s.entry(col).or_default().insert(key, (1, value)); },
+				Change::Remove(col, key) => { s.entry(col).or_default().remove(&key); },
+				Change::Store(col, hash, value) => {
+					s.entry(col).or_default().entry(hash.as_ref().to_vec())
+						.and_modify(|(c, _)| *c += 1)
+						.or_insert_with(|| (1, value));
+				},
+				Change::Reference(col, hash) => {
+					if let Entry::Occupied(mut entry) = s.entry(col).or_default().entry(hash.as_ref().to_vec()) {
+						entry.get_mut().0 += 1;
+					}
+				}
+				Change::Release(col, hash) => {
+					if let Entry::Occupied(mut entry) = s.entry(col).or_default().entry(hash.as_ref().to_vec()) {
+						entry.get_mut().0 -= 1;
+						if entry.get().0 == 0 {
+							entry.remove();
+						}
+					}
+				}
 			}
 		}
 
@@ -45,18 +60,11 @@ impl<H> Database<H> for MemDb<H>
 
 	fn get(&self, col: ColumnId, key: &[u8]) -> Option<Vec<u8>> {
 		let s = self.0.read();
-		s.0.get(&col).and_then(|c| c.get(key).cloned())
-	}
-
-	fn lookup(&self, hash: &H) -> Option<Vec<u8>> {
-		let s = self.0.read();
-		s.1.get(hash).cloned()
+		s.get(&col).and_then(|c| c.get(key).map(|(_, v)| v.clone()))
 	}
 }
 
-impl<H> MemDb<H>
-	where H: Clone + Send + Sync + Eq + PartialEq + Default + std::hash::Hash
-{
+impl MemDb {
 	/// Create a new instance
 	pub fn new() -> Self {
 		MemDb::default()
@@ -65,7 +73,7 @@ impl<H> MemDb<H>
 	/// Count number of values in a column
 	pub fn count(&self, col: ColumnId) -> usize {
 		let s = self.0.read();
-		s.0.get(&col).map(|c| c.len()).unwrap_or(0)
+		s.get(&col).map(|c| c.len()).unwrap_or(0)
 	}
 }