From 79cc8dc29b9528ad7d96a9a323b344c2b03ec347 Mon Sep 17 00:00:00 2001
From: sdong <siying.d@fb.com>
Date: Fri, 20 Dec 2019 21:42:19 -0800
Subject: [PATCH] db_stress: cover approximate size (#6213)

Summary:
db_stress to execute DB::GetApproximateSizes() with randomized keys and options. Return value is not validated but error will be reported.
Two ways to generate the range keys: (1) two random keys; (2) a small range.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/6213

Test Plan: (1) run "make crash_test" for a while; (2) hack the code to ingest some errors to see it is reported.

Differential Revision: D19204665

fbshipit-source-id: 652db36f13bcb5a3bd8fe4a10c0aa22a77a0bce2
---
 db_stress_tool/db_stress_common.h     |  1 +
 db_stress_tool/db_stress_gflags.cc    |  4 ++
 db_stress_tool/db_stress_test_base.cc | 53 +++++++++++++++++++++++++++
 db_stress_tool/db_stress_test_base.h  |  7 +++-
 4 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/db_stress_tool/db_stress_common.h b/db_stress_tool/db_stress_common.h
index fe13a571b..cca79031c 100644
--- a/db_stress_tool/db_stress_common.h
+++ b/db_stress_tool/db_stress_common.h
@@ -225,6 +225,7 @@ DECLARE_uint64(blob_db_file_size);
 DECLARE_bool(blob_db_enable_gc);
 DECLARE_double(blob_db_gc_cutoff);
 #endif  // !ROCKSDB_LITE
+DECLARE_int32(approximate_size_one_in);
 
 const long KB = 1024;
 const int kRandomValueMaxFactor = 3;
diff --git a/db_stress_tool/db_stress_gflags.cc b/db_stress_tool/db_stress_gflags.cc
index 02510f886..1c7390c3f 100644
--- a/db_stress_tool/db_stress_gflags.cc
+++ b/db_stress_tool/db_stress_gflags.cc
@@ -636,4 +636,8 @@ DEFINE_int32(verify_db_one_in, 0,
 DEFINE_int32(continuous_verification_interval, 1000,
              "While test is running, verify db every N milliseconds. 0 "
              "disables continuous verification.");
+
+DEFINE_int32(approximate_size_one_in, 64,
+             "If non-zero, DB::GetApproximateSizes() will be called against"
+             " random key ranges.");
 #endif  // GFLAGS
diff --git a/db_stress_tool/db_stress_test_base.cc b/db_stress_tool/db_stress_test_base.cc
index 3bc6e195a..be65b198b 100644
--- a/db_stress_tool/db_stress_test_base.cc
+++ b/db_stress_tool/db_stress_test_base.cc
@@ -630,6 +630,15 @@ void StressTest::OperateDb(ThreadState* thread) {
         }
       }
 
+#ifndef ROCKSDB_LITE
+      if (thread->rand.OneInOpt(FLAGS_approximate_size_one_in)) {
+        Status s =
+            TestApproximateSize(thread, i, rand_column_families, rand_keys);
+        if (!s.ok()) {
+          VerificationAbort(shared, "ApproximateSize Failed", s);
+        }
+      }
+#endif  // !ROCKSDB_LITE
       if (thread->rand.OneInOpt(FLAGS_acquire_snapshot_one_in)) {
         TestAcquireSnapshot(thread, rand_column_family, keystr, i);
       }
@@ -1205,6 +1214,50 @@ Status StressTest::TestBackupRestore(
   return s;
 }
 
+#ifndef ROCKSDB_LITE
+Status StressTest::TestApproximateSize(
+    ThreadState* thread, uint64_t iteration,
+    const std::vector<int>& rand_column_families,
+    const std::vector<int64_t>& rand_keys) {
+  // rand_keys likely only has one key. Just use the first one.
+  assert(!rand_keys.empty());
+  assert(!rand_column_families.empty());
+  int64_t key1 = rand_keys[0];
+  int64_t key2;
+  if (thread->rand.OneIn(2)) {
+    // Two totally random keys. This tends to cover large ranges.
+    key2 = GenerateOneKey(thread, iteration);
+    if (key2 < key1) {
+      std::swap(key1, key2);
+    }
+  } else {
+    // Unless users pass a very large FLAGS_max_key, it we should not worry
+    // about overflow. It is for testing, so we skip the overflow checking
+    // for simplicity.
+    key2 = key1 + static_cast<int64_t>(thread->rand.Uniform(1000));
+  }
+  std::string key1_str = Key(key1);
+  std::string key2_str = Key(key2);
+  Range range{Slice(key1_str), Slice(key2_str)};
+  SizeApproximationOptions sao;
+  sao.include_memtabtles = thread->rand.OneIn(2);
+  if (sao.include_memtabtles) {
+    sao.include_files = thread->rand.OneIn(2);
+  }
+  if (thread->rand.OneIn(2)) {
+    if (thread->rand.OneIn(2)) {
+      sao.files_size_error_margin = 0.0;
+    } else {
+      sao.files_size_error_margin =
+          static_cast<double>(thread->rand.Uniform(3));
+    }
+  }
+  uint64_t result;
+  return db_->GetApproximateSizes(
+      sao, column_families_[rand_column_families[0]], &range, 1, &result);
+}
+#endif  // ROCKSDB_LITE
+
 Status StressTest::TestCheckpoint(ThreadState* thread,
                                   const std::vector<int>& rand_column_families,
                                   const std::vector<int64_t>& rand_keys) {
diff --git a/db_stress_tool/db_stress_test_base.h b/db_stress_tool/db_stress_test_base.h
index 5463a911f..f16af3b0b 100644
--- a/db_stress_tool/db_stress_test_base.h
+++ b/db_stress_tool/db_stress_test_base.h
@@ -183,7 +183,12 @@ class StressTest {
   Status MaybeReleaseSnapshots(ThreadState* thread, uint64_t i);
 #ifndef ROCKSDB_LITE
   Status VerifyGetLiveAndWalFiles(ThreadState* thread);
-#endif    // !ROCKSDB_LITE
+  virtual Status TestApproximateSize(
+      ThreadState* thread, uint64_t iteration,
+      const std::vector<int>& rand_column_families,
+      const std::vector<int64_t>& rand_keys);
+#endif  // !ROCKSDB_LITE
+
   void VerificationAbort(SharedState* shared, std::string msg, Status s) const;
 
   void VerificationAbort(SharedState* shared, std::string msg, int cf,