[epoch] Revert effects that did not reach finality (MystenLabs#2633)

crates/sui-core/src/authority/authority_store.rs

@@ -1,3 +1,4 @@
+use std::iter;
 // Copyright (c) 2022, Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 use super::*;
@@ -966,6 +967,82 @@ impl<const ALL_OBJ_VER: bool, S: Eq + Serialize + for<'de> Deserialize<'de>>
         Ok(assigned_seq)
     }
 
+    /// This function is called at the end of epoch for each transaction that's
+    /// executed locally on the validator but didn't make to the last checkpoint.
+    /// The effects of the execution is reverted here.
+    /// The following things are reverted:
+    /// 1. Certificate and effects are deleted.
+    /// 2. Latest parent_sync entries for each mutated object are deleted.
+    /// 3. All new object states are deleted.
+    /// 4. owner_index table change is reverted.
+    pub fn revert_state_update(&self, tx_digest: &TransactionDigest) -> SuiResult {
+        let effects = self.get_effects(tx_digest)?;
+        let mut write_batch = self.certificates.batch();
+        write_batch = write_batch.delete_batch(&self.certificates, iter::once(tx_digest))?;
+        write_batch = write_batch.delete_batch(&self.effects, iter::once(tx_digest))?;
+
+        let all_new_refs = effects
+            .mutated
+            .iter()
+            .chain(effects.created.iter())
+            .chain(effects.unwrapped.iter())
+            .map(|(r, _)| r)
+            .chain(effects.deleted.iter())
+            .chain(effects.wrapped.iter());
+        write_batch = write_batch.delete_batch(&self.parent_sync, all_new_refs)?;
+
+        let all_new_object_keys = effects
+            .mutated
+            .iter()
+            .chain(effects.created.iter())
+            .chain(effects.unwrapped.iter())
+            .map(|((id, version, _), _)| ObjectKey(*id, *version));
+        write_batch = write_batch.delete_batch(&self.objects, all_new_object_keys)?;
+
+        // Reverting the change to the owner_index table is most complex.
+        // For each newly created (i.e. created and unwrapped) object, the entry in owner_index
+        // needs to be deleted; for each mutated object, we need to query the object state of
+        // the older version, and then rewrite the entry with the old object info.
+        // TODO: Validators should not need to maintain owner_index.
+        // This is dependent on https://github.com/MystenLabs/sui/issues/2629.
+        let owners_to_delete = effects
+            .created
+            .iter()
+            .chain(effects.unwrapped.iter())
+            .chain(effects.mutated.iter())
+            .map(|((id, _, _), owner)| (*owner, *id));
+        write_batch = write_batch.delete_batch(&self.owner_index, owners_to_delete)?;
+        let mutated_objects = effects
+            .mutated
+            .iter()
+            .map(|(r, _)| r)
+            .chain(effects.deleted.iter())
+            .chain(effects.wrapped.iter())
+            .map(|(id, version, _)| {
+                ObjectKey(
+                    *id,
+                    version
+                        .decrement()
+                        .expect("version revert should never fail"),
+                )
+            });
+        let old_objects = self
+            .objects
+            .multi_get(mutated_objects)?
+            .into_iter()
+            .map(|obj_opt| {
+                let obj = obj_opt.expect("Older object version not found");
+                (
+                    (obj.owner, obj.id()),
+                    ObjectInfo::new(&obj.compute_object_reference(), &obj),
+                )
+            });
+        write_batch = write_batch.insert_batch(&self.owner_index, old_objects)?;
+
+        write_batch.write()?;
+        Ok(())
+    }
+
     /// Returns the last entry we have for this object in the parents_sync index used
     /// to facilitate client and authority sync. In turn the latest entry provides the
     /// latest object_reference, and also the latest transaction that has interacted with

crates/sui-core/src/epoch/reconfiguration.rs

@@ -71,9 +71,12 @@ where
                 next_cp,
                 "finish_epoch_change called when there are still unprocessed transactions",
             );
-            if checkpoints.extra_transactions.iter().next().is_some() {
-                // TODO: Revert any tx that's executed but not in the checkpoint.
+            for (tx_digest, _) in checkpoints.extra_transactions.iter() {
+                self.state
+                    .database
+                    .revert_state_update(&tx_digest.transaction)?;
             }
+            checkpoints.extra_transactions.clear()?;
             // drop checkpoints lock
         } else {
             unreachable!();

crates/sui-core/src/unit_tests/authority_tests.rs

@@ -1491,6 +1491,75 @@ async fn test_transfer_sui_with_amount() {
     );
 }
 
+#[tokio::test]
+async fn test_store_revert_state_update() {
+    // This test checks the correctness of revert_state_update in SuiDataStore.
+    let (sender, sender_key) = get_key_pair();
+    let (recipient, _sender_key) = get_key_pair();
+    let gas_object_id = ObjectID::random();
+    let gas_object = Object::with_id_owner_for_testing(gas_object_id, sender);
+    let gas_object_ref = gas_object.compute_object_reference();
+    let authority_state = init_state_with_objects(vec![gas_object.clone()]).await;
+
+    let tx_data = TransactionData::new_transfer_sui(
+        recipient,
+        sender,
+        None,
+        gas_object.compute_object_reference(),
+        MAX_GAS,
+    );
+    let signature = Signature::new(&tx_data, &sender_key);
+    let transaction = Transaction::new(tx_data, signature);
+
+    let certificate = init_certified_transaction(transaction, &authority_state);
+    let tx_digest = *certificate.digest();
+    authority_state
+        .handle_confirmation_transaction(ConfirmationTransaction { certificate })
+        .await
+        .unwrap();
+
+    authority_state
+        .database
+        .revert_state_update(&tx_digest)
+        .unwrap();
+    assert_eq!(
+        authority_state
+            .database
+            .get_object(&gas_object_id)
+            .unwrap()
+            .unwrap()
+            .owner,
+        Owner::AddressOwner(sender),
+    );
+    assert_eq!(
+        authority_state
+            .database
+            .get_latest_parent_entry(gas_object_id)
+            .unwrap()
+            .unwrap(),
+        (gas_object_ref, TransactionDigest::genesis()),
+    );
+    assert!(authority_state
+        .database
+        .get_owner_objects(Owner::AddressOwner(recipient))
+        .unwrap()
+        .is_empty());
+    assert_eq!(
+        authority_state
+            .database
+            .get_owner_objects(Owner::AddressOwner(sender))
+            .unwrap()
+            .len(),
+        1,
+    );
+    assert!(authority_state
+        .database
+        .get_certified_transaction(&tx_digest)
+        .unwrap()
+        .is_none());
+    assert!(authority_state.database.get_effects(&tx_digest).is_err());
+}
+
 // helpers
 
 #[cfg(test)]