Make Flush() optionally wait for OnFlushCompleted (FlushOptions::listener_wait) (#14992)

Summary:
DB::Flush() with wait=true only waits until the flush result is committed (memtables removed from the immutable list). The EventListener::OnFlushCompleted callbacks run afterward, on the background flush thread, with the DB mutex released, and the waiter is only woken by a shared condition variable that many unrelated events also signal. As a result a Flush(wait=true) caller can return before -- or in the middle of -- its OnFlushCompleted callbacks. This surfaced as a flaky failure in the c_test table_filter_listener_metadata case, which asserts listener_state.flush_count == 3 immediately after Flush() returns (https://github.com/facebook/rocksdb/actions/runs/29976431449). It is also long-standing behavior that other tests have worked around with TEST_WaitForBackgroundWork().

This adds an opt-in FlushOptions::listener_wait (default false). When set together with wait=true, Flush() does not return until the OnFlushCompleted callbacks for the flushed memtables have finished. Default behavior is unchanged.

Implementation: MemTableList gains a per-column-family num_pending_flush_notifications_ counter (guarded by the DB mutex). The background flush thread reserves a pending notification at flush-commit time -- before the mutex is released to run the callbacks -- and releases it after they return, in both the single-CF (FlushMemTableToOutputFile) and atomic (AtomicFlushMemTablesToOutputFiles) paths. WaitForFlushMemTables treats a CF as finished only when its normal predicate holds and the counter is zero. Because the memtables are removed and the counter incremented under one continuous mutex hold (and, for atomic flush, all CFs are reserved before the per-CF notify loop starts releasing the mutex), a listener_wait waiter can never observe the flush as finished before the callbacks complete.

The counter is per-CF rather than per-flush, so a listener_wait waiter can wait somewhat longer than strictly necessary: under rapid flushing it may block until an unrelated, later flush's OnFlushCompleted on the same CF also drains (e.g. a waiter satisfied by GetEarliestMemTableID() advancing still waits for the whole CF's pending notifications). This can only over-wait, never return early, and is tolerated in favor of implementation simplicity; a per-memtable-id watermark could tighten it if it ever matters.

To keep the counter balanced against future edits, the release uses an RAII Defer (runs on every exit path) and ~MemTableList asserts the counter has drained to zero; SubPendingFlushNotifications asserts against underflow. A reservation is taken only when a callback will actually fire (committed flush info non-empty, not a mempurge, CF not dropped, listeners registered).

Also adds the C API rocksdb_flushoptions_set/get_listener_wait (regenerated via tools/c_api_gen) and a DBImpl::TEST_FlushMemTableWithListenerWait() helper.

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

Test Plan:
New db_flush_test cases ListenerWaitWaitsForOnFlushCompleted and ListenerWaitAtomicFlushWaitsForAllOnFlushCompleted use a listener that blocks inside OnFlushCompleted plus spurious bg_cv_ wakeups to deterministically verify that Flush(wait=true, listener_wait=true) does not return until the callback (all CFs, for atomic flush) has finished; the atomic case guards the per-CF mutex-release gap in the notify loop.

Converted existing racy/worked-around cases to the new option: c_test table_filter_listener_metadata, listener_test OnSingleDBFlushTest and MultiDBMultiListeners, and db_flush_test FireOnFlushCompletedAfterCommittedResult (the last previously relied on TEST_WaitForBackgroundWork()). The auto-generated C API roundtrip test covers the new flushoptions field.

db_stress StressTest::TestFlush now sets listener_wait on a hard-coded 1 in 4 of its manual flushes (both the atomic and non-atomic branches), so the crash test exercises the new path with the DbStressListener registered.

Reviewed By: jaykorean

Differential Revision: D113425407

Pulled By: pdillinger

fbshipit-source-id: 6fd0f5bc93fabaccd8010c2fead1d93c4c44ccfb

c_api_gen/c_generated_option_structs_auto.cc.inc

@@ -1678,6 +1678,16 @@ unsigned char rocksdb_flushoptions_get_force_atomic_flush(
   return opt->rep.force_atomic_flush;
 }
 
+void rocksdb_flushoptions_set_listener_wait(rocksdb_flushoptions_t* opt,
+                                            unsigned char v) {
+  opt->rep.listener_wait = v;
+}
+
+unsigned char rocksdb_flushoptions_get_listener_wait(
+    rocksdb_flushoptions_t* opt) {
+  return opt->rep.listener_wait;
+}
+
 /* FlushWALOptions */
 
 void rocksdb_flushwaloptions_set_sync(rocksdb_flushwaloptions_t* opt,

c_api_gen/c_generated_option_structs_auto.h.inc

@@ -545,6 +545,12 @@ extern ROCKSDB_LIBRARY_API void rocksdb_flushoptions_set_force_atomic_flush(
 extern ROCKSDB_LIBRARY_API unsigned char
 rocksdb_flushoptions_get_force_atomic_flush(rocksdb_flushoptions_t* opt);
 
+extern ROCKSDB_LIBRARY_API void rocksdb_flushoptions_set_listener_wait(
+    rocksdb_flushoptions_t* opt, unsigned char v);
+
+extern ROCKSDB_LIBRARY_API unsigned char rocksdb_flushoptions_get_listener_wait(
+    rocksdb_flushoptions_t* opt);
+
 /* FlushWALOptions */
 
 extern ROCKSDB_LIBRARY_API void rocksdb_flushwaloptions_set_sync(

c_api_gen/c_generated_roundtrip_tests.c.inc

@@ -463,6 +463,10 @@ static void rocksdb_roundtrip_test_rocksdb_flushoptions(void) {
   CheckCondition(rocksdb_flushoptions_get_force_atomic_flush(obj) == 1);
   rocksdb_flushoptions_set_force_atomic_flush(obj, 0);
   CheckCondition(rocksdb_flushoptions_get_force_atomic_flush(obj) == 0);
+  rocksdb_flushoptions_set_listener_wait(obj, 1);
+  CheckCondition(rocksdb_flushoptions_get_listener_wait(obj) == 1);
+  rocksdb_flushoptions_set_listener_wait(obj, 0);
+  CheckCondition(rocksdb_flushoptions_get_listener_wait(obj) == 0);
   rocksdb_flushoptions_set_wait(obj, 1);
   CheckCondition(rocksdb_flushoptions_get_wait(obj) == 1);
   rocksdb_flushoptions_set_wait(obj, 0);

db/c.cc

@@ -10644,6 +10644,16 @@ unsigned char rocksdb_flushoptions_get_force_atomic_flush(
   return opt->rep.force_atomic_flush;
 }
 
+void rocksdb_flushoptions_set_listener_wait(rocksdb_flushoptions_t* opt,
+                                            unsigned char v) {
+  opt->rep.listener_wait = v;
+}
+
+unsigned char rocksdb_flushoptions_get_listener_wait(
+    rocksdb_flushoptions_t* opt) {
+  return opt->rep.listener_wait;
+}
+
 /* FlushWALOptions */
 
 void rocksdb_flushwaloptions_set_sync(rocksdb_flushwaloptions_t* opt,

db/c_test.c

@@ -4531,6 +4531,11 @@ int main(int argc, char** argv) {
     CheckNoError(err);
 
     rocksdb_flushoptions_set_wait(flush_options, 1);
+    // Ensure the OnFlushCompleted callbacks (which update listener_state) have
+    // finished by the time rocksdb_flush() returns, so the flush_count /
+    // saw_flush_entries_* assertions below are not racy.
+    rocksdb_flushoptions_set_listener_wait(flush_options, 1);
+    CheckCondition(1 == rocksdb_flushoptions_get_listener_wait(flush_options));
 
     rocksdb_put(local_db, local_woptions, "a", 1, "1", 1, &err);
     CheckNoError(err);

db/db_flush_test.cc

@@ -8,7 +8,9 @@
 // found in the LICENSE file. See the AUTHORS file for names of contributors.
 
 #include <atomic>
+#include <condition_variable>
 #include <limits>
+#include <mutex>
 
 #include "db/db_impl/db_impl.h"
 #include "db/db_test_util.h"
@@ -2219,8 +2221,12 @@ TEST_F(DBFlushTest, FireOnFlushCompletedAfterCommittedResult) {
   listener->seq1 = db_->GetLatestSequenceNumber();
   // t1 will wait for the second flush complete before committing flush result.
   auto t1 = port::Thread([&]() {
-    // flush_opts.wait = true
-    ASSERT_OK(db_->Flush(FlushOptions()));
+    // flush_opts.wait = true, and listener_wait = true so that Flush() does not
+    // return until the OnFlushCompleted callbacks (which set completed1 and
+    // completed2) have finished running.
+    FlushOptions flush_opts;
+    flush_opts.listener_wait = true;
+    ASSERT_OK(db_->Flush(flush_opts));
   });
   // Wait for first flush started.
   TEST_SYNC_POINT(
@@ -2233,15 +2239,142 @@ TEST_F(DBFlushTest, FireOnFlushCompletedAfterCommittedResult) {
   flush_opts.wait = false;
   ASSERT_OK(db_->Flush(flush_opts));
   t1.join();
-  // Ensure background work is fully finished including listener callbacks
-  // before accessing listener state.
-  ASSERT_OK(dbfull()->TEST_WaitForBackgroundWork());
+  // listener_wait on t1's flush above guarantees both OnFlushCompleted
+  // callbacks have finished by the time Flush() returned; no need to
+  // additionally wait for background work here.
   ASSERT_TRUE(listener->completed1);
   ASSERT_TRUE(listener->completed2);
   SyncPoint::GetInstance()->DisableProcessing();
   SyncPoint::GetInstance()->ClearAllCallBacks();
 }
 
+namespace {
+// EventListener whose OnFlushCompleted blocks until explicitly released, used
+// to verify FlushOptions::listener_wait. Optionally only blocks on the first
+// callback (useful for atomic flush, which fires one callback per CF).
+class BlockingFlushListener : public EventListener {
+ public:
+  explicit BlockingFlushListener(bool block_only_first)
+      : block_only_first_(block_only_first) {}
+
+  void OnFlushCompleted(DB* /*db*/, const FlushJobInfo& /*info*/) override {
+    if (block_only_first_ && num_completed_.fetch_add(1) != 0) {
+      return;
+    }
+    std::unique_lock<std::mutex> lk(mu_);
+    entered_ = true;
+    cv_.notify_all();
+    cv_.wait(lk, [this] { return release_; });
+    returned_ = true;
+  }
+
+  void WaitUntilEntered() {
+    std::unique_lock<std::mutex> lk(mu_);
+    cv_.wait(lk, [this] { return entered_; });
+  }
+
+  void Release() {
+    {
+      std::lock_guard<std::mutex> lk(mu_);
+      release_ = true;
+    }
+    cv_.notify_all();
+  }
+
+  bool CallbackReturned() {
+    std::lock_guard<std::mutex> lk(mu_);
+    return returned_;
+  }
+
+ private:
+  const bool block_only_first_;
+  std::atomic<int> num_completed_{0};
+  std::mutex mu_;
+  std::condition_variable cv_;
+  bool entered_ = false;
+  bool release_ = false;
+  bool returned_ = false;
+};
+}  // anonymous namespace
+
+TEST_F(DBFlushTest, ListenerWaitWaitsForOnFlushCompleted) {
+  // With FlushOptions::listener_wait, Flush(wait=true) must not return until
+  // the OnFlushCompleted callback has finished, even if the flush result is
+  // already committed and bg_cv_ is signalled by unrelated background work.
+  auto listener = std::make_shared<BlockingFlushListener>(
+      /*block_only_first=*/false);
+  Options options = CurrentOptions();
+  options.listeners.push_back(listener);
+  Reopen(options);
+  ASSERT_OK(Put("k", "v"));
+
+  std::atomic<bool> flush_returned{false};
+  port::Thread flush_thread([&] {
+    // Exercises the TEST_ wrapper, which flushes with wait=true and
+    // listener_wait=true.
+    ASSERT_OK(dbfull()->TEST_FlushMemTableWithListenerWait());
+    flush_returned.store(true);
+  });
+
+  // Wait until we are inside OnFlushCompleted. At this point the memtable has
+  // already been committed/removed, which is what unblocks a plain
+  // Flush(wait=true).
+  listener->WaitUntilEntered();
+
+  // Pump spurious bg_cv_ wakeups. A waiter that ignored listener_wait would
+  // wake, observe the memtable already gone, and return here.
+  for (int i = 0; i < 10; ++i) {
+    db_->DisableManualCompaction();
+    db_->EnableManualCompaction();
+  }
+
+  // The OnFlushCompleted callback has not returned, so neither must Flush().
+  ASSERT_FALSE(listener->CallbackReturned());
+  ASSERT_FALSE(flush_returned.load());
+
+  listener->Release();
+  flush_thread.join();
+  ASSERT_TRUE(flush_returned.load());
+  ASSERT_TRUE(listener->CallbackReturned());
+}
+
+TEST_F(DBFlushTest, ListenerWaitAtomicFlushWaitsForAllOnFlushCompleted) {
+  // Same guarantee for atomic flush across multiple column families. The atomic
+  // flush removes all CFs' memtables together and then fires OnFlushCompleted
+  // per CF, releasing the DB mutex between CFs; listener_wait must still block
+  // Flush() until the (first, still-running) callback has finished.
+  auto listener = std::make_shared<BlockingFlushListener>(
+      /*block_only_first=*/true);
+  Options options = CurrentOptions();
+  options.atomic_flush = true;
+  options.listeners.push_back(listener);
+  CreateAndReopenWithCF({"pikachu"}, options);
+  ASSERT_OK(Put(0, "k0", "v0"));
+  ASSERT_OK(Put(1, "k1", "v1"));
+
+  std::atomic<bool> flush_returned{false};
+  port::Thread flush_thread([&] {
+    FlushOptions fo;
+    fo.wait = true;
+    fo.listener_wait = true;
+    ASSERT_OK(db_->Flush(fo, {handles_[0], handles_[1]}));
+    flush_returned.store(true);
+  });
+
+  listener->WaitUntilEntered();
+  for (int i = 0; i < 10; ++i) {
+    db_->DisableManualCompaction();
+    db_->EnableManualCompaction();
+  }
+  ASSERT_FALSE(listener->CallbackReturned());
+  ASSERT_FALSE(flush_returned.load());
+
+  listener->Release();
+  flush_thread.join();
+  ASSERT_TRUE(flush_returned.load());
+  ASSERT_TRUE(listener->CallbackReturned());
+}
+
 TEST_F(DBFlushTest, FlushWithBlob) {
   constexpr uint64_t min_blob_size = 10;
 

db/db_impl/db_impl.h

@@ -1203,6 +1203,13 @@ class DBImpl : public DB {
   Status TEST_FlushMemTable(ColumnFamilyData* cfd,
                             const FlushOptions& flush_opts);
 
+  // Like TEST_FlushMemTable(bool wait, ...), but also sets
+  // FlushOptions::listener_wait so the flush does not return until the
+  // registered EventListener::OnFlushCompleted callbacks for the flushed
+  // memtables have finished running.
+  Status TEST_FlushMemTableWithListenerWait(bool allow_write_stall = false,
+                                            ColumnFamilyHandle* cfh = nullptr);
+
   // Flush (multiple) ColumnFamilyData without using ColumnFamilyHandle. This
   // is because in certain cases, we can flush column families, wait for the
   // flush to complete, but delete the column family handle before the wait
@@ -2518,18 +2525,24 @@ class DBImpl : public DB {
   // memtable pending flush.
   // resuming_from_bg_err indicates whether the caller is attempting to resume
   // from background error.
+  // If wait_for_listener_notifications is true, also wait until the
+  // EventListener::OnFlushCompleted callbacks for the flushed memtables have
+  // finished running (see FlushOptions::listener_wait).
   Status WaitForFlushMemTable(
       ColumnFamilyData* cfd, const uint64_t* flush_memtable_id = nullptr,
       bool resuming_from_bg_err = false,
-      std::optional<FlushReason> flush_reason = std::nullopt) {
+      std::optional<FlushReason> flush_reason = std::nullopt,
+      bool wait_for_listener_notifications = false) {
     return WaitForFlushMemTables({cfd}, {flush_memtable_id},
-                                 resuming_from_bg_err, flush_reason);
+                                 resuming_from_bg_err, flush_reason,
+                                 wait_for_listener_notifications);
   }
   // Wait for memtables to be flushed for multiple column families.
   Status WaitForFlushMemTables(
       const autovector<ColumnFamilyData*>& cfds,
       const autovector<const uint64_t*>& flush_memtable_ids,
-      bool resuming_from_bg_err, std::optional<FlushReason> flush_reason);
+      bool resuming_from_bg_err, std::optional<FlushReason> flush_reason,
+      bool wait_for_listener_notifications = false);
 
   inline void WaitForPendingWrites() {
     mutex_.AssertHeld();

db/db_impl/db_impl_compaction_flush.cc

@@ -31,6 +31,7 @@
 #include "util/cast_util.h"
 #include "util/coding.h"
 #include "util/concurrent_task_limiter_impl.h"
+#include "util/defer.h"
 #include "util/udt_util.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -358,6 +359,34 @@ Status DBImpl::FlushMemTableToOutputFile(
     need_cancel = false;
   }
 
+  // Reserve a pending flush-completion notification before the DB mutex can be
+  // released, so that a Flush(wait=true, listener_wait=true) caller cannot
+  // observe this flush as finished until NotifyOnFlushCompleted() below has
+  // run. flush_job.Run() removed the memtables from the immutable list (which
+  // is what unblocks the waiter) and the DB mutex has been held continuously
+  // since, so there is no window in which the waiter could see the flush
+  // committed with a zero count. Reserved only when a notification will
+  // actually fire (matches the NotifyOnFlushCompleted call below): flush
+  // committed successfully, not a mempurge, listeners are registered, and this
+  // job committed at least one flush result to report.
+  const bool reserved_flush_notification =
+      s.ok() && !switched_to_mempurge &&
+      !immutable_db_options_.listeners.empty() &&
+      !flush_job.GetCommittedFlushJobsInfo()->empty();
+  if (reserved_flush_notification) {
+    cfd->imm()->AddPendingFlushNotifications(1);
+  }
+  // Release the reservation on every exit path (RAII) so a future early return
+  // cannot leak the count and hang listener_wait waiters. The DB mutex is held
+  // throughout this function, including at scope exit, which
+  // SubPendingFlushNotifications() requires. Runs after
+  // NotifyOnFlushCompleted() below has returned.
+  Defer release_flush_notification([&] {
+    if (reserved_flush_notification) {
+      cfd->imm()->SubPendingFlushNotifications(1);
+    }
+  });
+
   if (cfd->blob_partition_manager() != nullptr &&
       prepared_blob_generations > 0) {
     auto unconsumed_additions = flush_job.TakeExternalBlobFileAdditions();
@@ -900,6 +929,42 @@ Status DBImpl::AtomicFlushMemTablesToOutputFiles(
         directories_.GetDbDir(), log_buffer);
   }
 
+  // Reserve a pending flush-completion notification for every CF that the
+  // NotifyOnFlushCompleted loop below will service, while the DB mutex is still
+  // held continuously since InstallMemtableAtomicFlushResults removed the
+  // memtables. This closes the window where a
+  // Flush(wait=true, listener_wait=true) waiter could observe an atomic flush
+  // as finished before a later CF's OnFlushCompleted has run, since the notify
+  // loop releases and re-acquires the mutex between CFs. The reservations are
+  // released by a matching pass after the notify loop.
+  autovector<uint8_t> reserved_flush_notification;
+  reserved_flush_notification.reserve(num_cfs);
+  {
+    const bool has_listeners = !immutable_db_options_.listeners.empty();
+    for (int i = 0; i != num_cfs; ++i) {
+      const bool reserve = s.ok() && has_listeners &&
+                           !switched_to_mempurge[i] && !cfds[i]->IsDropped() &&
+                           !jobs[i]->GetCommittedFlushJobsInfo()->empty();
+      if (reserve) {
+        cfds[i]->imm()->AddPendingFlushNotifications(1);
+      }
+      reserved_flush_notification.push_back(reserve ? 1 : 0);
+    }
+  }
+  // Release the reservations on every exit path (RAII) so a future early return
+  // cannot leak the count and hang listener_wait waiters. Done independently of
+  // the notify loop's skips so the count stays balanced even if a CF was
+  // dropped after it was reserved. The DB mutex is held throughout this
+  // function, including at scope exit, which SubPendingFlushNotifications()
+  // requires.
+  Defer release_flush_notifications([&] {
+    for (int i = 0; i != num_cfs; ++i) {
+      if (reserved_flush_notification[i]) {
+        cfds[i]->imm()->SubPendingFlushNotifications(1);
+      }
+    }
+  });
+
   if (!s.ok()) {
     // If the atomic flush's combined MANIFEST write failed, the output files
     // are cached in the table cache (added by BuildTable during FlushJob::Run)
@@ -2794,7 +2859,7 @@ Status DBImpl::FlushMemTable(ColumnFamilyData* cfd,
     s = WaitForFlushMemTables(
         cfds, flush_memtable_ids,
         flush_reason == FlushReason::kErrorRecovery /* resuming_from_bg_err */,
-        flush_reason);
+        flush_reason, flush_options.listener_wait);
     InstrumentedMutexLock lock_guard(&mutex_);
     for (auto* tmp_cfd : cfds) {
       tmp_cfd->UnrefAndTryDelete();
@@ -2949,7 +3014,7 @@ Status DBImpl::AtomicFlushMemTables(
     s = WaitForFlushMemTables(
         cfds, flush_memtable_ids,
         flush_reason == FlushReason::kErrorRecovery /* resuming_from_bg_err */,
-        flush_reason);
+        flush_reason, flush_options.listener_wait);
     InstrumentedMutexLock lock_guard(&mutex_);
     for (auto* cfd : cfds) {
       cfd->UnrefAndTryDelete();
@@ -3118,7 +3183,8 @@ Status DBImpl::WaitUntilFlushWouldNotStallWrites(ColumnFamilyData* cfd,
 Status DBImpl::WaitForFlushMemTables(
     const autovector<ColumnFamilyData*>& cfds,
     const autovector<const uint64_t*>& flush_memtable_ids,
-    bool resuming_from_bg_err, std::optional<FlushReason> flush_reason) {
+    bool resuming_from_bg_err, std::optional<FlushReason> flush_reason,
+    bool wait_for_listener_notifications) {
   int num = static_cast<int>(cfds.size());
   // Wait until the compaction completes
   InstrumentedMutexLock l(&mutex_);
@@ -3163,7 +3229,13 @@ Status DBImpl::WaitForFlushMemTables(
             flush_reason.value() != FlushReason::kExternalFileIngestion ||
             cfds[i]->GetSuperVersion()->imm->GetID() ==
                 cfds[i]->imm()->current()->GetID()) {
-          ++num_finished;
+          // When the caller requested it, also wait for the flush-completion
+          // listener callbacks for this CF to finish before counting it as
+          // finished (see FlushOptions::listener_wait).
+          if (!wait_for_listener_notifications ||
+              cfds[i]->imm()->NumPendingFlushNotifications() == 0) {
+            ++num_finished;
+          }
         }
       }
     }

db/db_impl/db_impl_debug.cc

@@ -164,6 +164,22 @@ Status DBImpl::TEST_FlushMemTable(ColumnFamilyData* cfd,
   return FlushMemTable(cfd, flush_opts, FlushReason::kTest);
 }
 
+Status DBImpl::TEST_FlushMemTableWithListenerWait(bool allow_write_stall,
+                                                  ColumnFamilyHandle* cfh) {
+  FlushOptions fo;
+  fo.wait = true;
+  fo.listener_wait = true;
+  fo.allow_write_stall = allow_write_stall;
+  ColumnFamilyData* cfd;
+  if (cfh == nullptr) {
+    cfd = default_cf_handle_->cfd();
+  } else {
+    auto cfhi = static_cast_with_check<ColumnFamilyHandleImpl>(cfh);
+    cfd = cfhi->cfd();
+  }
+  return FlushMemTable(cfd, fo, FlushReason::kTest);
+}
+
 Status DBImpl::TEST_AtomicFlushMemTables(
     const autovector<ColumnFamilyData*>& provided_candidate_cfds,
     const FlushOptions& flush_opts) {

db/listener_test.cc

@@ -520,12 +520,13 @@ TEST_F(EventListenerTest, OnSingleDBFlushTest) {
   ASSERT_OK(Put(5, "nikitich", std::string(90000, 'n')));
   ASSERT_OK(Put(6, "alyosha", std::string(90000, 'a')));
   ASSERT_OK(Put(7, "popovich", std::string(90000, 'p')));
+  FlushOptions flush_opts;
+  flush_opts.listener_wait = true;
   for (int i = 1; i < 8; ++i) {
-    ASSERT_OK(Flush(i));
-    ASSERT_OK(dbfull()->TEST_WaitForFlushMemTable());
-    // Ensure background work is fully finished including listener callbacks
-    // before accessing listener state.
-    ASSERT_OK(dbfull()->TEST_WaitForBackgroundWork());
+    ASSERT_OK(db_->Flush(flush_opts, handles_[i]));
+    // listener_wait guarantees the OnFlushCompleted callbacks have finished
+    // running by the time Flush() returns, so the listener state below is up to
+    // date without separately waiting for background work.
     ASSERT_EQ(listener->flushed_dbs_.size(), i);
     ASSERT_EQ(listener->flushed_column_family_names_.size(), i);
   }
@@ -635,21 +636,17 @@ TEST_F(EventListenerTest, MultiDBMultiListeners) {
     }
   }
 
+  FlushOptions flush_opts;
+  flush_opts.listener_wait = true;
   for (size_t c = 0; c < cf_names.size(); ++c) {
     for (int d = 0; d < kNumDBs; ++d) {
-      ASSERT_OK(dbs[d]->Flush(FlushOptions(), vec_handles[d][c]));
-      ASSERT_OK(static_cast_with_check<DBImpl>(dbs[d].get())
-                    ->TEST_WaitForFlushMemTable());
+      // listener_wait ensures the OnFlushCompleted callbacks have finished by
+      // the time Flush() returns, so the listener state checked below is up to
+      // date without separately waiting for background work.
+      ASSERT_OK(dbs[d]->Flush(flush_opts, vec_handles[d][c]));
     }
   }
 
-  for (int d = 0; d < kNumDBs; ++d) {
-    // Ensure background work is fully finished including listener callbacks
-    // before accessing listener state.
-    ASSERT_OK(static_cast_with_check<DBImpl>(dbs[d].get())
-                  ->TEST_WaitForBackgroundWork());
-  }
-
   for (auto* listener : listeners) {
     int pos = 0;
     for (size_t c = 0; c < cf_names.size(); ++c) {

db/memtable_list.h

@@ -287,7 +287,12 @@ class MemTableList {
 
   // Should not delete MemTableList without making sure MemTableList::current()
   // is Unref()'d.
-  ~MemTableList() {}
+  ~MemTableList() {
+    // All flush-completion notifications must have drained before this column
+    // family's memtable list is destroyed; otherwise a reserve leaked (see
+    // num_pending_flush_notifications_).
+    assert(num_pending_flush_notifications_ == 0);
+  }
 
   MemTableListVersion* current() const { return current_; }
 
@@ -305,6 +310,20 @@ class MemTableList {
   // completely flushed and logged.
   int NumFlushed() const;
 
+  // Flush-completion listener notification accounting. All require the DB
+  // mutex to be held (this class is not internally synchronized). See
+  // num_pending_flush_notifications_ for details.
+  void AddPendingFlushNotifications(int n) {
+    num_pending_flush_notifications_ += n;
+  }
+  void SubPendingFlushNotifications(int n) {
+    num_pending_flush_notifications_ -= n;
+    assert(num_pending_flush_notifications_ >= 0);
+  }
+  int NumPendingFlushNotifications() const {
+    return num_pending_flush_notifications_;
+  }
+
   // Returns true if there is at least one memtable on which flush has
   // not yet started.
   bool IsFlushPending() const;
@@ -552,6 +571,15 @@ class MemTableList {
   // Last memtabe list version id, increase by 1 each time a new
   // MemtableListVersion is installed.
   uint64_t last_memtable_list_version_id_;
+
+  // Number of committed flush results for this column family whose
+  // EventListener::OnFlushCompleted callbacks have not finished running yet.
+  // Incremented (by the background flush thread) at flush-commit time, before
+  // the DB mutex is released to invoke the callbacks, and decremented after the
+  // callbacks return. Consulted by WaitForFlushMemTables() when the caller set
+  // FlushOptions::listener_wait, so that Flush(wait=true) does not return
+  // before the callbacks complete. Guarded by the DB mutex.
+  int num_pending_flush_notifications_ = 0;
 };
 
 // Installs memtable atomic flush results.

db_stress_tool/db_stress_test_base.cc

@@ -3905,6 +3905,11 @@ void StressTest::TestFlush(ThreadState* thread,
                            const std::vector<int>& rand_column_families) {
   FlushOptions flush_opts;
   assert(flush_opts.wait);
+  // Occasionally exercise the stronger guarantee that Flush() does not return
+  // until the OnFlushCompleted listener callbacks have finished.
+  if (thread->rand.OneIn(4)) {
+    flush_opts.listener_wait = true;
+  }
   Status status;
   if (FLAGS_atomic_flush) {
     status = db_->Flush(flush_opts, column_families_);

include/rocksdb/c.h

@@ -5483,6 +5483,12 @@ extern ROCKSDB_LIBRARY_API void rocksdb_flushoptions_set_force_atomic_flush(
 extern ROCKSDB_LIBRARY_API unsigned char
 rocksdb_flushoptions_get_force_atomic_flush(rocksdb_flushoptions_t* opt);
 
+extern ROCKSDB_LIBRARY_API void rocksdb_flushoptions_set_listener_wait(
+    rocksdb_flushoptions_t* opt, unsigned char v);
+
+extern ROCKSDB_LIBRARY_API unsigned char rocksdb_flushoptions_get_listener_wait(
+    rocksdb_flushoptions_t* opt);
+
 /* FlushWALOptions */
 
 extern ROCKSDB_LIBRARY_API void rocksdb_flushwaloptions_set_sync(

include/rocksdb/options.h

@@ -2609,8 +2609,22 @@ struct FlushOptions {
   // Default: false (uses DBOptions::atomic_flush setting).
   bool force_atomic_flush;
 
+  // If true (and `wait` is also true), Flush() will not return until the
+  // registered EventListener::OnFlushCompleted callbacks for the flushed
+  // memtables have finished running. By default (false), Flush(wait=true) may
+  // return as soon as the flush result is committed, which can be before (or
+  // while) the OnFlushCompleted callbacks execute on the background flush
+  // thread. Set this to true when the caller needs to observe the effects of
+  // its OnFlushCompleted listener(s) immediately after Flush() returns.
+  // Has no effect when `wait == false`.
+  // Default: false
+  bool listener_wait;
+
   FlushOptions()
-      : wait(true), allow_write_stall(false), force_atomic_flush(false) {}
+      : wait(true),
+        allow_write_stall(false),
+        force_atomic_flush(false),
+        listener_wait(false) {}
 };
 
 struct FlushWALOptions {

java/rocksjni/options.cc

@@ -8797,6 +8797,31 @@ jboolean Java_org_rocksdb_FlushOptions_allowWriteStall(JNIEnv*, jclass,
   return static_cast<jboolean>(flush_options->allow_write_stall);
 }
 
+/*
+ * Class:     org_rocksdb_FlushOptions
+ * Method:    setListenerWait
+ * Signature: (JZ)V
+ */
+void Java_org_rocksdb_FlushOptions_setListenerWait(JNIEnv*, jclass,
+                                                   jlong jhandle,
+                                                   jboolean jlistener_wait) {
+  auto* flush_options =
+      reinterpret_cast<ROCKSDB_NAMESPACE::FlushOptions*>(jhandle);
+  flush_options->listener_wait = jlistener_wait == JNI_TRUE;
+}
+
+/*
+ * Class:     org_rocksdb_FlushOptions
+ * Method:    listenerWait
+ * Signature: (J)Z
+ */
+jboolean Java_org_rocksdb_FlushOptions_listenerWait(JNIEnv*, jclass,
+                                                    jlong jhandle) {
+  auto* flush_options =
+      reinterpret_cast<ROCKSDB_NAMESPACE::FlushOptions*>(jhandle);
+  return static_cast<jboolean>(flush_options->listener_wait);
+}
+
 /*
  * Class:     org_rocksdb_FlushOptions
  * Method:    disposeInternal

java/src/main/java/org/rocksdb/FlushOptions.java

@@ -73,6 +73,42 @@ public class FlushOptions extends RocksObject {
     assert(isOwningHandle());
     return allowWriteStall(nativeHandle_);
   }
+
+  /**
+   * If true (and {@link #waitForFlush()} is also true), the flush operation will
+   * not return until the registered
+   * {@link org.rocksdb.AbstractEventListener}'s {@code onFlushCompleted}
+   * callbacks for the flushed memtables have finished running.
+   * <p>
+   * By default (false), a flush with wait may return as soon as the flush
+   * result is committed, which can be before (or while) the completion callbacks
+   * execute on the background flush thread. Has no effect when
+   * {@link #waitForFlush()} is false.
+   * <p>
+   * Default: false
+   *
+   * @param listenerWait true to wait for the flush-completion listener callbacks
+   *     to finish, false otherwise.
+   *
+   * @return instance of current FlushOptions.
+   */
+  public FlushOptions setListenerWait(final boolean listenerWait) {
+    assert (isOwningHandle());
+    setListenerWait(nativeHandle_, listenerWait);
+    return this;
+  }
+
+  /**
+   * Returns true if the flush operation waits for the flush-completion listener
+   * callbacks to finish, false otherwise.
+   *
+   * @return true if the flush waits for the flush-completion listener callbacks
+   */
+  public boolean listenerWait() {
+    assert (isOwningHandle());
+    return listenerWait(nativeHandle_);
+  }
+
   private static long newFlushOptionsInance() {
     RocksDB.loadLibrary();
     return newFlushOptions();
@@ -89,4 +125,6 @@ public class FlushOptions extends RocksObject {
   private static native boolean waitForFlush(final long handle);
   private static native void setAllowWriteStall(final long handle, final boolean allowWriteStall);
   private static native boolean allowWriteStall(final long handle);
+  private static native void setListenerWait(final long handle, final boolean listenerWait);
+  private static native boolean listenerWait(final long handle);
 }

java/src/test/java/org/rocksdb/FlushOptionsTest.java

@@ -28,4 +28,13 @@ public class FlushOptionsTest {
       assertThat(flushOptions.allowWriteStall()).isTrue();
     }
   }
+
+  @Test
+  public void listenerWait() {
+    try (final FlushOptions flushOptions = new FlushOptions()) {
+      assertThat(flushOptions.listenerWait()).isFalse();
+      flushOptions.setListenerWait(true);
+      assertThat(flushOptions.listenerWait()).isTrue();
+    }
+  }
 }

unreleased_history/public_api_changes/flush_listener_wait.md

@@ -0,0 +1 @@
+Added `FlushOptions::listener_wait` (default `false`). When set together with `FlushOptions::wait == true`, `DB::Flush()` will not return until the registered `EventListener::OnFlushCompleted` callbacks for the flushed memtables have finished running. By default (`false`), `Flush(wait=true)` may return as soon as the flush result is committed, which can be before (or while) the `OnFlushCompleted` callbacks execute on the background flush thread. Also added the corresponding C API `rocksdb_flushoptions_set_listener_wait()` / `rocksdb_flushoptions_get_listener_wait()`.