make replace_account methods fallible

solana-labs · Oct 3, 2023 · 56299f4 · 56299f4
1 parent 358d2e0
commit 56299f4
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 58 deletions.
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
@@ -8058,12 +8058,21 @@ impl Bank {
         }
 
         if new_feature_activations.contains(&feature_set::programify_feature_gate_program::id()) {
-            replace_account::replace_empty_account_with_upgradeable_program(
+            let datapoint_name = "bank-progamify_feature_gate_program";
+            if let Err(e) = replace_account::replace_empty_account_with_upgradeable_program(
                 self,
                 &feature::id(),
                 &inline_feature_gate_program::noop_program::id(),
-                "bank-apply_feature_gate_program",
-            );
+                datapoint_name,
+            ) {
+                warn!(
+                    "{}: Failed to replace empty account {} with upgradeable program: {}",
+                    datapoint_name,
+                    feature::id(),
+                    e
+                );
+                datapoint_warn!(datapoint_name, ("slot", self.slot(), i64),);
+            }
         }
     }
 

diff --git a/runtime/src/bank/replace_account.rs b/runtime/src/bank/replace_account.rs
@@ -8,8 +8,32 @@ use {
         pubkey::Pubkey,
     },
     std::sync::atomic::Ordering::Relaxed,
+    thiserror::Error,
 };
 
+/// Errors returned by `replace_account` methods
+#[derive(Debug, Error, PartialEq)]
+pub enum ReplaceAccountError {
+    /// Source account not found
+    #[error("Source account not found")]
+    SourceAccountNotFound,
+    /// Source data account not found
+    #[error("Source data account not found")]
+    SourceDataAccountNotFound,
+    /// Destination account not found
+    #[error("Destination account not found")]
+    DestinationAccountNotFound,
+    /// Destination account exists
+    #[error("Destination account exists")]
+    DestinationAccountExists,
+    /// Unable to serialize program account state
+    #[error("Unable to serialize program account state")]
+    UnableToSerializeProgramAccountState,
+    /// Not an upgradeable program
+    #[error("Not an upgradeable program")]
+    NotAnUpgradeableProgram,
+}
+
 /// Moves one account in place of another
 /// `source`: the account to replace with
 /// `destination`: the account to be replaced
@@ -55,7 +79,7 @@ pub(crate) fn replace_non_upgradeable_program_account(
     source_address: &Pubkey,
     destination_address: &Pubkey,
     datapoint_name: &'static str,
-) {
+) -> Result<(), ReplaceAccountError> {
     if let Some(destination_account) = bank.get_account_with_fixed_root(destination_address) {
         if let Some(source_account) = bank.get_account_with_fixed_root(source_address) {
             datapoint_info!(datapoint_name, ("slot", bank.slot, i64));
@@ -73,16 +97,13 @@ pub(crate) fn replace_non_upgradeable_program_account(
                 .write()
                 .unwrap()
                 .remove_programs([*destination_address].into_iter());
+
+            Ok(())
         } else {
-            warn!(
-                "Unable to find source program {}. Destination: {}",
-                source_address, destination_address
-            );
-            datapoint_warn!(datapoint_name, ("slot", bank.slot(), i64),);
+            Err(ReplaceAccountError::SourceAccountNotFound)
         }
     } else {
-        warn!("Unable to find destination program {}", destination_address);
-        datapoint_warn!(datapoint_name, ("slot", bank.slot(), i64),);
+        Err(ReplaceAccountError::DestinationAccountNotFound)
     }
 }
 
@@ -97,7 +118,7 @@ pub(crate) fn replace_empty_account_with_upgradeable_program(
     source_address: &Pubkey,
     destination_address: &Pubkey,
     datapoint_name: &'static str,
-) {
+) -> Result<(), ReplaceAccountError> {
     // Must be attempting to replace an empty account with a program
     // account _and_ data account
     if let Some(source_account) = bank.get_account_with_fixed_root(source_address) {
@@ -171,52 +192,24 @@ pub(crate) fn replace_empty_account_with_upgradeable_program(
                                 // Any remaining lamports in the source program account are burnt
                                 bank.capitalization
                                     .fetch_sub(change_in_capitalization, Relaxed);
+
+                                Ok(())
                             } else {
-                                error!(
-                                    "Unable to serialize program account state. \
-                                    Source program: {} Source data account: {} \
-                                    Destination program: {} Destination data account: {}",
-                                    source_address,
-                                    source_data_address,
-                                    destination_address,
-                                    destination_data_address,
-                                );
-                                datapoint_error!(datapoint_name, ("slot", bank.slot(), i64),);
+                                Err(ReplaceAccountError::UnableToSerializeProgramAccountState)
                             }
+                        } else {
+                            Err(ReplaceAccountError::DestinationAccountExists)
                         }
                     } else {
-                        warn!(
-                            "Unable to find data account for source program. \
-                            Source program: {} Source data account: {} \
-                            Destination program: {} Destination data account: {}",
-                            source_address,
-                            source_data_address,
-                            destination_address,
-                            destination_data_address,
-                        );
-                        datapoint_warn!(datapoint_name, ("slot", bank.slot(), i64),);
+                        Err(ReplaceAccountError::SourceDataAccountNotFound)
                     }
                 }
-                _ => {
-                    warn!(
-                        "Source program account is not an upgradeable program. \
-                        Source program: {} Source data account: {} \
-                        Destination program: {} Destination data account: {}",
-                        source_address,
-                        source_data_address,
-                        destination_address,
-                        destination_data_address,
-                    );
-                    datapoint_warn!(datapoint_name, ("slot", bank.slot(), i64),);
-                    return;
-                }
+                _ => Err(ReplaceAccountError::NotAnUpgradeableProgram),
             }
+        } else {
+            Err(ReplaceAccountError::NotAnUpgradeableProgram)
         }
     } else {
-        warn!(
-            "Unable to find source program {}. Destination: {}",
-            source_address, destination_address
-        );
-        datapoint_warn!(datapoint_name, ("slot", bank.slot(), i64),);
+        Err(ReplaceAccountError::SourceAccountNotFound)
     }
 }
diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
@@ -9,7 +9,8 @@ use {
     crate::{
         accounts_background_service::{PrunedBanksRequestHandler, SendDroppedBankCallback},
         bank::replace_account::{
-            replace_empty_account_with_upgradeable_program, replace_non_upgradeable_program_account,
+            replace_empty_account_with_upgradeable_program,
+            replace_non_upgradeable_program_account, ReplaceAccountError,
         },
         bank_client::BankClient,
         epoch_rewards_hasher::hash_rewards_into_partitions,
@@ -8071,7 +8072,8 @@ fn test_replace_non_upgradeable_program_account() {
         &source,
         &destination,
         "bank-apply_program_replacement",
-    );
+    )
+    .unwrap();
 
     // Destination program account balance is now the source program account's balance
     assert_eq!(bank.get_balance(&destination), source_lamports);
@@ -8188,7 +8190,8 @@ fn test_replace_empty_account_with_upgradeable_program_success(
         &source,
         &destination,
         "bank-apply_empty_account_replacement_for_program",
-    );
+    )
+    .unwrap();
 
     // Destination program account was created and funded to pay for minimum rent
     // for the PDA
@@ -8314,11 +8317,15 @@ fn test_replace_empty_account_with_upgradeable_program_fail_when_account_exists(
     let original_capitalization = bank.capitalization();
 
     // Attempt the replacement
-    replace_empty_account_with_upgradeable_program(
-        &mut bank,
-        &source,
-        &destination,
-        "bank-apply_empty_account_replacement_for_program",
+    assert_eq!(
+        replace_empty_account_with_upgradeable_program(
+            &mut bank,
+            &source,
+            &destination,
+            "bank-apply_empty_account_replacement_for_program",
+        )
+        .unwrap_err(),
+        ReplaceAccountError::DestinationAccountExists
     );
 
     // Everything should be unchanged