Fix WBM concurrency control, Add SetAllowStall(), Cleanup

HISTORY.md

@@ -1,5 +1,7 @@
 # Rocksdb Change Log
 ## Unreleased
+### New Features
+* Allow setting `WriteBufferManager::allow_stall` in runtime by `SetAllowStall()`, which was previously fixed after passed into `WriteBufferManager`'s constructor
 
 ## 8.1.0 (03/18/2023)
 ### Behavior changes

db/db_impl/db_impl.h

@@ -1200,17 +1200,13 @@ class DBImpl : public DB {
   // state_ is changed accordingly.
   class WBMStallInterface : public StallInterface {
    public:
-    enum State {
-      BLOCKED = 0,
-      RUNNING,
-    };
 
     WBMStallInterface() : state_cv_(&state_mutex_) {
       MutexLock lock(&state_mutex_);
       state_ = State::RUNNING;
     }
 
-    void SetState(State state) {
+    void SetState(State state) override {
       MutexLock lock(&state_mutex_);
       state_ = state;
     }
@@ -1871,9 +1867,9 @@ class DBImpl : public DB {
   Status DelayWrite(uint64_t num_bytes, WriteThread& write_thread,
                     const WriteOptions& write_options);
 
-  // Begin stalling of writes when memory usage increases beyond a certain
-  // threshold.
-  void WriteBufferManagerStallWrites();
+  // If stall conditions are met, begin stalling of writes with help of
+  // `WriteBufferManager`
+  void MaybeWriteBufferManagerStallWrites();
 
   Status ThrottleLowPriWritesIfNeeded(const WriteOptions& write_options,
                                       WriteBatch* my_batch);

db/db_impl/db_impl_write.cc

@@ -1218,7 +1218,7 @@ Status DBImpl::PreprocessWrite(const WriteOptions& write_options,
       status = Status::Incomplete("Write stall");
     } else {
       InstrumentedMutexLock l(&mutex_);
-      WriteBufferManagerStallWrites();
+      MaybeWriteBufferManagerStallWrites();
     }
   }
   InstrumentedMutexLock l(&log_write_mutex_);
@@ -1900,19 +1900,13 @@ Status DBImpl::DelayWrite(uint64_t num_bytes, WriteThread& write_thread,
 
 // REQUIRES: mutex_ is held
 // REQUIRES: this thread is currently at the front of the writer queue
-void DBImpl::WriteBufferManagerStallWrites() {
+void DBImpl::MaybeWriteBufferManagerStallWrites() {
   mutex_.AssertHeld();
   // First block future writer threads who want to add themselves to the queue
   // of WriteThread.
   write_thread_.BeginWriteStall();
   mutex_.Unlock();
-
-  // Change the state to State::Blocked.
-  static_cast<WBMStallInterface*>(wbm_stall_.get())
-      ->SetState(WBMStallInterface::State::BLOCKED);
-  // Then WriteBufferManager will add DB instance to its queue
-  // and block this thread by calling WBMStallInterface::Block().
-  write_buffer_manager_->BeginWriteStall(wbm_stall_.get());
+  write_buffer_manager_->MaybeBeginWriteStall(wbm_stall_.get());
   wbm_stall_->Block();
 
   mutex_.Lock();

db/db_write_buffer_manager_test.cc

@@ -846,6 +846,59 @@ TEST_P(DBWriteBufferManagerTest, StopSwitchingMemTablesOnceFlushing) {
   delete shared_wbm_db;
 }
 
+TEST_F(DBWriteBufferManagerTest, RuntimeChangeableThreadSafeParameters) {
+  for (std::string test_parameter : {"buffer_size", "allow_stall"}) {
+    Options options = CurrentOptions();
+    options.write_buffer_manager.reset(new WriteBufferManager(
+        10000, nullptr /* cahce */, true /* allow_stall */));
+    DestroyAndReopen(options);
+
+    // Pause flush thread so that the only way to exist write stall
+    // below is to change the `test_parameter` in the runtime successfully
+    std::unique_ptr<test::SleepingBackgroundTask> sleeping_task(
+        new test::SleepingBackgroundTask());
+    env_->SetBackgroundThreads(1, Env::HIGH);
+    env_->Schedule(&test::SleepingBackgroundTask::DoSleepTask,
+                   sleeping_task.get(), Env::Priority::HIGH);
+    sleeping_task->WaitUntilSleeping();
+
+    // After completing this write, any future write will be stalled by
+    // WriteBufferManager
+    ASSERT_OK(Put(Key(0), DummyString(10000)));
+
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->LoadDependency(
+        {{"WBMStallInterface::BlockDB",
+          "DBWriteBufferManagerTest::RuntimeChangeableThreadSafeParameters::"
+          "ChangeParameter"}});
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->EnableProcessing();
+
+    port::Thread thread1([&] { ASSERT_OK(Put(Key(0), DummyString(10000))); });
+
+    port::Thread thread2([&] {
+      TEST_SYNC_POINT(
+          "DBWriteBufferManagerTest::RuntimeChangeableThreadSafeParameters::"
+          "ChangeParameter");
+      if (test_parameter == "buffer_size") {
+        options.write_buffer_manager->SetBufferSize(
+            options.write_buffer_manager->buffer_size() * 200);
+      } else if (test_parameter == "allow_stall") {
+        options.write_buffer_manager->SetAllowStall(false);
+      }
+    });
+
+    // If `test_parameter` is successfully changed in thread2, the write stall
+    // encountered in thread1 will stop and the test will finish. Othwerwise,
+    // thread1 will hang forever.
+    thread1.join();
+    thread2.join();
+
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->ClearAllCallBacks();
+    ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->DisableProcessing();
+
+    sleeping_task->WakeUp();
+    sleeping_task->WaitUntilDone();
+  }
+}
 
 INSTANTIATE_TEST_CASE_P(DBWriteBufferManagerTest, DBWriteBufferManagerTest,
                         testing::Bool());

include/rocksdb/write_buffer_manager.h

@@ -27,18 +27,27 @@ class CacheReservationManager;
 // internal use only. Each DB instance contains ptr to StallInterface.
 class StallInterface {
  public:
+  enum State {
+    BLOCKED = 0,
+    RUNNING,
+  };
+
   virtual ~StallInterface() {}
 
   virtual void Block() = 0;
 
   virtual void Signal() = 0;
+
+  virtual void SetState(State /* state */) {}
 };
 
+// This class is thread-safe
 class WriteBufferManager final {
  public:
   // Parameters:
-  // _buffer_size: _buffer_size = 0 indicates no limit. Memory won't be capped.
-  // memory_usage() won't be valid and ShouldFlush() will always return true.
+  // _buffer_size: the memory limit. _buffer_size = 0 indicates no limit where
+  // memory won't be capped, memory_usage() won't be valid and ShouldFlush()
+  // will always return false.
   //
   // cache_: if `cache` is provided, we'll put dummy entries in the cache and
   // cost the memory allocated to the cache. It can be used even if _buffer_size
@@ -56,121 +65,134 @@ class WriteBufferManager final {
 
   ~WriteBufferManager();
 
-  // Returns true if buffer_limit is passed to limit the total memory usage and
-  // is greater than 0.
+  // Returns true if memory limit exixts (i.e, buffer size > 0);
+  //
+  // WARNING: If running without syncronization with `SetBufferSize()`, this
+  // function might not return the latest result changed by `SetBufferSize()`
+  // but an old result.
   bool enabled() const { return buffer_size() > 0; }
 
   // Returns true if pointer to cache is passed.
   bool cost_to_cache() const { return cache_res_mgr_ != nullptr; }
 
-  // Returns the total memory used by memtables.
-  // Only valid if enabled()
+  // Returns the total memory used by memtables.  Only valid if enabled()=true
+  //
+  // WARNING: If running without syncronization with `ReserveMem()/FreeMem()`,
+  // this function might not return the latest result changed by these functions
+  // but an old result.
   size_t memory_usage() const {
     return memory_used_.load(std::memory_order_relaxed);
   }
 
   // Returns the total memory used by active memtables.
+  //
+  // WARNING: If running without syncronization with
+  // `ReserveMem()/ScheduleFreeMem()`, this function might not return the latest
+  // result changed by these functions but an old result.
   size_t mutable_memtable_memory_usage() const {
     return memory_active_.load(std::memory_order_relaxed);
   }
 
+  // Return the number of dummy entries put in the cache used to cost the memory
+  // accounted by this WriteBufferManager to the cache
+  //
+  // WARNING: If running without syncronization with `ReserveMem()/FreeMem()`,
+  // this function might not return the latest result changed by these functions
+  // but an old result.
   size_t dummy_entries_in_cache_usage() const;
 
   // Returns the buffer_size.
+  //
+  // WARNING: If running without syncronization with `SetBufferSize()`, this
+  // function might not return the latest result changed by `SetBufferSize()`
+  // but an old result.
   size_t buffer_size() const {
     return buffer_size_.load(std::memory_order_relaxed);
   }
 
-  void SetBufferSize(size_t new_size) {
-    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.
-    MaybeEndWriteStall();
-  }
+  // REQUIRED: new_size != 0
+  void SetBufferSize(size_t new_size);
 
-  // Below functions should be called by RocksDB internally.
-
-  // Should only be called from write thread
-  bool ShouldFlush() const {
-    if (enabled()) {
-      if (mutable_memtable_memory_usage() >
-          mutable_limit_.load(std::memory_order_relaxed)) {
-        return true;
-      }
-      size_t local_size = buffer_size();
-      if (memory_usage() >= local_size &&
-          mutable_memtable_memory_usage() >= local_size / 2) {
-        // If the memory exceeds the buffer size, we trigger more aggressive
-        // flush. But if already more than half memory is being flushed,
-        // triggering more flush may not help. We will hold it instead.
-        return true;
-      }
-    }
-    return false;
-  }
+  void SetAllowStall(bool new_allow_stall);
 
-  // Returns true if total memory usage exceeded buffer_size.
-  // We stall the writes untill memory_usage drops below buffer_size. When the
-  // function returns true, all writer threads (including one checking this
-  // condition) across all DBs will be stalled. Stall is allowed only if user
-  // pass allow_stall = true during WriteBufferManager instance creation.
+  // WARNING: Should only be called from write thread
+  // WARNING: Should only be called by RocksDB internally
+  bool ShouldFlush() const;
+
+  // Returns true if stall conditions are met.
+  // Stall conditions: stall is allowed AND memory limit is set (i.e, buffer
+  // size > 0) AND total memory usage accounted by this WriteBufferManager
+  // exceeds the memory limit.
+  //
+  // WARNING: Should only be called by RocksDB internally.
   //
-  // Should only be called by RocksDB internally .
+  // WARNING: If running without syncronization with any functions that could
+  // change the stall conditions above, this function might not return the
+  // latest result changed by these functions but an old result.
   bool ShouldStall() const {
-    if (!allow_stall_ || !enabled()) {
+    if (allow_stall_.load(std::memory_order_relaxed) && enabled() &&
+        IsStallThresholdExceeded()) {
+      return true;
+    } else {
       return false;
     }
-
-    return IsStallActive() || IsStallThresholdExceeded();
-  }
-
-  // Returns true if stall is active.
-  bool IsStallActive() const {
-    return stall_active_.load(std::memory_order_relaxed);
-  }
-
-  // Returns true if stalling condition is met.
-  bool IsStallThresholdExceeded() const {
-    return memory_usage() >= buffer_size_;
   }
 
+  // WARNING: Should only be called by RocksDB internally.
   void ReserveMem(size_t mem);
 
   // We are in the process of freeing `mem` bytes, so it is not considered
   // when checking the soft limit.
+  //
+  // WARNING: Should only be called by RocksDB internally.
   void ScheduleFreeMem(size_t mem);
 
+  // WARNING: Should only be called by RocksDB internally.
   void FreeMem(size_t mem);
 
-  // Add the DB instance to the queue and block the DB.
-  // Should only be called by RocksDB internally.
-  void BeginWriteStall(StallInterface* wbm_stall);
+  // If stall conditions are met, WriteBufferManager
+  // will prepare for write stall (including changing `wbm_stall`'s state
+  // to be `State::Blocked`). Otherwise, this function does nothing.
+  //
+  // WARNING: Should only be called by RocksDB internally.
+  void MaybeBeginWriteStall(StallInterface* wbm_stall);
+
+  // WARNING: Should only be called by RocksDB internally.
+  void RemoveDBFromQueue(StallInterface* wbm_stall);
 
+ private:
   // If stall conditions have resolved, remove DB instances from queue and
   // signal them to continue.
+  //
+  // Called when stall conditions (see `ShouldStall()` API) might have been
+  // changed
+  //
+  // REQUIRED: wbm_mutex_ held
   void MaybeEndWriteStall();
 
-  void RemoveDBFromQueue(StallInterface* wbm_stall);
+  // REQUIRED: wbm_mutex_ held
+  bool IsStallThresholdExceeded() const {
+    return memory_usage() >= buffer_size_;
+  }
+
+  // WARNING:  Should only be called from write thread
+  // REQUIRED: wbm_mutex_ held
+  void ReserveMemWithCache(size_t mem);
+
+  // REQUIRED: wbm_mutex_ held
+  void FreeMemWithCache(size_t mem);
+
+  // Mutex used to protect WriteBufferManager's data variables.
+  mutable std::mutex wbm_mutex_;
 
- private:
   std::atomic<size_t> buffer_size_;
   std::atomic<size_t> mutable_limit_;
   std::atomic<size_t> memory_used_;
   // Memory that hasn't been scheduled to free.
   std::atomic<size_t> memory_active_;
   std::shared_ptr<CacheReservationManager> cache_res_mgr_;
-  // Protects cache_res_mgr_
-  std::mutex cache_res_mgr_mu_;
 
   std::list<StallInterface*> 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<bool> stall_active_;
-
-  void ReserveMemWithCache(size_t mem);
-  void FreeMemWithCache(size_t mem);
+  std::atomic<bool> allow_stall_;
 };
 }  // namespace ROCKSDB_NAMESPACE