kv/RocksDBStore: Improve estimate_prefix_size.

Modify to include both files and memtables.
Gives better prediction of data size, especially gives non-zero size
when all data is still in memtables.

Signed-off-by: Adam Kupczyk <akupczyk@ibm.com>

diff --git a/src/kv/RocksDBStore.cc b/src/kv/RocksDBStore.cc
index ca63ea064841480429b4b4a44ba11e892ba0249c..53da9cd047e1bd838a644debb9685067cdcc0cea 100644 (file)
--- a/src/kv/RocksDBStore.cc
+++ b/src/kv/RocksDBStore.cc
@@ -1394,6 +1394,13 @@ bool RocksDBStore::get_property(
 int64_t RocksDBStore::estimate_prefix_size(const string& prefix,
                                           const string& key_prefix)
 {
+  // The default mode is to derive estimates based on
+  // sst files alone (INCLUDE_FILES).
+  // This gives an irritating result when a batch of keys is
+  // just commited but estimate keeps showing 0.
+  rocksdb::DB::SizeApproximationFlags flags(
+    rocksdb::DB::SizeApproximationFlags::INCLUDE_FILES |
+    rocksdb::DB::SizeApproximationFlags::INCLUDE_MEMTABLES);
   uint64_t size = 0;
   auto p_iter = cf_handles.find(prefix);
   if (p_iter != cf_handles.end()) {
@@ -1402,14 +1409,14 @@ int64_t RocksDBStore::estimate_prefix_size(const string& prefix,
       string start = key_prefix + string(1, '\x00');
       string limit = key_prefix + string("\xff\xff\xff\xff");
       rocksdb::Range r(start, limit);
-      db->GetApproximateSizes(cf, &r, 1, &s);
+      db->GetApproximateSizes(cf, &r, 1, &s, flags);
       size += s;
     }
   } else {
     string start = combine_strings(prefix , key_prefix);
     string limit = combine_strings(prefix , key_prefix + "\xff\xff\xff\xff");
     rocksdb::Range r(start, limit);
-    db->GetApproximateSizes(default_cf, &r, 1, &size);
+    db->GetApproximateSizes(default_cf, &r, 1, &size, flags);
   }
   return size;
 }