Add standalone set head operations (#1600)

* Refactor head setting operation to a separate function

* Fix commit issue and implement set_head standalone in db

* Add standalone set head operations

* Address grumbles

* Change number_and_hash_to_lookup key in light mod to take reference

* Fix bug in set head commit

* Add a convenience fn

* Fix a deadlock

* Fix missing argument

substrate/core/client/src/backend.rs

@@ -44,6 +44,14 @@ impl NewBlockState {
 			NewBlockState::Normal => false,
 		}
 	}
+
+	/// Whether this block is considered final.
+	pub fn is_final(self) -> bool {
+		match self {
+			NewBlockState::Final => true,
+			NewBlockState::Best | NewBlockState::Normal => false,
+		}
+	}
 }
 
 /// Block insertion operation. Keeps hold if the inserted block state and data.
@@ -81,6 +89,8 @@ pub trait BlockImportOperation<Block, H> where
 		where I: IntoIterator<Item=(Vec<u8>, Option<Vec<u8>>)>;
 	/// Mark a block as finalized.
 	fn mark_finalized(&mut self, id: BlockId<Block>, justification: Option<Justification>) -> error::Result<()>;
+	/// Mark a block as new head. If both block import and set head are specified, set head overrides block import's best block rule.
+	fn mark_head(&mut self, id: BlockId<Block>) -> error::Result<()>;
 }
 
 /// Provides access to an auxiliary database.

substrate/core/client/src/client.rs

@@ -650,6 +650,25 @@ impl<B, E, Block, RA> Client<B, E, Block, RA> where
 		result
 	}
 
+	/// Set a block as best block.
+	pub fn set_head(
+		&self,
+		id: BlockId<Block>
+	) -> error::Result<()> {
+		self.lock_import_and_run(|operation| {
+			self.apply_head(operation, id)
+		})
+	}
+
+	/// Set a block as best block, and apply it to an operation.
+	pub fn apply_head(
+		&self,
+		operation: &mut ClientImportOperation<Block, Blake2Hasher, B>,
+		id: BlockId<Block>,
+	) -> error::Result<()> {
+		operation.op.mark_head(id)
+	}
+
 	/// Apply a checked and validated block to an operation. If a justification is provided
 	/// then `finalized` *must* be true.
 	pub fn apply_block(

substrate/core/client/src/in_mem.rs

@@ -168,55 +168,24 @@ impl<Block: BlockT> Blockchain<Block> {
 		new_state: NewBlockState,
 	) -> crate::error::Result<()> {
 		let number = header.number().clone();
-		let best_tree_route = match new_state.is_best() {
-			false => None,
-			true => {
-				let best_hash = self.storage.read().best_hash;
-				if &best_hash == header.parent_hash() {
-					None
-				} else {
-					let route = crate::blockchain::tree_route(
-						self,
-						BlockId::Hash(best_hash),
-						BlockId::Hash(*header.parent_hash()),
-					)?;
-					Some(route)
-				}
-			}
-		};
-
-		let mut storage = self.storage.write();
-
-		storage.leaves.import(hash.clone(), number.clone(), header.parent_hash().clone());
 
 		if new_state.is_best() {
-			if let Some(tree_route) = best_tree_route {
-				// apply retraction and enaction when reorganizing up to parent hash
-				let enacted = tree_route.enacted();
+			self.apply_head(&header)?;
+		}
 
-				for entry in enacted {
-					storage.hashes.insert(entry.number, entry.hash);
-				}
+		{
+			let mut storage = self.storage.write();
+			storage.leaves.import(hash.clone(), number.clone(), header.parent_hash().clone());
+			storage.blocks.insert(hash.clone(), StoredBlock::new(header, body, justification));
 
-				for entry in tree_route.retracted().iter().skip(enacted.len()) {
-					storage.hashes.remove(&entry.number);
-				}
+			if let NewBlockState::Final = new_state {
+				storage.finalized_hash = hash;
+				storage.finalized_number = number.clone();
 			}
 
-			storage.best_hash = hash.clone();
-			storage.best_number = number.clone();
-			storage.hashes.insert(number, hash.clone());
-		}
-
-		storage.blocks.insert(hash.clone(), StoredBlock::new(header, body, justification));
-
-		if let NewBlockState::Final = new_state {
-			storage.finalized_hash = hash;
-			storage.finalized_number = number.clone();
-		}
-
-		if number == Zero::zero() {
-			storage.genesis_hash = hash;
+			if number == Zero::zero() {
+				storage.genesis_hash = hash;
+			}
 		}
 
 		Ok(())
@@ -247,6 +216,58 @@ impl<Block: BlockT> Blockchain<Block> {
 		self.storage.write().header_cht_roots.insert(block, cht_root);
 	}
 
+	/// Set an existing block as head.
+	pub fn set_head(&self, id: BlockId<Block>) -> error::Result<()> {
+		let header = match self.header(id)? {
+			Some(h) => h,
+			None => return Err(error::ErrorKind::UnknownBlock(format!("{}", id)).into()),
+		};
+
+		self.apply_head(&header)
+	}
+
+	fn apply_head(&self, header: &<Block as BlockT>::Header) -> error::Result<()> {
+		let hash = header.hash();
+		let number = header.number();
+
+		// Note: this may lock storage, so it must happen before obtaining storage
+		// write lock.
+		let best_tree_route = {
+			let best_hash = self.storage.read().best_hash;
+			if &best_hash == header.parent_hash() {
+				None
+			} else {
+				let route = crate::blockchain::tree_route(
+					self,
+					BlockId::Hash(best_hash),
+					BlockId::Hash(*header.parent_hash()),
+				)?;
+				Some(route)
+			}
+		};
+
+		let mut storage = self.storage.write();
+
+		if let Some(tree_route) = best_tree_route {
+			// apply retraction and enaction when reorganizing up to parent hash
+			let enacted = tree_route.enacted();
+
+			for entry in enacted {
+				storage.hashes.insert(entry.number, entry.hash);
+			}
+
+			for entry in tree_route.retracted().iter().skip(enacted.len()) {
+				storage.hashes.remove(&entry.number);
+			}
+		}
+
+		storage.best_hash = hash.clone();
+		storage.best_number = number.clone();
+		storage.hashes.insert(number.clone(), hash.clone());
+
+		Ok(())
+	}
+
 	fn finalize_header(&self, id: BlockId<Block>, justification: Option<Justification>) -> error::Result<()> {
 		let hash = match self.header(id)? {
 			Some(h) => h.hash(),
@@ -388,6 +409,10 @@ impl<Block: BlockT> light::blockchain::Storage<Block> for Blockchain<Block>
 		Ok(())
 	}
 
+	fn set_head(&self, id: BlockId<Block>) -> error::Result<()> {
+		Blockchain::set_head(self, id)
+	}
+
 	fn last_finalized(&self) -> error::Result<Block::Hash> {
 		Ok(self.storage.read().finalized_hash.clone())
 	}
@@ -420,6 +445,7 @@ pub struct BlockImportOperation<Block: BlockT, H: Hasher> {
 	changes_trie_update: Option<MemoryDB<H>>,
 	aux: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 	finalized_blocks: Vec<(BlockId<Block>, Option<Justification>)>,
+	set_head: Option<BlockId<Block>>,
 }
 
 impl<Block, H> backend::BlockImportOperation<Block, H> for BlockImportOperation<Block, H>
@@ -500,6 +526,12 @@ where
 		self.finalized_blocks.push((block, justification));
 		Ok(())
 	}
+
+	fn mark_head(&mut self, block: BlockId<Block>) -> error::Result<()> {
+		assert!(self.pending_block.is_none(), "Only one set block per operation is allowed");
+		self.set_head = Some(block);
+		Ok(())
+	}
 }
 
 /// In-memory backend. Keeps all states and blocks in memory. Useful for testing.
@@ -572,6 +604,7 @@ where
 			changes_trie_update: None,
 			aux: Default::default(),
 			finalized_blocks: Default::default(),
+			set_head: None,
 		})
 	}
 
@@ -615,6 +648,10 @@ where
 			self.blockchain.write_aux(operation.aux);
 		}
 
+		if let Some(set_head) = operation.set_head {
+			self.blockchain.set_head(set_head)?;
+		}
+
 		Ok(())
 	}
 

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

@@ -50,6 +50,7 @@ pub struct ImportOperation<Block: BlockT, S, F, H> {
 	leaf_state: NewBlockState,
 	aux_ops: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 	finalized_blocks: Vec<BlockId<Block>>,
+	set_head: Option<BlockId<Block>>,
 	storage_update: Option<InMemoryState<H>>,
 	_phantom: ::std::marker::PhantomData<(S, F)>,
 }
@@ -120,6 +121,7 @@ impl<S, F, Block, H> ClientBackend<Block, H> for Backend<S, F, H> where
 			leaf_state: NewBlockState::Normal,
 			aux_ops: Vec::new(),
 			finalized_blocks: Vec::new(),
+			set_head: None,
 			storage_update: None,
 			_phantom: Default::default(),
 		})
@@ -165,6 +167,10 @@ impl<S, F, Block, H> ClientBackend<Block, H> for Backend<S, F, H> where
 			}
 		}
 
+		if let Some(set_head) = operation.set_head {
+			self.blockchain.storage().set_head(set_head)?;
+		}
+
 		Ok(())
 	}
 
@@ -294,6 +300,11 @@ where
 		self.finalized_blocks.push(block);
 		Ok(())
 	}
+
+	fn mark_head(&mut self, block: BlockId<Block>) -> ClientResult<()> {
+		self.set_head = Some(block);
+		Ok(())
+	}
 }
 
 impl<Block, S, F, H> StateBackend<H> for OnDemandState<Block, S, F>

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

@@ -45,6 +45,9 @@ pub trait Storage<Block: BlockT>: AuxStore + BlockchainHeaderBackend<Block> {
 		aux_ops: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 	) -> ClientResult<()>;
 
+	/// Set an existing block as new best block.
+	fn set_head(&self, block: BlockId<Block>) -> ClientResult<()>;
+
 	/// Mark historic header as finalized.
 	fn finalize_header(&self, block: BlockId<Block>) -> ClientResult<()>;
 
@@ -246,6 +249,10 @@ pub mod tests {
 			Ok(())
 		}
 
+		fn set_head(&self, _block: BlockId<Block>) -> ClientResult<()> {
+			Err(ClientErrorKind::Backend("Test error".into()).into())
+		}
+
 		fn finalize_header(&self, _block: BlockId<Block>) -> ClientResult<()> {
 			Err(ClientErrorKind::Backend("Test error".into()).into())
 		}