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
2026-06-12 22:51:13 +00:00 · 2019-02-12 17:21:25 +01:00
parent 6122f7d6b6
commit 22048dba60
7 changed files with 321 additions and 175 deletions
@@ -257,6 +257,7 @@ pub struct BlockImportOperation<Block: BlockT, H: Hasher> {
 	pending_block: Option<PendingBlock<Block>>,
 	aux_ops: Vec<(Vec<u8>, Option<Vec<u8>>)>,
 	finalized_blocks: Vec<(BlockId<Block>, Option<Justification>)>,
+	set_head: Option<BlockId<Block>>,
 }

 impl<Block: BlockT, H: Hasher> BlockImportOperation<Block, H> {
@@ -355,6 +356,12 @@ where Block: BlockT<Hash=H256>,
 		self.finalized_blocks.push((block, justification));
 		Ok(())
 	}
+
+	fn mark_head(&mut self, block: BlockId<Block>) -> Result<(), client::error::Error> {
+		assert!(self.set_head.is_none(), "Only one set head per operation is allowed");
+		self.set_head = Some(block);
+		Ok(())
+	}
 }

 struct StorageDb<Block: BlockT> {
@@ -563,6 +570,71 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 		})
 	}

+	/// Handle setting head within a transaction. `route_to` should be the last
+	/// block that existed in the database. `best_to` should be the best block
+	/// to be set.
+	///
+	/// In the case where the new best block is a block to be imported, `route_to`
+	/// should be the parent of `best_to`. In the case where we set an existing block
+	/// to be best, `route_to` should equal to `best_to`.
+	fn set_head_with_transaction(&self, transaction: &mut DBTransaction, route_to: Block::Hash, best_to: (NumberFor<Block>, Block::Hash)) -> Result<(Vec<Block::Hash>, Vec<Block::Hash>), client::error::Error> {
+		let mut enacted = Vec::default();
+		let mut retracted = Vec::default();
+
+		let meta = self.blockchain.meta.read();
+
+		// cannot find tree route with empty DB.
+		if meta.best_hash != Default::default() {
+			let tree_route = ::client::blockchain::tree_route(
+				&self.blockchain,
+				BlockId::Hash(meta.best_hash),
+				BlockId::Hash(route_to),
+			)?;
+
+			// uncanonicalize: check safety violations and ensure the numbers no longer
+			// point to these block hashes in the key mapping.
+			for r in tree_route.retracted() {
+				if r.hash == meta.finalized_hash {
+					warn!(
+						"Potential safety failure: reverting finalized block {:?}",
+						(&r.number, &r.hash)
+					);
+
+					return Err(::client::error::ErrorKind::NotInFinalizedChain.into());
+				}
+
+				retracted.push(r.hash.clone());
+				utils::remove_number_to_key_mapping(
+					transaction,
+					columns::KEY_LOOKUP,
+					r.number
+				);
+			}
+
+			// canonicalize: set the number lookup to map to this block's hash.
+			for e in tree_route.enacted() {
+				enacted.push(e.hash.clone());
+				utils::insert_number_to_key_mapping(
+					transaction,
+					columns::KEY_LOOKUP,
+					e.number,
+					e.hash
+				);
+			}
+		}
+
+		let lookup_key = utils::number_and_hash_to_lookup_key(best_to.0, &best_to.1);
+		transaction.put(columns::META, meta_keys::BEST_BLOCK, &lookup_key);
+		utils::insert_number_to_key_mapping(
+			transaction,
+			columns::KEY_LOOKUP,
+			best_to.0,
+			best_to.1,
+		);
+
+		Ok((enacted, retracted))
+	}
+
 	fn ensure_sequential_finalization(
 		&self,
 		header: &Block::Header,
@@ -592,7 +664,7 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 		if let Some(justification) = justification {
 			transaction.put(
 				columns::JUSTIFICATION,
-				&utils::number_and_hash_to_lookup_key(number, *hash),
+				&utils::number_and_hash_to_lookup_key(number, hash),
 				&justification.encode(),
 			);
 		}
@@ -657,7 +729,7 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 			}
 		}

-		if let Some(pending_block) = operation.pending_block {
+		let imported = if let Some(pending_block) = operation.pending_block {
 			let hash = pending_block.header.hash();
 			let parent_hash = *pending_block.header.parent_hash();
 			let number = pending_block.header.number().clone();
@@ -665,58 +737,11 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 			// blocks are keyed by number + hash.
 			let lookup_key = utils::number_and_hash_to_lookup_key(number, hash);

-			let mut enacted = Vec::default();
-			let mut retracted = Vec::default();
-
-			if pending_block.leaf_state.is_best() {
-				let meta = self.blockchain.meta.read();
-
-				// cannot find tree route with empty DB.
-				if meta.best_hash != Default::default() {
-					let tree_route = ::client::blockchain::tree_route(
-						&self.blockchain,
-						BlockId::Hash(meta.best_hash),
-						BlockId::Hash(parent_hash),
-					)?;
-
-					// uncanonicalize: check safety violations and ensure the numbers no longer
-					// point to these block hashes in the key mapping.
-					for r in tree_route.retracted() {
-						retracted.push(r.hash.clone());
-						if r.hash == meta.finalized_hash {
-							warn!("Potential safety failure: reverting finalized block {:?}",
-								(&r.number, &r.hash));
-
-							return Err(::client::error::ErrorKind::NotInFinalizedChain.into());
-						}
-
-						utils::remove_number_to_key_mapping(
-							&mut transaction,
-							columns::KEY_LOOKUP,
-							r.number
-						);
-					}
-
-					// canonicalize: set the number lookup to map to this block's hash.
-					for e in tree_route.enacted() {
-						enacted.push(e.hash.clone());
-						utils::insert_number_to_key_mapping(
-							&mut transaction,
-							columns::KEY_LOOKUP,
-							e.number,
-							e.hash
-						);
-					}
-				}
-
-				transaction.put(columns::META, meta_keys::BEST_BLOCK, &lookup_key);
-				utils::insert_number_to_key_mapping(
-					&mut transaction,
-					columns::KEY_LOOKUP,
-					number,
-					hash,
-				);
-			}
+			let (enacted, retracted) = if pending_block.leaf_state.is_best() {
+				self.set_head_with_transaction(&mut transaction, parent_hash, (number, hash))?
+			} else {
+				(Default::default(), Default::default())
+			};

 			utils::insert_hash_to_key_mapping(
 				&mut transaction,
@@ -752,15 +777,12 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {
 			apply_state_commit(&mut transaction, commit);

 			// Check if need to finalize. Genesis is always finalized instantly.
-			let finalized = number_u64 == 0 || match pending_block.leaf_state {
-				NewBlockState::Final => true,
-				_ => false,
-			};
+			let finalized = number_u64 == 0 || pending_block.leaf_state.is_final();

 			let header = &pending_block.header;
 			let is_best = pending_block.leaf_state.is_best();
 			let changes_trie_updates = operation.changes_trie_updates;
-		
+
 			self.changes_tries_storage.commit(&mut transaction, changes_trie_updates);

 			if finalized {
@@ -774,49 +796,60 @@ impl<Block: BlockT<Hash=H256>> Backend<Block> {

 			debug!(target: "db", "DB Commit {:?} ({}), best = {}", hash, number, is_best);

-			{
+			let displaced_leaf = {
 				let mut leaves = self.blockchain.leaves.write();
 				let displaced_leaf = leaves.import(hash, number, parent_hash);
 				leaves.prepare_transaction(&mut transaction, columns::META, meta_keys::LEAF_PREFIX);

-				let write_result = self.storage.db.write(transaction).map_err(db_err);
-				if let Err(e) = write_result {
-					// revert leaves set update, if there was one.
-					if let Some(displaced_leaf) = displaced_leaf {
-						leaves.undo(displaced_leaf);
-					}
-					return Err(e);
+				displaced_leaf
+			};
+
+			meta_updates.push((hash, number, pending_block.leaf_state.is_best(), finalized));
+
+			Some((number, hash, enacted, retracted, displaced_leaf, is_best))
+		} else {
+			None
+		};
+
+		if let Some(set_head) = operation.set_head {
+			if let Some(header) = ::client::blockchain::HeaderBackend::header(&self.blockchain, set_head)? {
+				let number = header.number();
+				let hash = header.hash();
+
+				self.set_head_with_transaction(
+					&mut transaction,
+					hash.clone(),
+					(number.clone(), hash.clone())
+				)?;
+			} else {
+				return Err(client::error::ErrorKind::UnknownBlock(format!("Cannot set head {:?}", set_head)).into())
+			}
+		}
+
+		let write_result = self.storage.db.write(transaction).map_err(db_err);
+
+		if let Some((number, hash, enacted, retracted, displaced_leaf, is_best)) = imported {
+			if let Err(e) = write_result {
+				if let Some(displaced_leaf) = displaced_leaf {
+					self.blockchain.leaves.write().undo(displaced_leaf);
 				}
-				drop(leaves);
+				return Err(e)
 			}

-			for (hash, number, is_best, is_finalized) in meta_updates {
-				self.blockchain.update_meta(hash, number, is_best, is_finalized);
-			}
-
-			self.blockchain.update_meta(
-				hash.clone(),
-				number.clone(),
-				is_best,
-				finalized,
-			);
-
-			// sync canonical state cache
 			operation.old_state.sync_cache(
 				&enacted,
 				&retracted,
 				operation.storage_updates,
 				Some(hash),
 				Some(number),
-				|| is_best
+				|| is_best,
 			);
-		} else {
-			// No pending block, just write the transaction and apply meta changes
-			self.storage.db.write(transaction).map_err(db_err)?;
-			for (hash, number, is_best, is_finalized) in meta_updates {
-				self.blockchain.update_meta(hash, number, is_best, is_finalized);
-			}
 		}
+
+		for (hash, number, is_best, is_finalized) in meta_updates {
+			self.blockchain.update_meta(hash, number, is_best, is_finalized);
+		}
+
 		Ok(())
 	}

@@ -911,6 +944,7 @@ impl<Block> client::backend::Backend<Block, Blake2Hasher> for Backend<Block> whe
 			changes_trie_updates: MemoryDB::default(),
 			aux_ops: Vec::new(),
 			finalized_blocks: Vec::new(),
+			set_head: None,
 		})
 	}

@@ -990,7 +1024,7 @@ impl<Block> client::backend::Backend<Block, Blake2Hasher> for Backend<Block> whe
 					let hash = self.blockchain.hash(best)?.ok_or_else(
 						|| client::error::ErrorKind::UnknownBlock(
 							format!("Error reverting to {}. Block hash not found.", best)))?;
-					let key = utils::number_and_hash_to_lookup_key(best.clone(), hash.clone());
+					let key = utils::number_and_hash_to_lookup_key(best.clone(), &hash);
 					transaction.put(columns::META, meta_keys::BEST_BLOCK, &key);
 					transaction.delete(columns::KEY_LOOKUP, removed.hash().as_ref());
 					self.storage.db.write(transaction).map_err(db_err)?;
@@ -1032,7 +1066,7 @@ impl<Block> client::backend::Backend<Block, Blake2Hasher> for Backend<Block> whe
 					Err(client::error::ErrorKind::UnknownBlock(format!("State already discarded for {:?}", block)).into())
 				}
 			},
-			Ok(None) => Err(client::error::ErrorKind::UnknownBlock(format!("{:?}", block)).into()),
+			Ok(None) => Err(client::error::ErrorKind::UnknownBlock(format!("Unknown state for block {:?}", block)).into()),
 			Err(e) => Err(e),
 		}
 	}
@@ -189,6 +189,61 @@ impl<Block: BlockT> LightStorage<Block> {
 					.cloned()))
 	}

+	/// Handle setting head within a transaction. `route_to` should be the last
+	/// block that existed in the database. `best_to` should be the best block
+	/// to be set.
+	///
+	/// In the case where the new best block is a block to be imported, `route_to`
+	/// should be the parent of `best_to`. In the case where we set an existing block
+	/// to be best, `route_to` should equal to `best_to`.
+	fn set_head_with_transaction(&self, transaction: &mut DBTransaction, route_to: Block::Hash, best_to: (NumberFor<Block>, Block::Hash)) -> Result<(), client::error::Error> {
+		let lookup_key = utils::number_and_hash_to_lookup_key(best_to.0, &best_to.1);
+
+		// handle reorg.
+		let meta = self.meta.read();
+		if meta.best_hash != Default::default() {
+			let tree_route = ::client::blockchain::tree_route(
+				self,
+				BlockId::Hash(meta.best_hash),
+				BlockId::Hash(route_to),
+			)?;
+
+			// update block number to hash lookup entries.
+			for retracted in tree_route.retracted() {
+				if retracted.hash == meta.finalized_hash {
+					// TODO: can we recover here?
+					warn!("Safety failure: reverting finalized block {:?}",
+						  (&retracted.number, &retracted.hash));
+				}
+
+				utils::remove_number_to_key_mapping(
+					transaction,
+					columns::KEY_LOOKUP,
+					retracted.number
+				);
+			}
+
+			for enacted in tree_route.enacted() {
+				utils::insert_number_to_key_mapping(
+					transaction,
+					columns::KEY_LOOKUP,
+					enacted.number,
+					enacted.hash
+				);
+			}
+		}
+
+		transaction.put(columns::META, meta_keys::BEST_BLOCK, &lookup_key);
+		utils::insert_number_to_key_mapping(
+			transaction,
+			columns::KEY_LOOKUP,
+			best_to.0,
+			best_to.1,
+		);
+
+		Ok(())
+	}
+
 	// Note that a block is finalized. Only call with child of last finalized block.
 	fn note_finalized(
 		&self,
@@ -325,51 +380,10 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 		}

 		// blocks are keyed by number + hash.
-		let lookup_key = utils::number_and_hash_to_lookup_key(number, hash);
+		let lookup_key = utils::number_and_hash_to_lookup_key(number, &hash);

 		if leaf_state.is_best() {
-			// handle reorg.
-			{
-				let meta = self.meta.read();
-				if meta.best_hash != Default::default() {
-					let tree_route = ::client::blockchain::tree_route(
-						self,
-						BlockId::Hash(meta.best_hash),
-						BlockId::Hash(parent_hash),
-					)?;
-
-					// update block number to hash lookup entries.
-					for retracted in tree_route.retracted() {
-						if retracted.hash == meta.finalized_hash {
-							warn!("Safety failure: reverting finalized block {:?}",
-								(&retracted.number, &retracted.hash));
-						}
-
-						utils::remove_number_to_key_mapping(
-							&mut transaction,
-							columns::KEY_LOOKUP,
-							retracted.number
-						);
-					}
-
-					for enacted in tree_route.enacted() {
-						utils::insert_number_to_key_mapping(
-							&mut transaction,
-							columns::KEY_LOOKUP,
-							enacted.number,
-							enacted.hash
-						);
-					}
-				}
-			}
-
-			transaction.put(columns::META, meta_keys::BEST_BLOCK, &lookup_key);
-			utils::insert_number_to_key_mapping(
-				&mut transaction,
-				columns::KEY_LOOKUP,
-				number,
-				hash,
-			);
+			self.set_head_with_transaction(&mut transaction, parent_hash, (number, hash))?;
 		}

 		utils::insert_hash_to_key_mapping(
@@ -426,6 +440,20 @@ impl<Block> LightBlockchainStorage<Block> for LightStorage<Block>
 		Ok(())
 	}

+	fn set_head(&self, id: BlockId<Block>) -> ClientResult<()> {
+		if let Some(header) = self.header(id)? {
+			let hash = header.hash();
+			let number = header.number();
+
+			let mut transaction = DBTransaction::new();
+			self.set_head_with_transaction(&mut transaction, hash.clone(), (number.clone(), hash.clone()))?;
+			self.db.write(transaction).map_err(db_err)?;
+			Ok(())
+		} else {
+			Err(ClientErrorKind::UnknownBlock(format!("Cannot set head {:?}", id)).into())
+		}
+	}
+
 	fn header_cht_root(&self, cht_size: u64, block: NumberFor<Block>) -> ClientResult<Block::Hash> {
 		self.read_cht_root(HEADER_CHT_PREFIX, cht_size, block)
 	}
@@ -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.
@@ -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(
@@ -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(())
 	}

@@ -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>
@@ -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())
 		}