set & get block cache capacity via cf_options (#198)

* change block cache capacity

librocksdb_sys/crocksdb/c.cc

@@ -1674,16 +1674,40 @@ void crocksdb_options_set_wal_bytes_per_sync(crocksdb_options_t *opt, uint64_t v
   opt->rep.wal_bytes_per_sync = v;
 }
 
-size_t crocksdb_options_get_block_cache_usage(crocksdb_options_t *opt) {
+static BlockBasedTableOptions* get_block_based_table_options(crocksdb_options_t *opt) {
   if (opt && opt->rep.table_factory != nullptr) {
     void* table_opt = opt->rep.table_factory->GetOptions();
     if (table_opt && strcmp(opt->rep.table_factory->Name(), block_base_table_str) == 0) {
-      auto opts = static_cast<BlockBasedTableOptions*>(table_opt);
-      if (opts->block_cache) {
-        return opts->block_cache->GetUsage();
-      }
+      return static_cast<BlockBasedTableOptions*>(table_opt);
     }
   }
+  return nullptr;
+}
+
+size_t crocksdb_options_get_block_cache_usage(crocksdb_options_t *opt) {
+  auto opts = get_block_based_table_options(opt);
+  if (opts && opts->block_cache) {
+    return opts->block_cache->GetUsage();
+  }
+  return 0;
+}
+
+void crocksdb_options_set_block_cache_capacity(crocksdb_options_t* opt, size_t capacity, char **errptr) {
+  Status s;
+  auto opts = get_block_based_table_options(opt);
+  if (opts && opts->block_cache) {
+    opts->block_cache->SetCapacity(capacity);
+  } else {
+    s = Status::InvalidArgument("failed to get block based table options");
+  }
+  SaveError(errptr, s);
+}
+
+size_t crocksdb_options_get_block_cache_capacity(crocksdb_options_t* opt) {
+  auto opts = get_block_based_table_options(opt);
+  if (opts && opts->block_cache) {
+    return opts->block_cache->GetCapacity();
+  }
   return 0;
 }
 

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -629,6 +629,10 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_block_based_table_factory
 
 extern C_ROCKSDB_LIBRARY_API size_t crocksdb_options_get_block_cache_usage(
     crocksdb_options_t *opt);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_block_cache_capacity(
+  crocksdb_options_t* opt, size_t capacity, char **errptr);
+extern C_ROCKSDB_LIBRARY_API size_t crocksdb_options_get_block_cache_capacity(
+  crocksdb_options_t* opt);
 
 /* Flush job info */
 

librocksdb_sys/src/lib.rs

@@ -539,6 +539,8 @@ extern "C" {
     pub fn crocksdb_options_set_ratelimiter(options: *mut Options, limiter: *mut DBRateLimiter);
     pub fn crocksdb_options_set_info_log(options: *mut Options, logger: *mut DBLogger);
     pub fn crocksdb_options_get_block_cache_usage(options: *const Options) -> usize;
+    pub fn crocksdb_options_set_block_cache_capacity(options: *const Options, capacity: usize, err: *mut *mut c_char);
+    pub fn crocksdb_options_get_block_cache_capacity(options: *const Options) -> usize;
     pub fn crocksdb_ratelimiter_create(
         rate_bytes_per_sec: i64,
         refill_period_us: i64,

src/rocksdb_options.rs

@@ -1325,6 +1325,20 @@ impl ColumnFamilyOptions {
         unsafe { crocksdb_ffi::crocksdb_options_get_block_cache_usage(self.inner) as u64 }
     }
 
+    pub fn set_block_cache_capacity(&self, capacity: u64) -> Result<(), String> {
+        unsafe {
+           ffi_try!(crocksdb_options_set_block_cache_capacity(
+                self.inner,
+                capacity as usize
+            ));
+            Ok(())
+        }
+    }
+
+    pub fn get_block_cache_capacity(&self) -> u64 {
+        unsafe { crocksdb_ffi::crocksdb_options_get_block_cache_capacity(self.inner) as u64 }
+    }
+
     pub fn set_fifo_compaction_options(&mut self, fifo_opts: FifoCompactionOptions) {
         unsafe {
             crocksdb_ffi::crocksdb_options_set_fifo_compaction_options(self.inner, fifo_opts.inner);

tests/test_rocksdb_options.rs

@@ -369,6 +369,36 @@ fn test_get_block_cache_usage() {
     assert!(db.get_options().get_block_cache_usage() > 0);
 }
 
+#[test]
+fn test_block_cache_capacity() {
+    let path = TempDir::new("_rust_rocksdb_set_and_get_block_cache_capacity").expect("");
+
+    let mut opts = DBOptions::new();
+    let mut cf_opts = ColumnFamilyOptions::new();
+    opts.create_if_missing(true);
+    let mut block_opts = BlockBasedOptions::new();
+    block_opts.set_lru_cache(16 * 1024 * 1024, -1, 0, 0.0);
+    cf_opts.set_block_based_table_factory(&block_opts);
+    let db = DB::open_cf(
+        opts,
+        path.path().to_str().unwrap(),
+        vec![("default", cf_opts)],
+    ).unwrap();
+
+    assert_eq!(
+        db.get_options().get_block_cache_capacity(),
+        16 * 1024 * 1024
+    );
+
+    let opt = db.get_options();
+    opt.set_block_cache_capacity(32 * 1024 * 1024).unwrap();
+
+    assert_eq!(
+        db.get_options().get_block_cache_capacity(),
+        32 * 1024 * 1024
+    );
+}
+
 #[test]
 fn test_disable_block_cache() {
     let mut cf_opts = ColumnFamilyOptions::new();