tlv-account-resolution: resolve extra metas without per-meta allocation

tlv-account-resolution/src/state.rs

@@ -214,6 +214,9 @@ impl ExtraAccountMetaList {
     ) -> Result<(), ProgramError> {
         let state = TlvStateBorrowed::unpack(data).unwrap();
         let extra_meta_list = ExtraAccountMetaList::unpack_with_tlv_state::<T>(&state)?;
+        if extra_meta_list.is_empty() {
+            return Ok(());
+        }
 
         let initial_accounts_len = account_infos.len() - extra_meta_list.len();
 
@@ -223,25 +226,25 @@ impl ExtraAccountMetaList {
             .map(account_info_to_meta)
             .collect::<Vec<_>>();
 
+        // Create a list of `Ref`s so we can reference account data in the
+        // resolution step. The list is built once, since rebuilding it for
+        // every meta makes cumulative heap usage quadratic under the
+        // never-freeing SBF bump allocator
+        let account_key_data_refs = account_infos
+            .iter()
+            .map(|info| {
+                let key = *info.key;
+                let data = info.try_borrow_data()?;
+                Ok((key, data))
+            })
+            .collect::<Result<Vec<_>, ProgramError>>()?;
+
         for (i, config) in extra_meta_list.iter().enumerate() {
-            let meta = {
-                // Create a list of `Ref`s so we can reference account data in the
-                // resolution step
-                let account_key_data_refs = account_infos
-                    .iter()
-                    .map(|info| {
-                        let key = *info.key;
-                        let data = info.try_borrow_data()?;
-                        Ok((key, data))
-                    })
-                    .collect::<Result<Vec<_>, ProgramError>>()?;
-
-                config.resolve(instruction_data, program_id, |usize| {
-                    account_key_data_refs
-                        .get(usize)
-                        .map(|(pubkey, opt_data)| (pubkey, Some(opt_data.as_ref())))
-                })?
-            };
+            let meta = config.resolve(instruction_data, program_id, |usize| {
+                account_key_data_refs
+                    .get(usize)
+                    .map(|(pubkey, opt_data)| (pubkey, Some(opt_data.as_ref())))
+            })?;
 
             // Ensure the account is in the correct position
             let expected_index = i
@@ -313,41 +316,54 @@ impl ExtraAccountMetaList {
         let state = TlvStateBorrowed::unpack(data)?;
         let bytes = state.get_first_bytes::<T>()?;
         let extra_account_metas = ListView::<ExtraAccountMeta>::unpack(bytes)?;
+        if extra_account_metas.is_empty() {
+            return Ok(());
+        }
+
+        // Create a list of `Ref`s so we can reference account data in the
+        // resolution step. The list is built once and each resolved account
+        // is appended to it, since rebuilding it for every meta makes
+        // cumulative heap usage quadratic under the never-freeing SBF bump
+        // allocator
+        let mut account_key_data_refs = cpi_account_infos
+            .iter()
+            .map(|info| {
+                let key = *info.key;
+                let data = info.try_borrow_data()?;
+                Ok((key, data))
+            })
+            .collect::<Result<Vec<_>, ProgramError>>()?;
+
+        // Collect resolved account infos separately, since
+        // `account_key_data_refs` holds borrows of `cpi_account_infos` until
+        // resolution completes
+        let mut resolved_account_infos = Vec::with_capacity(extra_account_metas.len());
 
         for extra_meta in extra_account_metas.iter() {
-            let mut meta = {
-                // Create a list of `Ref`s so we can reference account data in the
-                // resolution step
-                let account_key_data_refs = cpi_account_infos
-                    .iter()
-                    .map(|info| {
-                        let key = *info.key;
-                        let data = info.try_borrow_data()?;
-                        Ok((key, data))
-                    })
-                    .collect::<Result<Vec<_>, ProgramError>>()?;
-
-                extra_meta.resolve(
-                    &cpi_instruction.data,
-                    &cpi_instruction.program_id,
-                    |usize| {
-                        account_key_data_refs
-                            .get(usize)
-                            .map(|(pubkey, opt_data)| (pubkey, Some(opt_data.as_ref())))
-                    },
-                )?
-            };
+            let mut meta = extra_meta.resolve(
+                &cpi_instruction.data,
+                &cpi_instruction.program_id,
+                |usize| {
+                    account_key_data_refs
+                        .get(usize)
+                        .map(|(pubkey, opt_data)| (pubkey, Some(opt_data.as_ref())))
+                },
+            )?;
             de_escalate_account_meta(&mut meta, &cpi_instruction.accounts);
 
             let account_info = account_infos
                 .iter()
                 .find(|&x| *x.key == meta.pubkey)
-                .ok_or(AccountResolutionError::IncorrectAccount)?
-                .clone();
+                .ok_or(AccountResolutionError::IncorrectAccount)?;
 
+            account_key_data_refs.push((*account_info.key, account_info.try_borrow_data()?));
             cpi_instruction.accounts.push(meta);
-            cpi_account_infos.push(account_info);
+            resolved_account_infos.push(account_info.clone());
         }
+
+        // Release the data borrows before appending the resolved accounts
+        drop(account_key_data_refs);
+        cpi_account_infos.extend(resolved_account_infos);
         Ok(())
     }
 }

tlv-account-resolution/tests/allocation.rs

@@ -0,0 +1,163 @@
+#![allow(clippy::arithmetic_side_effects)]
+
+//! Guards against reintroducing per-meta list rebuilds during account
+//! resolution, which grow heap usage quadratically under Solana's
+//! never-freeing SBF bump allocator.
+
+use {
+    solana_account_info::AccountInfo,
+    solana_instruction::{AccountMeta, Instruction},
+    solana_pubkey::Pubkey,
+    spl_discriminator::{ArrayDiscriminator, SplDiscriminate},
+    spl_tlv_account_resolution::{account::ExtraAccountMeta, state::ExtraAccountMetaList},
+    std::{
+        alloc::{GlobalAlloc, Layout, System},
+        sync::atomic::{AtomicUsize, Ordering},
+    },
+};
+
+struct CountingAllocator;
+
+// Process-global and never decremented: a second `#[test]` in this file would
+// run on a parallel thread and pollute the measured deltas, so keep this file
+// to a single test.
+static ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0);
+
+unsafe impl GlobalAlloc for CountingAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        ALLOCATED_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
+        // SAFETY: forwards the caller's layout unchanged to the system
+        // allocator, adding only an atomic counter update
+        unsafe { System.alloc(layout) }
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // SAFETY: forwards a pointer and layout produced by the matching
+        // `alloc` above to the system allocator
+        unsafe { System.dealloc(ptr, layout) }
+    }
+}
+
+#[global_allocator]
+static GLOBAL: CountingAllocator = CountingAllocator;
+
+struct TestInstruction;
+impl SplDiscriminate for TestInstruction {
+    const SPL_DISCRIMINATOR: ArrayDiscriminator =
+        ArrayDiscriminator::new([1; ArrayDiscriminator::LENGTH]);
+}
+
+fn cpi_resolution_allocated_bytes(extra_meta_count: usize) -> usize {
+    let program_id = Pubkey::new_unique();
+    let owner = Pubkey::new_unique();
+
+    let keys: Vec<Pubkey> = (0..extra_meta_count + 2)
+        .map(|_| Pubkey::new_unique())
+        .collect();
+    let mut lamports: Vec<u64> = vec![0; keys.len()];
+    let mut datas: Vec<Vec<u8>> = vec![Vec::new(); keys.len()];
+
+    let account_infos: Vec<AccountInfo> = keys
+        .iter()
+        .zip(lamports.iter_mut())
+        .zip(datas.iter_mut())
+        .map(|((key, lamports), data)| {
+            AccountInfo::new(key, false, false, lamports, data, &owner, false)
+        })
+        .collect();
+
+    let extra_metas: Vec<ExtraAccountMeta> = keys[2..]
+        .iter()
+        .map(|key| ExtraAccountMeta::from(&AccountMeta::new_readonly(*key, false)))
+        .collect();
+    let mut buffer = vec![0; ExtraAccountMetaList::size_of(extra_metas.len()).unwrap()];
+    ExtraAccountMetaList::init::<TestInstruction>(&mut buffer, &extra_metas).unwrap();
+
+    let instruction_accounts = vec![
+        AccountMeta::new_readonly(keys[0], false),
+        AccountMeta::new_readonly(keys[1], false),
+    ];
+    let mut cpi_instruction = Instruction::new_with_bytes(program_id, &[], instruction_accounts);
+    let mut cpi_account_infos = account_infos[..2].to_vec();
+
+    let before = ALLOCATED_BYTES.load(Ordering::Relaxed);
+    ExtraAccountMetaList::add_to_cpi_instruction::<TestInstruction>(
+        &mut cpi_instruction,
+        &mut cpi_account_infos,
+        &buffer,
+        &account_infos,
+    )
+    .unwrap();
+    let after = ALLOCATED_BYTES.load(Ordering::Relaxed);
+
+    assert_eq!(cpi_instruction.accounts.len(), keys.len());
+    assert_eq!(cpi_account_infos.len(), keys.len());
+
+    after - before
+}
+
+fn check_resolution_allocated_bytes(extra_meta_count: usize) -> usize {
+    let program_id = Pubkey::new_unique();
+    let owner = Pubkey::new_unique();
+
+    let keys: Vec<Pubkey> = (0..extra_meta_count + 2)
+        .map(|_| Pubkey::new_unique())
+        .collect();
+    let mut lamports: Vec<u64> = vec![0; keys.len()];
+    let mut datas: Vec<Vec<u8>> = vec![Vec::new(); keys.len()];
+
+    let account_infos: Vec<AccountInfo> = keys
+        .iter()
+        .zip(lamports.iter_mut())
+        .zip(datas.iter_mut())
+        .map(|((key, lamports), data)| {
+            AccountInfo::new(key, false, false, lamports, data, &owner, false)
+        })
+        .collect();
+
+    let extra_metas: Vec<ExtraAccountMeta> = keys[2..]
+        .iter()
+        .map(|key| ExtraAccountMeta::from(&AccountMeta::new_readonly(*key, false)))
+        .collect();
+    let mut buffer = vec![0; ExtraAccountMetaList::size_of(extra_metas.len()).unwrap()];
+    ExtraAccountMetaList::init::<TestInstruction>(&mut buffer, &extra_metas).unwrap();
+
+    let before = ALLOCATED_BYTES.load(Ordering::Relaxed);
+    ExtraAccountMetaList::check_account_infos::<TestInstruction>(
+        &account_infos,
+        &[],
+        &program_id,
+        &buffer,
+    )
+    .unwrap();
+    let after = ALLOCATED_BYTES.load(Ordering::Relaxed);
+
+    after - before
+}
+
+fn assert_linear_growth(label: &str, small: usize, medium: usize, large: usize) {
+    let first_growth = medium.saturating_sub(small);
+    let second_growth = large.saturating_sub(medium);
+
+    assert!(
+        second_growth <= first_growth.saturating_mul(2).saturating_add(small),
+        "{label} allocation growth accelerated: 8 metas allocated {small} bytes, 16 allocated \
+         {medium}, and 32 allocated {large}; this indicates per-meta buffer rebuilds"
+    );
+}
+
+#[test]
+fn resolution_allocations_scale_linearly_with_meta_count() {
+    assert_linear_growth(
+        "add_to_cpi_instruction",
+        cpi_resolution_allocated_bytes(8),
+        cpi_resolution_allocated_bytes(16),
+        cpi_resolution_allocated_bytes(32),
+    );
+    assert_linear_growth(
+        "check_account_infos",
+        check_resolution_allocated_bytes(8),
+        check_resolution_allocated_bytes(16),
+        check_resolution_allocated_bytes(32),
+    );
+}