From c9cd404bcd0d8ae50227b8c4f6ec68d47341575a Mon Sep 17 00:00:00 2001
From: Venkatesh Radhakrishnan <rven@fb.com>
Date: Thu, 25 Jun 2015 14:43:25 -0700
Subject: [PATCH] Make flush check for shutdown

Summary:
Fixes task 7156865 where a compaction causes a hang in flush
memtable if CancelAllBackgroundWork was called prior to it.
Stack trace is in : https://phabricator.fb.com/P19848829
We end up waiting for a flush which will never happen because there are no background threads.

Test Plan: PreShutdownFlush

Reviewers: sdong, igor

Reviewed By: sdong, igor

Subscribers: dhruba, leveldb

Differential Revision: https://reviews.facebook.net/D40617
---
 db/db_impl.cc | 16 +++++++++-------
 db/db_test.cc | 11 +++++++++++
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/db/db_impl.cc b/db/db_impl.cc
index a89fdaffb3..e7465e0f99 100644
--- a/db/db_impl.cc
+++ b/db/db_impl.cc
@@ -270,18 +270,18 @@ DBImpl::DBImpl(const DBOptions& options, const std::string& dbname)
   LogFlush(db_options_.info_log);
 }
 
-// Will only lock the mutex_ and wait for completion if wait is true
+// Will lock the mutex_,  will wait for completion if wait is true
 void DBImpl::CancelAllBackgroundWork(bool wait) {
+  InstrumentedMutexLock l(&mutex_);
   shutting_down_.store(true, std::memory_order_release);
+  bg_cv_.SignalAll();
   if (!wait) {
     return;
   }
   // Wait for background work to finish
-  mutex_.Lock();
   while (bg_compaction_scheduled_ || bg_flush_scheduled_) {
     bg_cv_.Wait();
   }
-  mutex_.Unlock();
 }
 
 DBImpl::~DBImpl() {
@@ -299,12 +299,11 @@ DBImpl::~DBImpl() {
     }
     versions_->GetColumnFamilySet()->FreeDeadColumnFamilies();
   }
-  // CancelAllBackgroundWork called with false means we just set the
-  // shutdown marker, while holding the mutex_ here. After which we
-  // do a variant of the waiting after we release the lock and unschedule work
+  mutex_.Unlock();
+  // CancelAllBackgroundWork called with false means we just set the shutdown
+  // marker. After this we do a variant of the waiting and unschedule work
   // (to consider: moving all the waiting into CancelAllBackgroundWork(true))
   CancelAllBackgroundWork(false);
-  mutex_.Unlock();
   int compactions_unscheduled = env_->UnSchedule(this, Env::Priority::LOW);
   int flushes_unscheduled = env_->UnSchedule(this, Env::Priority::HIGH);
   mutex_.Lock();
@@ -2036,6 +2035,9 @@ Status DBImpl::WaitForFlushMemTable(ColumnFamilyData* cfd) {
   // Wait until the compaction completes
   InstrumentedMutexLock l(&mutex_);
   while (cfd->imm()->NumNotFlushed() > 0 && bg_error_.ok()) {
+    if (shutting_down_.load(std::memory_order_acquire)) {
+      return Status::ShutdownInProgress();
+    }
     bg_cv_.Wait();
   }
   if (!bg_error_.ok()) {
diff --git a/db/db_test.cc b/db/db_test.cc
index 20a4f872d8..f09578bb8a 100644
--- a/db/db_test.cc
+++ b/db/db_test.cc
@@ -11315,6 +11315,17 @@ TEST_F(DBTest, PreShutdownManualCompaction) {
   }
 }
 
+TEST_F(DBTest, PreShutdownFlush) {
+  Options options = CurrentOptions();
+  options.max_background_flushes = 0;
+  CreateAndReopenWithCF({"pikachu"}, options);
+  ASSERT_OK(Put(1, "key", "value"));
+  CancelAllBackgroundWork(db_);
+  Status s =
+      db_->CompactRange(CompactRangeOptions(), handles_[1], nullptr, nullptr);
+  ASSERT_TRUE(s.IsShutdownInProgress());
+}
+
 TEST_F(DBTest, PreShutdownMultipleCompaction) {
   const int kTestKeySize = 16;
   const int kTestValueSize = 984;