Add program heap bump instruction (backport solana-labs#20607) (solana-labs#20815)

* Add program heap bump instruction (solana-labs#20607)

(cherry picked from commit 5816451)

* nudge

Co-authored-by: Jack May <jack@solana.com>

Author: mergify[bot]

programs/bpf/benches/bpf_loader.rs

@@ -193,6 +193,54 @@ fn bench_program_execute_noop(bencher: &mut Bencher) {
     });
 }
 
+#[bench]
+fn bench_create_vm(bencher: &mut Bencher) {
+    const BUDGET: u64 = 200_000;
+    let loader_id = bpf_loader::id();
+
+    let accounts = [RefCell::new(AccountSharedData::new(
+        1,
+        10000001,
+        &solana_sdk::pubkey::new_rand(),
+    ))];
+    let keys = [solana_sdk::pubkey::new_rand()];
+    let keyed_accounts: Vec<_> = keys
+        .iter()
+        .zip(&accounts)
+        .map(|(key, account)| solana_sdk::keyed_account::KeyedAccount::new(&key, false, &account))
+        .collect();
+    let instruction_data = vec![0u8];
+
+    let mut invoke_context = MockInvokeContext::new(keyed_accounts);
+    invoke_context.compute_meter.remaining = BUDGET;
+
+    // Serialize account data
+    let keyed_accounts = invoke_context.get_keyed_accounts().unwrap();
+    let mut serialized = serialize_parameters(
+        &bpf_loader::id(),
+        &solana_sdk::pubkey::new_rand(),
+        keyed_accounts,
+        &instruction_data,
+    )
+    .unwrap();
+
+    let elf = load_elf("noop").unwrap();
+    let mut executable =
+        <dyn Executable<BpfError, ThisInstructionMeter>>::from_elf(&elf, None, Config::default())
+            .unwrap();
+    executable.set_syscall_registry(register_syscalls(&mut invoke_context).unwrap());
+
+    bencher.iter(|| {
+        let _ = create_vm(
+            &loader_id,
+            executable.as_ref(),
+            serialized.as_slice_mut(),
+            &mut invoke_context,
+        )
+        .unwrap();
+    });
+}
+
 #[bench]
 fn bench_instruction_count_tuner(_bencher: &mut Bencher) {
     const BUDGET: u64 = 200_000;

programs/bpf_loader/src/lib.rs

@@ -34,7 +34,7 @@ use solana_sdk::{
     entrypoint::{HEAP_LENGTH, SUCCESS},
     feature_set::{
         add_missing_program_error_mappings, close_upgradeable_program_accounts, fix_write_privs,
-        reduce_required_deploy_balance, upgradeable_close_instruction,
+        reduce_required_deploy_balance, requestable_heap_size, upgradeable_close_instruction,
     },
     ic_logger_msg, ic_msg,
     instruction::{AccountMeta, InstructionError},
@@ -150,6 +150,12 @@ pub fn create_vm<'a>(
     invoke_context: &'a mut dyn InvokeContext,
 ) -> Result<EbpfVm<'a, BpfError, ThisInstructionMeter>, EbpfError<BpfError>> {
     let bpf_compute_budget = invoke_context.get_bpf_compute_budget();
+    let heap_size = bpf_compute_budget.heap_size.unwrap_or(HEAP_LENGTH);
+    if invoke_context.is_feature_active(&requestable_heap_size::id()) {
+        let _ = invoke_context.get_compute_meter().borrow_mut().consume(
+            (heap_size as u64 / (32 * 1024)).saturating_sub(1) * bpf_compute_budget.heap_cost,
+        );
+    }
     let heap = AlignedMemory::new_with_size(
         bpf_compute_budget.heap_size.unwrap_or(HEAP_LENGTH),
         HOST_ALIGN,

runtime/src/bank.rs

@@ -3677,7 +3677,11 @@ impl Bank {
                         .unwrap_or_else(BpfComputeBudget::new);
 
                     let mut process_result = if feature_set.is_active(&tx_wide_compute_cap::id()) {
-                        compute_budget::process_request(&mut bpf_compute_budget, tx)
+                        compute_budget::process_request(
+                            &mut bpf_compute_budget,
+                            tx,
+                            feature_set.clone(),
+                        )
                     } else {
                         Ok(())
                     };
@@ -14253,6 +14257,7 @@ pub(crate) mod tests {
                 *compute_budget,
                 BpfComputeBudget {
                     max_units: 1,
+                    heap_size: Some(48 * 1024),
                     ..BpfComputeBudget::default()
                 }
             );
@@ -14264,6 +14269,7 @@ pub(crate) mod tests {
         let message = Message::new(
             &[
                 compute_budget::request_units(1),
+                compute_budget::request_heap_frame(48 * 1024),
                 Instruction::new_with_bincode(program_id, &0, vec![]),
             ],
             Some(&mint_keypair.pubkey()),

runtime/src/message_processor.rs

@@ -11,8 +11,8 @@ use solana_sdk::{
     bpf_loader_upgradeable::{self, UpgradeableLoaderState},
     feature_set::{
         demote_program_write_locks, fix_write_privs, instructions_sysvar_enabled,
-        neon_evm_compute_budget, remove_native_loader, tx_wide_compute_cap, updated_verify_policy,
-        FeatureSet,
+        neon_evm_compute_budget, remove_native_loader, requestable_heap_size, tx_wide_compute_cap,
+        updated_verify_policy, FeatureSet,
     },
     ic_logger_msg, ic_msg,
     instruction::{CompiledInstruction, Instruction, InstructionError},
@@ -1227,11 +1227,17 @@ impl MessageProcessor {
         let program_id = instruction.program_id(&message.account_keys);
 
         let mut bpf_compute_budget = bpf_compute_budget;
-        if feature_set.is_active(&neon_evm_compute_budget::id())
+        if !feature_set.is_active(&tx_wide_compute_cap::id())
+            && feature_set.is_active(&neon_evm_compute_budget::id())
             && *program_id == crate::neon_evm_program::id()
         {
             // Bump the compute budget for neon_evm
             bpf_compute_budget.max_units = bpf_compute_budget.max_units.max(500_000);
+        }
+        if !feature_set.is_active(&requestable_heap_size::id())
+            && feature_set.is_active(&neon_evm_compute_budget::id())
+            && *program_id == crate::neon_evm_program::id()
+        {
             bpf_compute_budget.heap_size = Some(256 * 1024);
         }