introduce ShrinkInProgress

runtime/src/account_storage.rs

@@ -80,6 +80,40 @@ 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 new_id = new_store.append_vec_id();
+ let mut storages = slot_storages.write().unwrap();
+ // 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 +126,43 @@ 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(),
+ "slot: {}, len: {}",
+ self.slot,
+ storages.len()
+ );
+ }
+}
+
+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,