tx wide compute budget (#18631)

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...