From bb77d1707221d17f4e76bae031a023f0b0c428e0 Mon Sep 17 00:00:00 2001 From: Peter Dillinger Date: Fri, 24 Jul 2026 09:36:04 -0700 Subject: [PATCH] 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_generated_option_structs_auto.cc.inc | 10 ++ .../c_generated_option_structs_auto.h.inc | 6 + c_api_gen/c_generated_roundtrip_tests.c.inc | 4 + db/c.cc | 10 ++ db/c_test.c | 5 + db/db_flush_test.cc | 143 +++++++++++++++++- db/db_impl/db_impl.h | 19 ++- db/db_impl/db_impl_compaction_flush.cc | 80 +++++++++- db/db_impl/db_impl_debug.cc | 16 ++ db/listener_test.cc | 27 ++-- db/memtable_list.h | 30 +++- db_stress_tool/db_stress_test_base.cc | 5 + include/rocksdb/c.h | 6 + include/rocksdb/options.h | 16 +- java/rocksjni/options.cc | 25 +++ .../main/java/org/rocksdb/FlushOptions.java | 38 +++++ .../java/org/rocksdb/FlushOptionsTest.java | 9 ++ .../public_api_changes/flush_listener_wait.md | 1 + 18 files changed, 421 insertions(+), 29 deletions(-) create mode 100644 unreleased_history/public_api_changes/flush_listener_wait.md diff --git a/c_api_gen/c_generated_option_structs_auto.cc.inc b/c_api_gen/c_generated_option_structs_auto.cc.inc index 1c68553831..eab2ece854 100644 --- a/c_api_gen/c_generated_option_structs_auto.cc.inc +++ b/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, diff --git a/c_api_gen/c_generated_option_structs_auto.h.inc b/c_api_gen/c_generated_option_structs_auto.h.inc index b5fce25c80..fef3d31baf 100644 --- a/c_api_gen/c_generated_option_structs_auto.h.inc +++ b/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( diff --git a/c_api_gen/c_generated_roundtrip_tests.c.inc b/c_api_gen/c_generated_roundtrip_tests.c.inc index b7c8330f09..e3ffa048e8 100644 --- a/c_api_gen/c_generated_roundtrip_tests.c.inc +++ b/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); diff --git a/db/c.cc b/db/c.cc index 43c779164b..2daf8a660c 100644 --- a/db/c.cc +++ b/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, diff --git a/db/c_test.c b/db/c_test.c index 8bb38e9d3a..8973393407 100644 --- a/db/c_test.c +++ b/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); diff --git a/db/db_flush_test.cc b/db/db_flush_test.cc index 3fe3d01266..83f823d14e 100644 --- a/db/db_flush_test.cc +++ b/db/db_flush_test.cc @@ -8,7 +8,9 @@ // found in the LICENSE file. See the AUTHORS file for names of contributors. #include +#include #include +#include #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 lk(mu_); + entered_ = true; + cv_.notify_all(); + cv_.wait(lk, [this] { return release_; }); + returned_ = true; + } + + void WaitUntilEntered() { + std::unique_lock lk(mu_); + cv_.wait(lk, [this] { return entered_; }); + } + + void Release() { + { + std::lock_guard lk(mu_); + release_ = true; + } + cv_.notify_all(); + } + + bool CallbackReturned() { + std::lock_guard lk(mu_); + return returned_; + } + + private: + const bool block_only_first_; + std::atomic 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( + /*block_only_first=*/false); + Options options = CurrentOptions(); + options.listeners.push_back(listener); + Reopen(options); + ASSERT_OK(Put("k", "v")); + + std::atomic 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( + /*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 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; diff --git a/db/db_impl/db_impl.h b/db/db_impl/db_impl.h index a9d8435aac..bd2f57ebd1 100644 --- a/db/db_impl/db_impl.h +++ b/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 flush_reason = std::nullopt) { + std::optional 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& cfds, const autovector& flush_memtable_ids, - bool resuming_from_bg_err, std::optional flush_reason); + bool resuming_from_bg_err, std::optional flush_reason, + bool wait_for_listener_notifications = false); inline void WaitForPendingWrites() { mutex_.AssertHeld(); diff --git a/db/db_impl/db_impl_compaction_flush.cc b/db/db_impl/db_impl_compaction_flush.cc index d383924cd6..c8ce2904f8 100644 --- a/db/db_impl/db_impl_compaction_flush.cc +++ b/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 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& cfds, const autovector& flush_memtable_ids, - bool resuming_from_bg_err, std::optional flush_reason) { + bool resuming_from_bg_err, std::optional flush_reason, + bool wait_for_listener_notifications) { int num = static_cast(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; + } } } } diff --git a/db/db_impl/db_impl_debug.cc b/db/db_impl/db_impl_debug.cc index 2ab541063e..12c6f39cff 100644 --- a/db/db_impl/db_impl_debug.cc +++ b/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(cfh); + cfd = cfhi->cfd(); + } + return FlushMemTable(cfd, fo, FlushReason::kTest); +} + Status DBImpl::TEST_AtomicFlushMemTables( const autovector& provided_candidate_cfds, const FlushOptions& flush_opts) { diff --git a/db/listener_test.cc b/db/listener_test.cc index ec1578bc08..26f0d349b1 100644 --- a/db/listener_test.cc +++ b/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(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(dbs[d].get()) - ->TEST_WaitForBackgroundWork()); - } - for (auto* listener : listeners) { int pos = 0; for (size_t c = 0; c < cf_names.size(); ++c) { diff --git a/db/memtable_list.h b/db/memtable_list.h index 468a1c33d0..c75cf51480 100644 --- a/db/memtable_list.h +++ b/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. diff --git a/db_stress_tool/db_stress_test_base.cc b/db_stress_tool/db_stress_test_base.cc index 9ab8fde026..2d4548ccac 100644 --- a/db_stress_tool/db_stress_test_base.cc +++ b/db_stress_tool/db_stress_test_base.cc @@ -3905,6 +3905,11 @@ void StressTest::TestFlush(ThreadState* thread, const std::vector& 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_); diff --git a/include/rocksdb/c.h b/include/rocksdb/c.h index 147e5945da..c7a29b4aa0 100644 --- a/include/rocksdb/c.h +++ b/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( diff --git a/include/rocksdb/options.h b/include/rocksdb/options.h index a8805da757..a62806ef43 100644 --- a/include/rocksdb/options.h +++ b/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 { diff --git a/java/rocksjni/options.cc b/java/rocksjni/options.cc index 13d4041e07..8cb9b17216 100644 --- a/java/rocksjni/options.cc +++ b/java/rocksjni/options.cc @@ -8797,6 +8797,31 @@ jboolean Java_org_rocksdb_FlushOptions_allowWriteStall(JNIEnv*, jclass, return static_cast(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(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(jhandle); + return static_cast(flush_options->listener_wait); +} + /* * Class: org_rocksdb_FlushOptions * Method: disposeInternal diff --git a/java/src/main/java/org/rocksdb/FlushOptions.java b/java/src/main/java/org/rocksdb/FlushOptions.java index cb723fb7db..61303621a1 100644 --- a/java/src/main/java/org/rocksdb/FlushOptions.java +++ b/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. + *

+ * 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. + *

+ * 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); } diff --git a/java/src/test/java/org/rocksdb/FlushOptionsTest.java b/java/src/test/java/org/rocksdb/FlushOptionsTest.java index f90ae911d8..740f2a5814 100644 --- a/java/src/test/java/org/rocksdb/FlushOptionsTest.java +++ b/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(); + } + } } diff --git a/unreleased_history/public_api_changes/flush_listener_wait.md b/unreleased_history/public_api_changes/flush_listener_wait.md new file mode 100644 index 0000000000..b94d110c94 --- /dev/null +++ b/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()`.