Optimize lookup_txos using rocksdb's MultiGet

See https://github.com/facebook/rocksdb/wiki/MultiGet-Performance
shesek · May 22, 2024 · 38c9576 · 38c9576
1 parent 5db7d2b
commit 38c9576
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 20 deletions.
diff --git a/src/new_index/db.rs b/src/new_index/db.rs
@@ -196,6 +196,14 @@ impl DB {
         self.db.get(key).unwrap().map(|v| v.to_vec())
     }
 
+    pub fn multi_get<K, I>(&self, keys: I) -> Vec<Result<Option<Vec<u8>>, rocksdb::Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = K>,
+    {
+        self.db.multi_get(keys)
+    }
+
     fn verify_compatibility(&self, config: &Config) {
         let mut compatibility_bytes = bincode::serialize_little(&DB_VERSION).unwrap();
 

diff --git a/src/new_index/schema.rs b/src/new_index/schema.rs
@@ -1038,26 +1038,19 @@ fn lookup_txos(
     outpoints: &BTreeSet<OutPoint>,
     allow_missing: bool,
 ) -> HashMap<OutPoint, TxOut> {
-    let pool = rayon::ThreadPoolBuilder::new()
-        .num_threads(16) // we need to saturate SSD IOPS
-        .thread_name(|i| format!("lookup-txo-{}", i))
-        .build()
-        .unwrap();
-    pool.install(|| {
-        outpoints
-            .par_iter()
-            .filter_map(|outpoint| {
-                lookup_txo(&txstore_db, &outpoint)
-                    .or_else(|| {
-                        if !allow_missing {
-                            panic!("missing txo {} in {:?}", outpoint, txstore_db);
-                        }
-                        None
-                    })
-                    .map(|txo| (*outpoint, txo))
-            })
-            .collect()
-    })
+    let mut remain_outpoints = outpoints.iter();
+    txstore_db
+        .multi_get(outpoints.iter().map(TxOutRow::key))
+        .into_iter()
+        .filter_map(|res| {
+            let outpoint = remain_outpoints.next().unwrap();
+            match res.unwrap() {
+                Some(txo) => Some((*outpoint, deserialize(&txo).expect("failed to parse TxOut"))),
+                None if allow_missing => None,
+                None => panic!("missing txo {}", outpoint),
+            }
+        })
+        .collect()
 }
 
 fn lookup_txo(txstore_db: &DB, outpoint: &OutPoint) -> Option<TxOut> {