From 0b0cb6f1a2f71eb4532416a959ebcf682ac9096b Mon Sep 17 00:00:00 2001
From: feilongliu <feilongliu@devvm662.vll1.facebook.com>
Date: Thu, 20 Jun 2019 13:04:13 -0700
Subject: [PATCH] Fix segfalut in ~DBWithTTLImpl() when called after Close()
 (#5485)

Summary:
~DBWithTTLImpl() fails after calling Close() function (will invoke the
Close() function of DBImpl), because the Close() function deletes
default_cf_handle_ which is used in the GetOptions() function called
in ~DBWithTTLImpl(), hence lead to segfault.

Fix by creating a Close() function for the DBWithTTLImpl class and do
the close and the work originally in ~DBWithTTLImpl(). If the Close()
function is not called, it will be called in the ~DBWithTTLImpl()
function.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/5485

Test Plan: make clean;  USE_CLANG=1 make all check -j

Differential Revision: D15924498

fbshipit-source-id: 567397fb972961059083a1ae0f9f99ff74872b78
---
 utilities/ttl/db_ttl_impl.cc | 21 ++++++++++++++----
 utilities/ttl/db_ttl_impl.h  |  6 +++++
 utilities/ttl/ttl_test.cc    | 43 ++++++++++++++++++++++++++++++++++--
 3 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/utilities/ttl/db_ttl_impl.cc b/utilities/ttl/db_ttl_impl.cc
index 47049a135..2c79d01ba 100644
--- a/utilities/ttl/db_ttl_impl.cc
+++ b/utilities/ttl/db_ttl_impl.cc
@@ -34,12 +34,25 @@ void DBWithTTLImpl::SanitizeOptions(int32_t ttl, ColumnFamilyOptions* options,
 }
 
 // Open the db inside DBWithTTLImpl because options needs pointer to its ttl
-DBWithTTLImpl::DBWithTTLImpl(DB* db) : DBWithTTL(db) {}
+DBWithTTLImpl::DBWithTTLImpl(DB* db) : DBWithTTL(db), closed_(false) {}
 
 DBWithTTLImpl::~DBWithTTLImpl() {
-  // Need to stop background compaction before getting rid of the filter
-  CancelAllBackgroundWork(db_, /* wait = */ true);
-  delete GetOptions().compaction_filter;
+  if (!closed_) {
+    Close();
+  }
+}
+
+Status DBWithTTLImpl::Close() {
+  Status ret = Status::OK();
+  if (!closed_) {
+    Options default_options = GetOptions();
+    // Need to stop background compaction before getting rid of the filter
+    CancelAllBackgroundWork(db_, /* wait = */ true);
+    ret = db_->Close();
+    delete default_options.compaction_filter;
+    closed_ = true;
+  }
+  return ret;
 }
 
 Status UtilityDB::OpenTtlDB(const Options& options, const std::string& dbname,
diff --git a/utilities/ttl/db_ttl_impl.h b/utilities/ttl/db_ttl_impl.h
index 593cd64a0..1111c13a7 100644
--- a/utilities/ttl/db_ttl_impl.h
+++ b/utilities/ttl/db_ttl_impl.h
@@ -35,6 +35,8 @@ class DBWithTTLImpl : public DBWithTTL {
 
   virtual ~DBWithTTLImpl();
 
+  virtual Status Close() override;
+
   Status CreateColumnFamilyWithTtl(const ColumnFamilyOptions& options,
                                    const std::string& column_family_name,
                                    ColumnFamilyHandle** handle,
@@ -99,6 +101,10 @@ class DBWithTTLImpl : public DBWithTTL {
   void SetTtl(int32_t ttl) override { SetTtl(DefaultColumnFamily(), ttl); }
 
   void SetTtl(ColumnFamilyHandle *h, int32_t ttl) override;
+
+ private:
+   // remember whether the Close completes or not
+   bool closed_;
 };
 
 class TtlIterator : public Iterator {
diff --git a/utilities/ttl/ttl_test.cc b/utilities/ttl/ttl_test.cc
index 38c6affab..61f5e6449 100644
--- a/utilities/ttl/ttl_test.cc
+++ b/utilities/ttl/ttl_test.cc
@@ -86,9 +86,24 @@ class TtlTest : public testing::Test {
     ASSERT_OK(DBWithTTL::Open(options_, dbname_, &db_ttl_, ttl, true));
   }
 
+  // Call db_ttl_->Close() before delete db_ttl_
   void CloseTtl() {
-    delete db_ttl_;
-    db_ttl_ = nullptr;
+    CloseTtlHelper(true);
+  }
+
+  // No db_ttl_->Close() before delete db_ttl_
+  void CloseTtlNoDBClose() {
+    CloseTtlHelper(false);
+  }
+
+  void CloseTtlHelper(bool close_db) {
+    if (db_ttl_ != nullptr) {
+      if (close_db) {
+        db_ttl_->Close();
+      }
+      delete db_ttl_;
+      db_ttl_ = nullptr;
+    }
   }
 
   // Populates and returns a kv-map
@@ -401,6 +416,30 @@ TEST_F(TtlTest, NoEffect) {
   CloseTtl();
 }
 
+
+// Rerun the NoEffect test with a different version of CloseTtl
+// function, where db is directly deleted without close.
+TEST_F(TtlTest, DestructWithoutClose) {
+  MakeKVMap(kSampleSize_);
+  int64_t boundary1 = kSampleSize_ / 3;
+  int64_t boundary2 = 2 * boundary1;
+
+  OpenTtl();
+  PutValues(0, boundary1);                       //T=0: Set1 never deleted
+  SleepCompactCheck(1, 0, boundary1);            //T=1: Set1 still there
+  CloseTtlNoDBClose();
+
+  OpenTtl(0);
+  PutValues(boundary1, boundary2 - boundary1);   //T=1: Set2 never deleted
+  SleepCompactCheck(1, 0, boundary2);            //T=2: Sets1 & 2 still there
+  CloseTtlNoDBClose();
+
+  OpenTtl(-1);
+  PutValues(boundary2, kSampleSize_ - boundary2); //T=3: Set3 never deleted
+  SleepCompactCheck(1, 0, kSampleSize_, true);    //T=4: Sets 1,2,3 still there
+  CloseTtlNoDBClose();
+}
+
 // Puts a set of values and checks its presence using Get during ttl
 TEST_F(TtlTest, PresentDuringTTL) {
   MakeKVMap(kSampleSize_);