babe: make it resilient to forks

client/consensus/babe/src/lib.rs

@@ -1526,10 +1526,11 @@ where
 				old_epoch_changes = Some((*epoch_changes).clone());
 
 				let mut viable_epoch = epoch_changes
-					.viable_epoch_mut(&epoch_descriptor, |slot| Epoch::genesis(&self.config, slot))
+					.viable_epoch(&epoch_descriptor, |slot| Epoch::genesis(&self.config, slot))
 					.ok_or_else(|| {
 						ConsensusError::ClientImport(Error::<Block>::FetchEpoch(parent_hash).into())
-					})?;
+					})?
+					.into_cloned();
 
 				let epoch_config = next_config_digest
 					.map(Into::into)
@@ -1552,14 +1553,16 @@ where
 
 					let original_epoch_index = epoch_data.epoch_index;
 
-					// NOTE: notice that we are only updating the `Epoch` from `EpochChanges` (that
-					// is stored in the map), and not the `EpochHeader` that is stored in the fork
-					// tree.  next time we search for an epoch for a given slot we will do it
-					// through the fork tree (which isn't updated), but the reason this works is
-					// because we will search in-depth in the tree with the predicate
-					// `epoch.start_slot <= slot` which will still match correctly without the
-					// updated `start_slot`. the new epoch that will get inserted below (after
-					// `increment`) will already use a correct `start_slot`.
+					// NOTE: notice that we are only updating a local copy of the `Epoch`, this
+					// makes it so that when we insert the new epoch into `EpochChanges` below
+					// (after incrementing it), it will use the correct epoch index and start slot.
+					// we do not update the original epoch that will be re-used because there might
+					// be other forks (that we haven't imported) where the epoch isn't skipped, and
+					// to import those forks we want to keep the original epoch data. not updating
+					// the original epoch works because when we search the tree for which epoch to
+					// use for a given slot, we will search in-depth with the predicate
+					// `epoch.start_slot <= slot` which will still match correctly without updating
+					// `start_slot` to the correct value as below.
 					epoch_data.epoch_index += skipped_epochs;
 					epoch_data.start_slot =
 						Slot::from(*epoch_data.start_slot + skipped_epochs * epoch_data.duration);

client/consensus/babe/src/tests.rs

@@ -1152,32 +1152,49 @@ fn allows_skipping_epochs() {
 		&mut block_import,
 	);
 
-	// we check the epoch data that was announced at block #7 (that we checked above)
-	let epoch3 = epoch_changes
+	// and the first block in epoch 3 (#8) announces epoch 4
+	let epoch4 = epoch_changes
 		.shared_data()
 		.epoch(&EpochIdentifier {
 			position: EpochIdentifierPosition::Regular,
-			hash: blocks[epoch_length as usize],
-			number: epoch_length + 1,
+			hash: block,
+			number: epoch_length + 2,
 		})
 		.unwrap()
 		.clone();
 
-	// and it now is updated to start at epoch 3 with an updated slot
-	assert_eq!(epoch3.epoch_index, 3);
-	assert_eq!(epoch3.start_slot, Slot::from(epoch_length * 3 + 1));
+	assert_eq!(epoch4.epoch_index, 4);
+	assert_eq!(epoch4.start_slot, Slot::from(epoch_length * 4 + 1));
 
-	// and the first block in epoch 3 (#8) announces epoch 4
-	let epoch = epoch_changes
+	// if we try to get the epoch data for a slot in epoch 3
+	let epoch3 = epoch_changes
 		.shared_data()
-		.epoch(&EpochIdentifier {
-			position: EpochIdentifierPosition::Regular,
-			hash: block,
-			number: epoch_length + 2,
-		})
+		.epoch_data_for_child_of(
+			descendent_query(&*client),
+			&block,
+			epoch_length + 2,
+			(epoch_length * 3 + 2).into(),
+			|slot| Epoch::genesis(&data.link.config, slot),
+		)
 		.unwrap()
-		.clone();
+		.unwrap();
+
+	// we get back the data for epoch 2
+	assert_eq!(epoch3, epoch2);
+
+	// but if we try to get the epoch data for a slot in epoch 4
+	let epoch4_ = epoch_changes
+		.shared_data()
+		.epoch_data_for_child_of(
+			descendent_query(&*client),
+			&block,
+			epoch_length + 2,
+			(epoch_length * 4 + 1).into(),
+			|slot| Epoch::genesis(&data.link.config, slot),
+		)
+		.unwrap()
+		.unwrap();
 
-	assert_eq!(epoch.epoch_index, 4);
-	assert_eq!(epoch.start_slot, Slot::from(epoch_length * 4 + 1));
+	// we get epoch 4 as expected
+	assert_eq!(epoch4, epoch4_);
 }