From c529ca48d190b9d9ebc327dbc7a39e12e29c6ac7 Mon Sep 17 00:00:00 2001
From: udi-speedb <106253580+udi-speedb@users.noreply.github.com>
Date: Sat, 19 Aug 2023 22:40:08 +0300
Subject: [PATCH] Proactive Flushes: Decrement number of flushes to init on
 non-wbm-initiated flush end (#632)

---
 HISTORY.md                       |  3 ++-
 memtable/write_buffer_manager.cc | 28 +++++++++++++++++++++-------
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/HISTORY.md b/HISTORY.md
index 3c1d1de807..8c0c5de704 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -19,7 +19,8 @@
 * build: Windows compilation fix (#568).
 * Logger: fix Block cache stats trace by spacing it from the last trace (#578).
 * WriteController: move the class to public interface which should have been done under #346.
-* unit tests: fix DBCompactionTest.DisableMultiManualCompaction by blocking all bg compaction threads which increased by default to 8 in #194. 
+* unit tests: fix DBCompactionTest.DisableMultiManualCompaction by blocking all bg compaction threads which increased by default to 8 in #194.
+* Proactive Flushes: fix accounting with non-WBM initiated flushes.
 
 ## Fig v2.5.0 (06/14/2023)
 Based on RocksDB 8.1.1
diff --git a/memtable/write_buffer_manager.cc b/memtable/write_buffer_manager.cc
index b228152e6b..3fef73ea52 100644
--- a/memtable/write_buffer_manager.cc
+++ b/memtable/write_buffer_manager.cc
@@ -658,6 +658,16 @@ void WriteBufferManager::InitiateFlushesThread() {
       while (num_flushes_to_initiate_ > 0U) {
         bool was_flush_initiated = false;
         {
+          // Below an initiator is requested to initate a flush. The initiator
+          // may call another WBM method that relies on these counters. The
+          // counters are updated here, while under the flushes_mu_ lock
+          // (released below) to ensure num_flushes_to_initiate_ can't become
+          // negative Not recalculating flush initiation size since the
+          // increment & decrement cancel each other with respect to the recalc.
+          ++num_running_flushes_;
+          assert(num_flushes_to_initiate_ > 0U);
+          --num_flushes_to_initiate_;
+
           // Unlocking the flushed_mu_ since flushing (via the initiator cb) may
           // call a WBM service (e.g., ReserveMem()), that, in turn, needs to
           // flushes_mu_lock the same mutex => will get stuck
@@ -669,19 +679,16 @@ void WriteBufferManager::InitiateFlushesThread() {
           // 2. Have all existing initiators failed to initiate a flush?
           if (flush_initiators_.empty() ||
               (num_repeated_failures_to_initiate >= flush_initiators_.size())) {
+            // No flush was initiated => undo the counters update
+            assert(num_running_flushes_ > 0U);
+            --num_running_flushes_;
+            ++num_flushes_to_initiate_;
             break;
           }
           assert(IsInitiatorIdxValid(next_candidate_initiator_idx_));
           auto& initiator = flush_initiators_[next_candidate_initiator_idx_];
           UpdateNextCandidateInitiatorIdx();
 
-          // Assuming initiator would flush (flushes_mu_lock is unlocked and
-          // called initiator may call another method that relies on these
-          // counters) Not recalculating flush initiation size since the
-          // increment & decrement cancel each other with respect to the recalc
-          ++num_running_flushes_;
-          --num_flushes_to_initiate_;
-
           // TODO: Use a weak-pointer for the registered initiators. That would
           // allow us to release the flushes_initiators_mu_ mutex before calling
           // the callback (which may take a long time).
@@ -690,6 +697,7 @@ void WriteBufferManager::InitiateFlushesThread() {
 
         if (!was_flush_initiated) {
           // No flush was initiated => undo the counters update
+          assert(num_running_flushes_ > 0U);
           --num_running_flushes_;
           ++num_flushes_to_initiate_;
           ++num_repeated_failures_to_initiate;
@@ -728,6 +736,12 @@ void WriteBufferManager::FlushStarted(bool wbm_initiated) {
   flushes_mu_->Lock();
 
   ++num_running_flushes_;
+  // Any number of non-wbm-initiated flushes may be initiated, so, we must not
+  // underflow num_flushes_to_initiate_
+  if (num_flushes_to_initiate_ > 0U) {
+    --num_flushes_to_initiate_;
+  }
+
   size_t curr_memory_used = memory_usage();
   RecalcFlushInitiationSize();
   ReevaluateNeedForMoreFlushesLockHeld(curr_memory_used);