From c6da00b418764ca8fec248b768f3744ca1e655ed Mon Sep 17 00:00:00 2001
From: Sage Weil <sage@redhat.com>
Date: Mon, 2 Oct 2017 22:17:54 -0500
Subject: [PATCH] kv: add per-prefix space estimation

Signed-off-by: Sage Weil <sage@redhat.com>
---
 src/kv/KeyValueDB.h    |  5 +++++
 src/kv/RocksDBStore.cc | 21 +++++++++++++++++++++
 src/kv/RocksDBStore.h  |  2 ++
 3 files changed, 28 insertions(+)

diff --git a/src/kv/KeyValueDB.h b/src/kv/KeyValueDB.h
index 097fe12827804..e8d7aced708a4 100644
--- a/src/kv/KeyValueDB.h
+++ b/src/kv/KeyValueDB.h
@@ -418,6 +418,11 @@ public:
 
   virtual ~KeyValueDB() {}
 
+  /// estimate space utilization for a prefix (in bytes)
+  virtual int64_t estimate_prefix_size(const string& prefix) {
+    return 0;
+  }
+
   /// compact the underlying store
   virtual void compact() {}
 
diff --git a/src/kv/RocksDBStore.cc b/src/kv/RocksDBStore.cc
index 928b4d915fb81..b59bb3dfeeae8 100644
--- a/src/kv/RocksDBStore.cc
+++ b/src/kv/RocksDBStore.cc
@@ -678,6 +678,27 @@ void RocksDBStore::split_stats(const std::string &s, char delim, std::vector<std
     }
 }
 
+int64_t RocksDBStore::estimate_prefix_size(const string& prefix)
+{
+  auto cf = get_cf_handle(prefix);
+  uint64_t size = 0;
+  uint8_t flags =
+    //rocksdb::DB::INCLUDE_MEMTABLES |  // do not include memtables...
+    rocksdb::DB::INCLUDE_FILES;
+  if (cf) {
+    string start(1, '\x00');
+    string limit("\xff\xff\xff\xff");
+    rocksdb::Range r(start, limit);
+    db->GetApproximateSizes(cf, &r, 1, &size, flags);
+  } else {
+    string limit = prefix + "\xff\xff\xff\xff";
+    rocksdb::Range r(prefix, limit);
+    db->GetApproximateSizes(default_cf,
+			    &r, 1, &size, flags);
+  }
+  return size;
+}
+
 void RocksDBStore::get_statistics(Formatter *f)
 {
   if (!g_conf()->rocksdb_perf)  {
diff --git a/src/kv/RocksDBStore.h b/src/kv/RocksDBStore.h
index 96058f04af9a2..08477dded929c 100644
--- a/src/kv/RocksDBStore.h
+++ b/src/kv/RocksDBStore.h
@@ -192,6 +192,8 @@ public:
     return logger;
   }
 
+  int64_t estimate_prefix_size(const string& prefix) override;
+
   struct  RocksWBHandler: public rocksdb::WriteBatch::Handler {
     std::string seen ;
     int num_seen = 0;
-- 
2.39.5