From 0ccaba2a05db5c5930018d0d78775f29ba622e16 Mon Sep 17 00:00:00 2001
From: Yi Wu <yiwu@fb.com>
Date: Tue, 23 May 2017 10:32:02 -0700
Subject: [PATCH] Fix rocksdb.estimate-num-keys DB property underflow

Summary:
rocksdb.estimate-num-keys is compute from `estimate_num_keys - 2 * estimate_num_deletes`. If  `2 * estimate_num_deletes > estimate_num_keys` it will underflow. Fixing it.
Closes https://github.com/facebook/rocksdb/pull/2348

Differential Revision: D5109272

Pulled By: yiwu-arbug

fbshipit-source-id: e1bfb91346a59b7282a282b615002507e9d7c246
---
 db/db_properties_test.cc | 12 ++++++++++++
 db/internal_stats.cc     | 14 ++++++++------
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/db/db_properties_test.cc b/db/db_properties_test.cc
index 5ce6fb8d7..41a4db82f 100644
--- a/db/db_properties_test.cc
+++ b/db/db_properties_test.cc
@@ -1300,6 +1300,18 @@ TEST_F(DBPropertiesTest, NeedCompactHintPersistentTest) {
     SetPerfLevel(kDisable);
   }
 }
+
+TEST_F(DBPropertiesTest, EstimateNumKeysUnderflow) {
+  Options options;
+  Reopen(options);
+  Put("foo", "bar");
+  Delete("foo");
+  Delete("foo");
+  uint64_t num_keys = 0;
+  ASSERT_TRUE(dbfull()->GetIntProperty("rocksdb.estimate-num-keys", &num_keys));
+  ASSERT_EQ(0, num_keys);
+}
+
 #endif  // ROCKSDB_LITE
 }  // namespace rocksdb
 
diff --git a/db/internal_stats.cc b/db/internal_stats.cc
index 5de370e18..faec981a5 100644
--- a/db/internal_stats.cc
+++ b/db/internal_stats.cc
@@ -677,12 +677,14 @@ bool InternalStats::HandleEstimateNumKeys(uint64_t* value, DBImpl* db,
   // Estimate number of entries in the column family:
   // Use estimated entries in tables + total entries in memtables.
   const auto* vstorage = cfd_->current()->storage_info();
-  *value = cfd_->mem()->num_entries() +
-           cfd_->imm()->current()->GetTotalNumEntries() -
-           (cfd_->mem()->num_deletes() +
-            cfd_->imm()->current()->GetTotalNumDeletes()) *
-               2 +
-           vstorage->GetEstimatedActiveKeys();
+  uint64_t estimate_keys = cfd_->mem()->num_entries() +
+                           cfd_->imm()->current()->GetTotalNumEntries() +
+                           vstorage->GetEstimatedActiveKeys();
+  uint64_t estimate_deletes =
+      cfd_->mem()->num_deletes() + cfd_->imm()->current()->GetTotalNumDeletes();
+  *value = estimate_keys > estimate_deletes * 2
+               ? estimate_keys - (estimate_deletes * 2)
+               : 0;
   return true;
 }