Only mmap file from snapshot once (#14815)

runtime/src/accounts_db.rs

@@ -381,13 +381,18 @@ impl AccountStorageEntry {
         }
     }
 
-    pub(crate) fn new_empty_map(id: AppendVecId, accounts_current_len: usize) -> Self {
+    pub(crate) fn new_existing(
+        slot: Slot,
+        id: AppendVecId,
+        accounts: AppendVec,
+        num_accounts: usize,
+    ) -> Self {
         Self {
             id: AtomicUsize::new(id),
-            slot: AtomicU64::new(0),
-            accounts: AppendVec::new_empty_map(accounts_current_len),
+            slot: AtomicU64::new(slot),
+            accounts,
             count_and_status: RwLock::new((0, AccountStorageStatus::Available)),
-            approx_store_count: AtomicUsize::new(0),
+            approx_store_count: AtomicUsize::new(num_accounts),
             alive_bytes: AtomicUsize::new(0),
         }
     }
@@ -521,13 +526,6 @@ impl AccountStorageEntry {
         count
     }
 
-    pub fn set_file<P: AsRef<Path>>(&mut self, path: P) -> IOResult<()> {
-        let num_accounts = self.accounts.set_file(path)?;
-        self.approx_store_count
-            .store(num_accounts, Ordering::Relaxed);
-        Ok(())
-    }
-
     pub fn get_relative_path(&self) -> Option<PathBuf> {
         AppendVec::get_relative_path(self.accounts.get_path())
     }

runtime/src/append_vec.rs

@@ -105,7 +105,7 @@ impl<'a> StoredAccountMeta<'a> {
 
     fn ref_executable_byte(&self) -> &u8 {
         // Use extra references to avoid value silently clamped to 1 (=true) and 0 (=false)
-        // Yes, this really happens; see test_set_file_crafted_executable
+        // Yes, this really happens; see test_new_from_file_crafted_executable
         let executable_bool: &bool = &self.account_meta.executable;
         // UNSAFE: Force to interpret mmap-backed bool as u8 to really read the actual memory content
         let executable_byte: &u8 = unsafe { &*(executable_bool as *const bool as *const u8) };
@@ -271,29 +271,28 @@ impl AppendVec {
     }
 
     #[allow(clippy::mutex_atomic)]
-    pub fn set_file<P: AsRef<Path>>(&mut self, path: P) -> io::Result<usize> {
-        // this AppendVec must not hold actual file;
-        assert_eq!(self.file_size, 0);
-
+    pub fn new_from_file<P: AsRef<Path>>(path: P, current_len: usize) -> io::Result<(Self, usize)> {
         let data = OpenOptions::new()
             .read(true)
             .write(true)
             .create(false)
             .open(&path)?;
 
-        let current_len = self.current_len.load(Ordering::Relaxed);
-        assert_eq!(current_len, *self.append_offset.lock().unwrap());
-
         let file_size = std::fs::metadata(&path)?.len();
         AppendVec::sanitize_len_and_size(current_len, file_size as usize)?;
 
         let map = unsafe { MmapMut::map_mut(&data)? };
 
-        self.file_size = file_size;
-        self.path = path.as_ref().to_path_buf();
-        self.map = map;
+        let new = AppendVec {
+            path: path.as_ref().to_path_buf(),
+            map,
+            append_offset: Mutex::new(current_len),
+            current_len: AtomicUsize::new(current_len),
+            file_size,
+            remove_on_drop: true,
+        };
 
-        let (sanitized, num_accounts) = self.sanitize_layout_and_length();
+        let (sanitized, num_accounts) = new.sanitize_layout_and_length();
         if !sanitized {
             return Err(std::io::Error::new(
                 std::io::ErrorKind::Other,
@@ -301,7 +300,7 @@ impl AppendVec {
             ));
         }
 
-        Ok(num_accounts)
+        Ok((new, num_accounts))
     }
 
     fn sanitize_layout_and_length(&self) -> (bool, usize) {
@@ -564,11 +563,9 @@ pub mod tests {
     }
 
     #[test]
-    fn test_append_vec_set_file_bad_size() {
-        let file = get_append_vec_path("test_append_vec_set_file_bad_size");
+    fn test_append_vec_new_from_file_bad_size() {
+        let file = get_append_vec_path("test_append_vec_new_from_file_bad_size");
         let path = &file.path;
-        let mut av = AppendVec::new_empty_map(0);
-        assert_eq!(av.accounts(0).len(), 0);
 
         let _data = OpenOptions::new()
             .read(true)
@@ -577,7 +574,7 @@ pub mod tests {
             .open(&path)
             .expect("create a test file for mmap");
 
-        let result = av.set_file(path);
+        let result = AppendVec::new_from_file(path, 0);
         assert_matches!(result, Err(ref message) if message.to_string() == *"too small file size 0 for AppendVec");
     }
 
@@ -693,10 +690,11 @@ pub mod tests {
     }
 
     #[test]
-    fn test_set_file_crafted_zero_lamport_account() {
+    fn test_new_from_file_crafted_zero_lamport_account() {
         let file = get_append_vec_path("test_append");
         let path = &file.path;
         let mut av = AppendVec::new(&path, true, 1024 * 1024);
+        av.set_no_remove_on_drop();
 
         let pubkey = solana_sdk::pubkey::new_rand();
         let owner = Pubkey::default();
@@ -713,16 +711,18 @@ pub mod tests {
         assert_eq!(av.get_account_test(index).unwrap(), account_with_meta);
 
         av.flush().unwrap();
-        av.file_size = 0;
-        let result = av.set_file(path);
+        let accounts_len = av.len();
+        drop(av);
+        let result = AppendVec::new_from_file(path, accounts_len);
         assert_matches!(result, Err(ref message) if message.to_string() == *"incorrect layout/length/data");
     }
 
     #[test]
-    fn test_set_file_crafted_data_len() {
-        let file = get_append_vec_path("test_set_file_crafted_data_len");
+    fn test_new_from_file_crafted_data_len() {
+        let file = get_append_vec_path("test_new_from_file_crafted_data_len");
         let path = &file.path;
         let mut av = AppendVec::new(&path, true, 1024 * 1024);
+        av.set_no_remove_on_drop();
 
         let crafted_data_len = 1;
 
@@ -739,16 +739,18 @@ pub mod tests {
         assert_eq!(account.meta.data_len, crafted_data_len);
 
         av.flush().unwrap();
-        av.file_size = 0;
-        let result = av.set_file(path);
+        let accounts_len = av.len();
+        drop(av);
+        let result = AppendVec::new_from_file(path, accounts_len);
         assert_matches!(result, Err(ref message) if message.to_string() == *"incorrect layout/length/data");
     }
 
     #[test]
-    fn test_set_file_too_large_data_len() {
-        let file = get_append_vec_path("test_set_file_too_large_data_len");
+    fn test_new_from_file_too_large_data_len() {
+        let file = get_append_vec_path("test_new_from_file_too_large_data_len");
         let path = &file.path;
         let mut av = AppendVec::new(&path, true, 1024 * 1024);
+        av.set_no_remove_on_drop();
 
         let too_large_data_len = u64::max_value();
         av.append_account_test(&create_test_account(10)).unwrap();
@@ -763,16 +765,18 @@ pub mod tests {
         assert_matches!(accounts.first(), None);
 
         av.flush().unwrap();
-        av.file_size = 0;
-        let result = av.set_file(path);
+        let accounts_len = av.len();
+        drop(av);
+        let result = AppendVec::new_from_file(path, accounts_len);
         assert_matches!(result, Err(ref message) if message.to_string() == *"incorrect layout/length/data");
     }
 
     #[test]
-    fn test_set_file_crafted_executable() {
-        let file = get_append_vec_path("test_set_file_crafted_executable");
+    fn test_new_from_file_crafted_executable() {
+        let file = get_append_vec_path("test_new_from_crafted_executable");
         let path = &file.path;
         let mut av = AppendVec::new(&path, true, 1024 * 1024);
+        av.set_no_remove_on_drop();
         av.append_account_test(&create_test_account(10)).unwrap();
         {
             let mut executable_account = create_test_account(10);
@@ -817,8 +821,9 @@ pub mod tests {
         }
 
         av.flush().unwrap();
-        av.file_size = 0;
-        let result = av.set_file(path);
+        let accounts_len = av.len();
+        drop(av);
+        let result = AppendVec::new_from_file(path, accounts_len);
         assert_matches!(result, Err(ref message) if message.to_string() == *"incorrect layout/length/data");
     }
 }

runtime/src/serde_snapshot.rs

@@ -9,6 +9,7 @@ use {
         epoch_stakes::EpochStakes,
         message_processor::MessageProcessor,
         rent_collector::RentCollector,
+        serde_snapshot::future::SerializableStorage,
         stakes::Stakes,
     },
     bincode,
@@ -70,7 +71,7 @@ trait TypeContext<'a> {
     type SerializableAccountStorageEntry: Serialize
         + DeserializeOwned
         + From<&'a AccountStorageEntry>
-        + Into<AccountStorageEntry>;
+        + SerializableStorage;
 
     fn serialize_bank_and_storage<S: serde::ser::Serializer>(
         serializer: S,
@@ -240,7 +241,7 @@ fn reconstruct_bank_from_fields<E, P>(
     caching_enabled: bool,
 ) -> Result<Bank, Error>
 where
-    E: Into<AccountStorageEntry>,
+    E: SerializableStorage,
     P: AsRef<Path>,
 {
     let mut accounts_db = reconstruct_accountsdb_from_fields(
@@ -274,7 +275,7 @@ fn reconstruct_accountsdb_from_fields<E, P>(
     caching_enabled: bool,
 ) -> Result<AccountsDB, Error>
 where
-    E: Into<AccountStorageEntry>,
+    E: SerializableStorage,
     P: AsRef<Path>,
 {
     let mut accounts_db = AccountsDB::new_with_config(
@@ -285,20 +286,6 @@ where
     );
     let AccountsDbFields(storage, version, slot, bank_hash_info) = accounts_db_fields;
 
-    // convert to two level map of slot -> id -> account storage entry
-    let storage = {
-        let mut map = HashMap::new();
-        for (slot, entries) in storage.into_iter() {
-            let sub_map = map.entry(slot).or_insert_with(HashMap::new);
-            for entry in entries.into_iter() {
-                let entry: AccountStorageEntry = entry.into();
-                entry.slot.store(slot, Ordering::Relaxed);
-                sub_map.insert(entry.append_vec_id(), Arc::new(entry));
-            }
-        }
-        map
-    };
-
     // Ensure all account paths exist
     for path in &accounts_db.paths {
         std::fs::create_dir_all(path)
@@ -320,14 +307,14 @@ where
             remaining_slots_to_process -= 1;
 
             let mut new_slot_storage = HashMap::new();
-            for (id, storage_entry) in slot_storage.drain() {
+            for storage_entry in slot_storage.drain(..) {
                 let path_index = thread_rng().gen_range(0, accounts_db.paths.len());
                 let local_dir = &accounts_db.paths[path_index];
 
                 // Move the corresponding AppendVec from the snapshot into the directory pointed
                 // at by `local_dir`
                 let append_vec_relative_path =
-                    AppendVec::new_relative_path(slot, storage_entry.append_vec_id());
+                    AppendVec::new_relative_path(slot, storage_entry.id());
                 let append_vec_abs_path = stream_append_vecs_path
                     .as_ref()
                     .join(&append_vec_relative_path);
@@ -342,9 +329,17 @@ where
 
                 // Notify the AppendVec of the new file location
                 let local_path = local_dir.join(append_vec_relative_path);
-                let mut u_storage_entry = Arc::try_unwrap(storage_entry).unwrap();
-                u_storage_entry.set_file(local_path)?;
-                new_slot_storage.insert(id, Arc::new(u_storage_entry));
+
+                let (accounts, num_accounts) =
+                    AppendVec::new_from_file(&local_path, storage_entry.current_len())?;
+                let u_storage_entry = AccountStorageEntry::new_existing(
+                    slot,
+                    storage_entry.id(),
+                    accounts,
+                    num_accounts,
+                );
+
+                new_slot_storage.insert(storage_entry.id(), Arc::new(u_storage_entry));
             }
             Ok((slot, new_slot_storage))
         })

runtime/src/serde_snapshot/future.rs

@@ -12,6 +12,20 @@ pub(super) struct SerializableAccountStorageEntry {
     accounts_current_len: usize,
 }
 
+pub trait SerializableStorage {
+    fn id(&self) -> AppendVecId;
+    fn current_len(&self) -> usize;
+}
+
+impl SerializableStorage for SerializableAccountStorageEntry {
+    fn id(&self) -> AppendVecId {
+        self.id
+    }
+    fn current_len(&self) -> usize {
+        self.accounts_current_len
+    }
+}
+
 #[cfg(all(test, RUSTC_WITH_SPECIALIZATION))]
 impl solana_frozen_abi::abi_example::IgnoreAsHelper for SerializableAccountStorageEntry {}
 
@@ -24,12 +38,6 @@ impl From<&AccountStorageEntry> for SerializableAccountStorageEntry {
     }
 }
 
-impl From<SerializableAccountStorageEntry> for AccountStorageEntry {
-    fn from(s: SerializableAccountStorageEntry) -> Self {
-        AccountStorageEntry::new_empty_map(s.id, s.accounts_current_len)
-    }
-}
-
 use std::sync::RwLock;
 // Deserializable version of Bank which need not be serializable,
 // because it's handled by SerializableVersionedBank.

runtime/store-tool/src/main.rs

@@ -25,15 +25,13 @@ fn main() {
 
     let file = value_t_or_exit!(matches, "file", String);
     let len = value_t_or_exit!(matches, "len", usize);
-    let mut store = AppendVec::new_empty_map(len);
+    let (mut store, num_accounts) = AppendVec::new_from_file(file, len).expect("should succeed");
     store.set_no_remove_on_drop();
-    store.set_file(file).expect("set_file failed");
-    let accounts = store.accounts(0);
     info!(
         "store: len: {} capacity: {} accounts: {}",
         store.len(),
         store.capacity(),
-        accounts.len()
+        num_accounts,
     );
     for account in store.accounts(0) {
         info!(