fix: fragment large allocations (#534)

this allows to efficiently reuse large free allocations, resulting from
realloc/remove of large accounts, without moving accountsdb storage
pointer to the end of the file

<!-- greptile_comment -->

## Greptile Summary

Updated On: 2025-09-09 14:49:25 UTC

This PR implements allocation splitting functionality for the accounts
database to improve memory efficiency. The changes allow large freed
allocations (from removed or reallocated accounts) to be split into
smaller chunks that can be reused by subsequent accounts, rather than
requiring the entire large allocation to be consumed by a single new
account.

The core implementation is in the `try_recycle_allocation` method in
`magicblock-accounts-db/src/index.rs`, which now checks if a recycled
allocation is larger than needed. When this occurs, it splits the
allocation - using exactly the requested amount and returning the
remainder to the deallocations index for future reuse. This is achieved
by calculating the remainder blocks, updating the offset for the unused
portion, and reinserting it into the recycleable allocations pool.

The changes include comprehensive test coverage with two new test
functions: `test_recycle_allocation_split` in the index tests and
`test_reallocation_split` in the main tests. These tests verify that
large allocations can be properly split and that memory pointers are
correctly managed during the splitting process. Helper methods have been
refactored to support creating accounts with custom sizes for testing
purposes.

This optimization addresses a significant inefficiency in the accounts
database where large freed allocations could only be reused by accounts
of equal or larger size, leading to memory fragmentation and unnecessary
growth of the storage file. The splitting functionality integrates
seamlessly with the existing allocation recycling system while
maintaining proper offset calculations and endianness handling.

## Confidence score: 4/5

- This PR appears safe to merge with careful attention to the memory
management logic
- Score reflects solid implementation with good test coverage, but
complexity in allocation splitting warrants review
- Pay close attention to the allocation splitting logic in
`magicblock-accounts-db/src/index.rs` lines 372-383

<!-- /greptile_comment -->

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

Author: bmuddha

magicblock-accounts-db/src/index.rs

@@ -250,8 +250,8 @@ impl AccountsDbIndex {
         Ok(())
     }
 
-    /// Ensures that current owner of account matches the one recorded in index
-    /// if not, the index cleanup will be performed and new entries inserted to
+    /// Ensures that current owner of account matches the one recorded in index,
+    /// if not, the index cleanup will be performed and new entry inserted to
     /// match the current state
     pub(crate) fn ensure_correct_owner(
         &self,
@@ -355,7 +355,7 @@ impl AccountsDbIndex {
 
     /// Check whether allocation of given size (in blocks) exists.
     /// These are the allocations which are leftovers from
-    /// accounts' reallocations due to their resizing
+    /// accounts' reallocations due to their resizing/removal
     pub(crate) fn try_recycle_allocation(
         &self,
         space: u32,
@@ -369,9 +369,18 @@ impl AccountsDbIndex {
         let (_, val) =
             cursor.get(Some(&space.to_le_bytes()), None, MDB_SET_RANGE_OP)?;
 
-        let (offset, blocks) = bytes!(#unpack, val, u32, u32);
+        let (offset, mut blocks) = bytes!(#unpack, val, u32, u32);
         // delete the allocation record from recycleable list
         cursor.del(WEMPTY)?;
+        // check whether the found allocation contains more space than necessary
+        let remainder = blocks.saturating_sub(space);
+        if remainder > 0 {
+            // split the allocation, to maximize the efficiency of block reuse
+            blocks = space;
+            let new_offset = offset.saturating_add(blocks);
+            let index_value = bytes!(#pack, new_offset, u32, remainder, u32);
+            cursor.put(&remainder.to_le_bytes(), &index_value, WEMPTY)?;
+        }
 
         drop(cursor);
         txn.commit()?;

magicblock-accounts-db/src/index/tests.rs

@@ -256,7 +256,7 @@ fn test_recycle_allocation_after_realloc() {
     assert_eq!(
         tenv.get_delloactions_count(),
         0,
-        "the number of deallocations should have decresed after recycling"
+        "the number of deallocations should have decreased after recycling"
     );
     let result = tenv.try_recycle_allocation(new_allocation.blocks);
     assert!(
@@ -278,7 +278,48 @@ fn test_recycle_allocation_after_realloc() {
     assert_eq!(
         tenv.get_delloactions_count(),
         0,
-        "the number of deallocations should have decresed after recycling"
+        "the number of deallocations should have decreased after recycling"
+    );
+}
+
+#[test]
+fn test_recycle_allocation_split() {
+    let tenv = setup();
+    let IndexAccount {
+        pubkey,
+        owner,
+        allocation,
+    } = tenv.account();
+
+    tenv.insert_account(&pubkey, &owner, allocation)
+        .expect("failed to insert account");
+    tenv.remove_account(&pubkey).unwrap();
+    assert_eq!(
+        tenv.get_delloactions_count(),
+        1,
+        "the number of deallocations should have increased after account removal"
+    );
+
+    let result = tenv
+        .try_recycle_allocation(allocation.blocks / 2)
+        .expect("failed to recycle allocation");
+    assert_eq!(result.blocks, allocation.blocks / 2);
+    assert_eq!(result.offset, allocation.offset);
+    let result = tenv
+        .try_recycle_allocation(allocation.blocks / 2)
+        .expect("failed to recycle allocation");
+    assert_eq!(result.blocks, allocation.blocks / 2);
+    assert!(result.offset > allocation.offset);
+
+    assert_eq!(
+        tenv.get_delloactions_count(),
+        0,
+        "the number of deallocations should have decreased after recycling"
+    );
+    let result = tenv.try_recycle_allocation(allocation.blocks);
+    assert!(
+        matches!(result, Err(AccountsDbError::NotFound)),
+        "deallocations index should have run out of existing allocations"
     );
 }
 

magicblock-accounts-db/src/tests.rs

@@ -541,6 +541,23 @@ fn test_many_insertions_to_accountsdb() {
     });
 }
 
+#[test]
+fn test_reallocation_split() {
+    let tenv = init_test_env();
+    const SIZE: usize = 1024;
+    tenv.account();
+    let account1 = tenv.account_with_size(SIZE * 2);
+    let data_ptr1 = account1.account.data().as_ptr();
+    let account2 = tenv.account_with_size(SIZE);
+    let data_ptr2 = account2.account.data().as_ptr();
+    tenv.remove_account(&account1.pubkey);
+    let account3 = tenv.account_with_size(SIZE / 4);
+    let account4 = tenv.account_with_size(SIZE / 4);
+    assert_eq!(account3.account.data().as_ptr(), data_ptr1);
+    assert!(account4.account.data().as_ptr() < data_ptr2);
+    assert!(account4.account.data().as_ptr() > data_ptr1);
+}
+
 // ==============================================================
 // ==============================================================
 //                      UTILITY CODE BELOW
@@ -580,8 +597,11 @@ fn init_test_env() -> AdbTestEnv {
 
 impl AdbTestEnv {
     fn account(&self) -> AccountWithPubkey {
+        Self::account_with_size(self, SPACE)
+    }
+    fn account_with_size(&self, size: usize) -> AccountWithPubkey {
         let pubkey = Pubkey::new_unique();
-        let mut account = AccountSharedData::new(LAMPORTS, SPACE, &OWNER);
+        let mut account = AccountSharedData::new(LAMPORTS, size, &OWNER);
         account.data_as_mut_slice()[..INIT_DATA_LEN]
             .copy_from_slice(ACCOUNT_DATA);
         self.adb.insert_account(&pubkey, &account);

magicblock-api/src/fund_account.rs

@@ -5,8 +5,12 @@ use magicblock_magic_program_api::{
     self, MAGIC_CONTEXT_PUBKEY, MAGIC_CONTEXT_SIZE,
 };
 use solana_sdk::{
-    account::Account, clock::Epoch, pubkey::Pubkey, signature::Keypair,
-    signer::Signer, system_program,
+    account::{Account, WritableAccount},
+    clock::Epoch,
+    pubkey::Pubkey,
+    signature::Keypair,
+    signer::Signer,
+    system_program,
 };
 
 use crate::{
@@ -15,17 +19,7 @@ use crate::{
 };
 
 pub(crate) fn fund_account(bank: &Bank, pubkey: &Pubkey, lamports: u64) {
-    bank.store_account(
-        *pubkey,
-        Account {
-            lamports,
-            data: vec![],
-            owner: system_program::id(),
-            executable: false,
-            rent_epoch: Epoch::MAX,
-        }
-        .into(),
-    );
+    fund_account_with_data(bank, pubkey, lamports, vec![]);
 }
 
 pub(crate) fn fund_account_with_data(
@@ -34,17 +28,23 @@ pub(crate) fn fund_account_with_data(
     lamports: u64,
     data: Vec<u8>,
 ) {
-    bank.store_account(
-        *pubkey,
-        Account {
-            lamports,
-            data,
-            owner: system_program::id(),
-            executable: false,
-            rent_epoch: Epoch::MAX,
-        }
-        .into(),
-    );
+    if let Some(mut acc) = bank.get_account(pubkey) {
+        acc.set_lamports(lamports);
+        acc.set_data(data);
+        bank.store_account(*pubkey, acc);
+    } else {
+        bank.store_account(
+            *pubkey,
+            Account {
+                lamports,
+                data,
+                owner: system_program::id(),
+                executable: false,
+                rent_epoch: Epoch::MAX,
+            }
+            .into(),
+        );
+    }
 }
 
 pub(crate) fn fund_validator_identity(bank: &Bank, validator_id: &Pubkey) {