apollo_consensus: Add functionality to consensus state to support reverts (#10248)

Author: guyf-starkware

crates/apollo_consensus/src/storage.rs

@@ -31,6 +31,9 @@ pub trait HeightVotedStorageTrait: Debug + Send + Sync {
     /// Sets the last height on which the node voted.
     fn set_prev_voted_height(&mut self, height: BlockNumber)
     -> Result<(), HeightVotedStorageError>;
+    /// If the last voted height is higher or equal to the height to revert to, set the last voted
+    /// height to be just before the height to revert.
+    fn revert_height(&mut self, height: BlockNumber) -> Result<(), HeightVotedStorageError>;
 }
 
 struct HeightVotedStorage {
@@ -40,7 +43,7 @@ struct HeightVotedStorage {
     storage_writer: StorageWriter,
 }
 
-/// Opens and returns the storage used for storing last voted height state.
+/// Returns a new concrete implementation of HeightVotedStorageTrait based on the configuration.
 pub fn get_voted_height_storage(config: StorageConfig) -> impl HeightVotedStorageTrait {
     let (storage_reader, storage_writer) = open_storage(config).expect("Failed to open storage");
     HeightVotedStorage { storage_reader, storage_writer }
@@ -86,6 +89,27 @@ impl HeightVotedStorageTrait for HeightVotedStorage {
             }
         }
     }
+
+    fn revert_height(&mut self, height: BlockNumber) -> Result<(), HeightVotedStorageError> {
+        let mut txn = self.storage_writer.begin_rw_txn()?;
+        let last_voted = txn.get_last_voted_marker()?;
+        match last_voted {
+            None => return Ok(()), // If none, nothing to revert.
+            Some(last_voted) => {
+                if last_voted.height >= height {
+                    if let Some(prev_height) = height.prev() {
+                        txn =
+                            txn.set_last_voted_marker(&LastVotedMarker { height: prev_height })?;
+                    } else {
+                        // We're reverting height 0, so clear the last voted height completely.
+                        txn = txn.clear_last_voted_marker()?;
+                    }
+                }
+            }
+        }
+        txn.commit()?;
+        Ok(())
+    }
 }
 
 // We must implement Debug for HeightVotedStorageImpl to allow it being used as a member in structs

crates/apollo_consensus/src/storage_test.rs

@@ -62,3 +62,54 @@ fn write_last_height_return_error_when_previous_last_height_is_equal() {
         Err(HeightVotedStorageError::InconsistentStorageState { error_msg: _ })
     );
 }
+
+#[test]
+fn revert_height_when_no_last_height_in_storage_does_nothing() {
+    let mut storage = get_voted_height_storage(get_new_storage_config());
+    assert!(storage.get_prev_voted_height().unwrap().is_none());
+    storage.revert_height(BlockNumber(1)).unwrap();
+    assert!(storage.get_prev_voted_height().unwrap().is_none());
+}
+
+#[test]
+fn revert_height_when_last_height_in_storage_is_lower_than_height_to_revert_to_does_nothing() {
+    const HEIGHT_TO_REVERT_TO: BlockNumber = BlockNumber(2);
+    // Storage has a lower height than what we revert (so should be a no-op)
+    let last_height_in_storage = HEIGHT_TO_REVERT_TO.prev().unwrap();
+
+    let mut storage = get_voted_height_storage(get_new_storage_config());
+    storage.set_prev_voted_height(last_height_in_storage).unwrap();
+    storage.revert_height(HEIGHT_TO_REVERT_TO).unwrap();
+    assert_eq!(storage.get_prev_voted_height().unwrap(), Some(last_height_in_storage));
+}
+
+#[test]
+fn revert_height_when_last_height_in_storage_is_higher_than_revert_height_reverts_the_given_height()
+{
+    const HEIGH_TO_REVERT_TO: BlockNumber = BlockNumber(2);
+    // Storage has a higher height than what we're reverting.
+    let last_height_in_storage = HEIGH_TO_REVERT_TO.unchecked_next().unchecked_next();
+
+    let mut storage = get_voted_height_storage(get_new_storage_config());
+    storage.set_prev_voted_height(last_height_in_storage).unwrap();
+    storage.revert_height(HEIGH_TO_REVERT_TO).unwrap();
+    assert_eq!(storage.get_prev_voted_height().unwrap(), Some(HEIGH_TO_REVERT_TO.prev().unwrap()));
+}
+
+#[test]
+fn revert_height_when_last_height_in_storage_is_equal_to_revert_height_reverts_the_given_height() {
+    const HEIGHT_TO_REVERT_TO: BlockNumber = BlockNumber(2);
+
+    let mut storage = get_voted_height_storage(get_new_storage_config());
+    storage.set_prev_voted_height(HEIGHT_TO_REVERT_TO).unwrap();
+    storage.revert_height(HEIGHT_TO_REVERT_TO).unwrap();
+    assert_eq!(storage.get_prev_voted_height().unwrap(), Some(HEIGHT_TO_REVERT_TO.prev().unwrap()));
+}
+
+#[test]
+fn revert_height_to_0_clears_the_last_height_in_storage() {
+    let mut storage = get_voted_height_storage(get_new_storage_config());
+    storage.set_prev_voted_height(BlockNumber(5)).unwrap();
+    storage.revert_height(BlockNumber(0)).unwrap();
+    assert!(storage.get_prev_voted_height().unwrap().is_none());
+}

crates/apollo_consensus/src/test_utils.rs

@@ -123,4 +123,7 @@ impl HeightVotedStorageTrait for NoOpHeightVotedStorage {
     ) -> Result<(), HeightVotedStorageError> {
         Ok(())
     }
+    fn revert_height(&mut self, _height: BlockNumber) -> Result<(), HeightVotedStorageError> {
+        Ok(())
+    }
 }

crates/apollo_storage/src/consensus.rs

@@ -61,6 +61,7 @@ where
     Self: Sized,
 {
     fn set_last_voted_marker(self, last_voted_marker: &LastVotedMarker) -> StorageResult<Self>;
+    fn clear_last_voted_marker(self) -> StorageResult<Self>;
 }
 
 impl<Mode: TransactionKind> ConsensusStorageReader for StorageTxn<'_, Mode> {
@@ -76,4 +77,10 @@ impl ConsensusStorageWriter for StorageTxn<'_, RW> {
         table.upsert(&self.txn, &(), last_voted_marker)?;
         Ok(self)
     }
+
+    fn clear_last_voted_marker(self) -> StorageResult<Self> {
+        let table = self.open_table(&self.tables.last_voted_marker)?;
+        table.delete(&self.txn, &())?;
+        Ok(self)
+    }
 }