Add RocksDB TPCC comparison path

Author: KKould

src/storage/rocksdb.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use crate::errors::DatabaseError;
-use crate::storage::table_codec::{Bytes, TableCodec};
+use crate::storage::table_codec::TableCodec;
 use crate::storage::{
     CheckpointableStorage, InnerIter, Storage, Transaction, TransactionIsolationLevel,
 };
@@ -665,6 +665,7 @@ macro_rules! impl_transaction {
                         read_opts.set_prefix_same_as_start(true);
                     }
                 }
+                set_iterate_upper_bound(&mut read_opts, max);
 
                 let mut iter = self.tx.raw_iterator_opt(read_opts);
                 match &min {
@@ -679,7 +680,6 @@ macro_rules! impl_transaction {
                 }
 
                 Ok($iter {
-                    upper: owned_bound(max),
                     iter,
                     advanced: false,
                     done: false,
@@ -698,7 +698,6 @@ impl_transaction!(RocksTransaction, RocksIter);
 impl_transaction!(OptimisticRocksTransaction, OptimisticRocksIter);
 
 pub struct OptimisticRocksIter<'txn, 'iter> {
-    upper: Bound<Bytes>,
     iter: DBRawIteratorWithThreadMode<'iter, rocksdb::Transaction<'txn, OptimisticTransactionDB>>,
     advanced: bool,
     done: bool,
@@ -707,17 +706,11 @@ pub struct OptimisticRocksIter<'txn, 'iter> {
 impl InnerIter for OptimisticRocksIter<'_, '_> {
     #[inline]
     fn try_next(&mut self) -> Result<Option<crate::storage::KeyValueRef<'_>>, DatabaseError> {
-        next(
-            &mut self.iter,
-            &self.upper,
-            &mut self.advanced,
-            &mut self.done,
-        )
+        next(&mut self.iter, &mut self.advanced, &mut self.done)
     }
 }
 
 pub struct RocksIter<'txn, 'iter> {
-    upper: Bound<Bytes>,
     iter: DBRawIteratorWithThreadMode<
         'iter,
         rocksdb::Transaction<'txn, TransactionDB<rocksdb::MultiThreaded>>,
@@ -729,19 +722,27 @@ pub struct RocksIter<'txn, 'iter> {
 impl InnerIter for RocksIter<'_, '_> {
     #[inline]
     fn try_next(&mut self) -> Result<Option<crate::storage::KeyValueRef<'_>>, DatabaseError> {
-        next(
-            &mut self.iter,
-            &self.upper,
-            &mut self.advanced,
-            &mut self.done,
-        )
+        next(&mut self.iter, &mut self.advanced, &mut self.done)
+    }
+}
+
+#[inline]
+fn set_iterate_upper_bound(read_opts: &mut ReadOptions, upper: Bound<&[u8]>) {
+    match upper {
+        Bound::Included(bytes) => {
+            let mut exclusive_upper = Vec::with_capacity(bytes.len() + 1);
+            exclusive_upper.extend_from_slice(bytes);
+            exclusive_upper.push(0);
+            read_opts.set_iterate_upper_bound(exclusive_upper);
+        }
+        Bound::Excluded(bytes) => read_opts.set_iterate_upper_bound(bytes),
+        Bound::Unbounded => {}
     }
 }
 
 #[inline]
 fn next<'a, D: rocksdb::DBAccess>(
     iter: &'a mut DBRawIteratorWithThreadMode<'_, D>,
-    upper: &Bound<Bytes>,
     advanced: &mut bool,
     done: &mut bool,
 ) -> Result<Option<crate::storage::KeyValueRef<'a>>, DatabaseError> {
@@ -762,28 +763,11 @@ fn next<'a, D: rocksdb::DBAccess>(
         iter.status()?;
         return Ok(None);
     };
-    let upper_bound_check = match upper {
-        Bound::Included(upper) => key <= upper.as_slice(),
-        Bound::Excluded(upper) => key < upper.as_slice(),
-        Bound::Unbounded => true,
-    };
-    if !upper_bound_check {
-        *done = true;
-        return Ok(None);
-    }
 
     *advanced = true;
     Ok(Some((key, value)))
 }
 
-fn owned_bound(bound: Bound<&[u8]>) -> Bound<Bytes> {
-    match bound {
-        Bound::Included(bytes) => Bound::Included(bytes.to_vec()),
-        Bound::Excluded(bytes) => Bound::Excluded(bytes.to_vec()),
-        Bound::Unbounded => Bound::Unbounded,
-    }
-}
-
 #[cfg(all(test, not(target_arch = "wasm32")))]
 mod test {
     use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef, TableName};
@@ -792,8 +776,8 @@ mod test {
     use crate::expression::range_detacher::Range;
     use crate::storage::rocksdb::RocksStorage;
     use crate::storage::{
-        IndexImplEnum, IndexImplParams, IndexIter, IndexIterState, IterBounds, PrimaryKeyIndexImpl,
-        Storage, Transaction,
+        IndexImplEnum, IndexImplParams, IndexIter, IndexIterState, InnerIter, IterBounds,
+        PrimaryKeyIndexImpl, Storage, Transaction,
     };
     use crate::types::index::{IndexMeta, IndexType};
     use crate::types::tuple::Tuple;
@@ -837,6 +821,34 @@ mod test {
         Ok(())
     }
 
+    #[test]
+    fn test_rocksdb_range_uses_iterate_upper_bound() -> Result<(), DatabaseError> {
+        let temp_dir = TempDir::new().expect("unable to create temporary working directory");
+        let storage = RocksStorage::new(temp_dir.path())?;
+        let mut transaction = storage.transaction()?;
+
+        transaction.set(b"a", b"1")?;
+        transaction.set(b"b", b"2")?;
+        transaction.set(b"b\0", b"3")?;
+        transaction.set(b"c", b"4")?;
+
+        let mut inclusive_iter = transaction.range(Bound::Included(b"a"), Bound::Included(b"b"))?;
+        let mut inclusive_keys = Vec::new();
+        while let Some((key, _)) = inclusive_iter.try_next()? {
+            inclusive_keys.push(key.to_vec());
+        }
+        assert_eq!(inclusive_keys, vec![b"a".to_vec(), b"b".to_vec()]);
+
+        let mut exclusive_iter = transaction.range(Bound::Included(b"a"), Bound::Excluded(b"b"))?;
+        let mut exclusive_keys = Vec::new();
+        while let Some((key, _)) = exclusive_iter.try_next()? {
+            exclusive_keys.push(key.to_vec());
+        }
+        assert_eq!(exclusive_keys, vec![b"a".to_vec()]);
+
+        Ok(())
+    }
+
     #[test]
     fn test_in_rocksdb_storage_works_with_data() -> Result<(), DatabaseError> {
         let temp_dir = TempDir::new().expect("unable to create temporary working directory");

tpcc/src/backend/kitesql_rocksdb.rs

@@ -17,12 +17,16 @@ use super::{
 };
 use crate::TpccError;
 use kite_sql::db::{prepare, DBTransaction, DataBaseBuilder, Database, TransactionIter};
-use kite_sql::storage::rocksdb::{RocksStorage, RocksTransaction};
-use kite_sql::storage::Transaction;
+use kite_sql::storage::rocksdb::{OptimisticRocksStorage, RocksStorage};
+use kite_sql::storage::{Storage, Transaction};
 use kite_sql::types::tuple::Tuple;
 
-pub struct KiteSqlRocksDbBackend {
-    database: Database<RocksStorage>,
+pub type KiteSqlRocksDbBackend = KiteSqlRocksBackend<RocksStorage>;
+pub type KiteSqlOptimisticRocksDbBackend = KiteSqlRocksBackend<OptimisticRocksStorage>;
+pub type KiteSqlRocksDbTransaction<'a> = KiteSqlRocksTransaction<'a, RocksStorage>;
+
+pub struct KiteSqlRocksBackend<S: Storage> {
+    database: Database<S>,
 }
 
 impl KiteSqlRocksDbBackend {
@@ -33,7 +37,19 @@ impl KiteSqlRocksDbBackend {
                 .build_rocksdb()?,
         })
     }
+}
 
+impl KiteSqlOptimisticRocksDbBackend {
+    pub fn new(path: &str, rocksdb_stats: bool) -> Result<Self, TpccError> {
+        Ok(Self {
+            database: DataBaseBuilder::path(path)
+                .storage_statistics(rocksdb_stats)
+                .build_optimistic()?,
+        })
+    }
+}
+
+impl<S: Storage> KiteSqlRocksBackend<S> {
     fn prepare_spec_groups(
         &self,
         specs: &[Vec<StatementSpec>],
@@ -53,16 +69,16 @@ impl KiteSqlRocksDbBackend {
         Ok(groups)
     }
 
-    fn start_transaction(&self) -> Result<KiteSqlRocksDbTransaction<'_>, TpccError> {
-        Ok(KiteSqlRocksDbTransaction {
+    fn start_transaction(&self) -> Result<KiteSqlRocksTransaction<'_, S>, TpccError> {
+        Ok(KiteSqlRocksTransaction {
             inner: self.database.new_transaction()?,
         })
     }
 }
 
-impl BackendControl for KiteSqlRocksDbBackend {
+impl<S: Storage> BackendControl for KiteSqlRocksBackend<S> {
     type Transaction<'a>
-        = KiteSqlRocksDbTransaction<'a>
+        = KiteSqlRocksTransaction<'a, S>
     where
         Self: 'a;
 
@@ -84,23 +100,23 @@ impl BackendControl for KiteSqlRocksDbBackend {
     }
 }
 
-impl SimpleExecutor for KiteSqlRocksDbBackend {
+impl<S: Storage> SimpleExecutor for KiteSqlRocksBackend<S> {
     fn execute_batch(&self, sql: &str) -> Result<(), TpccError> {
         self.database.run(sql)?.done()?;
         Ok(())
     }
 }
 
-pub struct KiteSqlRocksDbTransaction<'a> {
-    inner: DBTransaction<'a, RocksStorage>,
+pub struct KiteSqlRocksTransaction<'a, S: Storage> {
+    inner: DBTransaction<'a, S>,
 }
 
-impl<'a> KiteSqlRocksDbTransaction<'a> {
+impl<'a, S: Storage> KiteSqlRocksTransaction<'a, S> {
     pub(crate) fn execute_raw<'b>(
         &'b mut self,
         statement: &PreparedStatement,
         params: &[DbParam],
-    ) -> Result<RocksTxnResult<'b, 'a>, TpccError> {
+    ) -> Result<KiteSqlTxnResult<'b, S::TransactionType<'a>>, TpccError> {
         let PreparedStatement::KiteSql { statement, .. } = statement else {
             return Err(TpccError::InvalidBackend);
         };
@@ -110,7 +126,7 @@ impl<'a> KiteSqlRocksDbTransaction<'a> {
     }
 }
 
-impl<'a> BackendTransaction for KiteSqlRocksDbTransaction<'a> {
+impl<'a, S: Storage> BackendTransaction for KiteSqlRocksTransaction<'a, S> {
     fn query_one(
         &mut self,
         statement: &PreparedStatement,
@@ -202,5 +218,3 @@ impl<T: Transaction> Iterator for KiteSqlTxnResult<'_, T> {
         self.0.next().map(|item| item.map_err(TpccError::from))
     }
 }
-
-pub(crate) type RocksTxnResult<'a, 'txn> = KiteSqlTxnResult<'a, RocksTransaction<'txn>>;

tpcc/src/main.rs

@@ -14,7 +14,7 @@
 
 use crate::backend::dual::DualBackend;
 use crate::backend::kitesql_lmdb::KiteSqlLmdbBackend;
-use crate::backend::kitesql_rocksdb::KiteSqlRocksDbBackend;
+use crate::backend::kitesql_rocksdb::{KiteSqlOptimisticRocksDbBackend, KiteSqlRocksDbBackend};
 use crate::backend::sqlite::{SqliteBackend, SqliteProfile};
 use crate::backend::{
     BackendControl, BackendTransaction, ColumnType, PreparedStatement, StatementSpec,
@@ -36,6 +36,7 @@ use rand::prelude::ThreadRng;
 use rand::Rng;
 use std::fs;
 use std::path::Path;
+#[cfg(any(feature = "pprof", test))]
 use std::path::PathBuf;
 use std::time::{Duration, Instant};
 
@@ -125,6 +126,7 @@ struct Args {
 #[derive(Copy, Clone, Debug, ValueEnum)]
 enum BackendKind {
     KitesqlRocksdb,
+    KitesqlOptimisticRocksdb,
     KitesqlLmdb,
     Sqlite,
     Dual,
@@ -141,6 +143,11 @@ fn main() -> Result<(), TpccError> {
             let backend = KiteSqlRocksDbBackend::new(&args.path, args.rocksdb_stats)?;
             run_tpcc(&backend, &args, &mut rng)
         }
+        BackendKind::KitesqlOptimisticRocksdb => {
+            reset_db_path(Path::new(&args.path))?;
+            let backend = KiteSqlOptimisticRocksDbBackend::new(&args.path, args.rocksdb_stats)?;
+            run_tpcc(&backend, &args, &mut rng)
+        }
         BackendKind::KitesqlLmdb => {
             reset_db_path(Path::new(&args.path))?;
             let backend = KiteSqlLmdbBackend::new(&args.path)?;