fix: adb slot reuse

accountsdb/src/store/mod.rs

@@ -7,6 +7,7 @@ use std::sync::atomic::Ordering::{Acquire, Release};
 
 use solana_account::{
     AccountMode, AccountSharedData, BorrowedAccount, CoWAccount::*, DirtyMarkers, OwnedAccount,
+    StorageUnit,
 };
 use solana_pubkey::Pubkey;
 use tracing::{error, warn};
@@ -18,7 +19,7 @@ use crate::{
     metrics::{self, Operation},
     store::{
         index::{Index, OptRoTxn, OptRwTxn, OwnerIter, read_txn, write_txn},
-        kv::{Offset, OwnerAndOffset},
+        kv::{KeyTail, Offset, OwnerAndOffset},
         mmap::{DatabaseMeta, MappedStorage},
     },
 };
@@ -37,6 +38,18 @@ pub(crate) const VERSION: DatabaseVersion = 1;
 /// Version tag stored in the metadata header.
 pub(crate) type DatabaseVersion = u64;
 
+/// Snapshot of an in-place owned overwrite that is not yet durable.
+struct ReusedOverwrite {
+    /// Live image offset that was overwritten.
+    offset: Offset,
+    /// Full existing span, including headers, taken before serialize.
+    snapshot: Vec<StorageUnit>,
+    /// Account whose owner mapping may have been updated in the open txn.
+    pubkey: Pubkey,
+    /// Owner tag recorded in the index before this overwrite.
+    owner: KeyTail,
+}
+
 /// Persisted store backed by the mmap and LMDB index.
 pub(crate) struct PersistedStore {
     /// Mapped account storage.
@@ -84,8 +97,8 @@ impl PersistedStore {
     ///
     /// Borrowed accounts in authoritative modes are committed in place. Owned
     /// accounts in those modes are serialized into the mmap. Other modes delete
-    /// stale persisted entries. If the LMDB commit fails or database runs out of
-    /// space, the borrowed images are rolled back so in-memory state stays
+    /// stale persisted entries. If a later apply or the LMDB commit fails, borrowed
+    /// images and in-place owned reuse are rolled back so mmap state stays
     /// aligned with the durable index.
     pub(crate) fn upsert<'a, AC>(&self, accounts: AC) -> Result<()>
     where
@@ -94,26 +107,34 @@ impl PersistedStore {
         let mut applied = 0;
         let mut result = Ok(());
         let mut txn = None;
+        let mut reused = Vec::new();
         for entry in accounts.clone() {
-            result = self.apply(entry, &mut txn);
+            result = self.apply(entry, &mut txn, &mut reused);
             if result.is_err() {
                 break;
             }
             applied += 1;
         }
         // Commit once after the batch so the index and mmap stay in sync.
-        if let Some(txn) = txn
-            && result.is_ok()
+        if result.is_ok()
+            && let Some(txn) = txn.take()
         {
-            metrics::accounts(StoreKind::Persisted, self.index.accounts.len(&txn)?);
-            result = txn.commit().map_err(Into::into);
+            match self.index.accounts.len(&txn) {
+                Ok(count) => {
+                    metrics::accounts(StoreKind::Persisted, count);
+                    result = txn.commit().map_err(Into::into);
+                }
+                Err(error) => result = Err(error.into()),
+            }
         }
         if let Err(error) = &result {
             warn!(applied, ?error, "accounts persistence failed; rolling back");
-            // Only borrowed accounts need rollback here: owned inserts never
-            // mutate an existing borrowed image in place.
+            // Borrowed commits and in-place owned reuse write the mmap before
+            // the LMDB commit. Borrowed views undo via sequence; reused owned
+            // images restore the snapshot taken before serialize.
             let processed = accounts.into_iter().take(applied).map(|(_, a)| a);
             Self::rollback(processed);
+            self.restore_reused(&reused, txn.as_mut());
         }
 
         result
@@ -154,7 +175,12 @@ impl PersistedStore {
     }
 
     /// Applies one account state transition to the persisted backend.
-    fn apply<'e>(&'e self, acc: &AccountEntry, txn: OptRwTxn<'_, 'e>) -> Result<()> {
+    fn apply<'e>(
+        &'e self,
+        acc: &AccountEntry,
+        txn: OptRwTxn<'_, 'e>,
+        reused: &mut Vec<ReusedOverwrite>,
+    ) -> Result<()> {
         let (pubkey, account) = acc;
         // An account that has moved to a non-authoritative mode, or has been
         // closed, no longer belongs here, so drop any stale persisted entry.
@@ -169,7 +195,7 @@ impl PersistedStore {
         let markers = account.markers();
         match account.cow() {
             Borrowed(acc) => self.update(pubkey, acc, markers, txn),
-            Owned(acc) => self.insert(pubkey, acc, txn),
+            Owned(acc) => self.insert(pubkey, acc, txn, reused),
         }
     }
 
@@ -188,6 +214,25 @@ impl PersistedStore {
         }
     }
 
+    /// Restores overwritten owned spans and their prior owner mappings.
+    fn restore_reused(&self, reused: &[ReusedOverwrite], txn: Option<&mut heed::RwTxn<'_>>) {
+        for image in reused {
+            let ptr = self.storage.at(image.offset);
+            // SAFETY: `offset` is the live span we overwrote; `snapshot` is that
+            // span including headers, taken before serialize.
+            unsafe {
+                let dest = slice::from_raw_parts_mut(ptr.as_ptr(), image.snapshot.len());
+                dest.copy_from_slice(&image.snapshot);
+            }
+        }
+        let Some(txn) = txn else {
+            return;
+        };
+        for image in reused {
+            let _ = self.index.update_owner(&image.pubkey, image.owner, txn);
+        }
+    }
+
     /// Commits a borrowed image after updating its owner mapping if needed.
     fn update<'e>(
         &'e self,
@@ -210,15 +255,46 @@ impl PersistedStore {
     }
 
     /// Serializes an owned image into mapped storage and records its offset.
+    ///
+    /// Reuses the live slot when the new payload still fits. Otherwise
+    /// allocates a page-aligned span so small resizes do not append a copy.
     fn insert<'e>(
         &'e self,
         pubkey: &Pubkey,
         acc: &OwnedAccount,
         txn: OptRwTxn<'_, 'e>,
+        reused: &mut Vec<ReusedOverwrite>,
     ) -> Result<()> {
         let txn = write_txn(self.index.env(), txn)?;
-        let units = acc.units();
         let owner = acc.owner().into();
+        let live = self.index.accounts.get(txn, pubkey)?;
+        let existing = match live {
+            // SAFETY: live offsets come from the persisted index.
+            Some(data) => unsafe { BorrowedAccount::span(self.storage.at(data.offset)) },
+            None => 0,
+        };
+        let units = acc.units_at_least(existing);
+        if let Some(prior) = live
+            && existing >= acc.units()
+        {
+            let ptr = self.storage.at(prior.offset);
+            // SAFETY: `prior.offset` is the live image; `existing` is its full span.
+            let snapshot =
+                unsafe { slice::from_raw_parts(ptr.as_ptr(), existing as usize) }.to_vec();
+            self.index.update_owner(pubkey, owner, txn)?;
+            // SAFETY: `offset` is the live image and `existing` still fits.
+            unsafe {
+                let buffer = slice::from_raw_parts_mut(ptr.as_ptr(), existing as usize);
+                acc.serialize(buffer, pubkey);
+            }
+            reused.push(ReusedOverwrite {
+                offset: prior.offset,
+                snapshot,
+                pubkey: *pubkey,
+                owner: prior.owner,
+            });
+            return Ok(());
+        }
 
         let (ptr, offset) = if let Some(offset) = self.index.allocate(units, txn)? {
             let ptr = self.storage.at(offset);

accountsdb/src/tests.rs

@@ -759,3 +759,102 @@ fn test_large_accounts_growth_and_defrag() {
         assert!(acc.data().iter().all(|&b| b == *fill));
     }
 }
+
+// Same-size owned writes reuse the live slot instead of appending a copy.
+#[test]
+fn test_same_slot_reuse_when_size_matches() {
+    let (_dir, db) = db();
+    let owner = Pubkey::new_unique();
+    let key = Pubkey::new_unique();
+    store(&db, key, mutable_data(1, vec![1; 64], &owner));
+    let start = cursor(&db);
+    let live = offset(&db, &key);
+
+    store(&db, key, mutable_data(2, vec![2; 64], &owner));
+    assert!(offset(&db, &key) == live);
+    assert_eq!(cursor(&db), start);
+    assert_eq!(reload(&db, &key).data(), vec![2; 64]);
+}
+
+// Incremental +64-byte resizes reuse 32 KiB-rounded slack instead of
+// appending an exact-fit copy on every write.
+#[test]
+fn test_incremental_growth_reuses_slack_span() {
+    let (_dir, db) = db();
+    let owner = Pubkey::new_unique();
+    let key = Pubkey::new_unique();
+    store(&db, key, mutable_data(1, vec![0; 8], &owner));
+    let start = cursor(&db);
+
+    for i in 1..=2_000 {
+        store(&db, key, mutable_data(1, vec![i as u8; 8 + i * 64], &owner));
+    }
+    // Four slack growths at 32 KiB-rounded images (spans ~64, ~128, ~192,
+    // then ~256 KiB plus headers; 81940 units total), not 2,000 exact-fit copies.
+    assert!(
+        cursor(&db) - start < 120_000,
+        "incremental growth advanced the cursor by {}",
+        cursor(&db) - start
+    );
+    assert_eq!(reload(&db, &key).data().len(), 8 + 2_000 * 64);
+}
+
+// In-place reuse writes the mmap before the LMDB commit. A later apply
+// failure must restore the previous span (headers included) and the prior
+// owner index. A new allocation leaves the old image untouched.
+#[test]
+fn test_reused_image_restored_on_batch_failure() {
+    // Larger than the 64 MiB test map once serialized as a double image.
+    const OVERSIZE: usize = 40 << 20;
+    let owner = Pubkey::new_unique();
+    let next = Pubkey::new_unique();
+
+    {
+        let (_dir, db) = db();
+        let key = Pubkey::new_unique();
+        let original = vec![1; 64];
+        store(&db, key, mutable_data(1, original.clone(), &owner));
+        let live = offset(&db, &key);
+
+        let updated = mutable_data(2, vec![2; 64], &next);
+        let extra = Pubkey::new_unique();
+        assert!(
+            db.store(&[(key, updated), (extra, mutable_data(3, vec![0; OVERSIZE], &owner)),])
+                .is_err()
+        );
+
+        assert!(offset(&db, &key) == live);
+        let restored = reload(&db, &key);
+        assert_eq!(restored.lamports(), 1);
+        assert_eq!(restored.data(), original);
+        assert_eq!(restored.owner(), &owner);
+        assert_eq!(program(&db, &owner), vec![key]);
+        assert!(program(&db, &next).is_empty());
+        assert!(!in_persisted(&db, &extra));
+    }
+
+    {
+        let (_dir, db) = db();
+        let key = Pubkey::new_unique();
+        let original = vec![1; 8];
+        store(&db, key, mutable_data(1, original.clone(), &owner));
+        let live = offset(&db, &key);
+
+        let grown = mutable_data(2, vec![2; 64 * 1024], &next);
+        assert!(grown.owned().units() > reload(&db, &key).owned().units());
+        let extra = Pubkey::new_unique();
+        assert!(
+            db.store(&[(key, grown), (extra, mutable_data(3, vec![0; OVERSIZE], &owner)),])
+                .is_err()
+        );
+
+        assert!(offset(&db, &key) == live);
+        let restored = reload(&db, &key);
+        assert_eq!(restored.lamports(), 1);
+        assert_eq!(restored.data(), original);
+        assert_eq!(restored.owner(), &owner);
+        assert_eq!(program(&db, &owner), vec![key]);
+        assert!(program(&db, &next).is_empty());
+        assert!(!in_persisted(&db, &extra));
+    }
+}

solana/account/src/cow/owned.rs

@@ -23,21 +23,40 @@ impl OwnedAccount {
         self.allocation() * 2 + IMAGE_OFFSET as u32
     }
 
+    /// Span that can hold this account, reusing `existing_span` when it fits.
+    ///
+    /// New allocations round each image up to 32 KiB so +64-byte resizes
+    /// overwrite the live slot instead of appending an exact-fit copy.
+    pub fn units_at_least(&self, existing_span: u32) -> u32 {
+        let need = self.units();
+        if existing_span >= need {
+            return existing_span;
+        }
+        if existing_span == 0 {
+            return need;
+        }
+        const IMAGE_SLACK: u32 = (32 * 1024 / crate::STORAGE_UNIT) as u32;
+        let alloc = self.allocation().div_ceil(IMAGE_SLACK) * IMAGE_SLACK;
+        alloc * 2 + IMAGE_OFFSET as u32
+    }
+
     /// Returns the storage units needed for one image, rounded up to alignment.
     fn allocation(&self) -> u32 {
         (STATIC_SIZE + self.data.len()).div_ceil(ALIGNMENT) as u32
     }
 
-    /// Writes the account into a buffer sized by `units`.
+    /// Writes the account into a buffer sized by `units` or a larger slack span.
     ///
     /// # Safety
     ///
-    /// `buf` must be exactly `units()` storage units long.
-    /// `pubkey` is written into the image prefix so borrowed iteration can
-    /// recover the full account key without consulting the index.
+    /// `buf` must be at least `units()` storage units long and have even image
+    /// space after the header. `pubkey` is written into the image prefix so
+    /// borrowed iteration can recover the full account key without consulting
+    /// the index.
     pub unsafe fn serialize(&self, buf: &mut [StorageUnit], pubkey: &Pubkey) {
         let ptr = NonNull::new_unchecked(buf.as_mut_ptr());
-        debug_assert_eq!(self.units() as usize, buf.len());
+        debug_assert!(buf.len() >= self.units() as usize);
+        debug_assert_eq!((buf.len() - IMAGE_OFFSET) % 2, 0);
 
         fn write<U, T: Sized>(ptr: NonNull<U>, v: T) -> NonNull<T> {
             // SAFETY: `serialize` requires a buffer sized for the full layout.
@@ -47,7 +66,7 @@ impl OwnedAccount {
             }
         }
 
-        let allocation = self.allocation();
+        let allocation = ((buf.len() - IMAGE_OFFSET) / 2) as u32;
         let ptr = write(ptr, AccountHeader::new(allocation));
         // The image prefix stores the account pubkey for later iteration.
         let ptr = write(ptr, *pubkey);