rocksdb: add Env related API

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -960,6 +960,8 @@ crocksdb_options_set_compaction_priority(crocksdb_options_t *, unsigned char);
 
 extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_delayed_write_rate(
     crocksdb_options_t*, uint64_t);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_env(
+    crocksdb_options_t *, crocksdb_env_t *);
 
 /* RateLimiter */
 extern C_ROCKSDB_LIBRARY_API crocksdb_ratelimiter_t* crocksdb_ratelimiter_create(

librocksdb_sys/src/lib.rs

@@ -35,6 +35,7 @@ pub enum DBWriteBatch {}
 pub enum DBComparator {}
 pub enum DBFlushOptions {}
 pub enum DBCompactionFilter {}
+pub enum Env {}
 pub enum EnvOptions {}
 pub enum SstFileWriter {}
 pub enum IngestExternalFileOptions {}
@@ -216,6 +217,13 @@ pub enum DBTableProperty {
     CompressionName = 17,
 }
 
+#[repr(C)]
+pub enum Priority {
+    LOW = 0,
+    HIGH = 1,
+    TOTAL = 2,
+}
+
 pub fn error_message(ptr: *mut c_char) -> String {
     let c_str = unsafe { CStr::from_ptr(ptr) };
     let s = format!("{}", c_str.to_string_lossy());
@@ -393,6 +401,7 @@ extern "C" {
                                                                     v: u64);
     pub fn crocksdb_options_set_compaction_priority(options: *mut DBOptions,
                                                     v: CompactionPriority);
+    pub fn crocksdb_options_set_env(options: *mut DBOptions, env: *const Env);
     pub fn crocksdb_filterpolicy_create_bloom_full(bits_per_key: c_int) -> *mut DBFilterPolicy;
     pub fn crocksdb_filterpolicy_create_bloom(bits_per_key: c_int) -> *mut DBFilterPolicy;
     pub fn crocksdb_open(options: *mut DBOptions,
@@ -750,6 +759,11 @@ extern "C" {
                                                           ignore_snapshot: bool);
     pub fn crocksdb_compactionfilter_destroy(filter: *mut DBCompactionFilter);
 
+    // Env
+    pub fn crocksdb_create_default_env() -> *mut Env;
+    pub fn crocksdb_env_set_background_threads(opt: *mut Env, pri: Priority);
+    pub fn crocksdb_env_destroy(opt: *mut Env);
+    pub fn crocksdb_env_join_all_threads(opt: *mut Env);
     // EnvOptions
     pub fn crocksdb_envoptions_create() -> *mut EnvOptions;
     pub fn crocksdb_envoptions_destroy(opt: *mut EnvOptions);

src/lib.rs

@@ -40,7 +40,8 @@ pub use merge_operator::MergeOperands;
 pub use rocksdb::{DB, DBIterator, DBVector, Kv, SeekKey, Writable, WriteBatch, CFHandle, Range,
                   BackupEngine, SstFileWriter};
 pub use rocksdb_options::{BlockBasedOptions, Options, ReadOptions, WriteOptions, RestoreOptions,
-                          IngestExternalFileOptions, EnvOptions, HistogramData, CompactOptions};
+                          IngestExternalFileOptions, Env, EnvOptions, HistogramData,
+                          CompactOptions};
 pub use slice_transform::SliceTransform;
 pub use table_properties::{TableProperties, TablePropertiesCollection,
                            TablePropertiesCollectionView, UserCollectedProperties};

src/rocksdb_options.rs

@@ -19,7 +19,7 @@ use comparator::{self, ComparatorCallback, compare_callback};
 use crocksdb_ffi::{self, DBOptions, DBWriteOptions, DBBlockBasedTableOptions, DBReadOptions,
                    DBRestoreOptions, DBCompressionType, DBRecoveryMode, DBSnapshot, DBInstance,
                    DBFlushOptions, DBStatisticsTickerType, DBStatisticsHistogramType,
-                   DBRateLimiter, DBInfoLogLevel, DBCompactOptions};
+                   DBRateLimiter, DBInfoLogLevel, DBCompactOptions, Priority};
 use event_listener::{EventListener, new_event_listener};
 use libc::{self, c_int, size_t, c_void};
 use merge_operator::{self, MergeOperatorCallback, full_merge_callback, partial_merge_callback};
@@ -294,6 +294,31 @@ impl Drop for CompactOptions {
     }
 }
 
+pub struct Env {
+    pub inner: *mut crocksdb_ffi::Env,
+}
+
+impl Env {
+    pub fn new() -> Env {
+        unsafe { Env { inner: crocksdb_ffi::crocksdb_create_default_env() } }
+    }
+
+    pub fn set_background_threads(&mut self, n: Priority) {
+        unsafe {
+            crocksdb_ffi::crocksdb_env_set_background_threads(self.inner, n);
+        }
+    }
+}
+
+impl Drop for Env {
+    fn drop(&mut self) {
+        unsafe {
+            crocksdb_ffi::crocksdb_env_join_all_threads(self.inner);
+            crocksdb_ffi::crocksdb_env_destroy(self.inner)
+        }
+    }
+}
+
 pub struct Options {
     pub inner: *mut DBOptions,
     filter: Option<CompactionFilterHandle>,
@@ -949,6 +974,10 @@ impl Options {
             crocksdb_ffi::crocksdb_options_set_allow_concurrent_memtable_write(self.inner, v);
         }
     }
+
+    pub fn set_env(&mut self, env: &Env) {
+        unsafe { crocksdb_ffi::crocksdb_options_set_env(self.inner, env.inner) }
+    }
 }
 
 pub struct FlushOptions {

tests/test_rocksdb_options.rs

@@ -11,11 +11,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use rocksdb::{DB, Options, BlockBasedOptions, WriteOptions, SliceTransform, Writable,
+use rocksdb::{DB, Env, Options, BlockBasedOptions, WriteOptions, SliceTransform, Writable,
               CompactOptions};
 use rocksdb::crocksdb_ffi::{DBStatisticsHistogramType as HistogramType,
                             DBStatisticsTickerType as TickerType, DBInfoLogLevel as InfoLogLevel,
-                            CompactionPriority, DBCompressionType};
+                            CompactionPriority, DBCompressionType, Priority};
 use std::path::Path;
 use std::thread;
 use std::time::Duration;
@@ -406,3 +406,29 @@ fn test_clone_options() {
     let opts2 = opts.clone();
     assert_eq!(opts.get_compression(), opts2.get_compression());
 }
+
+#[test]
+fn test_set_background_threads() {
+    let path = TempDir::new("_rust_rocksdb_background_threads").expect("");
+    let mut env = Env::new();
+    env.set_background_threads(Priority::HIGH);
+    let mut opt = Options::new();
+    opt.create_if_missing(true);
+    opt.set_env(&env);
+    DB::open(opt, path.path().to_str().unwrap()).unwrap();
+}
+
+#[test]
+fn test_two_db_share_one_env() {
+    let path1 = TempDir::new("_rust_rocksdb_share_env1").expect("");
+    let path2 = TempDir::new("_rust_rocksdb_share_env2").expect("");
+    let env = Env::new();
+    let mut opt1 = Options::new();
+    opt1.create_if_missing(true);
+    opt1.set_env(&env);
+    let mut opt2 = Options::new();
+    opt2.create_if_missing(true);
+    opt2.set_env(&env);
+    DB::open(opt1, path1.path().to_str().unwrap()).unwrap();
+    DB::open(opt2, path2.path().to_str().unwrap()).unwrap();
+}