Merge pull request AleoNet#1658 from ljedrz/testnet3_storage_tweaks

A few storage-related tweaks
shliming · Mar 2, 2022 · 613f29e · 613f29e
2 parents e8c776f + f471fe2
commit 613f29e
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 29 deletions.
diff --git a/storage/src/state/ledger.rs b/storage/src/state/ledger.rs
@@ -231,11 +231,6 @@ impl<N: Network> LedgerState<N> {
         ledger.regenerate_ledger_tree()?;
         assert_eq!(ledger.ledger_tree.read().root(), latest_ledger_root);
 
-        // let value = storage.export()?;
-        // println!("{}", value);
-        // let storage_2 = S::open(".ledger_2", context)?;
-        // storage_2.import(value)?;
-
         info!("Ledger successfully loaded at block {}", ledger.latest_block_height());
         Ok(ledger)
     }
@@ -1279,8 +1274,8 @@ impl<N: Network> LedgerState<N> {
             let remove_ledger_root = self
                 .ledger_roots
                 .iter()
-                .filter(|(_, block_height)| current_block_height == *block_height)
-                .collect::<Vec<_>>();
+                .filter(|(_, block_height)| current_block_height == *block_height);
+
             for (ledger_root, _) in remove_ledger_root {
                 self.ledger_roots.remove(&ledger_root, Some(batch))?;
             }

diff --git a/storage/src/state/prover.rs b/storage/src/state/prover.rs
@@ -41,11 +41,6 @@ impl<N: Network> ProverState<N> {
             coinbase: CoinbaseState::open(storage)?,
         };
 
-        // let value = storage.export()?;
-        // println!("{}", value);
-        // let storage_2 = S::open(".ledger_2", context)?;
-        // storage_2.import(value)?;
-
         info!("Prover successfully initialized");
         Ok(prover)
     }

diff --git a/storage/src/storage/rocksdb/map.rs b/storage/src/storage/rocksdb/map.rs
@@ -47,15 +47,26 @@ pub struct DataMap<K: Serialize + DeserializeOwned, V: Serialize + DeserializeOw
 }
 
 impl<K: Serialize + DeserializeOwned, V: Serialize + DeserializeOwned> DataMap<K, V> {
-    fn get_raw<Q>(&self, key: &Q) -> Result<Option<Vec<u8>>>
+    #[inline]
+    fn create_prefixed_key<Q>(&self, key: &Q) -> Result<Vec<u8>>
     where
         K: Borrow<Q>,
         Q: Serialize + ?Sized,
     {
-        let mut key_buf = self.context.clone();
-        key_buf.reserve(bincode::serialized_size(&key)? as usize);
-        bincode::serialize_into(&mut key_buf, &key)?;
-        match self.storage.rocksdb.get(&key_buf)? {
+        let mut raw_key = self.context.clone();
+        raw_key.reserve(bincode::serialized_size(&key)? as usize);
+        bincode::serialize_into(&mut raw_key, &key)?;
+
+        Ok(raw_key)
+    }
+
+    fn get_raw<'a, Q>(&'a self, key: &Q) -> Result<Option<rocksdb::DBPinnableSlice<'a>>>
+    where
+        K: Borrow<Q>,
+        Q: Serialize + ?Sized,
+    {
+        let raw_key = self.create_prefixed_key(key)?;
+        match self.storage.rocksdb.get_pinned(&raw_key)? {
             Some(data) => Ok(Some(data)),
             None => Ok(None),
         }
@@ -109,15 +120,13 @@ impl<'a, K: Serialize + DeserializeOwned, V: Serialize + DeserializeOwned> Map<'
         K: Borrow<Q>,
         Q: Serialize + ?Sized,
     {
-        let mut key_buf = self.context.clone();
-        key_buf.reserve(bincode::serialized_size(&key)? as usize);
-        bincode::serialize_into(&mut key_buf, &key)?;
-        let value_buf = bincode::serialize(value)?;
+        let raw_key = self.create_prefixed_key(key)?;
+        let raw_value = bincode::serialize(value)?;
 
         if let Some(batch_id) = batch {
-            self.storage.batches.lock().entry(batch_id).or_default().put(&key_buf, &value_buf);
+            self.storage.batches.lock().entry(batch_id).or_default().put(&raw_key, &raw_value);
         } else {
-            self.storage.rocksdb.put(&key_buf, &value_buf)?;
+            self.storage.rocksdb.put(&raw_key, &raw_value)?;
         }
 
         Ok(())
@@ -133,14 +142,12 @@ impl<'a, K: Serialize + DeserializeOwned, V: Serialize + DeserializeOwned> Map<'
         K: Borrow<Q>,
         Q: Serialize + ?Sized,
     {
-        let mut key_buf = self.context.clone();
-        key_buf.reserve(bincode::serialized_size(&key)? as usize);
-        bincode::serialize_into(&mut key_buf, &key)?;
+        let raw_key = self.create_prefixed_key(key)?;
 
         if let Some(batch_id) = batch {
-            self.storage.batches.lock().entry(batch_id).or_default().delete(&key_buf);
+            self.storage.batches.lock().entry(batch_id).or_default().delete(&raw_key);
         } else {
-            self.storage.rocksdb.delete(&key_buf)?;
+            self.storage.rocksdb.delete(&raw_key)?;
         }
 
         Ok(())

diff --git a/storage/src/storage/rocksdb/tests.rs b/storage/src/storage/rocksdb/tests.rs
@@ -230,6 +230,6 @@ fn test_export_import() {
     let map = storage.open_map::<u32, String>(MapId::Test).expect("Failed to open data map");
 
     for i in 0..100 {
-        assert_eq!(map.get(&i).expect("Failed to insert"), Some(i.to_string()));
+        assert_eq!(map.get(&i).expect("Failed to get"), Some(i.to_string()));
     }
 }