Fix race in WriteBufferManager (#9009)

Summary:
EndWriteStall has a data race: `queue_.empty()` is checked outside of the
mutex, so once we enter the critical section another thread may already have
cleared the list, and accessing the `front()` is undefined behavior (and causes
interesting crashes under high concurrency).

This PR fixes the bug, and also rewrites the logic to make it easier to reason
about it. It also fixes another subtle bug: if some writers are stalled and
`SetBufferSize(0)` is called, which disables the WBM, the writer are not
unblocked because of an early `enabled()` check in `EndWriteStall()`.

It doesn't significantly change the locking behavior, as before writers won't
lock unless entering a stall condition, and `FreeMem` almost always locks if
stalling is allowed, but that is inevitable with the current design. Liveness is
guaranteed by the fact that if some writes are blocked, eventually all writes
will be blocked due to `stall_active_`, and eventually all memory is freed.

While at it, do a couple of optimizations:

- In `WBMStallInterface::Signal()` signal the CV only after releasing the
  lock. Signaling under the lock is a common pitfall, as it causes the woken-up
  thread to immediately go back to sleep because the mutex is still locked by
  the awaker.

- Move all allocations and deallocations outside of the lock.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/9009

Test Plan:
```
USE_CLANG=1 make -j64 all check
```

Reviewed By: akankshamahajan15

Differential Revision: D31550668

Pulled By: ot

fbshipit-source-id: 5125387c3dc7ecaaa2b8bbc736e58c4156698580

1 files changed, 61 insertions, 31 deletions

diff --git a/memtable/write_buffer_manager.cc b/memtable/write_buffer_manager.cc
index c599b658c..48278aaaf 100644
--- a/memtable/write_buffer_manager.cc
+++ b/memtable/write_buffer_manager.cc
@@ -39,7 +39,12 @@ WriteBufferManager::WriteBufferManager(size_t _buffer_size,
 #endif  // ROCKSDB_LITE
 }
 
-WriteBufferManager::~WriteBufferManager() = default;
+WriteBufferManager::~WriteBufferManager() {
+#ifndef NDEBUG
+  std::unique_lock<std::mutex> lock(mu_);
+  assert(queue_.empty());
+#endif
+}
 
 std::size_t WriteBufferManager::dummy_entries_in_cache_usage() const {
   if (cache_rev_mng_ != nullptr) {
@@ -98,9 +103,7 @@ void WriteBufferManager::FreeMem(size_t mem) {
     memory_used_.fetch_sub(mem, std::memory_order_relaxed);
   }
   // Check if stall is active and can be ended.
-  if (allow_stall_) {
-    EndWriteStall();
-  }
+  MaybeEndWriteStall();
 }
 
 void WriteBufferManager::FreeMemWithCache(size_t mem) {
@@ -127,47 +130,74 @@ void WriteBufferManager::FreeMemWithCache(size_t mem) {
 
 void WriteBufferManager::BeginWriteStall(StallInterface* wbm_stall) {
   assert(wbm_stall != nullptr);
-  if (wbm_stall) {
+  assert(allow_stall_);
+
+  // Allocate outside of the lock.
+  std::list<StallInterface*> new_node = {wbm_stall};
+
+  {
     std::unique_lock<std::mutex> lock(mu_);
-    queue_.push_back(wbm_stall);
+    // Verify if the stall conditions are stil active.
+    if (ShouldStall()) {
+      stall_active_.store(true, std::memory_order_relaxed);
+      queue_.splice(queue_.end(), std::move(new_node));
+    }
   }
-  // In case thread enqueue itself and memory got freed in parallel, end the
-  // stall.
-  if (!ShouldStall()) {
-    EndWriteStall();
+
+  // If the node was not consumed, the stall has ended already and we can signal
+  // the caller.
+  if (!new_node.empty()) {
+    new_node.front()->Signal();
   }
 }
 
-// Called when memory is freed in FreeMem.
-void WriteBufferManager::EndWriteStall() {
-  if (enabled() && !IsStallThresholdExceeded()) {
-    {
-      std::unique_lock<std::mutex> lock(mu_);
-      stall_active_.store(false, std::memory_order_relaxed);
-      if (queue_.empty()) {
-        return;
-      }
-    }
+// Called when memory is freed in FreeMem or the buffer size has changed.
+void WriteBufferManager::MaybeEndWriteStall() {
+  // Cannot early-exit on !enabled() because SetBufferSize(0) needs to unblock
+  // the writers.
+  if (!allow_stall_) {
+    return;
+  }
 
-    // Get the instances from the list and call WBMStallInterface::Signal to
-    // change the state to running and unblock the DB instances.
-    // Check ShouldStall() incase stall got active by other DBs.
-    while (!ShouldStall() && !queue_.empty()) {
-      std::unique_lock<std::mutex> lock(mu_);
-      StallInterface* wbm_stall = queue_.front();
-      queue_.pop_front();
-      wbm_stall->Signal();
-    }
+  if (IsStallThresholdExceeded()) {
+    return;  // Stall conditions have not resolved.
+  }
+
+  // Perform all deallocations outside of the lock.
+  std::list<StallInterface*> cleanup;
+
+  std::unique_lock<std::mutex> lock(mu_);
+  if (!stall_active_.load(std::memory_order_relaxed)) {
+    return;  // Nothing to do.
+  }
+
+  // Unblock new writers.
+  stall_active_.store(false, std::memory_order_relaxed);
+
+  // Unblock the writers in the queue.
+  for (StallInterface* wbm_stall : queue_) {
+    wbm_stall->Signal();
   }
+  cleanup = std::move(queue_);
 }
 
 void WriteBufferManager::RemoveDBFromQueue(StallInterface* wbm_stall) {
   assert(wbm_stall != nullptr);
+
+  // Deallocate the removed nodes outside of the lock.
+  std::list<StallInterface*> cleanup;
+
   if (enabled() && allow_stall_) {
     std::unique_lock<std::mutex> lock(mu_);
-    queue_.remove(wbm_stall);
-    wbm_stall->Signal();
+    for (auto it = queue_.begin(); it != queue_.end();) {
+      auto next = std::next(it);
+      if (*it == wbm_stall) {
+        cleanup.splice(cleanup.end(), queue_, std::move(it));
+      }
+      it = next;
+    }
   }
+  wbm_stall->Signal();
 }
 
 }  // namespace ROCKSDB_NAMESPACE