remove should_retain from mark_dirty_dead_stores

runtime/src/accounts_db.rs

@@ -3847,7 +3847,6 @@ impl AccountsDb {
         // Purge old, overwritten storage entries
         let (remaining_stores, dead_storages) = self.mark_dirty_dead_stores(
             slot,
-            |store| !shrink_collect.store_ids.contains(&store.append_vec_id()),
             // If all accounts are zero lamports, then we want to mark the entire OLD append vec as dirty.
             // otherwise, we'll call 'add_uncleaned_pubkeys_after_shrink' just on the unref'd keys below.
             shrink_collect.all_are_zero_lamports,
@@ -4003,15 +4002,13 @@ impl AccountsDb {
     }
 
     /// get stores for 'slot'
-    /// retain only the stores where 'should_retain(store)' == true
-    /// for stores not retained, insert in 'dead_storages' and optionally 'dirty_stores'
+    /// Drop 'shrink_in_progress', which will cause the old store to be removed from the storage map.
+    /// For 'shrink_in_progress'.'old_storage' which is not retained, insert in 'dead_storages' and optionally 'dirty_stores'
     /// This is the end of the life cycle of `shrink_in_progress`.
-    /// Dropping 'shrink_in_progress' here will cause the old store to be removed from the storage map.
     /// returns: (# of remaining stores for this slot, dead storages)
     pub(crate) fn mark_dirty_dead_stores(
         &self,
         slot: Slot,
-        should_retain: impl Fn(&AccountStorageEntry) -> bool,
         add_dirty_stores: bool,
         shrink_in_progress: Option<ShrinkInProgress>,
     ) -> (usize, SnapshotStorage) {
@@ -4027,21 +4024,19 @@ impl AccountsDb {
 
         let remaining_stores = if let Some(slot_stores) = self.storage.get_slot_stores(slot) {
             if let Some(shrink_in_progress) = shrink_in_progress {
-                let store = shrink_in_progress.old_storage().clone();
-                not_retaining_store(&store);
+                // shrink is in progress, so 1 new append vec to keep, 1 old one to throw away
+                let store = shrink_in_progress.old_storage();
+                not_retaining_store(store);
+                // drop removes the old append vec that was being shrunk from db's storage
                 drop(shrink_in_progress);
                 slot_stores.read().unwrap().len()
             } else {
+                // no shrink in progress, so all append vecs in this slot are dead
                 let mut list = slot_stores.write().unwrap();
-                list.retain(|_key, store| {
-                    if !should_retain(store) {
-                        not_retaining_store(store);
-                        false
-                    } else {
-                        true
-                    }
+                list.drain().for_each(|(_key, store)| {
+                    not_retaining_store(&store);
                 });
-                list.len()
+                0
             }
         } else {
             0
@@ -4506,7 +4501,8 @@ impl AccountsDb {
             return; // skipping slot with no useful accounts to write
         }
 
-        let (shrink_in_progress, time) = measure!(current_ancient.create_if_necessary(slot, self));
+        let (mut shrink_in_progress, time) =
+            measure!(current_ancient.create_if_necessary(slot, self));
         let mut create_and_insert_store_elapsed_us = time.as_us();
         let available_bytes = current_ancient.append_vec().accounts.remaining_bytes();
         // split accounts in 'slot' into:
@@ -4537,8 +4533,17 @@ impl AccountsDb {
             // Assert: it cannot be the case that we already had an ancient append vec at this slot and
             // yet that ancient append vec does not have room for the accounts stored at this slot currently
             assert_ne!(slot, current_ancient.slot());
-            let (_, time) = measure!(current_ancient.create_ancient_append_vec(slot, self));
+            let (shrink_in_progress_overflow, time) =
+                measure!(current_ancient.create_ancient_append_vec(slot, self));
             create_and_insert_store_elapsed_us += time.as_us();
+            // We cannot possibly be shrinking the original slot that created an ancient append vec
+            // AND not have enough room in the ancient append vec at that slot
+            // to hold all the contents of that slot.
+            // We need this new 'shrink_in_progress' to be used in 'remove_old_stores_shrink' below.
+            // All non-overflow accounts were put in a prior slot's ancient append vec. All overflow accounts
+            // are essentially being shrunk into a new ancient append vec in 'slot'.
+            assert!(shrink_in_progress.is_none());
+            shrink_in_progress = Some(shrink_in_progress_overflow);
 
             // write the overflow accounts to the next ancient storage
             let timing =
@@ -4581,7 +4586,7 @@ impl AccountsDb {
                 self.accounts_index
                     .clean_dead_slot(*slot, &mut AccountsIndexRootsStats::default());
                 self.bank_hashes.write().unwrap().remove(slot);
-                // all storages have been removed from here and recycled or dropped
+                // all storages have been removed from this slot and recycled or dropped
                 assert!(self
                     .storage
                     .remove(slot)
@@ -16496,56 +16501,71 @@ pub mod tests {
     }
 
     #[test]
-    fn test_mark_dirty_dead_stores() {
+    fn test_mark_dirty_dead_stores_empty() {
         let db = AccountsDb::new_single_for_tests();
         let slot = 0;
-        let called = AtomicUsize::default();
-        let add_dirty_stores = false;
-        let (remaining_stores, _) = db.mark_dirty_dead_stores(
-            slot,
-            |_store| {
-                called.fetch_add(1, Ordering::Relaxed);
-                false
-            },
-            add_dirty_stores,
-            None,
-        );
-        assert_eq!(0, called.load(Ordering::Relaxed));
-        assert_eq!(0, remaining_stores);
+        for add_dirty_stores in [false, true] {
+            let (remaining_stores, dead_storages) =
+                db.mark_dirty_dead_stores(slot, add_dirty_stores, None);
+            assert_eq!(remaining_stores, 0);
+            assert!(dead_storages.is_empty());
+            assert!(db.dirty_stores.is_empty());
+        }
+    }
 
-        let size = 1;
-        let inserted_store = db.create_and_insert_store(slot, size, "test");
-        let (remaining_stores, _) = db.mark_dirty_dead_stores(
-            slot,
-            |store| {
-                assert_eq!(store.append_vec_id(), inserted_store.append_vec_id());
-                called.fetch_add(1, Ordering::Relaxed);
-                true // retain
-            },
-            add_dirty_stores,
-            None,
-        );
-        assert_eq!(1, called.load(Ordering::Relaxed));
-        assert_eq!(1, remaining_stores);
+    #[test]
+    fn test_mark_dirty_dead_stores() {
+        let slot = 0;
 
-        let called = AtomicUsize::default();
-        let (remaining_stores, _) = db.mark_dirty_dead_stores(
-            slot,
-            |store| {
-                assert_eq!(store.append_vec_id(), inserted_store.append_vec_id());
-                called.fetch_add(1, Ordering::Relaxed);
-                false // don't retain
-            },
-            add_dirty_stores,
-            None,
-        );
-        assert_eq!(1, called.load(Ordering::Relaxed));
-        assert!(db
-            .get_storages_for_slot(slot)
-            .unwrap_or_default()
-            .is_empty());
-        assert_eq!(0, remaining_stores);
-        assert!(db.dirty_stores.is_empty());
+        // None for shrink_in_progress, 1 existing store at the slot
+        // contents should be the same
+        // this tests the case where there were no alive bytes during shrink so no-op
+        for add_dirty_stores in [false, true] {
+            let db = AccountsDb::new_single_for_tests();
+            let size = 1;
+            let existing_store = db.create_and_insert_store(slot, size, "test");
+            let old_id = existing_store.append_vec_id();
+            let (remaining_stores, dead_storages) =
+                db.mark_dirty_dead_stores(slot, add_dirty_stores, None);
+            assert_eq!(0, remaining_stores);
+            assert_eq!(dead_storages.len(), 1);
+            assert_eq!(dead_storages.first().unwrap().append_vec_id(), old_id);
+            if add_dirty_stores {
+                assert_eq!(1, db.dirty_stores.len());
+                let dirty_store = db.dirty_stores.get(&(slot, old_id)).unwrap();
+                assert_eq!(dirty_store.append_vec_id(), old_id);
+            } else {
+                assert!(db.dirty_stores.is_empty());
+            }
+            assert!(db.get_storages_for_slot(slot).unwrap().is_empty());
+        }
+
+        // use shrink_in_progress to cause us to drop the initial store
+        for add_dirty_stores in [false, true] {
+            let db = AccountsDb::new_single_for_tests();
+            let size = 1;
+            let old_store = db.create_and_insert_store(slot, size, "test");
+            let old_id = old_store.append_vec_id();
+            let shrink_in_progress = db.get_store_for_shrink(slot, 100);
+            let (remaining_stores, dead_storages) =
+                db.mark_dirty_dead_stores(slot, add_dirty_stores, Some(shrink_in_progress));
+            assert_eq!(1, remaining_stores);
+            assert_eq!(dead_storages.len(), 1);
+            assert_eq!(dead_storages.first().unwrap().append_vec_id(), old_id);
+            if add_dirty_stores {
+                assert_eq!(1, db.dirty_stores.len());
+                let dirty_store = db.dirty_stores.get(&(slot, old_id)).unwrap();
+                assert_eq!(dirty_store.append_vec_id(), old_id);
+            } else {
+                assert!(db.dirty_stores.is_empty());
+            }
+            assert_eq!(
+                1,
+                db.get_storages_for_slot(slot)
+                    .map(|storages| storages.len())
+                    .unwrap_or_default()
+            );
+        }
     }
 
     #[test]

runtime/src/snapshot_minimizer.rs

@@ -382,8 +382,7 @@ impl<'a> SnapshotMinimizer<'a> {
 
         let (_, mut dead_storages_this_time) = self.accounts_db().mark_dirty_dead_stores(
             slot,
-            |_store| true, /* ignored if shrink_in_progress is passed, otherwise retain all */
-            true,          // add_dirty_stores
+            true, // add_dirty_stores
             shrink_in_progress,
         );
         dead_storages