Make planning epoch config change in BABE easier

bin/node/runtime/src/lib.rs

@@ -93,7 +93,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	// and set impl_version to 0. If only runtime
 	// implementation changes and behavior does not, then leave spec_version as
 	// is and increment impl_version.
-	spec_version: 249,
+	spec_version: 250,
 	impl_version: 0,
 	apis: RUNTIME_API_VERSIONS,
 	transaction_version: 1,

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

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())
+	});
+}