Prune leaves on finality (#1871)

* support multiple leaves at same height in leaves set

* prune leaves on finalization

* test leaves behavior in client-db

* fix no-std compilation

substrate/core/client/db/src/light.rs

@@ -24,7 +24,8 @@ use kvdb::{KeyValueDB, DBTransaction};
 use client::backend::{AuxStore, NewBlockState};
 use client::blockchain::{BlockStatus, Cache as BlockchainCache,
 	HeaderBackend as BlockchainHeaderBackend, Info as BlockchainInfo};
-use client::{cht, LeafSet};
+use client::cht;
+use client::leaves::{LeafSet, FinalizationDisplaced};
 use client::error::{ErrorKind as ClientErrorKind, Result as ClientResult};
 use client::light::blockchain::Storage as LightBlockchainStorage;
 use parity_codec::{Decode, Encode};
@@ -250,6 +251,7 @@ impl<Block: BlockT> LightStorage<Block> {
 		transaction: &mut DBTransaction,
 		header: &Block::Header,
 		hash: Block::Hash,
+		displaced: &mut Option<FinalizationDisplaced<Block::Hash, NumberFor<Block>>>,
 	) -> ClientResult<()> {
 		let meta = self.meta.read();
 		if &meta.finalized_hash != header.parent_hash() {
@@ -311,6 +313,12 @@ impl<Block: BlockT> LightStorage<Block> {
 			}
 		}
 
+		let new_displaced = self.leaves.write().finalize_height(header.number().clone());
+		match displaced {
+			x @ &mut None => *x = Some(new_displaced),
+			&mut Some(ref mut displaced) => displaced.merge(new_displaced),
+		}
+
 		Ok(())
 	}
 
@@ -366,6 +374,7 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 		leaf_state: NewBlockState,
 		aux_ops: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 	) -> ClientResult<()> {
+		let mut finalization_displaced_leaves = None;
 		let mut transaction = DBTransaction::new();
 
 		let hash = header.hash();
@@ -405,7 +414,12 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 		};
 
 		if finalized {
-			self.note_finalized(&mut transaction, &header, hash)?;
+			self.note_finalized(
+				&mut transaction,
+				&header,
+				hash,
+				&mut finalization_displaced_leaves,
+			)?;
 		}
 
 		{
@@ -425,10 +439,18 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 			debug!("Light DB Commit {:?} ({})", hash, number);
 			let write_result = self.db.write(transaction).map_err(db_err);
 			if let Err(e) = write_result {
+				let mut leaves = self.leaves.write();
+				let mut undo = leaves.undo();
+
 				// revert leaves set update if there was one.
 				if let Some(displaced_leaf) = displaced_leaf {
-					leaves.undo(displaced_leaf);
+					undo.undo_import(displaced_leaf);
 				}
+
+				if let Some(finalization_displaced) = finalization_displaced_leaves {
+					undo.undo_finalization(finalization_displaced);
+				}
+
 				return Err(e);
 			}
 
@@ -464,10 +486,11 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 
 	fn finalize_header(&self, id: BlockId<Block>) -> ClientResult<()> {
 		if let Some(header) = self.header(id)? {
+			let mut displaced = None;
 			let mut transaction = DBTransaction::new();
 			let hash = header.hash();
 			let number = *header.number();
-			self.note_finalized(&mut transaction, &header, hash.clone())?;
+			self.note_finalized(&mut transaction, &header, hash.clone(), &mut displaced)?;
 			{
 				let mut cache = self.cache.0.write();
 				let cache_ops = cache.transaction(&mut transaction)
@@ -477,7 +500,12 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 					)?
 					.into_ops();
 
-				self.db.write(transaction).map_err(db_err)?;
+				if let Err(e) = self.db.write(transaction).map_err(db_err) {
+					if let Some(displaced) = displaced {
+						self.leaves.write().undo().undo_finalization(displaced);
+					}
+					return Err(e);
+				}
 				cache.commit(cache_ops);
 			}
 			self.update_meta(hash, header.number().clone(), false, true);
@@ -940,4 +968,28 @@ pub(crate) mod tests {
 		assert_eq!(db.get_aux(&[2]).unwrap(), Some(vec![102]));
 		assert_eq!(db.get_aux(&[3]).unwrap(), Some(vec![103]));
 	}
+
+	#[test]
+	fn test_leaves_pruned_on_finality() {
+		let db = LightStorage::<Block>::new_test();
+		let block0 = insert_final_block(&db, None, || default_header(&Default::default(), 0));
+
+		let block1_a = insert_block(&db, None, || default_header(&block0, 1));
+		let block1_b = insert_block(&db, None, || header_with_extrinsics_root(&block0, 1, [1; 32].into()));
+		let block1_c = insert_block(&db, None, || header_with_extrinsics_root(&block0, 1, [2; 32].into()));
+
+		assert_eq!(db.leaves.read().hashes(), vec![block1_a, block1_b, block1_c]);
+
+		let block2_a = insert_block(&db, None, || default_header(&block1_a, 2));
+		let block2_b = insert_block(&db, None, || header_with_extrinsics_root(&block1_b, 2, [1; 32].into()));
+		let block2_c = insert_block(&db, None, || header_with_extrinsics_root(&block1_b, 2, [2; 32].into()));
+
+		assert_eq!(db.leaves.read().hashes(), vec![block2_a, block2_b, block2_c, block1_c]);
+
+		db.finalize_header(BlockId::hash(block1_a)).unwrap();
+		db.finalize_header(BlockId::hash(block2_a)).unwrap();
+
+		// leaves at same height stay. Leaves at lower heights pruned.
+		assert_eq!(db.leaves.read().hashes(), vec![block2_a, block2_b, block2_c]);
+	}
 }