Fix compute budget bump (#21238)

program-runtime/src/instruction_processor.rs

@@ -362,38 +362,35 @@ impl InstructionProcessor {
         invoke_context: &mut dyn InvokeContext,
     ) -> Result<(), InstructionError> {
         let keyed_accounts = invoke_context.get_keyed_accounts()?;
-        if let Ok(root_account) = keyed_account_at_index(keyed_accounts, 0) {
+        let root_account = keyed_account_at_index(keyed_accounts, 0)?;
-            let root_id = root_account.unsigned_key();
+        let root_id = root_account.unsigned_key();
-            let owner_id = &root_account.owner()?;
+        let owner_id = &root_account.owner()?;
-            if solana_sdk::native_loader::check_id(owner_id) {
+        if solana_sdk::native_loader::check_id(owner_id) {
-                for (id, process_instruction) in &self.programs {
+            for (id, process_instruction) in &self.programs {
-                    if id == root_id {
+                if id == root_id {
-                        // Call the builtin program
+                    // Call the builtin program
-                        return process_instruction(
+                    return process_instruction(
-                            1, // root_id to be skipped
+                        1, // root_id to be skipped
-                            instruction_data,
-                            invoke_context,
-                        );
-                    }
-                }
-                if !invoke_context.is_feature_active(&remove_native_loader::id()) {
-                    // Call the program via the native loader
-                    return self.native_loader.process_instruction(
-                        0,
                         instruction_data,
                         invoke_context,
                     );
                 }
-            } else {
+            }
-                for (id, process_instruction) in &self.programs {
+            if !invoke_context.is_feature_active(&remove_native_loader::id()) {
-                    if id == owner_id {
+                // Call the program via the native loader
-                        // Call the program via a builtin loader
+                return self
-                        return process_instruction(
+                    .native_loader
-                            0, // no root_id was provided
+                    .process_instruction(0, instruction_data, invoke_context);
-                            instruction_data,
+            }
-                            invoke_context,
+        } else {
-                        );
+            for (id, process_instruction) in &self.programs {
-                    }
+                if id == owner_id {
+                    // Call the program via a builtin loader
+                    return process_instruction(
+                        0, // no root_id was provided
+                        instruction_data,
+                        invoke_context,
+                    );
                 }
             }
         }

program-runtime/src/invoke_context.rs

@@ -8,8 +8,8 @@ use solana_sdk::{
     account::{AccountSharedData, ReadableAccount},
     compute_budget::ComputeBudget,
     feature_set::{
-        demote_program_write_locks, do_support_realloc, remove_native_loader, tx_wide_compute_cap,
+        demote_program_write_locks, do_support_realloc, neon_evm_compute_budget,
-        FeatureSet,
+        remove_native_loader, requestable_heap_size, tx_wide_compute_cap, FeatureSet,
     },
     hash::Hash,
     ic_logger_msg, ic_msg,
@@ -103,6 +103,7 @@ pub struct ThisInvokeContext<'a> {
     sysvars: &'a [(Pubkey, Vec<u8>)],
     logger: Rc<RefCell<dyn Logger>>,
     compute_budget: ComputeBudget,
+    current_compute_budget: ComputeBudget,
     compute_meter: Rc<RefCell<dyn ComputeMeter>>,
     executors: Rc<RefCell<Executors>>,
     instruction_recorders: Option<&'a [InstructionRecorder]>,
@@ -137,6 +138,7 @@ impl<'a> ThisInvokeContext<'a> {
             programs,
             sysvars,
             logger: ThisLogger::new_ref(log_collector),
+            current_compute_budget: compute_budget,
             compute_budget,
             compute_meter,
             executors,
@@ -195,7 +197,7 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
             return Err(InstructionError::CallDepth);
         }
 
-        if let Some(index_of_program_id) = program_indices.last() {
+        let program_id = if let Some(index_of_program_id) = program_indices.last() {
             let program_id = &self.accounts[*index_of_program_id].0;
             let contains = self
                 .invoke_stack
@@ -210,11 +212,32 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
                 // Reentrancy not allowed unless caller is calling itself
                 return Err(InstructionError::ReentrancyNotAllowed);
             }
-        }
+            *program_id
+        } else {
+            return Err(InstructionError::UnsupportedProgramId);
+        };
 
         if self.invoke_stack.is_empty() {
+            let mut compute_budget = self.compute_budget;
+            if !self.is_feature_active(&tx_wide_compute_cap::id())
+                && self.is_feature_active(&neon_evm_compute_budget::id())
+                && program_id == crate::neon_evm_program::id()
+            {
+                // Bump the compute budget for neon_evm
+                compute_budget.max_units = compute_budget.max_units.max(500_000);
+            }
+            if !self.is_feature_active(&requestable_heap_size::id())
+                && self.is_feature_active(&neon_evm_compute_budget::id())
+                && program_id == crate::neon_evm_program::id()
+            {
+                // Bump the compute budget for neon_evm
+                compute_budget.heap_size = Some(256_usize.saturating_mul(1024));
+            }
+            self.current_compute_budget = compute_budget;
+
             if !self.feature_set.is_active(&tx_wide_compute_cap::id()) {
-                self.compute_meter = ThisComputeMeter::new_ref(self.compute_budget.max_units);
+                self.compute_meter =
+                    ThisComputeMeter::new_ref(self.current_compute_budget.max_units);
             }
 
             self.pre_accounts = Vec::with_capacity(instruction.accounts.len());
@@ -484,7 +507,7 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
         self.sysvars
     }
     fn get_compute_budget(&self) -> &ComputeBudget {
-        &self.compute_budget
+        &self.current_compute_budget
     }
     fn set_blockhash(&mut self, hash: Hash) {
         self.blockhash = hash;
@@ -1087,4 +1110,77 @@ mod tests {
             invoke_context.pop();
         }
     }
+
+    #[test]
+    fn test_invoke_context_compute_budget() {
+        let accounts = vec![
+            (
+                solana_sdk::pubkey::new_rand(),
+                Rc::new(RefCell::new(AccountSharedData::default())),
+            ),
+            (
+                crate::neon_evm_program::id(),
+                Rc::new(RefCell::new(AccountSharedData::default())),
+            ),
+        ];
+
+        let noop_message = Message::new(
+            &[Instruction::new_with_bincode(
+                accounts[0].0,
+                &MockInstruction::NoopSuccess,
+                vec![AccountMeta::new_readonly(accounts[0].0, false)],
+            )],
+            None,
+        );
+        let neon_message = Message::new(
+            &[Instruction::new_with_bincode(
+                crate::neon_evm_program::id(),
+                &MockInstruction::NoopSuccess,
+                vec![AccountMeta::new_readonly(accounts[0].0, false)],
+            )],
+            None,
+        );
+
+        let mut feature_set = FeatureSet::all_enabled();
+        feature_set.deactivate(&tx_wide_compute_cap::id());
+        feature_set.deactivate(&requestable_heap_size::id());
+        let mut invoke_context = ThisInvokeContext::new_mock_with_sysvars_and_features(
+            &accounts,
+            &[],
+            &[],
+            Arc::new(feature_set),
+        );
+
+        invoke_context
+            .push(&noop_message, &noop_message.instructions[0], &[0], None)
+            .unwrap();
+        assert_eq!(
+            *invoke_context.get_compute_budget(),
+            ComputeBudget::default()
+        );
+        invoke_context.pop();
+
+        invoke_context
+            .push(&neon_message, &neon_message.instructions[0], &[1], None)
+            .unwrap();
+        let expected_compute_budget = ComputeBudget {
+            max_units: 500_000,
+            heap_size: Some(256_usize.saturating_mul(1024)),
+            ..ComputeBudget::default()
+        };
+        assert_eq!(
+            *invoke_context.get_compute_budget(),
+            expected_compute_budget
+        );
+        invoke_context.pop();
+
+        invoke_context
+            .push(&noop_message, &noop_message.instructions[0], &[0], None)
+            .unwrap();
+        assert_eq!(
+            *invoke_context.get_compute_budget(),
+            ComputeBudget::default()
+        );
+        invoke_context.pop();
+    }
 }

program-runtime/src/lib.rs

@@ -5,3 +5,4 @@ pub mod instruction_recorder;
 pub mod invoke_context;
 pub mod log_collector;
 pub mod native_loader;
+pub mod neon_evm_program;

runtime/src/lib.rs

@@ -34,7 +34,6 @@ pub mod in_mem_accounts_index;
 pub mod inline_spl_token_v2_0;
 pub mod loader_utils;
 pub mod message_processor;
-pub mod neon_evm_program;
 pub mod non_circulating_supply;
 mod nonce_keyed_account;
 mod pubkey_bins;

runtime/src/message_processor.rs

@@ -9,10 +9,7 @@ use solana_program_runtime::{
 use solana_sdk::{
     account::{AccountSharedData, WritableAccount},
     compute_budget::ComputeBudget,
-    feature_set::{
+    feature_set::{prevent_calling_precompiles_as_programs, FeatureSet},
-        neon_evm_compute_budget, prevent_calling_precompiles_as_programs, requestable_heap_size,
-        tx_wide_compute_cap, FeatureSet,
-    },
     hash::Hash,
     message::Message,
     precompiles::is_precompile,
@@ -107,22 +104,6 @@ impl MessageProcessor {
                 }
             }
 
-            let mut compute_budget = compute_budget;
-            if !invoke_context.is_feature_active(&tx_wide_compute_cap::id())
-                && invoke_context.is_feature_active(&neon_evm_compute_budget::id())
-                && *program_id == crate::neon_evm_program::id()
-            {
-                // Bump the compute budget for neon_evm
-                compute_budget.max_units = compute_budget.max_units.max(500_000);
-            }
-            if !invoke_context.is_feature_active(&requestable_heap_size::id())
-                && invoke_context.is_feature_active(&neon_evm_compute_budget::id())
-                && *program_id == crate::neon_evm_program::id()
-            {
-                // Bump the compute budget for neon_evm
-                compute_budget.heap_size = Some(256 * 1024);
-            }
-
             invoke_context.set_instruction_index(instruction_index);
             let result = invoke_context
                 .push(message, instruction, program_indices, None)

sdk/src/feature_set.rs

@@ -373,6 +373,18 @@ impl FeatureSet {
             inactive: HashSet::new(),
         }
     }
+
+    /// Activate a feature
+    pub fn activate(&mut self, feature_id: &Pubkey, slot: u64) {
+        self.inactive.remove(feature_id);
+        self.active.insert(*feature_id, slot);
+    }
+
+    /// Deactivate a feature
+    pub fn deactivate(&mut self, feature_id: &Pubkey) {
+        self.active.remove(feature_id);
+        self.inactive.insert(*feature_id);
+    }
 }
 
 #[cfg(test)]
@@ -433,4 +445,16 @@ mod test {
                 .collect()
         );
     }
+
+    #[test]
+    fn test_feature_set_activate_deactivate() {
+        let mut feature_set = FeatureSet::default();
+
+        let feature = Pubkey::new_unique();
+        assert!(!feature_set.is_active(&feature));
+        feature_set.activate(&feature, 0);
+        assert!(feature_set.is_active(&feature));
+        feature_set.deactivate(&feature);
+        assert!(!feature_set.is_active(&feature));
+    }
 }