Add an option to disable write stall and expose num_input_files of compaction job info (#649)

* add an option to disable write stall
Signed-off-by: Connor1996 <zbk602423539@gmail.com>

librocksdb_sys/crocksdb/c.cc

@@ -2154,6 +2154,11 @@ int crocksdb_compactionjobinfo_output_level(
   return info->rep.output_level;
 }
 
+size_t crocksdb_compactionjobinfo_num_input_files(
+    const crocksdb_compactionjobinfo_t* info) {
+  return info->rep.stats.num_input_files;
+}
+
 size_t crocksdb_compactionjobinfo_num_input_files_at_output_level(
     const crocksdb_compactionjobinfo_t* info) {
   return info->rep.stats.num_input_files_at_output_level;
@@ -2970,6 +2975,16 @@ int crocksdb_options_get_max_background_compactions(
   return opt->rep.max_background_compactions;
 }
 
+void crocksdb_options_set_base_background_compactions(crocksdb_options_t* opt,
+                                                      int n) {
+  opt->rep.base_background_compactions = n;
+}
+
+int crocksdb_options_get_base_background_compactions(
+    const crocksdb_options_t* opt) {
+  return opt->rep.base_background_compactions;
+}
+
 void crocksdb_options_set_max_background_flushes(crocksdb_options_t* opt,
                                                  int n) {
   opt->rep.max_background_flushes = n;
@@ -3059,6 +3074,16 @@ int crocksdb_options_get_disable_auto_compactions(
   return opt->rep.disable_auto_compactions;
 }
 
+void crocksdb_options_set_disable_write_stall(crocksdb_options_t* opt,
+                                              unsigned char disable) {
+  opt->rep.disable_write_stall = disable;
+}
+
+unsigned char crocksdb_options_get_disable_write_stall(
+    const crocksdb_options_t* opt) {
+  return opt->rep.disable_write_stall;
+}
+
 void crocksdb_options_set_delete_obsolete_files_period_micros(
     crocksdb_options_t* opt, uint64_t v) {
   opt->rep.delete_obsolete_files_period_micros = v;

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -843,6 +843,8 @@ crocksdb_compactionjobinfo_total_input_bytes(
 extern C_ROCKSDB_LIBRARY_API uint64_t
 crocksdb_compactionjobinfo_total_output_bytes(
     const crocksdb_compactionjobinfo_t*);
+extern C_ROCKSDB_LIBRARY_API size_t crocksdb_compactionjobinfo_num_input_files(
+    const crocksdb_compactionjobinfo_t* info);
 extern C_ROCKSDB_LIBRARY_API size_t
 crocksdb_compactionjobinfo_num_input_files_at_output_level(
     const crocksdb_compactionjobinfo_t* info);
@@ -1120,6 +1122,10 @@ extern C_ROCKSDB_LIBRARY_API void
 crocksdb_options_set_max_background_compactions(crocksdb_options_t*, int);
 extern C_ROCKSDB_LIBRARY_API int
 crocksdb_options_get_max_background_compactions(const crocksdb_options_t*);
+extern C_ROCKSDB_LIBRARY_API void
+crocksdb_options_set_base_background_compactions(crocksdb_options_t*, int);
+extern C_ROCKSDB_LIBRARY_API int
+crocksdb_options_get_base_background_compactions(const crocksdb_options_t*);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_max_background_flushes(
     crocksdb_options_t*, int);
 extern C_ROCKSDB_LIBRARY_API int crocksdb_options_get_max_background_flushes(
@@ -1223,6 +1229,10 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_disable_auto_compactions(
     crocksdb_options_t*, int);
 extern C_ROCKSDB_LIBRARY_API int crocksdb_options_get_disable_auto_compactions(
     const crocksdb_options_t*);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_disable_write_stall(
+    crocksdb_options_t*, unsigned char);
+extern C_ROCKSDB_LIBRARY_API unsigned char
+crocksdb_options_get_disable_write_stall(const crocksdb_options_t*);
 extern C_ROCKSDB_LIBRARY_API void
 crocksdb_options_set_delete_obsolete_files_period_micros(crocksdb_options_t*,
                                                          uint64_t);

rocksdb @ f95aaf8d

-Subproject commit 2aba25166553d4410564f0f41b55e772e672434d
+Subproject commit f95aaf8dac667334bb6314bd06d80749f63177bb

librocksdb_sys/src/lib.rs

@@ -760,6 +760,11 @@ extern "C" {
         max_bg_compactions: c_int,
     );
     pub fn crocksdb_options_get_max_background_compactions(options: *const Options) -> c_int;
+    pub fn crocksdb_options_set_base_background_compactions(
+        options: *mut Options,
+        base_bg_compactions: c_int,
+    );
+    pub fn crocksdb_options_get_base_background_compactions(options: *const Options) -> c_int;
     pub fn crocksdb_options_set_max_background_flushes(
         options: *mut Options,
         max_bg_flushes: c_int,
@@ -767,6 +772,8 @@ extern "C" {
     pub fn crocksdb_options_get_max_background_flushes(options: *const Options) -> c_int;
     pub fn crocksdb_options_set_disable_auto_compactions(options: *mut Options, v: c_int);
     pub fn crocksdb_options_get_disable_auto_compactions(options: *const Options) -> c_int;
+    pub fn crocksdb_options_set_disable_write_stall(options: *mut Options, v: bool);
+    pub fn crocksdb_options_get_disable_write_stall(options: *const Options) -> bool;
     pub fn crocksdb_options_set_report_bg_io_stats(options: *mut Options, v: c_int);
     pub fn crocksdb_options_set_compaction_readahead_size(options: *mut Options, v: size_t);
     pub fn crocksdb_options_set_wal_recovery_mode(options: *mut Options, mode: DBRecoveryMode);
@@ -2177,6 +2184,7 @@ extern "C" {
     pub fn crocksdb_compactionjobinfo_output_records(info: *const DBCompactionJobInfo) -> u64;
     pub fn crocksdb_compactionjobinfo_total_input_bytes(info: *const DBCompactionJobInfo) -> u64;
     pub fn crocksdb_compactionjobinfo_total_output_bytes(info: *const DBCompactionJobInfo) -> u64;
+    pub fn crocksdb_compactionjobinfo_num_input_files(info: *const DBCompactionJobInfo) -> size_t;
     pub fn crocksdb_compactionjobinfo_num_input_files_at_output_level(
         info: *const DBCompactionJobInfo,
     ) -> size_t;

src/rocksdb_options.rs

@@ -834,6 +834,16 @@ impl DBOptions {
         unsafe { crocksdb_ffi::crocksdb_options_get_max_background_compactions(self.inner) as i32 }
     }
 
+    pub fn set_base_background_compactions(&mut self, n: c_int) {
+        unsafe {
+            crocksdb_ffi::crocksdb_options_set_base_background_compactions(self.inner, n);
+        }
+    }
+
+    pub fn get_base_background_compactions(&self) -> i32 {
+        unsafe { crocksdb_ffi::crocksdb_options_get_base_background_compactions(self.inner) as i32 }
+    }
+
     pub fn set_max_background_flushes(&mut self, n: c_int) {
         unsafe {
             crocksdb_ffi::crocksdb_options_set_max_background_flushes(self.inner, n);
@@ -1741,6 +1751,14 @@ impl ColumnFamilyOptions {
         unsafe { crocksdb_ffi::crocksdb_options_get_disable_auto_compactions(self.inner) == 1 }
     }
 
+    pub fn set_disable_write_stall(&mut self, disable: bool) {
+        unsafe { crocksdb_ffi::crocksdb_options_set_disable_write_stall(self.inner, disable) }
+    }
+
+    pub fn get_disable_write_stall(&self) -> bool {
+        unsafe { crocksdb_ffi::crocksdb_options_get_disable_write_stall(self.inner) }
+    }
+
     pub fn set_block_based_table_factory(&mut self, factory: &BlockBasedOptions) {
         unsafe {
             crocksdb_ffi::crocksdb_options_set_block_based_table_factory(self.inner, factory.inner);

tests/cases/test_rocksdb_options.rs

@@ -594,9 +594,29 @@ fn test_set_max_background_compactions_and_flushes() {
     opts.create_if_missing(true);
     opts.set_max_background_compactions(4);
     opts.set_max_background_flushes(1);
+    opts.set_base_background_compactions(2);
+    assert_eq!(opts.get_max_background_compactions(), 4);
+    assert_eq!(opts.get_max_background_flushes(), 1);
+    assert_eq!(opts.get_base_background_compactions(), 2);
     DB::open(opts, path.path().to_str().unwrap()).unwrap();
 }
 
+#[test]
+fn test_set_disable_write_stall() {
+    let path = tempdir_with_prefix("_rust_rocksdb_disable_write_stall");
+    let mut opts = DBOptions::new();
+    opts.create_if_missing(true);
+    let mut cfopts = ColumnFamilyOptions::new();
+    cfopts.set_disable_write_stall(true);
+    assert_eq!(cfopts.get_disable_write_stall(), true);
+    DB::open_cf(
+        opts.clone(),
+        path.path().to_str().unwrap(),
+        vec![("default", cfopts)],
+    )
+    .unwrap();
+}
+
 #[test]
 fn test_set_compaction_pri() {
     let path = tempdir_with_prefix("_rust_rocksdb_compaction_pri");