From 064f163b7071504a49b8fe234850eba08bd7ecd1 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Fri, 13 Jan 2023 15:11:25 -0600
Subject: [PATCH] get_account_storage_entry handles missing shrink in progress
 entry correctly (#29682)

---
 runtime/src/account_storage.rs | 45 +++++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 3 deletions(-)

diff --git a/runtime/src/account_storage.rs b/runtime/src/account_storage.rs
index 80035719c..4d33227e6 100644
--- a/runtime/src/account_storage.rs
+++ b/runtime/src/account_storage.rs
@@ -42,9 +42,9 @@ impl AccountStorage {
         self.get_slot_stores_shrinking_in_progress_ok(slot)
             .and_then(|storage_map| storage_map.read().unwrap().get(&store_id).cloned())
             .or_else(|| {
-                self.shrink_in_progress_map
-                    .get(&slot)
-                    .map(|entry| Arc::clone(entry.value()))
+                self.shrink_in_progress_map.get(&slot).and_then(|entry| {
+                    (entry.value().append_vec_id() == store_id).then(|| Arc::clone(entry.value()))
+                })
             })
     }
 
@@ -485,4 +485,43 @@ pub(crate) mod tests {
         let sample = storage.get_test_storage();
         storage.shrinking_in_progress(0, sample);
     }
+
+    #[test]
+    fn test_missing() {
+        // already called 'shrink_in_progress' on this slot, but it finished, so we succeed
+        // verify data structures during and after shrink and then with subsequent shrink call
+        let storage = AccountStorage::default();
+        let sample = storage.get_test_storage();
+        let id = sample.append_vec_id();
+        let slot_stores = SlotStores::default();
+        slot_stores.write().unwrap().insert(id, sample.clone());
+        let missing_id = 9999;
+        let slot = sample.slot();
+        // id is missing since not in maps at all
+        assert!(storage.get_account_storage_entry(slot, id).is_none());
+        // missing should always be missing
+        assert!(storage
+            .get_account_storage_entry(slot, missing_id)
+            .is_none());
+        storage.map.insert(slot, slot_stores.clone());
+        // id is found in map
+        assert!(storage.get_account_storage_entry(slot, id).is_some());
+        assert!(storage
+            .get_account_storage_entry(slot, missing_id)
+            .is_none());
+        storage
+            .shrink_in_progress_map
+            .insert(slot, Arc::clone(&sample));
+        // id is found in map
+        assert!(storage
+            .get_account_storage_entry(slot, missing_id)
+            .is_none());
+        assert!(storage.get_account_storage_entry(slot, id).is_some());
+        slot_stores.write().unwrap().clear();
+        // id is found in shrink_in_progress_map
+        assert!(storage
+            .get_account_storage_entry(slot, missing_id)
+            .is_none());
+        assert!(storage.get_account_storage_entry(slot, id).is_some());
+    }
 }