add approximate memtable stats (#110)

* add approximate memtable stats

librocksdb_sys/crocksdb/c.cc

@@ -1010,6 +1010,28 @@ void crocksdb_approximate_sizes_cf(
   delete[] ranges;
 }
 
+void crocksdb_approximate_memtable_stats(
+    const crocksdb_t* db,
+    const char* range_start_key, size_t range_start_key_len,
+    const char* range_limit_key, size_t range_limit_key_len,
+    uint64_t* count, uint64_t* size) {
+  auto start = Slice(range_start_key, range_start_key_len);
+  auto limit = Slice(range_limit_key, range_limit_key_len);
+  Range range(start, limit);
+  db->rep->GetApproximateMemTableStats(range, count, size);
+}
+
+void crocksdb_approximate_memtable_stats_cf(
+    const crocksdb_t* db, const crocksdb_column_family_handle_t* cf,
+    const char* range_start_key, size_t range_start_key_len,
+    const char* range_limit_key, size_t range_limit_key_len,
+    uint64_t* count, uint64_t* size) {
+  auto start = Slice(range_start_key, range_start_key_len);
+  auto limit = Slice(range_limit_key, range_limit_key_len);
+  Range range(start, limit);
+  db->rep->GetApproximateMemTableStats(cf->rep, range, count, size);
+}
+
 void crocksdb_delete_file(
     crocksdb_t* db,
     const char* name) {

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -370,6 +370,18 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_approximate_sizes_cf(
     const size_t* range_start_key_len, const char* const* range_limit_key,
     const size_t* range_limit_key_len, uint64_t* sizes);
 
+extern C_ROCKSDB_LIBRARY_API void crocksdb_approximate_memtable_stats(
+    const crocksdb_t* db,
+    const char* range_start_key, size_t range_start_key_len,
+    const char* range_limit_key, size_t range_limit_key_len,
+    uint64_t* count, uint64_t* size);
+
+extern C_ROCKSDB_LIBRARY_API void crocksdb_approximate_memtable_stats_cf(
+    const crocksdb_t* db, const crocksdb_column_family_handle_t* cf,
+    const char* range_start_key, size_t range_start_key_len,
+    const char* range_limit_key, size_t range_limit_key_len,
+    uint64_t* count, uint64_t* size);
+
 extern C_ROCKSDB_LIBRARY_API void crocksdb_compact_range(crocksdb_t* db,
                                                       const char* start_key,
                                                       size_t start_key_len,

librocksdb_sys/src/lib.rs

@@ -688,6 +688,21 @@ extern "C" {
                                          range_limit_key: *const *const u8,
                                          range_limit_key_len: *const size_t,
                                          sizes: *mut uint64_t);
+    pub fn crocksdb_approximate_memtable_stats(db: *const DBInstance,
+                                               range_start_key: *const u8,
+                                               range_start_key_len: size_t,
+                                               range_limit_key: *const u8,
+                                               range_limit_key_len: size_t,
+                                               count: *mut uint64_t,
+                                               size: *mut uint64_t);
+    pub fn crocksdb_approximate_memtable_stats_cf(db: *const DBInstance,
+                                                  cf: *const DBCFHandle,
+                                                  range_start_key: *const u8,
+                                                  range_start_key_len: size_t,
+                                                  range_limit_key: *const u8,
+                                                  range_limit_key_len: size_t,
+                                                  count: *mut uint64_t,
+                                                  size: *mut uint64_t);
     pub fn crocksdb_compactoptions_create() -> *mut DBCompactOptions;
     pub fn crocksdb_compactoptions_destroy(opt: *mut DBCompactOptions);
     pub fn crocksdb_compactoptions_set_exclusive_manual_compaction(opt: *mut DBCompactOptions,

src/rocksdb.rs

@@ -820,6 +820,37 @@ impl DB {
         sizes
     }
 
+    // Return the approximate number of records and size in the range of memtables.
+    pub fn get_approximate_memtable_stats(&self, range: &Range) -> (u64, u64) {
+        let (mut count, mut size) = (0, 0);
+        unsafe {
+            crocksdb_ffi::crocksdb_approximate_memtable_stats(self.inner,
+                                                              range.start_key.as_ptr(),
+                                                              range.start_key.len(),
+                                                              range.end_key.as_ptr(),
+                                                              range.end_key.len(),
+                                                              &mut count,
+                                                              &mut size);
+        }
+        (count, size)
+    }
+
+    // Return the approximate number of records and size in the range of memtables of the cf.
+    pub fn get_approximate_memtable_stats_cf(&self, cf: &CFHandle, range: &Range) -> (u64, u64) {
+        let (mut count, mut size) = (0, 0);
+        unsafe {
+            crocksdb_ffi::crocksdb_approximate_memtable_stats_cf(self.inner,
+                                                                 cf.inner,
+                                                                 range.start_key.as_ptr(),
+                                                                 range.start_key.len(),
+                                                                 range.end_key.as_ptr(),
+                                                                 range.end_key.len(),
+                                                                 &mut count,
+                                                                 &mut size);
+        }
+        (count, size)
+    }
+
     pub fn compact_range(&self, start_key: Option<&[u8]>, end_key: Option<&[u8]>) {
         unsafe {
             let (start, s_len) = start_key.map_or((ptr::null(), 0), |k| (k.as_ptr(), k.len()));
@@ -1927,4 +1958,32 @@ mod test {
         assert_eq!(key_versions[1].value, "value3");
         assert_eq!(key_versions[1].seq, 3);
     }
+
+    #[test]
+    fn test_get_approximate_memtable_stats() {
+        let mut opts = DBOptions::new();
+        opts.create_if_missing(true);
+        let path = TempDir::new("_rust_rocksdb_get_approximate_memtable_stats").expect("");
+        let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
+
+        let samples = [(b"key1", b"value1"),
+                       (b"key2", b"value2"),
+                       (b"key3", b"value3"),
+                       (b"key4", b"value4")];
+
+        for &(k, v) in &samples {
+            db.put(k, v).unwrap();
+        }
+
+        let range = Range::new(b"a", b"z");
+
+        let (count, size) = db.get_approximate_memtable_stats(&range);
+        assert!(count > 0);
+        assert!(size > 0);
+
+        let cf = db.cf_handle("default").unwrap();
+        let (count, size) = db.get_approximate_memtable_stats_cf(cf, &range);
+        assert!(count > 0);
+        assert!(size > 0);
+    }
 }