Add program heap bump instruction (#20607)

programs/bpf/benches/bpf_loader.rs

@@ -198,6 +198,58 @@ 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(&loader_id, keyed_accounts);
+    invoke_context.compute_meter.remaining = BUDGET;
+
+    // Serialize account data
+    let keyed_accounts = invoke_context.get_keyed_accounts().unwrap();
+    let (mut serialized, account_lengths) = serialize_parameters(
+        &loader_id,
+        &solana_sdk::pubkey::new_rand(),
+        keyed_accounts,
+        &instruction_data,
+    )
+    .unwrap();
+
+    let elf = load_elf("noop").unwrap();
+    let executable = <dyn Executable<BpfError, ThisInstructionMeter>>::from_elf(
+        &elf,
+        None,
+        Config::default(),
+        register_syscalls(&mut invoke_context).unwrap(),
+    )
+    .unwrap();
+
+    bencher.iter(|| {
+        let _ = create_vm(
+            &loader_id,
+            executable.as_ref(),
+            serialized.as_slice_mut(),
+            &mut invoke_context,
+            &account_lengths,
+        )
+        .unwrap();
+    });
+}
+
 #[bench]
 fn bench_instruction_count_tuner(_bencher: &mut Bencher) {
     const BUDGET: u64 = 200_000;

programs/bpf_loader/src/lib.rs

@@ -32,7 +32,8 @@ use solana_sdk::{
     entrypoint::{HEAP_LENGTH, SUCCESS},
     feature_set::{
         add_missing_program_error_mappings, close_upgradeable_program_accounts, do_support_realloc,
-        fix_write_privs, reduce_required_deploy_balance, stop_verify_mul64_imm_nonzero,
+        fix_write_privs, reduce_required_deploy_balance, requestable_heap_size,
+        stop_verify_mul64_imm_nonzero,
     },
     ic_logger_msg, ic_msg,
     instruction::{AccountMeta, InstructionError},
@@ -151,6 +152,13 @@ pub fn create_vm<'a>(
     orig_data_lens: &'a [usize],
 ) -> Result<EbpfVm<'a, BpfError, ThisInstructionMeter>, EbpfError<BpfError>> {
     let compute_budget = invoke_context.get_compute_budget();
+    let heap_size = 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) * compute_budget.heap_cost);
+    }
     let mut heap =
         AlignedMemory::new_with_size(compute_budget.heap_size.unwrap_or(HEAP_LENGTH), HOST_ALIGN);
     let mut vm = EbpfVm::new(program, heap.as_slice_mut(), parameter_bytes)?;

runtime/src/bank.rs

@@ -3809,7 +3809,7 @@ impl Bank {
                     let mut compute_budget = self.compute_budget.unwrap_or_else(ComputeBudget::new);
 
                     let mut process_result = if feature_set.is_active(&tx_wide_compute_cap::id()) {
-                        compute_budget.process_transaction(tx)
+                        compute_budget.process_transaction(tx, feature_set.clone())
                     } else {
                         Ok(())
                     };
@@ -14810,6 +14810,7 @@ pub(crate) mod tests {
                 *compute_budget,
                 ComputeBudget {
                     max_units: 1,
+                    heap_size: Some(48 * 1024),
                     ..ComputeBudget::default()
                 }
             );
@@ -14821,6 +14822,7 @@ pub(crate) mod tests {
         let message = Message::new(
             &[
                 ComputeBudgetInstruction::request_units(1),
+                ComputeBudgetInstruction::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::{
     compute_budget::ComputeBudget,
     feature_set::{
         demote_program_write_locks, do_support_realloc, neon_evm_compute_budget,
-        prevent_calling_precompiles_as_programs, remove_native_loader, tx_wide_compute_cap,
+        prevent_calling_precompiles_as_programs, remove_native_loader, requestable_heap_size,
-        FeatureSet,
+        tx_wide_compute_cap, FeatureSet,
     },
     fee_calculator::FeeCalculator,
     hash::Hash,
@@ -538,11 +538,18 @@ impl MessageProcessor {
             }
 
             let mut compute_budget = compute_budget;
-            if invoke_context.is_feature_active(&neon_evm_compute_budget::id())
+            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);
             }
 

sdk/src/compute_budget.rs

@@ -3,15 +3,19 @@
 use {
     crate::{
         borsh::try_from_slice_unchecked,
+        entrypoint::HEAP_LENGTH as MIN_HEAP_FRAME_BYTES,
+        feature_set::{requestable_heap_size, FeatureSet},
         instruction::{Instruction, InstructionError},
         transaction::{SanitizedTransaction, TransactionError},
     },
     borsh::{BorshDeserialize, BorshSchema, BorshSerialize},
+    std::sync::Arc,
 };
 
 crate::declare_id!("ComputeBudget111111111111111111111111111111");
 
 const MAX_UNITS: u32 = 1_000_000;
+const MAX_HEAP_FRAME_BYTES: u32 = 256 * 1024;
 
 /// Compute Budget Instructions
 #[derive(
@@ -30,12 +34,24 @@ pub enum ComputeBudgetInstruction {
     /// Request a specific maximum number of compute units the transaction is
     /// allowed to consume.
     RequestUnits(u32),
+    /// Request a specific transaction-wide program heap frame size in bytes.
+    /// The value requested must be a multiple of 1024. This new heap frame size
+    /// applies to each program executed, including all calls to CPIs.
+    RequestHeapFrame(u32),
 }
 impl ComputeBudgetInstruction {
     /// Create a `ComputeBudgetInstruction::RequestUnits` `Instruction`
     pub fn request_units(units: u32) -> Instruction {
         Instruction::new_with_borsh(id(), &ComputeBudgetInstruction::RequestUnits(units), vec![])
     }
+    /// Create a `ComputeBudgetInstruction::RequestHeapFrame` `Instruction`
+    pub fn request_heap_frame(bytes: u32) -> Instruction {
+        Instruction::new_with_borsh(
+            id(),
+            &ComputeBudgetInstruction::RequestHeapFrame(bytes),
+            vec![],
+        )
+    }
 }
 
 #[derive(Clone, Copy, Debug, AbiExample, PartialEq)]
@@ -74,6 +90,9 @@ pub struct ComputeBudget {
     pub syscall_base_cost: u64,
     /// Optional program heap region size, if `None` then loader default
     pub heap_size: Option<usize>,
+    /// Number of compute units per additional 32k heap above the default (~.5
+    /// us per 32k at 15 units/us rounded up)
+    pub heap_cost: u64,
 }
 impl Default for ComputeBudget {
     fn default() -> Self {
@@ -99,24 +118,39 @@ impl ComputeBudget {
             secp256k1_recover_cost: 25_000,
             syscall_base_cost: 100,
             heap_size: None,
+            heap_cost: 8,
         }
     }
     pub fn process_transaction(
         &mut self,
         tx: &SanitizedTransaction,
+        feature_set: Arc<FeatureSet>,
     ) -> Result<(), TransactionError> {
         let error = TransactionError::InstructionError(0, InstructionError::InvalidInstructionData);
-        // Compute budget instruction must be in 1st or 2nd instruction (avoid nonce marker)
+        // Compute budget instruction must be in the 1st 3 instructions (avoid
-        for (program_id, instruction) in tx.message().program_instructions_iter().take(2) {
+        // nonce marker), otherwise ignored
+        for (program_id, instruction) in tx.message().program_instructions_iter().take(3) {
             if check_id(program_id) {
-                let ComputeBudgetInstruction::RequestUnits(units) =
+                match try_from_slice_unchecked(&instruction.data) {
-                    try_from_slice_unchecked::<ComputeBudgetInstruction>(&instruction.data)
+                    Ok(ComputeBudgetInstruction::RequestUnits(units)) => {
-                        .map_err(|_| error.clone())?;
                         if units > MAX_UNITS {
                             return Err(error);
                         }
                         self.max_units = units as u64;
                     }
+                    Ok(ComputeBudgetInstruction::RequestHeapFrame(bytes)) => {
+                        if !feature_set.is_active(&requestable_heap_size::id())
+                            || bytes > MAX_HEAP_FRAME_BYTES
+                            || bytes < MIN_HEAP_FRAME_BYTES as u32
+                            || bytes % 1024 != 0
+                        {
+                            return Err(error);
+                        }
+                        self.heap_size = Some(bytes as usize);
+                    }
+                    _ => return Err(error),
+                }
+            }
         }
         Ok(())
     }
@@ -135,77 +169,151 @@ mod tests {
         tx.try_into().unwrap()
     }
 
-    #[test]
+    macro_rules! test {
-    fn test_process_transaction() {
+        ( $instructions: expr, $expected_error: expr, $expected_budget: expr ) => {
             let payer_keypair = Keypair::new();
-        let mut compute_budget = ComputeBudget::default();
-
             let tx = sanitize_tx(Transaction::new(
                 &[&payer_keypair],
-            Message::new(&[], Some(&payer_keypair.pubkey())),
+                Message::new($instructions, Some(&payer_keypair.pubkey())),
                 Hash::default(),
             ));
-        compute_budget.process_transaction(&tx).unwrap();
+            let feature_set = Arc::new(FeatureSet::all_enabled());
-        assert_eq!(compute_budget, ComputeBudget::default());
+            let mut compute_budget = ComputeBudget::default();
+            let result = compute_budget.process_transaction(&tx, feature_set);
+            assert_eq!($expected_error as Result<(), TransactionError>, result);
+            assert_eq!(compute_budget, $expected_budget);
+        };
+    }
 
-        let tx = sanitize_tx(Transaction::new(
+    #[test]
-            &[&payer_keypair],
+    fn test_process_transaction() {
-            Message::new(
+        // Units
+        test!(&[], Ok(()), ComputeBudget::default());
+        test!(
             &[
                 ComputeBudgetInstruction::request_units(1),
                 Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
             ],
-                Some(&payer_keypair.pubkey()),
+            Ok(()),
-            ),
-            Hash::default(),
-        ));
-        compute_budget.process_transaction(&tx).unwrap();
-        assert_eq!(
-            compute_budget,
             ComputeBudget {
                 max_units: 1,
                 ..ComputeBudget::default()
             }
         );
-
+        test!(
-        let tx = sanitize_tx(Transaction::new(
-            &[&payer_keypair],
-            Message::new(
             &[
                 ComputeBudgetInstruction::request_units(MAX_UNITS + 1),
                 Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
             ],
-                Some(&payer_keypair.pubkey()),
-            ),
-            Hash::default(),
-        ));
-        let result = compute_budget.process_transaction(&tx);
-        assert_eq!(
-            result,
             Err(TransactionError::InstructionError(
                 0,
-                InstructionError::InvalidInstructionData
+                InstructionError::InvalidInstructionData,
-            ))
+            )),
+            ComputeBudget::default()
         );
-
+        test!(
-        let tx = sanitize_tx(Transaction::new(
-            &[&payer_keypair],
-            Message::new(
             &[
                 Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
                 ComputeBudgetInstruction::request_units(MAX_UNITS),
             ],
-                Some(&payer_keypair.pubkey()),
+            Ok(()),
-            ),
-            Hash::default(),
-        ));
-        compute_budget.process_transaction(&tx).unwrap();
-        assert_eq!(
-            compute_budget,
             ComputeBudget {
                 max_units: MAX_UNITS as u64,
                 ..ComputeBudget::default()
             }
         );
+        test!(
+            &[
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                ComputeBudgetInstruction::request_units(1),
+            ],
+            Ok(()),
+            ComputeBudget::default()
+        );
+
+        // HeapFrame
+        test!(&[], Ok(()), ComputeBudget::default());
+        test!(
+            &[
+                ComputeBudgetInstruction::request_heap_frame(40 * 1024),
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+            ],
+            Ok(()),
+            ComputeBudget {
+                heap_size: Some(40 * 1024),
+                ..ComputeBudget::default()
+            }
+        );
+        test!(
+            &[
+                ComputeBudgetInstruction::request_heap_frame(40 * 1024 + 1),
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+            ],
+            Err(TransactionError::InstructionError(
+                0,
+                InstructionError::InvalidInstructionData,
+            )),
+            ComputeBudget::default()
+        );
+        test!(
+            &[
+                ComputeBudgetInstruction::request_heap_frame(31 * 1024),
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+            ],
+            Err(TransactionError::InstructionError(
+                0,
+                InstructionError::InvalidInstructionData,
+            )),
+            ComputeBudget::default()
+        );
+        test!(
+            &[
+                ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES + 1),
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+            ],
+            Err(TransactionError::InstructionError(
+                0,
+                InstructionError::InvalidInstructionData,
+            )),
+            ComputeBudget::default()
+        );
+        test!(
+            &[
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
+            ],
+            Ok(()),
+            ComputeBudget {
+                heap_size: Some(MAX_HEAP_FRAME_BYTES as usize),
+                ..ComputeBudget::default()
+            }
+        );
+        test!(
+            &[
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                ComputeBudgetInstruction::request_heap_frame(1), // ignored
+            ],
+            Ok(()),
+            ComputeBudget::default()
+        );
+
+        // Combined
+        test!(
+            &[
+                Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
+                ComputeBudgetInstruction::request_units(MAX_UNITS),
+            ],
+            Ok(()),
+            ComputeBudget {
+                max_units: MAX_UNITS as u64,
+                heap_size: Some(MAX_HEAP_FRAME_BYTES as usize),
+                ..ComputeBudget::default()
+            }
+        );
     }
 }

sdk/src/feature_set.rs

@@ -245,6 +245,10 @@ pub mod turbine_peers_shuffle {
     solana_sdk::declare_id!("4VvpgRD6UsHvkXwpuQhtR5NG1G4esMaExeWuSEpsYRUa");
 }
 
+pub mod requestable_heap_size {
+    solana_sdk::declare_id!("CCu4boMmfLuqcmfTLPHQiUo22ZdUsXjgzPAURYaWt1Bw");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -301,6 +305,7 @@ lazy_static! {
         (remove_native_loader::id(), "Remove support for the native loader"),
         (send_to_tpu_vote_port::id(), "Send votes to the tpu vote port"),
         (turbine_peers_shuffle::id(), "turbine peers shuffle patch"),
+        (requestable_heap_size::id(), "Requestable heap frame size"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()