From e04f237152d2e55dd55631bcca8bce7cfaecbeba Mon Sep 17 00:00:00 2001
From: Wei Tang <hi@that.world>
Date: Wed, 20 May 2020 22:08:20 +0200
Subject: [PATCH] Make planning epoch config change in BABE easier (#5776)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Make planning epoch config change in BABE easier

* Bump node runtime version

* Update frame/babe/src/lib.rs

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>

* BabeEpochConfiguration -> NextConfigDescriptor

* Add tests for babe config changes

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>
Co-authored-by: André Silva <andre.beat@gmail.com>
---
 substrate/frame/babe/src/lib.rs   | 21 +++++++++++++++++---
 substrate/frame/babe/src/tests.rs | 33 +++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/substrate/frame/babe/src/lib.rs b/substrate/frame/babe/src/lib.rs
index 6d32c222493..153ff0e992d 100644
--- a/substrate/frame/babe/src/lib.rs
+++ b/substrate/frame/babe/src/lib.rs
@@ -42,7 +42,7 @@ use sp_inherents::{InherentIdentifier, InherentData, ProvideInherent, MakeFatalE
 use sp_consensus_babe::{
 	BABE_ENGINE_ID, ConsensusLog, BabeAuthorityWeight, SlotNumber,
 	inherents::{INHERENT_IDENTIFIER, BabeInherentData},
-	digests::{NextEpochDescriptor, PreDigest},
+	digests::{NextEpochDescriptor, NextConfigDescriptor, PreDigest},
 };
 use sp_consensus_vrf::schnorrkel;
 pub use sp_consensus_babe::{AuthorityId, VRF_OUTPUT_LENGTH, RANDOMNESS_LENGTH, PUBLIC_KEY_LENGTH};
@@ -136,6 +136,9 @@ decl_storage! {
 		// variable to its underlying value.
 		pub Randomness get(fn randomness): schnorrkel::Randomness;
 
+		/// Next epoch configuration, if changed.
+		NextEpochConfig: Option<NextConfigDescriptor>;
+
 		/// Next epoch randomness.
 		NextRandomness: schnorrkel::Randomness;
 
@@ -364,6 +367,15 @@ impl<T: Trait> Module<T> {
 			})
 	}
 
+	/// Plan an epoch config change. The epoch config change is recorded and will be enacted on the
+	/// next call to `enact_epoch_change`. The config will be activated one epoch after. Multiple calls to this
+	/// method will replace any existing planned config change that had not been enacted yet.
+	pub fn plan_config_change(
+		config: NextConfigDescriptor,
+	) {
+		NextEpochConfig::put(config);
+	}
+
 	/// DANGEROUS: Enact an epoch change. Should be done on every block where `should_epoch_change` has returned `true`,
 	/// and the caller is the only caller of this function.
 	///
@@ -399,12 +411,15 @@ impl<T: Trait> Module<T> {
 		// so that nodes can track changes.
 		let next_randomness = NextRandomness::get();
 
-		let next = NextEpochDescriptor {
+		let next_epoch = NextEpochDescriptor {
 			authorities: next_authorities,
 			randomness: next_randomness,
 		};
+		Self::deposit_consensus(ConsensusLog::NextEpochData(next_epoch));
 
-		Self::deposit_consensus(ConsensusLog::NextEpochData(next))
+		if let Some(next_config) = NextEpochConfig::take() {
+			Self::deposit_consensus(ConsensusLog::NextConfigData(next_config));
+		}
 	}
 
 	// finds the start slot of the current epoch. only guaranteed to
diff --git a/substrate/frame/babe/src/tests.rs b/substrate/frame/babe/src/tests.rs
index be2d3ed036e..ecb3639fc57 100644
--- a/substrate/frame/babe/src/tests.rs
+++ b/substrate/frame/babe/src/tests.rs
@@ -22,6 +22,7 @@ use mock::*;
 use frame_support::traits::OnFinalize;
 use pallet_session::ShouldEndSession;
 use sp_core::crypto::IsWrappedBy;
+use sp_consensus_babe::AllowedSlots;
 use sp_consensus_vrf::schnorrkel::{VRFOutput, VRFProof};
 
 const EMPTY_RANDOMNESS: [u8; 32] = [
@@ -150,3 +151,35 @@ fn can_predict_next_epoch_change() {
 		assert_eq!(Babe::next_expected_epoch_change(System::block_number()), Some(5 + 2));
 	})
 }
+
+#[test]
+fn can_enact_next_config() {
+	new_test_ext(0).1.execute_with(|| {
+		assert_eq!(<Test as Trait>::EpochDuration::get(), 3);
+		// this sets the genesis slot to 6;
+		go_to_block(1, 6);
+		assert_eq!(Babe::genesis_slot(), 6);
+		assert_eq!(Babe::current_slot(), 6);
+		assert_eq!(Babe::epoch_index(), 0);
+		go_to_block(2, 7);
+
+		Babe::plan_config_change(NextConfigDescriptor::V1 {
+			c: (1, 4),
+			allowed_slots: AllowedSlots::PrimarySlots,
+		});
+
+		progress_to_block(4);
+		Babe::on_finalize(9);
+		let header = System::finalize();
+
+		let consensus_log = sp_consensus_babe::ConsensusLog::NextConfigData(
+			sp_consensus_babe::digests::NextConfigDescriptor::V1 {
+				c: (1, 4),
+				allowed_slots: AllowedSlots::PrimarySlots,
+			}
+		);
+		let consensus_digest = DigestItem::Consensus(BABE_ENGINE_ID, consensus_log.encode());
+
+		assert_eq!(header.digest.logs[2], consensus_digest.clone())
+	});
+}
-- 
GitLab