tx wide compute budget (#18631)

Cargo.lock

@@ -4544,6 +4544,13 @@ dependencies = [
  "url 2.2.2",
 ]
 
+[[package]]
+name = "solana-compute-budget-program"
+version = "1.8.0"
+dependencies = [
+ "solana-sdk",
+]
+
 [[package]]
 name = "solana-config-program"
 version = "1.8.0"
@@ -5450,6 +5457,7 @@ dependencies = [
  "rustc_version 0.4.0",
  "serde",
  "serde_derive",
+ "solana-compute-budget-program",
  "solana-config-program",
  "solana-frozen-abi 1.8.0",
  "solana-frozen-abi-macro 1.8.0",
@@ -5483,6 +5491,8 @@ version = "1.8.0"
 dependencies = [
  "assert_matches",
  "bincode",
+ "borsh 0.9.1",
+ "borsh-derive 0.9.1",
  "bs58 0.4.0",
  "bv",
  "byteorder",

Cargo.toml

@@ -43,6 +43,7 @@ members = [
     "poh-bench",
     "program-test",
     "programs/bpf_loader",
+    "programs/compute-budget",
     "programs/config",
     "programs/failure",
     "programs/noop",

programs/bpf/Cargo.lock

@@ -3048,6 +3048,13 @@ dependencies = [
  "url",
 ]
 
+[[package]]
+name = "solana-compute-budget-program"
+version = "1.8.0"
+dependencies = [
+ "solana-sdk",
+]
+
 [[package]]
 name = "solana-config-program"
 version = "1.8.0"
@@ -3374,6 +3381,7 @@ dependencies = [
  "rustc_version 0.4.0",
  "serde",
  "serde_derive",
+ "solana-compute-budget-program",
  "solana-config-program",
  "solana-frozen-abi 1.8.0",
  "solana-frozen-abi-macro 1.8.0",
@@ -3398,6 +3406,8 @@ version = "1.8.0"
 dependencies = [
  "assert_matches",
  "bincode",
+ "borsh 0.9.1",
+ "borsh-derive 0.9.1",
  "bs58 0.4.0",
  "bv",
  "byteorder 1.4.3",

programs/bpf/tests/programs.rs

@@ -34,6 +34,7 @@ use solana_sdk::{
     bpf_loader, bpf_loader_deprecated, bpf_loader_upgradeable,
     client::SyncClient,
     clock::MAX_PROCESSING_AGE,
+    compute_budget,
     entrypoint::{MAX_PERMITTED_DATA_INCREASE, SUCCESS},
     instruction::{AccountMeta, CompiledInstruction, Instruction, InstructionError},
     keyed_account::KeyedAccount,
@@ -1234,8 +1235,6 @@ fn test_program_bpf_call_depth() {
 
     solana_logger::setup();
 
-    println!("Test program: solana_bpf_rust_call_depth");
-
     let GenesisConfigInfo {
         genesis_config,
         mint_keypair,
@@ -1269,6 +1268,40 @@ fn test_program_bpf_call_depth() {
     assert!(result.is_err());
 }
 
+#[cfg(feature = "bpf_rust")]
+#[test]
+fn test_program_bpf_compute_budget() {
+    solana_logger::setup();
+
+    let GenesisConfigInfo {
+        genesis_config,
+        mint_keypair,
+        ..
+    } = create_genesis_config(50);
+    let mut bank = Bank::new(&genesis_config);
+    let (name, id, entrypoint) = solana_bpf_loader_program!();
+    bank.add_builtin(&name, id, entrypoint);
+    let bank_client = BankClient::new(bank);
+    let program_id = load_bpf_program(
+        &bank_client,
+        &bpf_loader::id(),
+        &mint_keypair,
+        "solana_bpf_rust_noop",
+    );
+    let message = Message::new(
+        &[
+            compute_budget::request_units(1),
+            Instruction::new_with_bincode(program_id, &0, vec![]),
+        ],
+        Some(&mint_keypair.pubkey()),
+    );
+    let result = bank_client.send_and_confirm_message(&[&mint_keypair], message);
+    assert_eq!(
+        result.unwrap_err().unwrap(),
+        TransactionError::InstructionError(1, InstructionError::ProgramFailedToComplete),
+    );
+}
+
 #[test]
 fn assert_instruction_count() {
     solana_logger::setup();

programs/compute-budget/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "solana-compute-budget-program"
+description = "Solana Compute Budget program"
+version = "1.8.0"
+homepage = "https://solana.com/"
+documentation = "https://docs.rs/solana-compute-budget-program"
+repository = "https://github.com/solana-labs/solana"
+authors = ["Solana Maintainers <maintainers@solana.foundation>"]
+license = "Apache-2.0"
+edition = "2018"
+
+[dependencies]
+solana-sdk = { path = "../../sdk", version = "=1.8.0" }
+
+[lib]
+crate-type = ["lib"]
+name = "solana_compute_budget_program"
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]

programs/compute-budget/src/lib.rs

@@ -0,0 +1,12 @@
+use solana_sdk::{
+    instruction::InstructionError, process_instruction::InvokeContext, pubkey::Pubkey,
+};
+
+pub fn process_instruction(
+    _program_id: &Pubkey,
+    _data: &[u8],
+    _invoke_context: &mut dyn InvokeContext,
+) -> Result<(), InstructionError> {
+    // Do nothing, compute budget instructions handled by the runtime
+    Ok(())
+}

runtime/Cargo.toml

@@ -37,6 +37,7 @@ regex = "1.5.4"
 serde = { version = "1.0.126", features = ["rc"] }
 serde_derive = "1.0.103"
 solana-config-program = { path = "../programs/config", version = "=1.8.0" }
+solana-compute-budget-program = { path = "../programs/compute-budget", version = "=1.8.0" }
 solana-frozen-abi = { path = "../frozen-abi", version = "=1.8.0" }
 solana-frozen-abi-macro = { path = "../frozen-abi/macro", version = "=1.8.0" }
 solana-logger = { path = "../logger", version = "=1.8.0" }

runtime/src/bank.rs

@@ -76,23 +76,26 @@ use solana_sdk::{
         INITIAL_RENT_EPOCH, MAX_PROCESSING_AGE, MAX_RECENT_BLOCKHASHES,
         MAX_TRANSACTION_FORWARDING_DELAY, SECONDS_PER_DAY,
     },
+    compute_budget,
     epoch_info::EpochInfo,
     epoch_schedule::EpochSchedule,
     feature,
-    feature_set::{self, FeatureSet},
+    feature_set::{self, tx_wide_compute_cap, FeatureSet},
     fee_calculator::{FeeCalculator, FeeRateGovernor},
     genesis_config::{ClusterType, GenesisConfig},
     hard_forks::HardForks,
     hash::{extend_and_hash, hashv, Hash},
     incinerator,
     inflation::Inflation,
-    instruction::CompiledInstruction,
+    instruction::{CompiledInstruction, InstructionError},
     lamports::LamportsError,
     message::Message,
     native_loader,
     native_token::sol_to_lamports,
     nonce, nonce_account,
-    process_instruction::{BpfComputeBudget, Executor, ProcessInstructionWithContext},
+    process_instruction::{
+        BpfComputeBudget, ComputeMeter, Executor, ProcessInstructionWithContext,
+    },
     program_utils::limited_deserialize,
     pubkey::Pubkey,
     recent_blockhashes_account,
@@ -412,6 +415,28 @@ impl CachedExecutors {
     }
 }
 
+pub struct TxComputeMeter {
+    remaining: u64,
+}
+impl TxComputeMeter {
+    pub fn new(cap: u64) -> Self {
+        Self { remaining: cap }
+    }
+}
+impl ComputeMeter for TxComputeMeter {
+    fn consume(&mut self, amount: u64) -> std::result::Result<(), InstructionError> {
+        let exceeded = self.remaining < amount;
+        self.remaining = self.remaining.saturating_sub(amount);
+        if exceeded {
+            return Err(InstructionError::ComputationalBudgetExceeded);
+        }
+        Ok(())
+    }
+    fn get_remaining(&self) -> u64 {
+        self.remaining
+    }
+}
+
 #[derive(Default, Debug)]
 pub struct BankRc {
     /// where all the Accounts are stored
@@ -3172,76 +3197,96 @@ impl Bank {
             Vec::with_capacity(sanitized_txs.len());
         let mut transaction_log_messages: Vec<Option<Vec<String>>> =
             Vec::with_capacity(sanitized_txs.len());
-        let bpf_compute_budget = self
-            .bpf_compute_budget
-            .unwrap_or_else(BpfComputeBudget::new);
 
         let executed: Vec<TransactionExecutionResult> = loaded_txs
             .iter_mut()
             .zip(sanitized_txs.as_transactions_iter())
             .map(|(accs, tx)| match accs {
                 (Err(e), _nonce_rollback) => {
-                    inner_instructions.push(None);
                     transaction_log_messages.push(None);
+                    inner_instructions.push(None);
                     (Err(e.clone()), None)
                 }
                 (Ok(loaded_transaction), nonce_rollback) => {
+                    let feature_set = self.feature_set.clone();
                     signature_count += u64::from(tx.message().header.num_required_signatures);
-                    let executors = self.get_executors(&tx.message, &loaded_transaction.loaders);
 
-                    let (account_refcells, loader_refcells) = Self::accounts_to_refcells(
+                    let mut bpf_compute_budget = self
-                        &mut loaded_transaction.accounts,
+                        .bpf_compute_budget
-                        &mut loaded_transaction.loaders,
+                        .unwrap_or_else(BpfComputeBudget::new);
-                    );
 
-                    let instruction_recorders = if enable_cpi_recording {
+                    let mut process_result = if feature_set.is_active(&tx_wide_compute_cap::id()) {
-                        let ix_count = tx.message.instructions.len();
+                        compute_budget::process_request(&mut bpf_compute_budget, tx)
-                        let mut recorders = Vec::with_capacity(ix_count);
-                        recorders.resize_with(ix_count, InstructionRecorder::default);
-                        Some(recorders)
                     } else {
-                        None
+                        Ok(())
                     };
 
-                    let log_collector = if enable_log_recording {
-                        Some(Rc::new(LogCollector::default()))
-                    } else {
-                        None
-                    };
-
-                    let mut process_result = self.message_processor.process_message(
-                        tx.message(),
-                        &loader_refcells,
-                        &account_refcells,
-                        &self.rent_collector,
-                        log_collector.clone(),
-                        executors.clone(),
-                        instruction_recorders.as_deref(),
-                        self.feature_set.clone(),
-                        bpf_compute_budget,
-                        &mut timings.details,
-                        self.rc.accounts.clone(),
-                        &self.ancestors,
-                    );
-
-                    transaction_log_messages.push(Self::collect_log_messages(log_collector));
-                    inner_instructions.push(Self::compile_recorded_instructions(
-                        instruction_recorders,
-                        &tx.message,
-                    ));
-
-                    if let Err(e) = Self::refcells_to_accounts(
-                        &mut loaded_transaction.accounts,
-                        &mut loaded_transaction.loaders,
-                        account_refcells,
-                        loader_refcells,
-                    ) {
-                        warn!("Account lifetime mismanagement");
-                        process_result = Err(e);
-                    }
-
                     if process_result.is_ok() {
-                        self.update_executors(executors);
+                        let executors =
+                            self.get_executors(&tx.message, &loaded_transaction.loaders);
+
+                        let (account_refcells, loader_refcells) = Self::accounts_to_refcells(
+                            &mut loaded_transaction.accounts,
+                            &mut loaded_transaction.loaders,
+                        );
+
+                        let instruction_recorders = if enable_cpi_recording {
+                            let ix_count = tx.message.instructions.len();
+                            let mut recorders = Vec::with_capacity(ix_count);
+                            recorders.resize_with(ix_count, InstructionRecorder::default);
+                            Some(recorders)
+                        } else {
+                            None
+                        };
+
+                        let log_collector = if enable_log_recording {
+                            Some(Rc::new(LogCollector::default()))
+                        } else {
+                            None
+                        };
+
+                        let compute_meter = Rc::new(RefCell::new(TxComputeMeter::new(
+                            bpf_compute_budget.max_units,
+                        )));
+
+                        process_result = self.message_processor.process_message(
+                            tx.message(),
+                            &loader_refcells,
+                            &account_refcells,
+                            &self.rent_collector,
+                            log_collector.clone(),
+                            executors.clone(),
+                            instruction_recorders.as_deref(),
+                            feature_set,
+                            bpf_compute_budget,
+                            compute_meter,
+                            &mut timings.details,
+                            self.rc.accounts.clone(),
+                            &self.ancestors,
+                        );
+
+                        transaction_log_messages.push(Self::collect_log_messages(log_collector));
+                        inner_instructions.push(Self::compile_recorded_instructions(
+                            instruction_recorders,
+                            &tx.message,
+                        ));
+
+                        if let Err(e) = Self::refcells_to_accounts(
+                            &mut loaded_transaction.accounts,
+                            &mut loaded_transaction.loaders,
+                            account_refcells,
+                            loader_refcells,
+                        ) {
+                            warn!("Account lifetime mismanagement");
+                            process_result = Err(e);
+                        }
+
+                        if process_result.is_ok() {
+                            self.update_executors(executors);
+                        }
+                    } else {
+                        transaction_log_messages.push(None);
+                        inner_instructions.push(None);
                     }
 
                     let nonce_rollback =
@@ -5490,6 +5535,7 @@ pub(crate) mod tests {
     use solana_sdk::{
         account::Account,
         clock::{DEFAULT_SLOTS_PER_EPOCH, DEFAULT_TICKS_PER_SLOT},
+        compute_budget,
         epoch_schedule::MINIMUM_SLOTS_PER_EPOCH,
         feature::Feature,
         genesis_config::create_genesis_config,
@@ -13708,4 +13754,47 @@ pub(crate) mod tests {
         rent_debits.push(&Pubkey::default(), i64::MAX as u64, 0);
         assert_eq!(rent_debits.0.len(), 2);
     }
+
+    #[test]
+    fn test_compute_request_instruction() {
+        solana_logger::setup();
+        let GenesisConfigInfo {
+            genesis_config,
+            mint_keypair,
+            ..
+        } = create_genesis_config_with_leader(
+            1_000_000_000_000_000,
+            &Pubkey::new_unique(),
+            bootstrap_validator_stake_lamports(),
+        );
+        let mut bank = Bank::new(&genesis_config);
+
+        fn mock_ix_processor(
+            _pubkey: &Pubkey,
+            _data: &[u8],
+            invoke_context: &mut dyn InvokeContext,
+        ) -> std::result::Result<(), InstructionError> {
+            let compute_budget = invoke_context.get_bpf_compute_budget();
+            assert_eq!(
+                *compute_budget,
+                BpfComputeBudget {
+                    max_units: 1,
+                    ..BpfComputeBudget::default()
+                }
+            );
+            Ok(())
+        }
+        let program_id = solana_sdk::pubkey::new_rand();
+        bank.add_builtin("mock_program", program_id, mock_ix_processor);
+
+        let message = Message::new(
+            &[
+                compute_budget::request_units(1),
+                Instruction::new_with_bincode(program_id, &0, vec![]),
+            ],
+            Some(&mint_keypair.pubkey()),
+        );
+        let tx = Transaction::new(&[&mint_keypair], message, bank.last_blockhash());
+        bank.process_transaction(&tx).unwrap();
+    }
 }

runtime/src/builtins.rs

@@ -3,6 +3,7 @@ use crate::{
     system_instruction_processor,
 };
 use solana_sdk::{
+    feature_set,
     instruction::InstructionError,
     process_instruction::{stable_log, InvokeContext, ProcessInstructionWithContext},
     pubkey::Pubkey,
@@ -86,7 +87,15 @@ pub enum ActivationType {
 /// normal child Bank creation.
 /// https://github.com/solana-labs/solana/blob/84b139cc94b5be7c9e0c18c2ad91743231b85a0d/runtime/src/bank.rs#L1723
 fn feature_builtins() -> Vec<(Builtin, Pubkey, ActivationType)> {
-    vec![]
+    vec![(
+        Builtin::new(
+            "compute_budget_program",
+            solana_sdk::compute_budget::id(),
+            solana_compute_budget_program::process_instruction,
+        ),
+        feature_set::tx_wide_compute_cap::id(),
+        ActivationType::NewProgram,
+    )]
 }
 
 pub(crate) fn get() -> Builtins {

runtime/src/message_processor.rs

@@ -10,7 +10,8 @@ use solana_sdk::{
     account_utils::StateMut,
     bpf_loader_upgradeable::{self, UpgradeableLoaderState},
     feature_set::{
-        instructions_sysvar_enabled, neon_evm_compute_budget, updated_verify_policy, FeatureSet,
+        instructions_sysvar_enabled, neon_evm_compute_budget, tx_wide_compute_cap,
+        updated_verify_policy, FeatureSet,
     },
     ic_logger_msg, ic_msg,
     instruction::{CompiledInstruction, Instruction, InstructionError},
@@ -301,6 +302,7 @@ impl<'a> ThisInvokeContext<'a> {
         programs: &'a [(Pubkey, ProcessInstructionWithContext)],
         log_collector: Option<Rc<LogCollector>>,
         bpf_compute_budget: BpfComputeBudget,
+        compute_meter: Rc<RefCell<dyn ComputeMeter>>,
         executors: Rc<RefCell<Executors>>,
         instruction_recorder: Option<InstructionRecorder>,
         feature_set: Arc<FeatureSet>,
@@ -314,6 +316,13 @@ impl<'a> ThisInvokeContext<'a> {
             executable_accounts,
             accounts,
         );
+        let compute_meter = if feature_set.is_active(&tx_wide_compute_cap::id()) {
+            compute_meter
+        } else {
+            Rc::new(RefCell::new(ThisComputeMeter {
+                remaining: bpf_compute_budget.max_units,
+            }))
+        };
         let mut invoke_context = Self {
             invoke_stack: Vec::with_capacity(bpf_compute_budget.max_invoke_depth),
             rent,
@@ -322,9 +331,7 @@ impl<'a> ThisInvokeContext<'a> {
             programs,
             logger: Rc::new(RefCell::new(ThisLogger { log_collector })),
             bpf_compute_budget,
-            compute_meter: Rc::new(RefCell::new(ThisComputeMeter {
+            compute_meter,
-                remaining: bpf_compute_budget.max_units,
-            })),
             executors,
             instruction_recorder,
             feature_set,
@@ -1138,6 +1145,7 @@ impl MessageProcessor {
         instruction_index: usize,
         feature_set: Arc<FeatureSet>,
         bpf_compute_budget: BpfComputeBudget,
+        compute_meter: Rc<RefCell<dyn ComputeMeter>>,
         timings: &mut ExecuteDetailsTimings,
         account_db: Arc<Accounts>,
         ancestors: &Ancestors,
@@ -1164,7 +1172,7 @@ impl MessageProcessor {
             && *program_id == crate::neon_evm_program::id()
         {
             // Bump the compute budget for neon_evm
-            bpf_compute_budget.max_units = 500_000;
+            bpf_compute_budget.max_units = bpf_compute_budget.max_units.max(500_000);
             bpf_compute_budget.heap_size = Some(256 * 1024);
         }
 
@@ -1178,6 +1186,7 @@ impl MessageProcessor {
             &self.programs,
             log_collector,
             bpf_compute_budget,
+            compute_meter,
             executors,
             instruction_recorder,
             feature_set,
@@ -1218,6 +1227,7 @@ impl MessageProcessor {
         instruction_recorders: Option<&[InstructionRecorder]>,
         feature_set: Arc<FeatureSet>,
         bpf_compute_budget: BpfComputeBudget,
+        compute_meter: Rc<RefCell<dyn ComputeMeter>>,
         timings: &mut ExecuteDetailsTimings,
         account_db: Arc<Accounts>,
         ancestors: &Ancestors,
@@ -1240,6 +1250,7 @@ impl MessageProcessor {
                     instruction_index,
                     feature_set.clone(),
                     bpf_compute_budget,
+                    compute_meter.clone(),
                     timings,
                     account_db.clone(),
                     ancestors,
@@ -1266,6 +1277,7 @@ mod tests {
         instruction::{AccountMeta, Instruction, InstructionError},
         message::Message,
         native_loader::create_loadable_account_for_test,
+        process_instruction::MockComputeMeter,
     };
 
     #[test]
@@ -1313,6 +1325,7 @@ mod tests {
             &[],
             None,
             BpfComputeBudget::default(),
+            Rc::new(RefCell::new(MockComputeMeter::default())),
             Rc::new(RefCell::new(Executors::default())),
             None,
             Arc::new(FeatureSet::all_enabled()),
@@ -1926,6 +1939,7 @@ mod tests {
             None,
             Arc::new(FeatureSet::all_enabled()),
             BpfComputeBudget::new(),
+            Rc::new(RefCell::new(MockComputeMeter::default())),
             &mut ExecuteDetailsTimings::default(),
             Arc::new(Accounts::default()),
             &ancestors,
@@ -1953,6 +1967,7 @@ mod tests {
             None,
             Arc::new(FeatureSet::all_enabled()),
             BpfComputeBudget::new(),
+            Rc::new(RefCell::new(MockComputeMeter::default())),
             &mut ExecuteDetailsTimings::default(),
             Arc::new(Accounts::default()),
             &ancestors,
@@ -1984,6 +1999,7 @@ mod tests {
             None,
             Arc::new(FeatureSet::all_enabled()),
             BpfComputeBudget::new(),
+            Rc::new(RefCell::new(MockComputeMeter::default())),
             &mut ExecuteDetailsTimings::default(),
             Arc::new(Accounts::default()),
             &ancestors,
@@ -2107,6 +2123,7 @@ mod tests {
             None,
             Arc::new(FeatureSet::all_enabled()),
             BpfComputeBudget::new(),
+            Rc::new(RefCell::new(MockComputeMeter::default())),
             &mut ExecuteDetailsTimings::default(),
             Arc::new(Accounts::default()),
             &ancestors,
@@ -2138,6 +2155,7 @@ mod tests {
             None,
             Arc::new(FeatureSet::all_enabled()),
             BpfComputeBudget::new(),
+            Rc::new(RefCell::new(MockComputeMeter::default())),
             &mut ExecuteDetailsTimings::default(),
             Arc::new(Accounts::default()),
             &ancestors,
@@ -2167,6 +2185,7 @@ mod tests {
             None,
             Arc::new(FeatureSet::all_enabled()),
             BpfComputeBudget::new(),
+            Rc::new(RefCell::new(MockComputeMeter::default())),
             &mut ExecuteDetailsTimings::default(),
             Arc::new(Accounts::default()),
             &ancestors,
@@ -2268,6 +2287,7 @@ mod tests {
             programs.as_slice(),
             None,
             BpfComputeBudget::default(),
+            Rc::new(RefCell::new(MockComputeMeter::default())),
             Rc::new(RefCell::new(Executors::default())),
             None,
             Arc::new(FeatureSet::all_enabled()),
@@ -2324,6 +2344,7 @@ mod tests {
                 programs.as_slice(),
                 None,
                 BpfComputeBudget::default(),
+                Rc::new(RefCell::new(MockComputeMeter::default())),
                 Rc::new(RefCell::new(Executors::default())),
                 None,
                 Arc::new(FeatureSet::all_enabled()),

sdk/Cargo.toml

@@ -40,6 +40,8 @@ full = [
 [dependencies]
 assert_matches = { version = "1.5.0", optional = true }
 bincode = "1.3.3"
+borsh = "0.9.0"
+borsh-derive = "0.9.0"
 bs58 = "0.4.0"
 bv = { version = "0.11.1", features = ["serde"] }
 byteorder = { version = "1.4.3", optional = true }

sdk/src/compute_budget.rs

@@ -0,0 +1,142 @@
+#![cfg(feature = "full")]
+
+use crate::{
+    process_instruction::BpfComputeBudget,
+    transaction::{Transaction, TransactionError},
+};
+use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
+use solana_sdk::{
+    borsh::try_from_slice_unchecked,
+    instruction::{Instruction, InstructionError},
+};
+
+crate::declare_id!("ComputeBudget111111111111111111111111111111");
+
+const MAX_UNITS: u64 = 1_000_000;
+
+/// Compute Budget Instructions
+#[derive(
+    Serialize,
+    Deserialize,
+    BorshSerialize,
+    BorshDeserialize,
+    BorshSchema,
+    Debug,
+    Clone,
+    PartialEq,
+    AbiExample,
+    AbiEnumVisitor,
+)]
+pub enum ComputeBudgetInstruction {
+    /// Request a specific maximum number of compute units the transaction is
+    /// allowed to consume.
+    RequestUnits(u64),
+}
+
+/// Create a `ComputeBudgetInstruction::RequestUnits` `Instruction`
+pub fn request_units(units: u64) -> Instruction {
+    Instruction::new_with_borsh(id(), &ComputeBudgetInstruction::RequestUnits(units), vec![])
+}
+
+pub fn process_request(
+    compute_budget: &mut BpfComputeBudget,
+    tx: &Transaction,
+) -> Result<(), TransactionError> {
+    let error = TransactionError::InstructionError(0, InstructionError::InvalidInstructionData);
+    // Compute budget instruction must be in 1st or 2nd instruction (avoid nonce marker)
+    for instruction in tx.message().instructions.iter().take(2) {
+        if check_id(instruction.program_id(&tx.message().account_keys)) {
+            let ComputeBudgetInstruction::RequestUnits(units) =
+                try_from_slice_unchecked::<ComputeBudgetInstruction>(&instruction.data)
+                    .map_err(|_| error.clone())?;
+            if units > MAX_UNITS {
+                return Err(error);
+            }
+            compute_budget.max_units = units;
+        }
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        compute_budget, hash::Hash, message::Message, pubkey::Pubkey, signature::Keypair,
+        signer::Signer,
+    };
+
+    #[test]
+    fn test_process_request() {
+        let payer_keypair = Keypair::new();
+        let mut compute_budget = BpfComputeBudget::default();
+
+        let tx = Transaction::new(
+            &[&payer_keypair],
+            Message::new(&[], Some(&payer_keypair.pubkey())),
+            Hash::default(),
+        );
+        process_request(&mut compute_budget, &tx).unwrap();
+        assert_eq!(compute_budget, BpfComputeBudget::default());
+
+        let tx = Transaction::new(
+            &[&payer_keypair],
+            Message::new(
+                &[
+                    compute_budget::request_units(1),
+                    Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                ],
+                Some(&payer_keypair.pubkey()),
+            ),
+            Hash::default(),
+        );
+        process_request(&mut compute_budget, &tx).unwrap();
+        assert_eq!(
+            compute_budget,
+            BpfComputeBudget {
+                max_units: 1,
+                ..BpfComputeBudget::default()
+            }
+        );
+
+        let tx = Transaction::new(
+            &[&payer_keypair],
+            Message::new(
+                &[
+                    compute_budget::request_units(MAX_UNITS + 1),
+                    Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                ],
+                Some(&payer_keypair.pubkey()),
+            ),
+            Hash::default(),
+        );
+        let result = process_request(&mut compute_budget, &tx);
+        assert_eq!(
+            result,
+            Err(TransactionError::InstructionError(
+                0,
+                InstructionError::InvalidInstructionData
+            ))
+        );
+
+        let tx = Transaction::new(
+            &[&payer_keypair],
+            Message::new(
+                &[
+                    Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
+                    compute_budget::request_units(MAX_UNITS),
+                ],
+                Some(&payer_keypair.pubkey()),
+            ),
+            Hash::default(),
+        );
+        process_request(&mut compute_budget, &tx).unwrap();
+        assert_eq!(
+            compute_budget,
+            BpfComputeBudget {
+                max_units: MAX_UNITS,
+                ..BpfComputeBudget::default()
+            }
+        );
+    }
+}

sdk/src/feature_set.rs

@@ -171,6 +171,10 @@ pub mod libsecp256k1_0_5_upgrade_enabled {
     solana_sdk::declare_id!("DhsYfRjxfnh2g7HKJYSzT79r74Afa1wbHkAgHndrA1oy");
 }
 
+pub mod tx_wide_compute_cap {
+    solana_sdk::declare_id!("5ekBxc8itEnPv4NzGJtr8BVVQLNMQuLMNQQj7pHoLNZ9");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -213,6 +217,7 @@ lazy_static! {
         (neon_evm_compute_budget::id(), "bump neon_evm's compute budget"),
         (rent_for_sysvars::id(), "collect rent from accounts owned by sysvars"),
         (libsecp256k1_0_5_upgrade_enabled::id(), "upgrade libsecp256k1 to v0.5.0"),
+        (tx_wide_compute_cap::id(), "Transaction wide compute cap"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()

sdk/src/lib.rs

@@ -14,6 +14,7 @@ pub mod account_utils;
 pub mod builtins;
 pub mod client;
 pub mod commitment_config;
+pub mod compute_budget;
 pub mod derivation_path;
 pub mod deserialize_utils;
 pub mod entrypoint;

sdk/src/process_instruction.rs

@@ -147,7 +147,7 @@ pub fn get_sysvar<T: Sysvar>(
     })
 }
 
-#[derive(Clone, Copy, Debug, AbiExample)]
+#[derive(Clone, Copy, Debug, AbiExample, PartialEq)]
 pub struct BpfComputeBudget {
     /// Number of compute units that an instruction is allowed.  Compute units
     /// are consumed by program execution, resources they use, etc...