From 2a93957a36397f2efa66d16336b8b3cc5968f0db Mon Sep 17 00:00:00 2001
From: Yi Wu <yiwu@fb.com>
Date: Tue, 16 May 2017 16:27:41 -0700
Subject: [PATCH] Fix DBTest.FlushSchedule

---
 db/flush_scheduler.cc | 31 +++++++++++++++++--------------
 db/write_thread.h     |  2 +-
 2 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/db/flush_scheduler.cc b/db/flush_scheduler.cc
index 05de92b7e14..eb24efb9cf5 100644
--- a/db/flush_scheduler.cc
+++ b/db/flush_scheduler.cc
@@ -15,11 +15,9 @@ namespace rocksdb {
 
 void FlushScheduler::ScheduleFlush(ColumnFamilyData* cfd) {
 #ifndef NDEBUG
-  {
-    std::lock_guard<std::mutex> lock(checking_mutex_);
-    assert(checking_set_.count(cfd) == 0);
-    checking_set_.insert(cfd);
-  }
+  std::lock_guard<std::mutex> lock(checking_mutex_);
+  assert(checking_set_.count(cfd) == 0);
+  checking_set_.insert(cfd);
 #endif  // NDEBUG
   cfd->Ref();
 // Suppress false positive clang analyzer warnings.
@@ -36,8 +34,11 @@ void FlushScheduler::ScheduleFlush(ColumnFamilyData* cfd) {
 }
 
 ColumnFamilyData* FlushScheduler::TakeNextColumnFamily() {
+#ifndef NDEBUG
+  std::lock_guard<std::mutex> lock(checking_mutex_);
+#endif  // NDEBUG
   while (true) {
-    if (Empty()) {
+    if (head_.load(std::memory_order_relaxed) == nullptr) {
       return nullptr;
     }
 
@@ -48,12 +49,9 @@ ColumnFamilyData* FlushScheduler::TakeNextColumnFamily() {
     delete node;
 
 #ifndef NDEBUG
-    {
-      std::lock_guard<std::mutex> lock(checking_mutex_);
-      auto iter = checking_set_.find(cfd);
-      assert(iter != checking_set_.end());
-      checking_set_.erase(iter);
-    }
+    auto iter = checking_set_.find(cfd);
+    assert(iter != checking_set_.end());
+    checking_set_.erase(iter);
 #endif  // NDEBUG
 
     if (!cfd->IsDropped()) {
@@ -69,8 +67,13 @@ ColumnFamilyData* FlushScheduler::TakeNextColumnFamily() {
 }
 
 bool FlushScheduler::Empty() {
+#ifndef NDEBUG
+  std::lock_guard<std::mutex> lock(checking_mutex_);
+#endif  // NDEBUG
   auto rv = head_.load(std::memory_order_relaxed) == nullptr;
-  // assert(rv == checking_set_.empty());
+#ifndef NDEBUG
+  assert(rv == checking_set_.empty());
+#endif  // NDEBUG
   return rv;
 }
 
@@ -81,7 +84,7 @@ void FlushScheduler::Clear() {
       delete cfd;
     }
   }
-  assert(Empty());
+  assert(head_.load(std::memory_order_relaxed) == nullptr);
 }
 
 }  // namespace rocksdb
diff --git a/db/write_thread.h b/db/write_thread.h
index 795eac67b84..13439d64305 100644
--- a/db/write_thread.h
+++ b/db/write_thread.h
@@ -192,7 +192,7 @@ class WriteThread {
     Status FinalStatus() {
       if (!status.ok()) {
         // a non-ok memtable write status takes presidence
-        // assert(callback == nullptr || callback_status.ok());
+        assert(callback == nullptr || callback_status.ok());
         return status;
       } else if (!callback_status.ok()) {
         // if the callback failed then that is the status we want