Get block cache usage size (#34)

librocksdb_sys/crocksdb/c.cc

@@ -85,6 +85,8 @@ using std::shared_ptr;
 
 extern "C" {
 
+const char* block_base_table_str = "BlockBasedTable";
+
 struct crocksdb_t                 { DB*               rep; };
 struct crocksdb_backup_engine_t   { BackupEngine*     rep; };
 struct crocksdb_backup_engine_info_t { std::vector<BackupInfo> rep; };
@@ -1053,12 +1055,6 @@ void crocksdb_enable_file_deletions(
   SaveError(errptr, db->rep->EnableFileDeletions(force));
 }
 
-crocksdb_options_t* crocksdb_get_options(const crocksdb_t* db) {
-  crocksdb_options_t* options = new crocksdb_options_t;
-  options->rep = db->rep->GetOptions();
-  return options;
-}
-
 crocksdb_options_t* crocksdb_get_options_cf(
     const crocksdb_t* db,
     crocksdb_column_family_handle_t* column_family) {
@@ -1493,6 +1489,16 @@ void crocksdb_options_set_max_subcompactions(crocksdb_options_t *opt, size_t v)
   opt->rep.max_subcompactions = v;
 }
 
+size_t crocksdb_options_get_block_cache_usage(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) {
+      return static_cast<BlockBasedTableOptions*>(table_opt)->block_cache->GetUsage();
+    }
+  }
+  return 0;
+}
+
 crocksdb_cuckoo_table_options_t*
 crocksdb_cuckoo_options_create() {
   return new crocksdb_cuckoo_table_options_t;

librocksdb_sys/crocksdb/rocksdb/c.h

@@ -362,9 +362,6 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_disable_file_deletions(crocksdb_t* db
 extern C_ROCKSDB_LIBRARY_API void crocksdb_enable_file_deletions(
     crocksdb_t* db, unsigned char force, char** errptr);
 
-extern C_ROCKSDB_LIBRARY_API crocksdb_options_t* crocksdb_get_options(
-    const crocksdb_t* db);
-
 extern C_ROCKSDB_LIBRARY_API crocksdb_options_t* crocksdb_get_options_cf(
     const crocksdb_t* db, crocksdb_column_family_handle_t* column_family);
 
@@ -533,6 +530,9 @@ crocksdb_block_based_options_set_pin_l0_filter_and_index_blocks_in_cache(
 extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_block_based_table_factory(
     crocksdb_options_t* opt, crocksdb_block_based_table_options_t* table_options);
 
+extern C_ROCKSDB_LIBRARY_API size_t crocksdb_options_get_block_cache_usage(
+    crocksdb_options_t *opt);
+
 /* Cuckoo table options */
 
 extern C_ROCKSDB_LIBRARY_API crocksdb_cuckoo_table_options_t*

librocksdb_sys/src/lib.rs

@@ -72,8 +72,8 @@ pub enum DBCompactionStyle {
 
 #[repr(C)]
 pub enum DBUniversalCompactionStyle {
-    crocksdb_similar_size_compaction_stop_style = 0,
-    crocksdb_total_size_compaction_stop_style = 1,
+    SimilarSize = 0,
+    TotalSize = 1,
 }
 
 #[derive(Copy, Clone, PartialEq)]
@@ -171,7 +171,6 @@ macro_rules! ffi_try {
 // TODO audit the use of boolean arguments, b/c I think they need to be u8
 // instead...
 extern "C" {
-    pub fn crocksdb_get_options(db: *mut DBInstance) -> *mut DBOptions;
     pub fn crocksdb_get_options_cf(db: *mut DBInstance, cf: *mut DBCFHandle) -> *mut DBOptions;
     pub fn crocksdb_options_create() -> *mut DBOptions;
     pub fn crocksdb_options_destroy(opts: *mut DBOptions);
@@ -291,6 +290,7 @@ extern "C" {
                                                                  ratio: c_double);
     pub fn crocksdb_options_set_ratelimiter(options: *mut DBOptions, limiter: *mut DBRateLimiter);
     pub fn crocksdb_options_set_info_log(options: *mut DBOptions, logger: *mut DBLogger);
+    pub fn crocksdb_options_get_block_cache_usage(options: *const DBOptions) -> usize;
     pub fn crocksdb_ratelimiter_create(rate_bytes_per_sec: i64,
                                        refill_period_us: i64,
                                        fairness: i32)

src/rocksdb.rs

@@ -917,8 +917,9 @@ impl DB {
     }
 
     pub fn get_options(&self) -> Options {
+        let cf = self.cf_handle("default").unwrap();
         unsafe {
-            let inner = crocksdb_ffi::crocksdb_get_options(self.inner);
+            let inner = crocksdb_ffi::crocksdb_get_options_cf(self.inner, cf.inner);
             Options::from_raw(inner)
         }
     }
@@ -1000,6 +1001,14 @@ impl DB {
 
         DB::open_default(restore_db_path)
     }
+
+    pub fn get_block_cache_usage(&self) -> u64 {
+        self.get_options().get_block_cache_usage()
+    }
+
+    pub fn get_block_cache_usage_cf(&self, cf: &CFHandle) -> u64 {
+        self.get_options_cf(cf).get_block_cache_usage()
+    }
 }
 
 impl Writable for DB {
@@ -1745,4 +1754,22 @@ mod test {
         let opts = Options::new();
         assert!(DB::destroy(&opts, path).is_ok());
     }
+
+    #[test]
+    fn block_cache_usage() {
+        let path = TempDir::new("_rust_rocksdb_block_cache_usage").expect("");
+        let db = DB::open_default(path.path().to_str().unwrap()).unwrap();
+
+        for i in 0..200 {
+            db.put(format!("k_{}", i).as_bytes(), b"v").unwrap();
+        }
+        db.flush(true).unwrap();
+        for i in 0..200 {
+            db.get(format!("k_{}", i).as_bytes()).unwrap();
+        }
+
+        assert!(db.get_block_cache_usage() > 0);
+        let cf_handle = db.cf_handle("default").unwrap();
+        assert!(db.get_block_cache_usage_cf(cf_handle) > 0);
+    }
 }

src/rocksdb_options.rs

@@ -792,6 +792,12 @@ impl Options {
 
         Ok(())
     }
+
+    pub fn get_block_cache_usage(&self) -> u64 {
+        unsafe {
+            crocksdb_ffi::crocksdb_options_get_block_cache_usage(self.inner) as u64
+        }
+    }
 }
 
 pub struct FlushOptions {

tests/test_rocksdb_options.rs

@@ -191,8 +191,7 @@ fn test_set_pin_l0_filter_and_index_blocks_in_cache() {
     let mut block_opts = BlockBasedOptions::new();
     block_opts.set_pin_l0_filter_and_index_blocks_in_cache(true);
     opts.set_block_based_table_factory(&block_opts);
-    let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
-    drop(db);
+    DB::open(opts, path.path().to_str().unwrap()).unwrap();
 }
 #[test]
 fn test_pending_compaction_bytes_limit() {
@@ -201,7 +200,7 @@ fn test_pending_compaction_bytes_limit() {
     opts.create_if_missing(true);
     opts.set_soft_pending_compaction_bytes_limit(64 * 1024 * 1024 * 1024);
     opts.set_hard_pending_compaction_bytes_limit(256 * 1024 * 1024 * 1024);
-    let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
+    DB::open(opts, path.path().to_str().unwrap()).unwrap();
 }
 
 #[test]
@@ -212,3 +211,27 @@ fn test_set_max_subcompactions() {
     opts.set_max_subcompactions(4);
     DB::open(opts, path.path().to_str().unwrap()).unwrap();
 }
+
+#[test]
+fn test_get_block_cache_usage() {
+    let path = TempDir::new("_rust_rocksdb_set_cache_and_index").expect("");
+
+    let mut opts = Options::new();
+    assert_eq!(opts.get_block_cache_usage(), 0);
+
+    opts.create_if_missing(true);
+    let mut block_opts = BlockBasedOptions::new();
+    block_opts.set_lru_cache(16 * 1024 * 1024);
+    opts.set_block_based_table_factory(&block_opts);
+    let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
+
+    for i in 0..200 {
+        db.put(format!("k_{}", i).as_bytes(), b"v").unwrap();
+    }
+    db.flush(true).unwrap();
+    for i in 0..200 {
+        db.get(format!("k_{}", i).as_bytes()).unwrap();
+    }
+
+    assert!(db.get_options().get_block_cache_usage() > 0);
+}