Replicates `AccountsDataMeter` in `TransactionContext` (#26438)

Replicates AccountsDataMeter in TransactionContext.

cli/src/program.rs

@@ -2074,7 +2074,7 @@ fn read_and_verify_elf(program_location: &str) -> Result<Vec<u8>, Box<dyn std::e
     let mut program_data = Vec::new();
     file.read_to_end(&mut program_data)
         .map_err(|err| format!("Unable to read program file: {}", err))?;
-    let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1);
+    let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1, 0);
     let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
 
     // Verify the program

program-runtime/src/invoke_context.rs

@@ -2,7 +2,7 @@
 use solana_sdk::keyed_account::{create_keyed_accounts_unified, KeyedAccount};
 use {
     crate::{
-        accounts_data_meter::AccountsDataMeter,
+        accounts_data_meter::{AccountsDataMeter, MAX_ACCOUNTS_DATA_LEN},
         compute_budget::ComputeBudget,
         ic_logger_msg, ic_msg,
         log_collector::LogCollector,
@@ -252,7 +252,7 @@ impl<'a> InvokeContext<'a> {
         feature_set: Arc<FeatureSet>,
         blockhash: Hash,
         lamports_per_signature: u64,
-        initial_accounts_data_len: u64,
+        prev_accounts_data_len: u64,
     ) -> Self {
         Self {
             transaction_context,
@@ -265,7 +265,7 @@ impl<'a> InvokeContext<'a> {
             current_compute_budget: compute_budget,
             compute_budget,
             compute_meter: ComputeMeter::new_ref(compute_budget.compute_unit_limit),
-            accounts_data_meter: AccountsDataMeter::new(initial_accounts_data_len),
+            accounts_data_meter: AccountsDataMeter::new(prev_accounts_data_len),
             executors,
             feature_set,
             timings: ExecuteDetailsTimings::default(),
@@ -1152,6 +1152,7 @@ pub fn with_mock_invoke_context<R, F: FnMut(&mut InvokeContext) -> R>(
         preparation.transaction_accounts,
         ComputeBudget::default().max_invoke_depth.saturating_add(1),
         1,
+        MAX_ACCOUNTS_DATA_LEN,
     );
     let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
     invoke_context
@@ -1182,6 +1183,7 @@ pub fn mock_process_instruction(
         preparation.transaction_accounts,
         ComputeBudget::default().max_invoke_depth.saturating_add(1),
         1,
+        MAX_ACCOUNTS_DATA_LEN,
     );
     let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
     if let Some(sysvar_cache) = sysvar_cache_override {
@@ -1226,7 +1228,7 @@ mod tests {
             desired_result: Result<(), InstructionError>,
         },
         Resize {
-            new_len: usize,
+            new_len: u64,
         },
     }
 
@@ -1312,7 +1314,7 @@ mod tests {
                 }
                 MockInstruction::Resize { new_len } => instruction_context
                     .try_borrow_instruction_account(transaction_context, 0)?
-                    .set_data(&vec![0; new_len])
+                    .set_data(&vec![0; new_len as usize])
                     .unwrap(),
             }
         } else {
@@ -1355,7 +1357,7 @@ mod tests {
             });
         }
         let mut transaction_context =
-            TransactionContext::new(accounts, ComputeBudget::default().max_invoke_depth, 1);
+            TransactionContext::new(accounts, ComputeBudget::default().max_invoke_depth, 1, 0);
         let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
 
         // Check call depth increases and has a limit
@@ -1463,7 +1465,7 @@ mod tests {
         let accounts = vec![(solana_sdk::pubkey::new_rand(), AccountSharedData::default())];
         let instruction_accounts = vec![];
         let program_indices = vec![0];
-        let mut transaction_context = TransactionContext::new(accounts, 1, 1);
+        let mut transaction_context = TransactionContext::new(accounts, 1, 1, 0);
         let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         invoke_context
             .push(&instruction_accounts, &program_indices, &[])
@@ -1508,7 +1510,7 @@ mod tests {
                 is_writable: instruction_account_index < 2,
             })
             .collect::<Vec<_>>();
-        let mut transaction_context = TransactionContext::new(accounts, 2, 8);
+        let mut transaction_context = TransactionContext::new(accounts, 2, 8, 0);
         let mut invoke_context =
             InvokeContext::new_mock(&mut transaction_context, builtin_programs);
 
@@ -1640,7 +1642,7 @@ mod tests {
     fn test_invoke_context_compute_budget() {
         let accounts = vec![(solana_sdk::pubkey::new_rand(), AccountSharedData::default())];
 
-        let mut transaction_context = TransactionContext::new(accounts, 1, 3);
+        let mut transaction_context = TransactionContext::new(accounts, 1, 3, 0);
         let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         invoke_context.compute_budget =
             ComputeBudget::new(compute_budget::DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT as u64);
@@ -1658,8 +1660,9 @@ mod tests {
         solana_logger::setup();
 
         let program_key = Pubkey::new_unique();
-        let user_account_data_len = 123;
-        let user_account = AccountSharedData::new(100, user_account_data_len, &program_key);
+        let user_account_data_len = 123u64;
+        let user_account =
+            AccountSharedData::new(100, user_account_data_len as usize, &program_key);
         let dummy_account = AccountSharedData::new(10, 0, &program_key);
         let mut program_account = AccountSharedData::new(500, 500, &native_loader::id());
         program_account.set_executable(true);
@@ -1674,7 +1677,8 @@ mod tests {
             process_instruction: mock_process_instruction,
         }];
 
-        let mut transaction_context = TransactionContext::new(accounts, 1, 3);
+        let mut transaction_context =
+            TransactionContext::new(accounts, 1, 3, user_account_data_len * 2);
         let mut invoke_context =
             InvokeContext::new_mock(&mut transaction_context, &builtin_programs);
 
@@ -1684,7 +1688,13 @@ mod tests {
         invoke_context
             .accounts_data_meter
             .set_maximum(user_account_data_len as u64 * 3);
-        let remaining_account_data_len = invoke_context.accounts_data_meter.remaining() as usize;
+        let remaining_account_data_len = invoke_context
+            .transaction_context
+            .get_total_resize_remaining();
+        assert_eq!(
+            remaining_account_data_len,
+            invoke_context.accounts_data_meter.remaining(),
+        );
 
         let instruction_accounts = [
             InstructionAccount {
@@ -1719,6 +1729,12 @@ mod tests {
 
             assert!(result.is_ok());
             assert_eq!(invoke_context.accounts_data_meter.remaining(), 0);
+            assert_eq!(
+                invoke_context
+                    .transaction_context
+                    .get_total_resize_remaining(),
+                0
+            );
         }
 
         // Test 2: Resize the account to *the same size*, so not consuming any additional size; this must succeed
@@ -1737,6 +1753,12 @@ mod tests {
 
             assert!(result.is_ok());
             assert_eq!(invoke_context.accounts_data_meter.remaining(), 0);
+            assert_eq!(
+                invoke_context
+                    .transaction_context
+                    .get_total_resize_remaining(),
+                0
+            );
         }
 
         // Test 3: Resize the account to exceed the budget; this must fail
@@ -1759,6 +1781,12 @@ mod tests {
                 Err(solana_sdk::instruction::InstructionError::MaxAccountsDataSizeExceeded)
             ));
             assert_eq!(invoke_context.accounts_data_meter.remaining(), 0);
+            assert_eq!(
+                invoke_context
+                    .transaction_context
+                    .get_total_resize_remaining(),
+                0
+            );
         }
     }
 }

programs/bpf_loader/benches/serialization.rs

@@ -101,7 +101,7 @@ fn create_inputs() -> TransactionContext {
             },
         )
         .collect::<Vec<_>>();
-    let mut transaction_context = TransactionContext::new(transaction_accounts, 1, 1);
+    let mut transaction_context = TransactionContext::new(transaction_accounts, 1, 1, 0);
     let instruction_data = vec![1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
     transaction_context
         .push(&[0], &instruction_accounts, &instruction_data, true)

programs/bpf_loader/src/serialization.rs

@@ -454,7 +454,7 @@ mod tests {
             &program_indices,
         );
         let mut transaction_context =
-            TransactionContext::new(preparation.transaction_accounts, 1, 1);
+            TransactionContext::new(preparation.transaction_accounts, 1, 1, 0);
         let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         invoke_context
             .push(

programs/bpf_loader/src/syscalls.rs

@@ -3366,17 +3366,14 @@ mod tests {
          $program_key:ident,
          $loader_key:expr $(,)?) => {
             let $program_key = Pubkey::new_unique();
-            let mut $transaction_context = TransactionContext::new(
-                vec![
-                    (
-                        $loader_key,
-                        AccountSharedData::new(0, 0, &native_loader::id()),
-                    ),
-                    ($program_key, AccountSharedData::new(0, 0, &$loader_key)),
-                ],
-                1,
-                1,
-            );
+            let transaction_accounts = vec![
+                (
+                    $loader_key,
+                    AccountSharedData::new(0, 0, &native_loader::id()),
+                ),
+                ($program_key, AccountSharedData::new(0, 0, &$loader_key)),
+            ];
+            let mut $transaction_context = TransactionContext::new(transaction_accounts, 1, 1, 0);
             let mut $invoke_context = InvokeContext::new_mock(&mut $transaction_context, &[]);
             $invoke_context.push(&[], &[0, 1], &[]).unwrap();
         };

programs/stake/src/stake_state.rs

@@ -2785,6 +2785,7 @@ mod tests {
             )],
             1,
             1,
+            0,
         )
     }
 
@@ -2894,7 +2895,7 @@ mod tests {
 
     #[test]
     fn test_things_can_merge() {
-        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1);
+        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1, 0);
         let invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         let good_stake = Stake {
             credits_observed: 4242,
@@ -2993,7 +2994,7 @@ mod tests {
 
     #[test]
     fn test_metas_can_merge() {
-        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1);
+        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1, 0);
         let invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         // Identical Metas can merge
         assert!(MergeKind::metas_can_merge(
@@ -3140,7 +3141,7 @@ mod tests {
 
     #[test]
     fn test_merge_kind_get_if_mergeable() {
-        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1);
+        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1, 0);
         let invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         let authority_pubkey = Pubkey::new_unique();
         let initial_lamports = 4242424242;
@@ -3379,7 +3380,7 @@ mod tests {
 
     #[test]
     fn test_merge_kind_merge() {
-        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1);
+        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1, 0);
         let invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         let clock = Clock::default();
         let lamports = 424242;

programs/vote/benches/process_vote.rs

@@ -107,7 +107,8 @@ fn bench_process_vote_instruction(
     instruction_data: Vec<u8>,
 ) {
     bencher.iter(|| {
-        let mut transaction_context = TransactionContext::new(transaction_accounts.clone(), 1, 1);
+        let mut transaction_context =
+            TransactionContext::new(transaction_accounts.clone(), 1, 1, 0);
         let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         invoke_context
             .push(&instruction_accounts, &[0], &instruction_data)

rbpf-cli/src/main.rs

@@ -216,7 +216,8 @@ native machine code before execting it in the virtual machine.",
     let program_indices = [0, 1];
     let preparation =
         prepare_mock_invoke_context(transaction_accounts, instruction_accounts, &program_indices);
-    let mut transaction_context = TransactionContext::new(preparation.transaction_accounts, 1, 1);
+    let mut transaction_context =
+        TransactionContext::new(preparation.transaction_accounts, 1, 1, 0);
     let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
     invoke_context
         .push(

runtime/src/bank.rs

@@ -4283,12 +4283,14 @@ impl Bank {
             get_executors_time.as_us()
         );
 
+        let prev_accounts_data_len = self.load_accounts_data_size();
         let mut transaction_accounts = Vec::new();
         std::mem::swap(&mut loaded_transaction.accounts, &mut transaction_accounts);
         let mut transaction_context = TransactionContext::new(
             transaction_accounts,
             compute_budget.max_invoke_depth.saturating_add(1),
             tx.message().instructions().len(),
+            MAX_ACCOUNTS_DATA_LEN.saturating_sub(prev_accounts_data_len),
         );
 
         let pre_account_state_info =
@@ -4319,7 +4321,7 @@ impl Bank {
             &*self.sysvar_cache.read().unwrap(),
             blockhash,
             lamports_per_signature,
-            self.load_accounts_data_size(),
+            prev_accounts_data_len,
             &mut executed_units,
         );
         process_message_time.stop();
@@ -4376,6 +4378,7 @@ impl Bank {
             accounts,
             instruction_trace,
             mut return_data,
+            ..
         } = transaction_context.into();
         loaded_transaction.accounts = accounts;
 
@@ -18574,6 +18577,7 @@ pub(crate) mod tests {
             loaded_txs[0].0.as_ref().unwrap().accounts.clone(),
             compute_budget.max_invoke_depth.saturating_add(1),
             number_of_instructions_at_transaction_level,
+            0,
         );
 
         assert_eq!(

runtime/src/message_processor.rs

@@ -282,7 +282,7 @@ mod tests {
                 create_loadable_account_for_test("mock_system_program"),
             ),
         ];
-        let mut transaction_context = TransactionContext::new(accounts, 1, 3);
+        let mut transaction_context = TransactionContext::new(accounts, 1, 3, 0);
         let program_indices = vec![vec![2]];
         let executors = Rc::new(RefCell::new(Executors::default()));
         let account_keys = transaction_context.get_keys_of_accounts().to_vec();
@@ -502,7 +502,7 @@ mod tests {
                 create_loadable_account_for_test("mock_system_program"),
             ),
         ];
-        let mut transaction_context = TransactionContext::new(accounts, 1, 3);
+        let mut transaction_context = TransactionContext::new(accounts, 1, 3, 0);
         let program_indices = vec![vec![2]];
         let executors = Rc::new(RefCell::new(Executors::default()));
         let account_metas = vec![
@@ -661,7 +661,7 @@ mod tests {
             (secp256k1_program::id(), secp256k1_account),
             (mock_program_id, mock_program_account),
         ];
-        let mut transaction_context = TransactionContext::new(accounts, 1, 2);
+        let mut transaction_context = TransactionContext::new(accounts, 1, 2, 0);
 
         let message = SanitizedMessage::Legacy(Message::new(
             &[

runtime/src/nonce_keyed_account.rs

@@ -344,7 +344,7 @@ mod test {
                     is_writable: true,
                 },
             ];
-            let mut transaction_context = TransactionContext::new(accounts, 1, 2);
+            let mut transaction_context = TransactionContext::new(accounts, 1, 2, 0);
             let mut $invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         };
     }

runtime/src/system_instruction_processor.rs

@@ -786,7 +786,7 @@ mod tests {
 
     #[test]
     fn test_address_create_with_seed_mismatch() {
-        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1);
+        let mut transaction_context = TransactionContext::new(Vec::new(), 1, 1, 0);
         let invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
         let from = Pubkey::new_unique();
         let seed = "dull boy";

sdk/src/transaction_context.rs

@@ -48,6 +48,8 @@ pub struct TransactionContext {
     number_of_instructions_at_transaction_level: usize,
     instruction_trace: InstructionTrace,
     return_data: TransactionReturnData,
+    total_resize_limit: u64,
+    total_resize_delta: RefCell<i64>,
 }
 
 impl TransactionContext {
@@ -56,6 +58,7 @@ impl TransactionContext {
         transaction_accounts: Vec<TransactionAccount>,
         instruction_context_capacity: usize,
         number_of_instructions_at_transaction_level: usize,
+        total_resize_limit: u64,
     ) -> Self {
         let (account_keys, accounts): (Vec<Pubkey>, Vec<RefCell<AccountSharedData>>) =
             transaction_accounts
@@ -70,6 +73,8 @@ impl TransactionContext {
             number_of_instructions_at_transaction_level,
             instruction_trace: Vec::with_capacity(number_of_instructions_at_transaction_level),
             return_data: TransactionReturnData::default(),
+            total_resize_limit,
+            total_resize_delta: RefCell::new(0),
         }
     }
 
@@ -249,6 +254,18 @@ impl TransactionContext {
     pub fn get_instruction_trace(&self) -> &InstructionTrace {
         &self.instruction_trace
     }
+
+    /// Returns (in bytes) how much data can still be allocated
+    pub fn get_total_resize_remaining(&self) -> u64 {
+        let total_resize_delta = *self.total_resize_delta.borrow();
+        if total_resize_delta >= 0 {
+            self.total_resize_limit
+                .saturating_sub(total_resize_delta as u64)
+        } else {
+            self.total_resize_limit
+                .saturating_add(total_resize_delta.saturating_neg() as u64)
+        }
+    }
 }
 
 /// Return data at the end of a transaction
@@ -586,6 +603,9 @@ impl<'a> BorrowedAccount<'a> {
         if data.len() == self.account.data().len() {
             self.account.data_as_mut_slice().copy_from_slice(data);
         } else {
+            let mut total_resize_delta = self.transaction_context.total_resize_delta.borrow_mut();
+            *total_resize_delta = total_resize_delta
+                .saturating_add((data.len() as i64).saturating_sub(self.get_data().len() as i64));
             self.account.set_data_from_slice(data);
         }
         Ok(())
@@ -594,8 +614,11 @@ impl<'a> BorrowedAccount<'a> {
     /// Resizes the account data (transaction wide)
     ///
     /// Fills it with zeros at the end if is extended or truncates at the end otherwise.
-    pub fn set_data_length(&mut self, new_len: usize) -> Result<(), InstructionError> {
-        self.account.data_mut().resize(new_len, 0);
+    pub fn set_data_length(&mut self, new_length: usize) -> Result<(), InstructionError> {
+        let mut total_resize_delta = self.transaction_context.total_resize_delta.borrow_mut();
+        *total_resize_delta = total_resize_delta
+            .saturating_add((new_length as i64).saturating_sub(self.get_data().len() as i64));
+        self.account.data_mut().resize(new_length, 0);
         Ok(())
     }
 
@@ -666,7 +689,9 @@ pub struct ExecutionRecord {
     pub accounts: Vec<TransactionAccount>,
     pub instruction_trace: InstructionTrace,
     pub return_data: TransactionReturnData,
+    pub total_resize_delta: i64,
 }
+
 /// Used by the bank in the runtime to write back the processed accounts and recorded instructions
 impl From<TransactionContext> for ExecutionRecord {
     fn from(context: TransactionContext) -> Self {
@@ -681,6 +706,7 @@ impl From<TransactionContext> for ExecutionRecord {
                 .collect(),
             instruction_trace: context.instruction_trace,
             return_data: context.return_data,
+            total_resize_delta: RefCell::into_inner(context.total_resize_delta),
         }
     }
 }