diff --git a/src/database/kvtree.rs b/src/database/kvtree.rs index beda110a52cbf804c2ac68cd352c8de918bb7dc1..52e5b146ce67b7809c7a7d9d46f8e84b281a44e7 100644 --- a/src/database/kvtree.rs +++ b/src/database/kvtree.rs @@ -1,4 +1,4 @@ -use std::{future::Future, pin::Pin, sync::Arc}; +use std::{future::Future, pin::Pin}; use crate::Result; @@ -6,11 +6,13 @@ pub(crate) trait KvTree: Send + Sync { fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>>; #[allow(dead_code)] - #[cfg(feature = "rocksdb")] - fn multi_get( - &self, _iter: Vec<(&Arc<rust_rocksdb::BoundColumnFamily<'_>>, Vec<u8>)>, - ) -> Vec<Result<Option<Vec<u8>>, rust_rocksdb::Error>> { - unimplemented!() + fn multi_get(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>> { + let mut ret: Vec<Option<Vec<u8>>> = Vec::with_capacity(keys.len()); + for key in keys { + ret.push(self.get(key)?); + } + + Ok(ret) } fn insert(&self, key: &[u8], value: &[u8]) -> Result<()>; diff --git a/src/database/rocksdb/kvtree.rs b/src/database/rocksdb/kvtree.rs index 4b30b552ebb8c788fc208e401836c664350615a8..1d9adaa4f85b28ccd1030ad5a7a87aee236a2f2e 100644 --- a/src/database/rocksdb/kvtree.rs +++ b/src/database/rocksdb/kvtree.rs @@ -23,13 +23,27 @@ fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> { Ok(self.db.rocks.get_cf_opt(&self.cf(), key, &readoptions)?) } - fn multi_get( - &self, iter: Vec<(&Arc<rust_rocksdb::BoundColumnFamily<'_>>, Vec<u8>)>, - ) -> Vec<Result<Option<Vec<u8>>, rust_rocksdb::Error>> { + fn multi_get(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>> { let mut readoptions = rust_rocksdb::ReadOptions::default(); readoptions.set_total_order_seek(true); - self.db.rocks.multi_get_cf_opt(iter, &readoptions) + // Optimization can be `true` if key vector is pre-sorted **by the column + // comparator**. + const SORTED: bool = false; + + let mut ret: Vec<Option<Vec<u8>>> = Vec::with_capacity(keys.len()); + for res in self + .db + .rocks + .batched_multi_get_cf_opt(&self.cf(), keys, SORTED, &readoptions) + { + match res? { + Some(res) => ret.push(Some((*res).to_vec())), + None => ret.push(None), + } + } + + Ok(ret) } fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> {