introduce ShrinkInProgress

runtime/src/account_storage.rs

@@ -80,6 +80,42 @@ impl AccountStorage {
             .insert(store.append_vec_id(), store)
             .is_none());
     }
+
+    /// called when shrinking begins on a slot and append vec.
+    /// When 'ShrinkInProgress' is dropped by caller, the old store will be removed from the storage map.
+    pub(crate) fn shrinking_in_progress(
+        &self,
+        slot: Slot,
+        new_store: Arc<AccountStorageEntry>,
+    ) -> ShrinkInProgress<'_> {
+        let slot_storages: SlotStores = self.get_slot_stores(slot).unwrap_or_else(||
+            // DashMap entry.or_insert() returns a RefMut, essentially a write lock,
+            // which is dropped after this block ends, minimizing time held by the lock.
+            // However, we still want to persist the reference to the `SlotStores` behind
+            // the lock, hence we clone it out, (`SlotStores` is an Arc so is cheap to clone).
+            self
+                .map
+                .entry(slot)
+                .or_insert(Arc::new(RwLock::new(HashMap::new())))
+                .clone());
+
+        let shrinking_store = Arc::clone(slot_storages.read().unwrap().iter().next().unwrap().1);
+
+        let previous_id = shrinking_store.append_vec_id();
+        let new_id = new_store.append_vec_id();
+        let mut storages = slot_storages.write().unwrap();
+        assert!(storages.remove(&previous_id).is_some());
+        // insert 'new_store' into 'map'
+        assert!(storages.insert(new_id, Arc::clone(&new_store)).is_none());
+
+        ShrinkInProgress {
+            storage: self,
+            slot,
+            new_store,
+            old_store: shrinking_store,
+        }
+    }
+
     #[cfg(test)]
     pub(crate) fn insert_empty_at_slot(&self, slot: Slot) {
         self.map
@@ -92,6 +128,38 @@ impl AccountStorage {
     }
 }
 
+/// exists while there is a shrink in progress
+/// keeps track of the 'new_store' being created and the 'old_store' being replaced.
+pub(crate) struct ShrinkInProgress<'a> {
+    storage: &'a AccountStorage,
+    /// newly shrunk store with a subset of contents from 'old_store'
+    new_store: Arc<AccountStorageEntry>,
+    /// old store which will be shrunk and replaced
+    old_store: Arc<AccountStorageEntry>,
+    slot: Slot,
+}
+
+/// called when the shrink is no longer in progress. This means we can release the old append vec and update the map of slot -> append vec
+impl<'a> Drop for ShrinkInProgress<'a> {
+    fn drop(&mut self) {
+        // the slot must be in the map
+        let slot_storages: SlotStores = self.storage.get_slot_stores(self.slot).unwrap();
+
+        let mut storages = slot_storages.write().unwrap();
+        // the id must be in the hashmap
+        assert!(storages.remove(&self.old_store.append_vec_id()).is_some());
+    }
+}
+
+impl<'a> ShrinkInProgress<'a> {
+    pub(crate) fn new_storage(&self) -> &Arc<AccountStorageEntry> {
+        &self.new_store
+    }
+    pub(crate) fn old_storage(&self) -> &Arc<AccountStorageEntry> {
+        &self.old_store
+    }
+}
+
 #[derive(Debug, Eq, PartialEq, Copy, Clone, Deserialize, Serialize, AbiExample, AbiEnumVisitor)]
 pub enum AccountStorageStatus {
     Available = 0,