More engine abstraction, work toward Peekable and Iterable for Snapshot

cmd/src/server.rs

@@ -12,6 +12,7 @@ use kvproto::deadlock::create_deadlock;
 use kvproto::debugpb::create_debug;
 use kvproto::import_sstpb::create_import_sst;
 use pd_client::{PdClient, RpcClient};
+use std::convert::TryFrom;
 use std::fs::File;
 use std::path::Path;
 use std::sync::{Arc, Mutex};
@@ -205,9 +206,12 @@ fn run_raft_server(pd_client: RpcClient, cfg: &TiKvConfig, security_mgr: Arc<Sec
     )
     .unwrap_or_else(|e| fatal!("failed to create raft storage: {}", e));
 
+    let bps = i64::try_from(cfg.server.snap_max_write_bytes_per_sec.0)
+        .unwrap_or_else(|_| fatal!("snap_max_write_bytes_per_sec > i64::max_value"));
+
     // Create snapshot manager, server.
     let snap_mgr = SnapManagerBuilder::default()
-        .max_write_bytes_per_sec(cfg.server.snap_max_write_bytes_per_sec.0)
+        .max_write_bytes_per_sec(bps)
         .max_total_size(cfg.server.snap_max_total_size.0)
         .build(
             snap_path.as_path().to_str().unwrap().to_owned(),

components/backup/src/endpoint.rs

@@ -7,8 +7,9 @@ use std::sync::atomic::*;
 use std::sync::*;
 use std::time::*;
 
-use engine::rocks::util::io_limiter::IOLimiter;
 use engine::DB;
+use engine_rocks::RocksIOLimiter;
+use engine_traits::IOLimiter;
 use external_storage::*;
 use futures::lazy;
 use futures::prelude::Future;
@@ -64,7 +65,7 @@ impl fmt::Debug for Task {
 
 #[derive(Clone)]
 struct LimitedStorage {
-    limiter: Option<Arc<IOLimiter>>,
+    limiter: Option<Arc<RocksIOLimiter>>,
     storage: Arc<dyn ExternalStorage>,
 }
 
@@ -77,7 +78,7 @@ impl Task {
         let cancel = Arc::new(AtomicBool::new(false));
 
         let limiter = if req.get_rate_limit() != 0 {
-            Some(Arc::new(IOLimiter::new(req.get_rate_limit() as _)))
+            Some(Arc::new(RocksIOLimiter::new(req.get_rate_limit() as _)))
         } else {
             None
         };
@@ -839,7 +840,7 @@ pub mod tests {
         }
 
         // TODO: check key number for each snapshot.
-        let limiter = Arc::new(IOLimiter::new(10 * 1024 * 1024 /* 10 MB/s */));
+        let limiter = Arc::new(RocksIOLimiter::new(10 * 1024 * 1024 /* 10 MB/s */));
         for (ts, len) in backup_tss {
             let mut req = BackupRequest::new();
             req.set_start_key(vec![]);

components/backup/src/writer.rs

@@ -3,9 +3,10 @@
 use std::sync::Arc;
 use std::time::Instant;
 
-use engine::rocks::util::io_limiter::{IOLimiter, LimitReader};
 use engine::{CF_DEFAULT, CF_WRITE, DB};
+use engine_rocks::RocksIOLimiter;
 use engine_rocks::{RocksEngine, RocksSstWriter, RocksSstWriterBuilder};
+use engine_traits::LimitReader;
 use engine_traits::{SstWriter, SstWriterBuilder};
 use external_storage::ExternalStorage;
 use kvproto::backup::File;
@@ -60,7 +61,7 @@ impl Writer {
         name: &str,
         cf: &'static str,
         buf: &mut Vec<u8>,
-        limiter: Option<Arc<IOLimiter>>,
+        limiter: Option<Arc<RocksIOLimiter>>,
         storage: &dyn ExternalStorage,
     ) -> Result<File> {
         buf.reserve(self.writer.file_size() as _);
@@ -93,12 +94,16 @@ pub struct BackupWriter {
     name: String,
     default: Writer,
     write: Writer,
-    limiter: Option<Arc<IOLimiter>>,
+    limiter: Option<Arc<RocksIOLimiter>>,
 }
 
 impl BackupWriter {
     /// Create a new BackupWriter.
-    pub fn new(db: Arc<DB>, name: &str, limiter: Option<Arc<IOLimiter>>) -> Result<BackupWriter> {
+    pub fn new(
+        db: Arc<DB>,
+        name: &str,
+        limiter: Option<Arc<RocksIOLimiter>>,
+    ) -> Result<BackupWriter> {
         let default = RocksSstWriterBuilder::new()
             .set_in_memory(true)
             .set_cf(CF_DEFAULT)

components/engine/Cargo.toml

@@ -26,6 +26,7 @@ serde_derive = "1.0"
 toml = "0.4"
 hex = "0.3"
 tikv_util = { path = "../tikv_util" }
+engine_traits = { path = "../engine_traits" }
 
 [dependencies.prometheus]
 git = "https://github.com/pingcap/rust-prometheus.git"

components/engine/src/iterable.rs

@@ -5,141 +5,36 @@ pub use crate::rocks::{DBIterator, ReadOptions, DB};
 use crate::Result;
 use tikv_util::keybuilder::KeyBuilder;
 
-#[derive(Clone, PartialEq)]
-enum SeekMode {
-    TotalOrder,
-    Prefix,
-}
+pub use engine_traits::IterOptions as IterOption;
+pub use engine_traits::SeekMode;
 
-pub struct IterOption {
-    lower_bound: Option<KeyBuilder>,
-    upper_bound: Option<KeyBuilder>,
-    prefix_same_as_start: bool,
-    fill_cache: bool,
-    // only supported when Titan enabled, otherwise it doesn't take effect.
-    titan_key_only: bool,
-    seek_mode: SeekMode,
+pub trait IterOptionsExt {
+    fn build_read_opts(self) -> ReadOptions;
 }
 
-impl IterOption {
-    pub fn new(
-        lower_bound: Option<KeyBuilder>,
-        upper_bound: Option<KeyBuilder>,
-        fill_cache: bool,
-    ) -> IterOption {
-        IterOption {
-            lower_bound,
-            upper_bound,
-            prefix_same_as_start: false,
-            fill_cache,
-            titan_key_only: false,
-            seek_mode: SeekMode::TotalOrder,
-        }
-    }
-
-    #[inline]
-    pub fn use_prefix_seek(mut self) -> IterOption {
-        self.seek_mode = SeekMode::Prefix;
-        self
-    }
-
-    #[inline]
-    pub fn total_order_seek_used(&self) -> bool {
-        self.seek_mode == SeekMode::TotalOrder
-    }
-
-    #[inline]
-    pub fn fill_cache(&mut self, v: bool) {
-        self.fill_cache = v;
-    }
-
-    #[inline]
-    pub fn titan_key_only(&mut self, v: bool) {
-        self.titan_key_only = v;
-    }
-
-    #[inline]
-    pub fn lower_bound(&self) -> Option<&[u8]> {
-        self.lower_bound.as_ref().map(|v| v.as_slice())
-    }
-
-    #[inline]
-    pub fn set_lower_bound(&mut self, bound: &[u8], reserved_prefix_len: usize) {
-        let builder = KeyBuilder::from_slice(bound, reserved_prefix_len, 0);
-        self.lower_bound = Some(builder);
-    }
-
-    pub fn set_vec_lower_bound(&mut self, bound: Vec<u8>) {
-        self.lower_bound = Some(KeyBuilder::from_vec(bound, 0, 0));
-    }
-
-    pub fn set_lower_bound_prefix(&mut self, prefix: &[u8]) {
-        if let Some(ref mut builder) = self.lower_bound {
-            builder.set_prefix(prefix);
-        }
-    }
-
-    #[inline]
-    pub fn upper_bound(&self) -> Option<&[u8]> {
-        self.upper_bound.as_ref().map(|v| v.as_slice())
-    }
-
-    #[inline]
-    pub fn set_upper_bound(&mut self, bound: &[u8], reserved_prefix_len: usize) {
-        let builder = KeyBuilder::from_slice(bound, reserved_prefix_len, 0);
-        self.upper_bound = Some(builder);
-    }
-
-    pub fn set_vec_upper_bound(&mut self, bound: Vec<u8>) {
-        self.upper_bound = Some(KeyBuilder::from_vec(bound, 0, 0));
-    }
-
-    pub fn set_upper_bound_prefix(&mut self, prefix: &[u8]) {
-        if let Some(ref mut builder) = self.upper_bound {
-            builder.set_prefix(prefix);
-        }
-    }
-
-    #[inline]
-    pub fn set_prefix_same_as_start(mut self, enable: bool) -> IterOption {
-        self.prefix_same_as_start = enable;
-        self
-    }
-
-    pub fn build_read_opts(self) -> ReadOptions {
+impl IterOptionsExt for IterOption {
+    fn build_read_opts(self) -> ReadOptions {
         let mut opts = ReadOptions::new();
-        opts.fill_cache(self.fill_cache);
-        if self.titan_key_only {
+        opts.fill_cache(self.fill_cache());
+        if self.titan_key_only() {
             opts.set_titan_key_only(true);
         }
         if self.total_order_seek_used() {
             opts.set_total_order_seek(true);
-        } else if self.prefix_same_as_start {
+        } else if self.prefix_same_as_start() {
             opts.set_prefix_same_as_start(true);
         }
-        if let Some(builder) = self.lower_bound {
-            opts.set_iterate_lower_bound(builder.build());
+        let (lower, upper) = self.build_bounds();
+        if let Some(lower) = lower {
+            opts.set_iterate_lower_bound(lower);
         }
-        if let Some(builder) = self.upper_bound {
-            opts.set_iterate_upper_bound(builder.build());
+        if let Some(upper) = upper {
+            opts.set_iterate_upper_bound(upper);
         }
         opts
     }
 }
 
-impl Default for IterOption {
-    fn default() -> IterOption {
-        IterOption {
-            lower_bound: None,
-            upper_bound: None,
-            prefix_same_as_start: false,
-            fill_cache: true,
-            titan_key_only: false,
-            seek_mode: SeekMode::TotalOrder,
-        }
-    }
-}
-
 // TODO: refactor this trait into rocksdb trait.
 pub trait Iterable {
     fn new_iterator(&self, iter_opt: IterOption) -> DBIterator<&DB>;

components/engine/src/rocks/db.rs

@@ -3,6 +3,7 @@
 use std::option::Option;
 
 use super::{util, DBIterator, DBVector, WriteBatch, DB};
+use crate::iterable::IterOptionsExt;
 use crate::{IterOption, Iterable, Mutable, Peekable, Result};
 
 impl Peekable for DB {

components/engine/src/rocks/snapshot.rs

@@ -6,8 +6,10 @@ use std::option::Option;
 use std::sync::Arc;
 
 use super::{CFHandle, DBVector, ReadOptions, UnsafeSnap, DB};
+use crate::iterable::IterOptionsExt;
 use crate::{DBIterator, Error, IterOption, Iterable, Peekable, Result};
 
+#[repr(C)] // Guarantee same representation as in engine_rocks
 pub struct Snapshot {
     db: Arc<DB>,
     snap: UnsafeSnap,
@@ -43,23 +45,6 @@ impl Snapshot {
     pub fn get_db(&self) -> Arc<DB> {
         Arc::clone(&self.db)
     }
-
-    pub fn db_iterator(&self, iter_opt: IterOption) -> DBIterator<Arc<DB>> {
-        let mut opt = iter_opt.build_read_opts();
-        unsafe {
-            opt.set_snapshot(&self.snap);
-        }
-        DBIterator::new(Arc::clone(&self.db), opt)
-    }
-
-    pub fn db_iterator_cf(&self, cf: &str, iter_opt: IterOption) -> Result<DBIterator<Arc<DB>>> {
-        let handle = super::util::get_cf_handle(&self.db, cf)?;
-        let mut opt = iter_opt.build_read_opts();
-        unsafe {
-            opt.set_snapshot(&self.snap);
-        }
-        Ok(DBIterator::new_cf(Arc::clone(&self.db), handle, opt))
-    }
 }
 
 impl Debug for Snapshot {