Merge pull request #21 from BusyJay/busyjay/introduce-unsafe-snap

introduce UnsafeSnap

src/main.rs

@@ -143,7 +143,7 @@ fn main() {
 
 #[cfg(test)]
 mod tests {
-    use rocksdb::{BlockBasedOptions, DB, Options, DBCompressionType};
+    use rocksdb::{BlockBasedOptions, DB, DBCompressionType, Options};
     use rocksdb::DBCompactionStyle::DBUniversal;
 
     #[allow(dead_code)]
@@ -151,13 +151,14 @@ mod tests {
                                      opts: &mut Options,
                                      blockopts: &mut BlockBasedOptions)
                                      -> DB {
-        let per_level_compression: [DBCompressionType; 7] = [DBCompressionType::DBNo,
-                                                DBCompressionType::DBNo,
-                                                DBCompressionType::DBNo,
-                                                DBCompressionType::DBLz4,
-                                                DBCompressionType::DBLz4,
-                                                DBCompressionType::DBLz4,
-                                                DBCompressionType::DBLz4];
+        let per_level_compression: [DBCompressionType; 7] =
+            [DBCompressionType::DBNo,
+             DBCompressionType::DBNo,
+             DBCompressionType::DBNo,
+             DBCompressionType::DBLz4,
+             DBCompressionType::DBLz4,
+             DBCompressionType::DBLz4,
+             DBCompressionType::DBLz4];
 
         opts.create_if_missing(true);
         opts.set_max_open_files(10000);

src/rocksdb.rs

@@ -44,9 +44,18 @@ pub struct ReadOptions {
     inner: rocksdb_ffi::DBReadOptions,
 }
 
+/// The UnsafeSnap must be destroyed by db, it maybe be leaked
+/// if not using it properly, hence named as unsafe.
+///
+/// This object is convenient for wrapping snapshot by yourself. In most
+/// cases, using `Snapshot` is enough.
+pub struct UnsafeSnap {
+    inner: rocksdb_ffi::DBSnapshot,
+}
+
 pub struct Snapshot<'a> {
     db: &'a DB,
-    inner: rocksdb_ffi::DBSnapshot,
+    snap: UnsafeSnap,
 }
 
 // We need to find a better way to add a lifetime in here.
@@ -69,7 +78,7 @@ impl<'a> From<&'a [u8]> for SeekKey<'a> {
 }
 
 impl<'a> DBIterator<'a> {
-    fn new(db: &'a DB, readopts: &ReadOptions) -> DBIterator<'a> {
+    pub fn new(db: &'a DB, readopts: &ReadOptions) -> DBIterator<'a> {
         unsafe {
             let iterator = rocksdb_ffi::rocksdb_create_iterator(db.inner,
                                                                 readopts.inner);
@@ -195,11 +204,11 @@ impl<'a> Drop for DBIterator<'a> {
 
 impl<'a> Snapshot<'a> {
     pub fn new(db: &DB) -> Snapshot {
-        let snapshot =
-            unsafe { rocksdb_ffi::rocksdb_create_snapshot(db.inner) };
-        Snapshot {
-            db: db,
-            inner: snapshot,
+        unsafe {
+            Snapshot {
+                db: db,
+                snap: db.unsafe_snap(),
+            }
         }
     }
 
@@ -209,13 +218,17 @@ impl<'a> Snapshot<'a> {
     }
 
     pub fn iter_opt(&self, mut opt: ReadOptions) -> DBIterator {
-        opt.set_snapshot(self);
+        unsafe {
+            opt.set_snapshot(&self.snap);
+        }
         DBIterator::new(self.db, &opt)
     }
 
     pub fn get(&self, key: &[u8]) -> Result<Option<DBVector>, String> {
         let mut readopts = ReadOptions::new();
-        readopts.set_snapshot(self);
+        unsafe {
+            readopts.set_snapshot(&self.snap);
+        }
         self.db.get_opt(key, &readopts)
     }
 
@@ -224,16 +237,16 @@ impl<'a> Snapshot<'a> {
                   key: &[u8])
                   -> Result<Option<DBVector>, String> {
         let mut readopts = ReadOptions::new();
-        readopts.set_snapshot(self);
+        unsafe {
+            readopts.set_snapshot(&self.snap);
+        }
         self.db.get_cf_opt(cf, key, &readopts)
     }
 }
 
 impl<'a> Drop for Snapshot<'a> {
     fn drop(&mut self) {
-        unsafe {
-            rocksdb_ffi::rocksdb_release_snapshot(self.db.inner, self.inner);
-        }
+        unsafe { self.db.release_snap(&self.snap) }
     }
 }
 
@@ -606,6 +619,14 @@ impl DB {
         Snapshot::new(self)
     }
 
+    pub unsafe fn unsafe_snap(&self) -> UnsafeSnap {
+        UnsafeSnap { inner: rocksdb_ffi::rocksdb_create_snapshot(self.inner) }
+    }
+
+    pub unsafe fn release_snap(&self, snap: &UnsafeSnap) {
+        rocksdb_ffi::rocksdb_release_snapshot(self.inner, snap.inner)
+    }
+
     pub fn put_opt(&self,
                    key: &[u8],
                    value: &[u8],
@@ -1061,11 +1082,9 @@ impl ReadOptions {
         }
     }
 
-    fn set_snapshot(&mut self, snapshot: &Snapshot) {
-        unsafe {
-            rocksdb_ffi::rocksdb_readoptions_set_snapshot(self.inner,
-                                                          snapshot.inner);
-        }
+    pub unsafe fn set_snapshot(&mut self, snapshot: &UnsafeSnap) {
+        rocksdb_ffi::rocksdb_readoptions_set_snapshot(self.inner,
+                                                      snapshot.inner);
     }
 }
 

src/rocksdb_options.rs

@@ -162,7 +162,8 @@ impl Options {
         }
     }
 
-    pub fn compression_per_level(&mut self, level_types: &[DBCompressionType]) {
+    pub fn compression_per_level(&mut self,
+                                 level_types: &[DBCompressionType]) {
         unsafe {
             rocksdb_ffi::rocksdb_options_set_compression_per_level(self.inner,
                                                                 level_types.as_ptr(),
@@ -371,9 +372,11 @@ impl Options {
     pub fn set_report_bg_io_stats(&mut self, enable: bool) {
         unsafe {
             if enable {
-                rocksdb_ffi::rocksdb_options_set_report_bg_io_stats(self.inner, 1);
+                rocksdb_ffi::rocksdb_options_set_report_bg_io_stats(self.inner,
+                                                                    1);
             } else {
-                rocksdb_ffi::rocksdb_options_set_report_bg_io_stats(self.inner, 0);
+                rocksdb_ffi::rocksdb_options_set_report_bg_io_stats(self.inner,
+                                                                    0);
             }
         }
     }