Fix light client synchronization on master (#3301)

* value ranges in consensus cache

* skip values in cache

* read epoch0 + epoch1 data from genesis in babe

* sync authorities + session validators at genesis

* removed some debug printlns

* fixed cache encoding

* Revert "skip values in cache"

This reverts commit ce451c32823aaa4b67d99ca5b58f1bf3984df4db.

* Revert "value ranges in consensus cache"

This reverts commit 9062f9434cddd14a01275ddbfcd904b04282e63b.

* get rid of cache::AUTHORITIES in Babe

* cleaning up

* cleaning up

* update spec version

* lost changes

* fixed tests

* Update node/runtime/src/lib.rs

Co-Authored-By: DemiMarie-parity <48690212+DemiMarie-parity@users.noreply.github.com>

* fix once-per-block condition

* fix standalone babe + temp_storage in BuildGenesis

* fix benhes compilation

* fixed comment

* re-added light nodes to integration tests

* finalize_with_ancestors from extra_requests

* post-merge fix

* aaand removed debug code

* (another one)

* fix warn in logs (do not call ForkTree::finalize twice for the same block)

* sync digest.next_authorities with actual next authorities

* more docs

* reverting all commits affecting storage

* also remove keys from babe trait

* fixed warnings

* post-merge fixes

* reverted some redundant changes

* reverted more changes

substrate/CODEOWNERS

@@ -8,7 +8,7 @@
 #   can be everywhere.
 # - Multiple owners are supported.
 # - Either handle (e.g, @pepyakin) or email can be used. Keep in mind, that handles might work better because they
-#   are more recognizable on GitHub, you can use them for mentioning unlike an email. 
+#   are more recognizable on GitHub, you can use them for mentioning unlike an email.
 # - The latest matching rule, if multiple, takes precedence.
 
 /srml/contracts/ @pepyakin

substrate/core/consensus/aura/src/lib.rs

@@ -43,25 +43,24 @@ use client::{
 	runtime_api::ApiExt, error::Result as CResult, backend::AuxStore, BlockOf,
 };
 
-use sr_primitives::{generic::{self, BlockId, OpaqueDigestItemId}, Justification};
+use sr_primitives::{generic::{BlockId, OpaqueDigestItemId}, Justification};
 use sr_primitives::traits::{Block as BlockT, Header, DigestItemFor, ProvideRuntimeApi, Zero, Member};
 
 use primitives::crypto::Pair;
 use inherents::{InherentDataProviders, InherentData};
 
-use futures::{prelude::*, future};
+use futures::prelude::*;
 use parking_lot::Mutex;
-use futures_timer::Delay;
-use log::{error, warn, debug, info, trace};
+use log::{debug, info, trace};
 
 use srml_aura::{
 	InherentType as AuraInherent, AuraInherentData,
 	timestamp::{TimestampInherentData, InherentType as TimestampInherent, InherentError as TIError}
 };
-use substrate_telemetry::{telemetry, CONSENSUS_TRACE, CONSENSUS_DEBUG, CONSENSUS_WARN, CONSENSUS_INFO};
+use substrate_telemetry::{telemetry, CONSENSUS_TRACE, CONSENSUS_DEBUG, CONSENSUS_INFO};
 
 use slots::{CheckedHeader, SlotData, SlotWorker, SlotInfo, SlotCompatible};
-use slots::{SignedDuration, check_equivocation};
+use slots::check_equivocation;
 
 use keystore::KeyStorePtr;
 
@@ -335,7 +334,7 @@ fn find_pre_digest<B: BlockT, P: Pair>(header: &B::Header) -> Result<u64, String
 ///
 /// This digest item will always return `Some` when used with `as_aura_seal`.
 //
-// FIXME #1018 needs misbehavior types. The `transaction_pool` parameter will be 
+// FIXME #1018 needs misbehavior types. The `transaction_pool` parameter will be
 // used to submit such misbehavior reports.
 fn check_header<C, B: BlockT, P: Pair, T>(
 	client: &C,

substrate/core/consensus/babe/src/lib.rs

@@ -21,20 +21,19 @@
 #![forbid(unsafe_code, missing_docs)]
 pub use babe_primitives::*;
 pub use consensus_common::SyncOracle;
-use std::{collections::HashMap, sync::Arc, u64, fmt::{Debug, Display}, pin::Pin, time::{Instant, Duration}};
+use std::{collections::HashMap, sync::Arc, u64, pin::Pin, time::{Instant, Duration}};
 use babe_primitives;
 use consensus_common::ImportResult;
 use consensus_common::import_queue::{
 	BoxJustificationImport, BoxFinalityProofImport,
 };
 use consensus_common::well_known_cache_keys::Id as CacheKeyId;
-use sr_primitives::{generic, generic::{BlockId, OpaqueDigestItemId}, Justification};
+use sr_primitives::{generic::{BlockId, OpaqueDigestItemId}, Justification};
 use sr_primitives::traits::{
 	Block as BlockT, Header, DigestItemFor, NumberFor, ProvideRuntimeApi,
-	SimpleBitOps, Zero,
+	Zero,
 };
 use keystore::KeyStorePtr;
-use runtime_support::serde::{Serialize, Deserialize};
 use codec::{Decode, Encode};
 use parking_lot::{Mutex, MutexGuard};
 use primitives::{Blake2Hasher, H256, Pair, Public};
@@ -44,8 +43,6 @@ use substrate_telemetry::{
 	telemetry,
 	CONSENSUS_TRACE,
 	CONSENSUS_DEBUG,
-	CONSENSUS_WARN,
-	CONSENSUS_INFO,
 };
 use schnorrkel::{
 	keys::Keypair,
@@ -72,12 +69,11 @@ use client::{
 };
 use fork_tree::ForkTree;
 use slots::{CheckedHeader, check_equivocation};
-use futures::{prelude::*, future};
+use futures::prelude::*;
 use futures01::Stream as _;
-use futures_timer::Delay;
 use log::{error, warn, debug, info, trace};
 
-use slots::{SlotWorker, SlotData, SlotInfo, SlotCompatible, SignedDuration};
+use slots::{SlotWorker, SlotData, SlotInfo, SlotCompatible};
 
 mod aux_schema;
 #[cfg(test)]
@@ -256,7 +252,8 @@ impl<H, B, C, E, I, Error, SO> slots::SimpleSlotWorker<B> for BabeWorker<C, E, I
 	}
 
 	fn epoch_data(&self, block: &B::Hash) -> Result<Self::EpochData, consensus_common::Error> {
-		epoch(self.client.as_ref(), &BlockId::Hash(*block))
+		epoch_from_runtime(self.client.as_ref(), &BlockId::Hash(*block))
+			.ok_or(consensus_common::Error::InvalidAuthoritiesSet)
 	}
 
 	fn authorities_len(&self, epoch_data: &Self::EpochData) -> usize {
@@ -397,7 +394,7 @@ fn find_next_epoch_digest<B: BlockT>(header: &B::Header) -> Result<Option<Epoch>
 /// unsigned.  This is required for security and must not be changed.
 ///
 /// This digest item will always return `Some` when used with `as_babe_pre_digest`.
-// FIXME #1018 needs misbehavior types. The `transaction_pool` parameter will be 
+// FIXME #1018 needs misbehavior types. The `transaction_pool` parameter will be
 // used to submit such misbehavior reports.
 fn check_header<B: BlockT + Sized, C: AuxStore, T>(
 	client: &C,
@@ -595,24 +592,51 @@ impl<B: BlockT, C, T> Verifier<B> for BabeVerifier<C, T> where
 
 		let hash = header.hash();
 		let parent_hash = *header.parent_hash();
-		let Epoch { authorities, randomness, epoch_index, .. } =
-			epoch(self.api.as_ref(), &BlockId::Hash(parent_hash))
-				.map_err(|e| format!("Could not fetch epoch at {:?}: {:?}", parent_hash, e))?;
+
+		let epoch = epoch(self.api.as_ref(), &BlockId::Hash(parent_hash))
+			.map_err(|e| format!("Could not fetch epoch at {:?}: {:?}", parent_hash, e))?;
+		let (epoch, maybe_next_epoch) = epoch.deconstruct();
+		let Epoch { authorities, randomness, epoch_index, .. } = epoch;
 
 		// We add one to allow for some small drift.
 		// FIXME #1019 in the future, alter this queue to allow deferring of headers
-		let checked_header = check_header::<B, C, T>(
+		let mut checked_header = check_header::<B, C, T>(
 			&self.api,
 			slot_now + 1,
-			header,
+			header.clone(),
 			hash,
 			&authorities,
 			randomness,
 			epoch_index,
 			self.config.c(),
 			self.transaction_pool.as_ref().map(|x| &**x),
-		)?;
+		);
 
+		// if we have failed to check header using (presumably) current epoch AND we're probably in the next epoch
+		// => check using next epoch
+		// (this is only possible on the light client at epoch#0)
+		if epoch_index == 0 && checked_header.is_err() {
+			if let Some(Epoch { authorities, randomness, epoch_index, .. }) = maybe_next_epoch {
+				let checked_header_next = check_header::<B, C, T>(
+					&self.api,
+					slot_now + 1,
+					header,
+					hash,
+					&authorities,
+					randomness,
+					epoch_index,
+					self.config.c(),
+					self.transaction_pool.as_ref().map(|x| &**x),
+				);
+
+				match checked_header_next {
+					Ok(checked_header_next) => checked_header = Ok(checked_header_next),
+					Err(_) => (),
+				}
+			}
+		}
+
+		let checked_header = checked_header?;
 		match checked_header {
 			CheckedHeader::Checked(pre_header, (pre_digest, seal)) => {
 				let BabePreDigest { slot_number, .. } = pre_digest.as_babe_pre_digest()
@@ -665,31 +689,75 @@ impl<B: BlockT, C, T> Verifier<B> for BabeVerifier<C, T> where
 	}
 }
 
+/// Regular BABE epoch or spanned genesis epoch.
+#[derive(Debug, Decode, Encode)]
+enum MaybeSpanEpoch {
+	/// Genesis entry. Has the data for epoch#0 and epoch#1.
+	Genesis(Epoch, Epoch),
+	/// Regular entry. Has the data for the epoch after next (i.e. current epoch + 2).
+	Regular(Epoch),
+}
+
+impl MaybeSpanEpoch {
+	pub fn deconstruct(self) -> (Epoch, Option<Epoch>) {
+		match self {
+			MaybeSpanEpoch::Genesis(epoch0, epoch1) => (epoch0, Some(epoch1)),
+			MaybeSpanEpoch::Regular(epoch) => (epoch, None),
+		}
+	}
+
+	#[cfg(test)]
+	pub fn into_regular(self) -> Option<Epoch> {
+		match self {
+			MaybeSpanEpoch::Regular(epoch) => Some(epoch),
+			_ => None,
+		}
+	}
+}
+
 /// Extract current epoch data from cache and fallback to querying the runtime
 /// if the cache isn't populated.
-fn epoch<B, C>(client: &C, at: &BlockId<B>) -> Result<Epoch, ConsensusError> where
+fn epoch<B, C>(client: &C, at: &BlockId<B>) -> Result<MaybeSpanEpoch, ConsensusError> where
 	B: BlockT,
 	C: ProvideRuntimeApi + ProvideCache<B>,
 	C::Api: BabeApi<B>,
 {
-	client
-		.cache()
-		.and_then(|cache| cache.get_at(&well_known_cache_keys::EPOCH, at)
+	epoch_from_cache(client, at)
+		.or_else(|| epoch_from_runtime(client, at).map(MaybeSpanEpoch::Regular))
+		.ok_or(consensus_common::Error::InvalidAuthoritiesSet)
+}
+
+/// Extract current epoch data from cache.
+fn epoch_from_cache<B, C>(client: &C, at: &BlockId<B>) -> Option<MaybeSpanEpoch> where
+	B: BlockT,
+	C: ProvideCache<B>,
+{
+	// the epoch that is BABE-valid at the block is not the epoch that is cache-valid at the block
+	// we need to go back for maximum two steps
+	client.cache()
+		.and_then(|cache| cache
+			.get_at(&well_known_cache_keys::EPOCH, at)
 			.and_then(|v| Decode::decode(&mut &v[..]).ok()))
-		.or_else(|| {
-			if client.runtime_api().has_api::<dyn BabeApi<B>>(at).unwrap_or(false) {
-				let s = BabeApi::epoch(&*client.runtime_api(), at).ok()?;
-				if s.authorities.is_empty() {
-					error!("No authorities!");
-					None
-				} else {
-					Some(s)
-				}
-			} else {
-				error!("bad api!");
-				None
-			}
-		}).ok_or(consensus_common::Error::InvalidAuthoritiesSet)
+}
+
+/// Extract current epoch from runtime.
+fn epoch_from_runtime<B, C>(client: &C, at: &BlockId<B>) -> Option<Epoch> where
+	B: BlockT,
+	C: ProvideRuntimeApi,
+	C::Api: BabeApi<B>,
+{
+	if client.runtime_api().has_api::<dyn BabeApi<B>>(at).unwrap_or(false) {
+		let s = BabeApi::epoch(&*client.runtime_api(), at).ok()?;
+		if s.authorities.is_empty() {
+			error!("No authorities!");
+			None
+		} else {
+			Some(s)
+		}
+	} else {
+		error!("bad api!");
+		None
+	}
 }
 
 /// The BABE import queue type.
@@ -801,7 +869,7 @@ fn initialize_authorities_cache<B, C>(client: &C) -> Result<(), ConsensusError>
 
 	// check if we already have initialized the cache
 	let genesis_id = BlockId::Number(Zero::zero());
-	let genesis_epoch: Option<Epoch> = cache
+	let genesis_epoch: Option<MaybeSpanEpoch> = cache
 		.get_at(&well_known_cache_keys::EPOCH, &genesis_id)
 		.and_then(|v| Decode::decode(&mut &v[..]).ok());
 	if genesis_epoch.is_some() {
@@ -814,7 +882,11 @@ fn initialize_authorities_cache<B, C>(client: &C) -> Result<(), ConsensusError>
 			error,
 		)));
 
-	let genesis_epoch = epoch(client, &genesis_id)?;
+	let epoch0 = epoch_from_runtime(client, &genesis_id).ok_or(consensus_common::Error::InvalidAuthoritiesSet)?;
+	let mut epoch1 = epoch0.clone();
+	epoch1.epoch_index = 1;
+
+	let genesis_epoch = MaybeSpanEpoch::Genesis(epoch0, epoch1);
 	cache.initialize(&well_known_cache_keys::EPOCH, genesis_epoch.encode())
 		.map_err(map_err)
 }
@@ -990,6 +1062,16 @@ impl<B, E, Block, I, RA, PRA> BlockImport<Block> for BabeBlockImport<B, E, Block
 		// this way we can revert it if there's any error
 		let mut old_epoch_changes = None;
 
+		if let Some(enacted_epoch) = enacted_epoch.as_ref() {
+			let enacted_epoch = &enacted_epoch.data;
+
+			// update the current epoch in the client cache
+			new_cache.insert(
+				well_known_cache_keys::EPOCH,
+				MaybeSpanEpoch::Regular(enacted_epoch.clone()).encode(),
+			);
+		}
+
 		if let Some(next_epoch) = next_epoch_digest {
 			if let Some(enacted_epoch) = enacted_epoch {
 				let enacted_epoch = &enacted_epoch.data;
@@ -1000,27 +1082,6 @@ impl<B, E, Block, I, RA, PRA> BlockImport<Block> for BabeBlockImport<B, E, Block
 						next_epoch.epoch_index,
 					)));
 				}
-
-				// update the current epoch in the client cache
-				new_cache.insert(
-					well_known_cache_keys::EPOCH,
-					enacted_epoch.encode(),
-				);
-
-				let current_epoch = epoch(&*self.api, &BlockId::Hash(parent_hash))?;
-
-				// if the authorities have changed then we populate the
-				// `AUTHORITIES` key with the enacted epoch, so that the inner
-				// `ImportBlock` can process it (`EPOCH` is specific to BABE).
-				// e.g. in the case of GRANDPA it would require a justification
-				// for the block, expecting that the authorities actually
-				// changed.
-				if current_epoch.authorities != enacted_epoch.authorities {
-					new_cache.insert(
-						well_known_cache_keys::AUTHORITIES,
-						enacted_epoch.encode(),
-					);
-				}
 			}
 
 			old_epoch_changes = Some(epoch_changes.clone());
@@ -1158,7 +1219,10 @@ pub mod test_helpers {
 		C: ProvideRuntimeApi + ProvideCache<B>,
 		C::Api: BabeApi<B>,
 	{
-		let epoch = epoch(client, at).unwrap();
+		let epoch = match epoch(client, at).unwrap() {
+			MaybeSpanEpoch::Regular(epoch) => epoch,
+			_ => unreachable!("it is always Regular epoch on full nodes"),
+		};
 
 		super::claim_slot(
 			slot_number,

substrate/core/consensus/babe/src/tests.rs

@@ -20,7 +20,7 @@
 // https://github.com/paritytech/substrate/issues/2532
 #![allow(deprecated)]
 use super::*;
-use super::generic::DigestItem;
+use sr_primitives::generic::{self, DigestItem};
 
 use babe_primitives::AuthorityPair;
 use client::{LongestChain, block_builder::BlockBuilder};
@@ -341,7 +341,7 @@ fn authorities_call_works() {
 	let client = test_client::new();
 
 	assert_eq!(client.info().chain.best_number, 0);
-	assert_eq!(epoch(&client, &BlockId::Number(0)).unwrap().authorities, vec![
+	assert_eq!(epoch(&client, &BlockId::Number(0)).unwrap().into_regular().unwrap().authorities, vec![
 		(Keyring::Alice.public().into(), 1),
 		(Keyring::Bob.public().into(), 1),
 		(Keyring::Charlie.public().into(), 1),

substrate/core/finality-grandpa/src/import.rs

@@ -425,7 +425,7 @@ impl<B, E, Block: BlockT<Hash=H256>, RA, PRA, SC> BlockImport<Block>
 
 		// we don't want to finalize on `inner.import_block`
 		let mut justification = block.justification.take();
-		let enacts_consensus_change = new_cache.contains_key(&well_known_cache_keys::AUTHORITIES);
+		let enacts_consensus_change = !new_cache.is_empty();
 		let import_result = (&*self.inner).import_block(block, new_cache);
 
 		let mut imported_aux = {

substrate/core/finality-grandpa/src/light_import.rs

@@ -246,7 +246,7 @@ fn do_import_block<B, E, Block: BlockT<Hash=H256>, RA, J>(
 
 	// we don't want to finalize on `inner.import_block`
 	let justification = block.justification.take();
-	let enacts_consensus_change = new_cache.contains_key(&well_known_cache_keys::AUTHORITIES);
+	let enacts_consensus_change = !new_cache.is_empty();
 	let import_result = BlockImport::import_block(&mut client, block, new_cache);
 
 	let mut imported_aux = match import_result {

substrate/core/network/src/protocol/sync/extra_requests.rs

@@ -19,7 +19,7 @@ use crate::protocol::sync::{PeerSync, PeerSyncState};
 use fork_tree::ForkTree;
 use libp2p::PeerId;
 use log::warn;
-use sr_primitives::traits::{Block as BlockT, NumberFor};
+use sr_primitives::traits::{Block as BlockT, NumberFor, Zero};
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::time::{Duration, Instant};
 
@@ -38,6 +38,8 @@ pub(crate) type ExtraRequest<B> = (<B as BlockT>::Hash, NumberFor<B>);
 #[derive(Debug)]
 pub(crate) struct ExtraRequests<B: BlockT> {
 	tree: ForkTree<B::Hash, NumberFor<B>, ()>,
+	/// best finalized block number that we have seen since restart
+	best_seen_finalized_number: NumberFor<B>,
 	/// requests which have been queued for later processing
 	pending_requests: VecDeque<ExtraRequest<B>>,
 	/// requests which are currently underway to some peer
@@ -52,6 +54,7 @@ impl<B: BlockT> ExtraRequests<B> {
 	pub(crate) fn new() -> Self {
 		ExtraRequests {
 			tree: ForkTree::new(),
+			best_seen_finalized_number: Zero::zero(),
 			pending_requests: VecDeque::new(),
 			active_requests: HashMap::new(),
 			failed_requests: HashMap::new(),
@@ -80,7 +83,7 @@ impl<B: BlockT> ExtraRequests<B> {
 		match self.tree.import(request.0, request.1, (), &is_descendent_of) {
 			Ok(true) => {
 				// this is a new root so we add it to the current `pending_requests`
-				self.pending_requests.push_back((request.0, request.1))
+				self.pending_requests.push_back((request.0, request.1));
 			}
 			Err(err) => {
 				warn!(target: "sync", "Failed to insert request {:?} into tree: {:?}", request, err);
@@ -93,7 +96,7 @@ impl<B: BlockT> ExtraRequests<B> {
 	/// Retry any pending request if a peer disconnected.
 	pub(crate) fn peer_disconnected(&mut self, who: &PeerId) {
 		if let Some(request) = self.active_requests.remove(who) {
-			self.pending_requests.push_front(request)
+			self.pending_requests.push_front(request);
 		}
 	}
 
@@ -128,7 +131,10 @@ impl<B: BlockT> ExtraRequests<B> {
 			return Ok(())
 		}
 
-		self.tree.finalize(best_finalized_hash, best_finalized_number, &is_descendent_of)?;
+		if best_finalized_number > self.best_seen_finalized_number {
+			self.tree.finalize_with_ancestors(best_finalized_hash, best_finalized_number, &is_descendent_of)?;
+			self.best_seen_finalized_number = best_finalized_number;
+		}
 
 		let roots = self.tree.roots().collect::<HashSet<_>>();
 
@@ -176,6 +182,7 @@ impl<B: BlockT> ExtraRequests<B> {
 		self.active_requests.clear();
 		self.pending_requests.clear();
 		self.pending_requests.extend(self.tree.roots().map(|(&h, &n, _)| (h, n)));
+		self.best_seen_finalized_number = finalized_number;
 
 		true
 	}

substrate/core/service/src/chain_spec.rs

@@ -81,6 +81,7 @@ impl<'a, G: RuntimeGenesis> BuildStorage for &'a ChainSpec<G> {
 			)),
 		}
 	}
+
 	fn assimilate_storage(self, _: &mut (StorageOverlay, ChildrenStorageOverlay)) -> Result<(), String> {
 		Err("`assimilate_storage` not implemented for `ChainSpec`.".into())
 	}

substrate/core/service/test/src/lib.rs

@@ -354,7 +354,7 @@ pub fn sync<F, B, E>(spec: FactoryChainSpec<F>, mut block_factory: B, mut extrin
 {
 	const NUM_FULL_NODES: usize = 10;
 	// FIXME: BABE light client support is currently not working.
-	const NUM_LIGHT_NODES: usize = 0;
+	const NUM_LIGHT_NODES: usize = 10;
 	const NUM_BLOCKS: usize = 512;
 	let temp = TempDir::new("substrate-sync-test").expect("Error creating test dir");
 	let mut network = TestNet::<F>::new(
@@ -410,7 +410,7 @@ pub fn consensus<F>(spec: FactoryChainSpec<F>, authorities: Vec<String>) where
 	F::LightService: Future<Item=(), Error=service::Error>,
 {
 	const NUM_FULL_NODES: usize = 10;
-	const NUM_LIGHT_NODES: usize = 0;
+	const NUM_LIGHT_NODES: usize = 10;
 	const NUM_BLOCKS: usize = 10; // 10 * 2 sec block production time = ~20 seconds
 	let temp = TempDir::new("substrate-conensus-test").expect("Error creating test dir");
 	let mut network = TestNet::<F>::new(

substrate/core/utils/fork-tree/src/lib.rs

@@ -240,14 +240,16 @@ impl<H, N, V> ForkTree<H, N, V> where
 	/// with the given hash exists. All other roots are pruned, and the children
 	/// of the finalized node become the new roots.
 	pub fn finalize_root(&mut self, hash: &H) -> Option<V> {
-		if let Some(position) = self.roots.iter().position(|node| node.hash == *hash) {
-			let node = self.roots.swap_remove(position);
-			self.roots = node.children;
-			self.best_finalized_number = Some(node.number);
-			return Some(node.data);
-		}
+		self.roots.iter().position(|node| node.hash == *hash)
+			.map(|position| self.finalize_root_at(position))
+	}
 
-		None
+	/// Finalize root at given positiion. See `finalize_root` comment for details.
+	fn finalize_root_at(&mut self, position: usize) -> V {
+		let node = self.roots.swap_remove(position);
+		self.roots = node.children;
+		self.best_finalized_number = Some(node.number);
+		return node.data;
 	}
 
 	/// Finalize a node in the tree. This method will make sure that the node
@@ -305,6 +307,79 @@ impl<H, N, V> ForkTree<H, N, V> where
 		}
 	}
 
+	/// Finalize a node in the tree and all its ancestors. The given function
+	/// `is_descendent_of` should return `true` if the second hash (target) is
+	// a descendent of the first hash (base).
+	pub fn finalize_with_ancestors<F, E>(
+		&mut self,
+		hash: &H,
+		number: N,
+		is_descendent_of: &F,
+	) -> Result<FinalizationResult<V>, Error<E>>
+		where E: std::error::Error,
+				F: Fn(&H, &H) -> Result<bool, E>
+	{
+		if let Some(ref best_finalized_number) = self.best_finalized_number {
+			if number <= *best_finalized_number {
+				return Err(Error::Revert);
+			}
+		}
+
+		// check if one of the current roots is being finalized
+		if let Some(root) = self.finalize_root(hash) {
+			return Ok(FinalizationResult::Changed(Some(root)));
+		}
+
+		// we need to:
+		// 1) remove all roots that are not ancestors AND not descendants of finalized block;
+		// 2) if node is descendant - just leave it;
+		// 3) if node is ancestor - 'open it'
+		let mut changed = false;
+		let mut idx = 0;
+		while idx != self.roots.len() {
+			let (is_finalized, is_descendant, is_ancestor) = {
+				let root = &self.roots[idx];
+				let is_finalized = root.hash == *hash;
+				let is_descendant = !is_finalized
+					&& root.number > number && is_descendent_of(hash, &root.hash).unwrap_or(false);
+				let is_ancestor = !is_finalized && !is_descendant
+					&& root.number < number && is_descendent_of(&root.hash, hash).unwrap_or(false);
+				(is_finalized, is_descendant, is_ancestor)
+			};
+
+			// if we have met finalized root - open it and return
+			if is_finalized {
+				return Ok(FinalizationResult::Changed(Some(self.finalize_root_at(idx))));
+			}
+
+			// if node is descendant of finalized block - just leave it as is
+			if is_descendant {
+				idx += 1;
+				continue;
+			}
+
+			// if node is ancestor of finalized block - remove it and continue with children
+			if is_ancestor {
+				let root = self.roots.swap_remove(idx);
+				self.roots.extend(root.children);
+				changed = true;
+				continue;
+			}
+
+			// if node is neither ancestor, nor descendant of the finalized block - remove it
+			self.roots.swap_remove(idx);
+			changed = true;
+		}
+
+		self.best_finalized_number = Some(number);
+
+		if changed {
+			Ok(FinalizationResult::Changed(None))
+		} else {
+			Ok(FinalizationResult::Unchanged)
+		}
+	}
+
 	/// Checks if any node in the tree is finalized by either finalizing the
 	/// node itself or a child node that's not in the tree, guaranteeing that
 	/// the node being finalized isn't a descendent of any of the node's
@@ -580,23 +655,32 @@ mod test {
 		//   /   - G
 		//  /   /
 		// A - F - H - I
+		//          \
+		//           - L - M - N
+		//              \
+		//               - O
 		//  \
 		//   — J - K
 		//
+		// (where N is not a part of fork tree)
 		let is_descendent_of = |base: &&str, block: &&str| -> Result<bool, TestError> {
-			let letters = vec!["B", "C", "D", "E", "F", "G", "H", "I", "J", "K"];
+			let letters = vec!["B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L"];
 			match (*base, *block) {
 				("A", b) => Ok(letters.into_iter().any(|n| n == b)),
 				("B", b) => Ok(b == "C" || b == "D" || b == "E"),
 				("C", b) => Ok(b == "D" || b == "E"),
 				("D", b) => Ok(b == "E"),
 				("E", _) => Ok(false),
-				("F", b) => Ok(b == "G" || b == "H" || b == "I"),
+				("F", b) => Ok(b == "G" || b == "H" || b == "I" || b == "L" || b == "M" || b == "N" || b == "O"),
 				("G", _) => Ok(false),
-				("H", b) => Ok(b == "I"),
+				("H", b) => Ok(b == "I" || b == "L" || b == "M" || b == "O"),
 				("I", _) => Ok(false),
 				("J", b) => Ok(b == "K"),
 				("K", _) => Ok(false),
+				("L", b) => Ok(b == "M" || b == "O" || b == "N"),
+				("M", b) => Ok(b == "N"),
+				("N", _) => Ok(false),
+				("O", _) => Ok(false),
 				("0", _) => Ok(true),
 				_ => Ok(false),
 			}
@@ -614,6 +698,9 @@ mod test {
 
 		tree.import("H", 3, (), &is_descendent_of).unwrap();
 		tree.import("I", 4, (), &is_descendent_of).unwrap();
+		tree.import("L", 4, (), &is_descendent_of).unwrap();
+		tree.import("M", 5, (), &is_descendent_of).unwrap();
+		tree.import("O", 5, (), &is_descendent_of).unwrap();
 
 		tree.import("J", 2, (), &is_descendent_of).unwrap();
 		tree.import("K", 3, (), &is_descendent_of).unwrap();
@@ -770,7 +857,7 @@ mod test {
 
 		assert_eq!(
 			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
-			vec![("I", 4)],
+			vec![("I", 4), ("L", 4)],
 		);
 
 		// finalizing a node from another fork that isn't part of the tree clears the tree
@@ -782,6 +869,71 @@ mod test {
 		assert!(tree.roots.is_empty());
 	}
 
+	#[test]
+	fn finalize_with_ancestor_works() {
+		let (mut tree, is_descendent_of) = test_fork_tree();
+
+		let original_roots = tree.roots.clone();
+
+		// finalizing a block prior to any in the node doesn't change the tree
+		assert_eq!(
+			tree.finalize_with_ancestors(&"0", 0, &is_descendent_of),
+			Ok(FinalizationResult::Unchanged),
+		);
+
+		assert_eq!(tree.roots, original_roots);
+
+		// finalizing "A" opens up three possible forks
+		assert_eq!(
+			tree.finalize_with_ancestors(&"A", 1, &is_descendent_of),
+			Ok(FinalizationResult::Changed(Some(()))),
+		);
+
+		assert_eq!(
+			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
+			vec![("B", 2), ("F", 2), ("J", 2)],
+		);
+
+		// finalizing H:
+		// 1) removes roots that are not ancestors/descendants of H (B, J)
+		// 2) opens root that is ancestor of H (F -> G+H)
+		// 3) finalizes the just opened root H (H -> I + L)
+		assert_eq!(
+			tree.finalize_with_ancestors(&"H", 3, &is_descendent_of),
+			Ok(FinalizationResult::Changed(Some(()))),
+		);
+
+		assert_eq!(
+			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
+			vec![("I", 4), ("L", 4)],
+		);
+
+		assert_eq!(
+			tree.best_finalized_number,
+			Some(3),
+		);
+
+		// finalizing N (which is not a part of the tree):
+		// 1) removes roots that are not ancestors/descendants of N (I)
+		// 2) opens root that is ancestor of N (L -> M+O)
+		// 3) removes roots that are not ancestors/descendants of N (O)
+		// 4) opens root that is ancestor of N (M -> {})
+		assert_eq!(
+			tree.finalize_with_ancestors(&"N", 6, &is_descendent_of),
+			Ok(FinalizationResult::Changed(None)),
+		);
+
+		assert_eq!(
+			tree.roots().map(|(h, n, _)| (h.clone(), n.clone())).collect::<Vec<_>>(),
+			vec![],
+		);
+
+		assert_eq!(
+			tree.best_finalized_number,
+			Some(6),
+		);
+	}
+
 	#[test]
 	fn finalize_with_descendent_works() {
 		#[derive(Debug, PartialEq)]
@@ -927,7 +1079,9 @@ mod test {
 			vec![
 				("A", 1),
 				("J", 2), ("K", 3),
-				("F", 2), ("H", 3), ("I", 4),
+				("F", 2), ("H", 3), ("L", 4), ("O", 5),
+				("M", 5),
+				("I", 4),
 				("G", 3),
 				("B", 2), ("C", 3), ("D", 4), ("E", 5),
 			],

substrate/node/runtime/src/lib.rs

@@ -80,8 +80,8 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	// and set impl_version to equal spec_version. If only runtime
 	// implementation changes and behavior does not, then leave spec_version as
 	// is and increment impl_version.
-	spec_version: 143,
-	impl_version: 143,
+	spec_version: 144,
+	impl_version: 144,
 	apis: RUNTIME_API_VERSIONS,
 };
 

substrate/srml/babe/src/lib.rs

@@ -162,6 +162,10 @@ decl_storage! {
 		/// epoch.
 		SegmentIndex build(|_| 0): u32;
 		UnderConstruction: map u32 => Vec<[u8; 32 /* VRF_OUTPUT_LENGTH */]>;
+
+		/// Temporary value (cleared at block finalization) which is true
+		/// if per-block initialization has already been called for current block.
+		Initialized get(initialized): Option<bool>;
 	}
 	add_extra_genesis {
 		config(authorities): Vec<(AuthorityId, BabeWeight)>;
@@ -193,25 +197,12 @@ decl_module! {
 
 		/// Initialization
 		fn on_initialize() {
-			for digest in Self::get_inherent_digests()
-				.logs
-				.iter()
-				.filter_map(|s| s.as_pre_runtime())
-				.filter_map(|(id, mut data)| if id == BABE_ENGINE_ID {
-					RawBabePreDigest::decode(&mut data).ok()
-				} else {
-					None
-				})
-			{
-				if EpochStartSlot::get() == 0 {
-					EpochStartSlot::put(digest.slot_number);
-				}
+			Self::do_initialize();
+		}
 
-				CurrentSlot::put(digest.slot_number);
-				Self::deposit_vrf_output(&digest.vrf_output);
-
-				return;
-			}
+		/// Block finalization
+		fn on_finalize() {
+			Initialized::kill();
 		}
 	}
 }
@@ -248,6 +239,12 @@ impl<T: Trait> IsMember<AuthorityId> for Module<T> {
 
 impl<T: Trait> session::ShouldEndSession<T::BlockNumber> for Module<T> {
 	fn should_end_session(_: T::BlockNumber) -> bool {
+		// it might be (and it is in current implementation) that session module is calling
+		// should_end_session() from it's own on_initialize() handler
+		// => because session on_initialize() is called earlier than ours, let's ensure
+		// that we have synced with digest before checking if session should be ended
+		Self::do_initialize();
+
 		let diff = CurrentSlot::get().saturating_sub(EpochStartSlot::get());
 		diff >= T::EpochDuration::get()
 	}
@@ -285,6 +282,36 @@ impl<T: Trait> Module<T> {
 		}
 	}
 
+	fn do_initialize() {
+		// since do_initialize can be called twice (if session module is present)
+		// => let's ensure that we only modify the storage once per block
+		let initialized = Self::initialized().unwrap_or(false);
+		if initialized {
+			return;
+		}
+
+		Initialized::put(true);
+		for digest in Self::get_inherent_digests()
+			.logs
+			.iter()
+			.filter_map(|s| s.as_pre_runtime())
+			.filter_map(|(id, mut data)| if id == BABE_ENGINE_ID {
+				RawBabePreDigest::decode(&mut data).ok()
+			} else {
+				None
+			})
+		{
+			if EpochStartSlot::get() == 0 {
+				EpochStartSlot::put(digest.slot_number);
+			}
+
+			CurrentSlot::put(digest.slot_number);
+			Self::deposit_vrf_output(&digest.vrf_output);
+
+			return;
+		}
+	}
+
 	/// Call this function exactly once when an epoch changes, to update the
 	/// randomness. Returns the new randomness.
 	fn randomness_change_epoch(next_epoch_index: u64) -> [u8; RANDOMNESS_LENGTH] {

substrate/srml/support/test/tests/genesisconfig.rs

@@ -15,18 +15,18 @@
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
 pub trait Trait {
-    type BlockNumber: codec::Codec + Default;
-    type Origin;
+	type BlockNumber: codec::Codec + Default;
+	type Origin;
 }
 
 srml_support::decl_module! {
-    pub struct Module<T: Trait> for enum Call where origin: T::Origin {}
+	pub struct Module<T: Trait> for enum Call where origin: T::Origin {}
 }
 
 srml_support::decl_storage! {
-    trait Store for Module<T: Trait> as Example {
-            pub AppendableDM config(t): double_map u32, blake2_256(T::BlockNumber) => Vec<u32>;
-    }
+	trait Store for Module<T: Trait> as Example {
+		pub AppendableDM config(t): double_map u32, blake2_256(T::BlockNumber) => Vec<u32>;
+	}
 }
 
 struct Test;