From bb355521f0a81c01f2fa660752b325d8b6fd225c Mon Sep 17 00:00:00 2001
From: Matthew Von-Maszewski <matthewvon@users.noreply.github.com>
Date: Fri, 31 Jan 2025 09:39:53 -0500
Subject: [PATCH] PLAT-7494:  rocksdb stall on manual Flush() and
 CompactRange() (#57)

---
 db/db_impl/db_impl_compaction_flush.cc |  8 ++++++++
 db/memtable_list.cc                    |  3 ++-
 db/memtable_list.h                     | 10 ++++++++++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/db/db_impl/db_impl_compaction_flush.cc b/db/db_impl/db_impl_compaction_flush.cc
index f552310e8..948bf2e48 100644
--- a/db/db_impl/db_impl_compaction_flush.cc
+++ b/db/db_impl/db_impl_compaction_flush.cc
@@ -1989,6 +1989,10 @@ Status DBImpl::FlushMemTable(ColumnFamilyData* cfd,
     }
   }
 
+  if (flush_reason == FlushReason::kManualCompaction ||
+      flush_reason == FlushReason::kManualFlush) {
+    cfd->imm()->BeginManualOperation();
+  }
   autovector<FlushRequest> flush_reqs;
   autovector<uint64_t> memtable_ids_to_wait;
   {
@@ -2107,6 +2111,10 @@ Status DBImpl::FlushMemTable(ColumnFamilyData* cfd,
       tmp_cfd->UnrefAndTryDelete();
     }
   }
+  if (flush_reason == FlushReason::kManualCompaction ||
+      flush_reason == FlushReason::kManualFlush) {
+    cfd->imm()->CompleteManualOperation();
+  }
   TEST_SYNC_POINT("DBImpl::FlushMemTable:FlushMemTableFinished");
   return s;
 }
diff --git a/db/memtable_list.cc b/db/memtable_list.cc
index f447ee735..0740a62a7 100644
--- a/db/memtable_list.cc
+++ b/db/memtable_list.cc
@@ -329,7 +329,8 @@ bool MemTableListVersion::TrimHistory(autovector<MemTable*>* to_delete,
 // not yet started.
 bool MemTableList::IsFlushPending() const {
   if ((flush_requested_ && num_flush_not_started_ > 0) ||
-      (num_flush_not_started_ >= min_write_buffer_number_to_merge_)) {
+      (num_flush_not_started_ >= min_write_buffer_number_to_merge_) ||
+      (active_manuals_ && num_flush_not_started_)) {
     assert(imm_flush_needed.load(std::memory_order_relaxed));
     return true;
   }
diff --git a/db/memtable_list.h b/db/memtable_list.h
index 866ecccb6..7219c7970 100644
--- a/db/memtable_list.h
+++ b/db/memtable_list.h
@@ -391,6 +391,13 @@ class MemTableList {
   void RemoveOldMemTables(uint64_t log_number,
                           autovector<MemTable*>* to_delete);
 
+  void BeginManualOperation() { ++active_manuals_; };
+
+  void CompleteManualOperation() {
+    assert(active_manuals_ >= 1);
+    --active_manuals_;
+  };
+
  private:
   friend Status InstallMemtableAtomicFlushResults(
       const autovector<MemTableList*>* imm_lists,
@@ -436,6 +443,9 @@ class MemTableList {
 
   // Cached value of current_->HasHistory().
   std::atomic<bool> current_has_history_;
+
+  // count of manual flush/compactions active
+  std::atomic<int> active_manuals_{0};
 };
 
 // Installs memtable atomic flush results.