add compute budget instruction to set loaded accounts data size limit

docs/src/developing/programming-model/runtime.md

@@ -60,8 +60,9 @@ to.
 As the transaction is processed compute units are consumed by its
 instruction's programs performing operations such as executing SBF instructions,
 calling syscalls, etc... When the transaction consumes its entire budget, or
-exceeds a bound such as attempting a call stack that is too deep, the runtime
-halts the transaction processing and returns an error.
+exceeds a bound such as attempting a call stack that is too deep, or loaded
+account data size exceeds limit, the runtime halts the transaction processing and
+returns an error.
 
 The following operations incur a compute cost:
 
@@ -150,6 +151,21 @@ let instruction = ComputeBudgetInstruction::set_compute_unit_limit(300_000);
 let instruction = ComputeBudgetInstruction::set_compute_unit_price(1);
 ```
 
+### Accounts data size limit
+
+A transaction should request the maximum bytes of accounts data it is
+allowed to load by including a `SetLoadedAccountsDataSizeLimit` instruction, requested
+limit is capped by `MAX_ACCOUNTS_DATA_SIZE_BYTES`. If no
+`SetLoadedAccountsDataSizeLimit` is provided, the transaction is defaulted to
+have limit of `DEFAULT_ACCOUNTS_DATA_SIZE_BYTES`.
+
+The `ComputeBudgetInstruction::set_accounts_data_size_limit` function can be used
+to create this instruction:
+
+```rust
+let instruction = ComputeBudgetInstruction::set_accounts_data_size_limit(100_000);
+```
+
 ## New Features
 
 As Solana evolves, new features or patches may be introduced that changes the

program-runtime/src/compute_budget.rs

@@ -10,6 +10,10 @@ use {
     },
 };
 
+/// The total accounts data a transaction can load is limited to 64MiB to not break
+/// anyone in Mainnet-beta today. It can be set by set_loaded_accounts_data_size_limit instruction
+pub const MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES: usize = 64 * 1024 * 1024;
+
 pub const DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT: u32 = 200_000;
 pub const MAX_COMPUTE_UNIT_LIMIT: u32 = 1_400_000;
 const MAX_HEAP_FRAME_BYTES: u32 = 256 * 1024;
@@ -109,6 +113,9 @@ pub struct ComputeBudget {
     pub alt_bn128_pairing_one_pair_cost_other: u64,
     /// Big integer modular exponentiation cost
     pub big_modular_exponentiation_cost: u64,
+    /// Maximum accounts data size, in bytes, that a transaction is allowed to load; The
+    /// value is capped by MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES to prevent overuse of memory.
+    pub loaded_accounts_data_size_limit: usize,
 }
 
 impl Default for ComputeBudget {
@@ -157,6 +164,7 @@ impl ComputeBudget {
             alt_bn128_pairing_one_pair_cost_first: 36_364,
             alt_bn128_pairing_one_pair_cost_other: 12_121,
             big_modular_exponentiation_cost: 33,
+            loaded_accounts_data_size_limit: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
         }
     }
 
@@ -166,11 +174,13 @@ impl ComputeBudget {
         default_units_per_instruction: bool,
         support_request_units_deprecated: bool,
         enable_request_heap_frame_ix: bool,
+        support_set_loaded_accounts_data_size_limit_ix: bool,
     ) -> Result<PrioritizationFeeDetails, TransactionError> {
         let mut num_non_compute_budget_instructions: usize = 0;
         let mut updated_compute_unit_limit = None;
         let mut requested_heap_size = None;
         let mut prioritization_fee = None;
+        let mut updated_loaded_accounts_data_size_limit = None;
 
         for (i, (program_id, instruction)) in instructions.enumerate() {
             if compute_budget::check_id(program_id) {
@@ -214,6 +224,14 @@ impl ComputeBudget {
                         prioritization_fee =
                             Some(PrioritizationFeeType::ComputeUnitPrice(micro_lamports));
                     }
+                    Ok(ComputeBudgetInstruction::SetLoadedAccountsDataSizeLimit(bytes))
+                        if support_set_loaded_accounts_data_size_limit_ix =>
+                    {
+                        if updated_loaded_accounts_data_size_limit.is_some() {
+                            return Err(duplicate_instruction_error);
+                        }
+                        updated_loaded_accounts_data_size_limit = Some(bytes as usize);
+                    }
                     _ => return Err(invalid_instruction_data_error),
                 }
             } else {
@@ -250,6 +268,10 @@ impl ComputeBudget {
         .unwrap_or(MAX_COMPUTE_UNIT_LIMIT)
         .min(MAX_COMPUTE_UNIT_LIMIT) as u64;
 
+        self.loaded_accounts_data_size_limit = updated_loaded_accounts_data_size_limit
+            .unwrap_or(MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES)
+            .min(MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES);
+
         Ok(prioritization_fee
             .map(|fee_type| PrioritizationFeeDetails::new(fee_type, self.compute_unit_limit))
             .unwrap_or_default())
@@ -272,7 +294,7 @@ mod tests {
     };
 
     macro_rules! test {
-        ( $instructions: expr, $expected_result: expr, $expected_budget: expr, $enable_request_heap_frame_ix: expr ) => {
+        ( $instructions: expr, $expected_result: expr, $expected_budget: expr, $enable_request_heap_frame_ix: expr, $support_set_loaded_accounts_data_size_limit_ix: expr ) => {
             let payer_keypair = Keypair::new();
             let tx = SanitizedTransaction::from_transaction_for_tests(Transaction::new(
                 &[&payer_keypair],
@@ -285,12 +307,19 @@ mod tests {
                 true,
                 false, /*not support request_units_deprecated*/
                 $enable_request_heap_frame_ix,
+                $support_set_loaded_accounts_data_size_limit_ix,
             );
             assert_eq!($expected_result, result);
             assert_eq!(compute_budget, $expected_budget);
         };
         ( $instructions: expr, $expected_result: expr, $expected_budget: expr) => {
-            test!($instructions, $expected_result, $expected_budget, true);
+            test!(
+                $instructions,
+                $expected_result,
+                $expected_budget,
+                true,
+                false
+            );
         };
     }
 
@@ -561,6 +590,7 @@ mod tests {
                 compute_unit_limit: 0,
                 ..ComputeBudget::default()
             },
+            false,
             false
         );
 
@@ -575,6 +605,7 @@ mod tests {
                 InstructionError::InvalidInstructionData
             )),
             ComputeBudget::default(),
+            false,
             false
         );
         test!(
@@ -587,6 +618,7 @@ mod tests {
                 InstructionError::InvalidInstructionData,
             )),
             ComputeBudget::default(),
+            false,
             false
         );
         test!(
@@ -601,6 +633,7 @@ mod tests {
                 InstructionError::InvalidInstructionData,
             )),
             ComputeBudget::default(),
+            false,
             false
         );
         test!(
@@ -615,6 +648,7 @@ mod tests {
                 InstructionError::InvalidInstructionData,
             )),
             ComputeBudget::default(),
+            false,
             false
         );
 
@@ -635,7 +669,170 @@ mod tests {
                 compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT as u64 * 7,
                 ..ComputeBudget::default()
             },
+            false,
             false
         );
     }
+
+    #[test]
+    fn test_process_loaded_accounts_data_size_limit_instruction() {
+        let enable_request_heap_frame_ix: bool = true;
+
+        // Assert for empty instructions, change value of support_set_loaded_accounts_data_size_limit_ix
+        // will not change results, which should all be default
+        for support_set_loaded_accounts_data_size_limit_ix in [true, false] {
+            test!(
+                &[],
+                Ok(PrioritizationFeeDetails::default()),
+                ComputeBudget {
+                    compute_unit_limit: 0,
+                    ..ComputeBudget::default()
+                },
+                enable_request_heap_frame_ix,
+                support_set_loaded_accounts_data_size_limit_ix
+            );
+        }
+
+        // Assert when set_loaded_accounts_data_size_limit presents,
+        // if support_set_loaded_accounts_data_size_limit_ix then
+        //     budget is set with data_size
+        // else
+        //     return InstructionError
+        let data_size: usize = 1;
+        for support_set_loaded_accounts_data_size_limit_ix in [true, false] {
+            let (expected_result, expected_budget) =
+                if support_set_loaded_accounts_data_size_limit_ix {
+                    (
+                        Ok(PrioritizationFeeDetails::default()),
+                        ComputeBudget {
+                            compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT as u64,
+                            loaded_accounts_data_size_limit: data_size,
+                            ..ComputeBudget::default()
+                        },
+                    )
+                } else {
+                    (
+                        Err(TransactionError::InstructionError(
+                            0,
+                            InstructionError::InvalidInstructionData,
+                        )),
+                        ComputeBudget::default(),
+                    )
+                };
+
+            test!(
+                &[
+                    ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size as u32),
+                    Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
+                ],
+                expected_result,
+                expected_budget,
+                enable_request_heap_frame_ix,
+                support_set_loaded_accounts_data_size_limit_ix
+            );
+        }
+
+        // Assert when set_loaded_accounts_data_size_limit presents, with greater than max value
+        // if support_set_loaded_accounts_data_size_limit_ix then
+        //     budget is set to max data size
+        // else
+        //     return InstructionError
+        let data_size: usize = MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES + 1;
+        for support_set_loaded_accounts_data_size_limit_ix in [true, false] {
+            let (expected_result, expected_budget) =
+                if support_set_loaded_accounts_data_size_limit_ix {
+                    (
+                        Ok(PrioritizationFeeDetails::default()),
+                        ComputeBudget {
+                            compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT as u64,
+                            loaded_accounts_data_size_limit: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
+                            ..ComputeBudget::default()
+                        },
+                    )
+                } else {
+                    (
+                        Err(TransactionError::InstructionError(
+                            0,
+                            InstructionError::InvalidInstructionData,
+                        )),
+                        ComputeBudget::default(),
+                    )
+                };
+
+            test!(
+                &[
+                    ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size as u32),
+                    Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
+                ],
+                expected_result,
+                expected_budget,
+                enable_request_heap_frame_ix,
+                support_set_loaded_accounts_data_size_limit_ix
+            );
+        }
+
+        // Assert when set_loaded_accounts_data_size_limit is not presented
+        // if support_set_loaded_accounts_data_size_limit_ix then
+        //     budget is set to default data size
+        // else
+        //     return
+        for support_set_loaded_accounts_data_size_limit_ix in [true, false] {
+            let (expected_result, expected_budget) = (
+                Ok(PrioritizationFeeDetails::default()),
+                ComputeBudget {
+                    compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT as u64,
+                    loaded_accounts_data_size_limit: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
+                    ..ComputeBudget::default()
+                },
+            );
+
+            test!(
+                &[Instruction::new_with_bincode(
+                    Pubkey::new_unique(),
+                    &0_u8,
+                    vec![]
+                ),],
+                expected_result,
+                expected_budget,
+                enable_request_heap_frame_ix,
+                support_set_loaded_accounts_data_size_limit_ix
+            );
+        }
+
+        // Assert when set_loaded_accounts_data_size_limit presents more than once,
+        // if support_set_loaded_accounts_data_size_limit_ix then
+        //     return DuplicateInstruction
+        // else
+        //     return InstructionError
+        let data_size: usize = MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES;
+        for support_set_loaded_accounts_data_size_limit_ix in [true, false] {
+            let (expected_result, expected_budget) =
+                if support_set_loaded_accounts_data_size_limit_ix {
+                    (
+                        Err(TransactionError::DuplicateInstruction(2)),
+                        ComputeBudget::default(),
+                    )
+                } else {
+                    (
+                        Err(TransactionError::InstructionError(
+                            1,
+                            InstructionError::InvalidInstructionData,
+                        )),
+                        ComputeBudget::default(),
+                    )
+                };
+
+            test!(
+                &[
+                    Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
+                    ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size as u32),
+                    ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size as u32),
+                ],
+                expected_result,
+                expected_budget,
+                enable_request_heap_frame_ix,
+                support_set_loaded_accounts_data_size_limit_ix
+            );
+        }
+    }
 }

programs/sbf/tests/programs.rs

@@ -3590,6 +3590,7 @@ fn test_program_fees() {
         false,
         true,
         true,
+        true,
     );
     bank_client
         .send_and_confirm_message(&[&mint_keypair], message)
@@ -3614,6 +3615,7 @@ fn test_program_fees() {
         false,
         true,
         true,
+        true,
     );
     assert!(expected_normal_fee < expected_prioritized_fee);
 

runtime/src/accounts.rs

@@ -31,8 +31,9 @@ use {
         bpf_loader_upgradeable::{self, UpgradeableLoaderState},
         clock::{BankId, Slot},
         feature_set::{
-            self, enable_request_heap_frame_ix, remove_congestion_multiplier_from_fee_calculation,
-            remove_deprecated_request_unit_ix, use_default_units_in_fee_calculation, FeatureSet,
+            self, add_set_tx_loaded_accounts_data_size_instruction, enable_request_heap_frame_ix,
+            remove_congestion_multiplier_from_fee_calculation, remove_deprecated_request_unit_ix,
+            use_default_units_in_fee_calculation, FeatureSet,
         },
         fee::FeeStructure,
         genesis_config::ClusterType,
@@ -646,6 +647,7 @@ impl Accounts {
                             !feature_set.is_active(&remove_deprecated_request_unit_ix::id()),
                             feature_set.is_active(&remove_congestion_multiplier_from_fee_calculation::id()),
                             feature_set.is_active(&enable_request_heap_frame_ix::id()) || self.accounts_db.expected_cluster_type() != ClusterType::MainnetBeta,
+                            feature_set.is_active(&add_set_tx_loaded_accounts_data_size_instruction::id()),
                         )
                     } else {
                         return (Err(TransactionError::BlockhashNotFound), None);
@@ -1673,6 +1675,7 @@ mod tests {
             false,
             true,
             true,
+            true,
         );
         assert_eq!(fee, lamports_per_signature);