From 22d4dc5066fbfc1a9613c206ee9d8978929d7d5b Mon Sep 17 00:00:00 2001 From: Giuseppe Ottaviano Date: Tue, 12 Oct 2021 00:14:41 -0700 Subject: [PATCH] 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 --- db/db_impl/db_impl.h | 6 +- include/rocksdb/write_buffer_manager.h | 32 ++++----- memtable/write_buffer_manager.cc | 94 +++++++++++++++++--------- 3 files changed, 81 insertions(+), 51 deletions(-) diff --git a/db/db_impl/db_impl.h b/db/db_impl/db_impl.h index ca227c99e..c8a18562f 100644 --- a/db/db_impl/db_impl.h +++ b/db/db_impl/db_impl.h @@ -1098,8 +1098,10 @@ class DBImpl : public DB { // Called from WriteBufferManager. This function changes the state_ // to State::RUNNING indicating the stall is cleared and DB can proceed. void Signal() override { - MutexLock lock(&state_mutex_); - state_ = State::RUNNING; + { + MutexLock lock(&state_mutex_); + state_ = State::RUNNING; + } state_cv_.Signal(); } diff --git a/include/rocksdb/write_buffer_manager.h b/include/rocksdb/write_buffer_manager.h index add957d84..a5471aabd 100644 --- a/include/rocksdb/write_buffer_manager.h +++ b/include/rocksdb/write_buffer_manager.h @@ -85,9 +85,7 @@ class WriteBufferManager { buffer_size_.store(new_size, std::memory_order_relaxed); mutable_limit_.store(new_size * 7 / 8, std::memory_order_relaxed); // Check if stall is active and can be ended. - if (allow_stall_) { - EndWriteStall(); - } + MaybeEndWriteStall(); } // Below functions should be called by RocksDB internally. @@ -118,17 +116,12 @@ class WriteBufferManager { // pass allow_stall = true during WriteBufferManager instance creation. // // Should only be called by RocksDB internally . - bool ShouldStall() { - if (allow_stall_ && enabled()) { - if (IsStallActive()) { - return true; - } - if (IsStallThresholdExceeded()) { - stall_active_.store(true, std::memory_order_relaxed); - return true; - } + bool ShouldStall() const { + if (!allow_stall_ || !enabled()) { + return false; } - return false; + + return IsStallActive() || IsStallThresholdExceeded(); } // Returns true if stall is active. @@ -137,7 +130,9 @@ class WriteBufferManager { } // Returns true if stalling condition is met. - bool IsStallThresholdExceeded() { return memory_usage() >= buffer_size_; } + bool IsStallThresholdExceeded() const { + return memory_usage() >= buffer_size_; + } void ReserveMem(size_t mem); @@ -151,8 +146,9 @@ class WriteBufferManager { // Should only be called by RocksDB internally. void BeginWriteStall(StallInterface* wbm_stall); - // Remove DB instances from queue and signal them to continue. - void EndWriteStall(); + // If stall conditions have resolved, remove DB instances from queue and + // signal them to continue. + void MaybeEndWriteStall(); void RemoveDBFromQueue(StallInterface* wbm_stall); @@ -167,9 +163,11 @@ class WriteBufferManager { std::mutex cache_rev_mng_mu_; std::list queue_; - // Protects the queue_ + // Protects the queue_ and stall_active_. std::mutex mu_; bool allow_stall_; + // Value should only be changed by BeginWriteStall() and MaybeEndWriteStall() + // while holding mu_, but it can be read without a lock. std::atomic stall_active_; void ReserveMemWithCache(size_t mem); 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 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 new_node = {wbm_stall}; + + { std::unique_lock 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 lock(mu_); - stall_active_.store(false, std::memory_order_relaxed); - if (queue_.empty()) { - 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 lock(mu_); - StallInterface* wbm_stall = queue_.front(); - queue_.pop_front(); - wbm_stall->Signal(); - } +// 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; } + + if (IsStallThresholdExceeded()) { + return; // Stall conditions have not resolved. + } + + // Perform all deallocations outside of the lock. + std::list cleanup; + + std::unique_lock 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 cleanup; + if (enabled() && allow_stall_) { std::unique_lock 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