Notify listeners before DB shutdown begins (#14769)

Summary:
### Notify listeners before DB shutdown begins

db_stress listener bookkeeping correlates compaction callbacks with file deletion callbacks. DBImpl intentionally skips some compaction listener callbacks once shutdown starts, but file deletion callbacks can still be delivered. That mismatch is normally tolerable for production listeners, but db_stress uses listener-local tracking to detect callback consistency and can report a false positive when shutdown interrupts a compaction callback sequence.

The failed-open case is the important gap. DB::Open() can create a DBImpl, recover or flush files, schedule background compaction, and then fail during late open work such as persisting OPTIONS or waiting for open-time compaction under fault injection. At that point DB::Open() tears down the internal DBImpl before returning an error to StressTest::Open(). If OnCompactionBegin already recorded an input file or job in DbStressListener, DBImpl shutdown can suppress the later OnCompactionPreCommit/OnCompactionCompleted callbacks that would normally clear that state. Since control has not returned to the stress harness yet, StressTest::CleanUp()/Reopen() cannot notify the listener in time. A later shutdown-time callback, or the next open retry reusing the same listener, can then observe stale tracking and abort even though RocksDB did not compact the same SST concurrently.

Add EventListener::OnDBShutdownBegin and fire it once from DBImpl::CancelAllBackgroundWork() before publishing shutting_down_. The callback also covers cleanup of a failed DB::Open() attempt, where the DB pointer refers to the internal DBImpl that was never returned to the caller. Track shutdown_notification_sent_ separately from shutting_down_ because listeners are invoked with mutex_ released, and a concurrent or reentrant cancellation must not deliver the callback twice.

Update DbStressListener to consume the DBImpl-driven shutdown notification instead of relying on StressTest::CleanUp()/Reopen() to manually notify it. This lets db_stress mark itself as shutting down before DBImpl starts skipping shutdown-sensitive compaction notifications, including during failed-open cleanup.

### Also keep listener state scoped to one db_stress open attempt.

Initialize listeners at the top of each non-transactional open retry before enabling open fault injection, so retry attempts get fresh listener state without widening open fault injection to listener construction. Factor the open fault setup into a small helper to keep the retry loop readable. The transaction open path still initializes listeners once because it does not use this open-fault retry loop.

### Also fix multi-ops transaction listener checks during shutdown

A TSAN race was reported where MultiOpsTxnsStressListener::OnCompactionCompleted called VerifyPkSkFast on a background compaction thread while the main thread was destroying the transaction DB wrapper in StressTest::CleanUp(). The reported read was a virtual call through the DB object and the write was the WritePreparedTxnDB/WriteUnpreparedTxnDB destructor updating the vptr.

The vulnerable ordering is that DBImpl can already be inside NotifyOnCompactionCompleted before shutdown is requested. It unlocks db mutex while iterating listeners; another listener such as DbStressListener can spend time in its callback, giving the main thread time to enter CleanUp()/Close(). The DB object is still shutting down, but MultiOpsTxnsStressListener may be invoked later in the same callback iteration and call VerifyPkSkFast through stress_test_->db_aptr_, racing with DB wrapper destruction.

With DBImpl-owned EventListener::OnDBShutdownBegin callback, we have MultiOpsTxnsStressListener consume that callback directly. Once DBImpl begins shutdown, the listener skips both flush-completed and compaction-completed verification callbacks, avoiding DB access during teardown. This also covers failed-open cleanup without name-based downcasts in StressTest.

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

Test Plan: CI

Reviewed By: pdillinger

Differential Revision: D106011452

Pulled By: xingbowang

fbshipit-source-id: 768838ddcd9910de5d1b5204c990a4d88dbc850c

db/db_impl/db_impl.cc

@@ -495,6 +495,21 @@ void DBImpl::WaitForAsyncFileOpen() {
   }
 }
 
+void DBImpl::NotifyOnDBShutdownBegin() {
+  mutex_.AssertHeld();
+  if (shutdown_notification_sent_ || immutable_db_options_.listeners.empty()) {
+    return;
+  }
+  shutdown_notification_sent_ = true;
+
+  // release lock while notifying events
+  mutex_.Unlock();
+  for (const auto& listener : immutable_db_options_.listeners) {
+    listener->OnDBShutdownBegin(this);
+  }
+  mutex_.Lock();
+}
+
 // Will lock the mutex_,  will wait for completion if wait is true
 void DBImpl::CancelAllBackgroundWork(bool wait) {
   ROCKS_LOG_INFO(immutable_db_options_.info_log,
@@ -541,6 +556,9 @@ void DBImpl::CancelAllBackgroundWork(bool wait) {
     immutable_db_options_.compaction_service->CancelAwaitingJobs();
   }
 
+  if (!already_shutting_down) {
+    NotifyOnDBShutdownBegin();
+  }
   shutting_down_.store(true, std::memory_order_release);
   bg_cv_.SignalAll();
   if (!wait) {

db/db_impl/db_impl.h

@@ -1455,6 +1455,13 @@ class DBImpl : public DB {
   std::atomic<int> next_job_id_ = 1;
 
   std::atomic<bool> shutting_down_ = false;
+  // Protected by mutex_. This is separate from shutting_down_ because
+  // OnDBShutdownBegin must fire before shutting_down_ is published, and the
+  // notification releases mutex_ while invoking listeners. Marking that the
+  // notification has started prevents a concurrent or reentrant
+  // CancelAllBackgroundWork() call from firing it again during that unlocked
+  // callback window.
+  bool shutdown_notification_sent_ = false;
 
   // No new background jobs can be queued if true. This is used to prevent new
   // background jobs from being queued after WaitForCompact() completes waiting
@@ -1538,6 +1545,7 @@ class DBImpl : public DB {
                                    const Status& st,
                                    const CompactionJobStats& job_stats,
                                    int job_id);
+  void NotifyOnDBShutdownBegin();
   void NotifyOnMemTableSealed(ColumnFamilyData* cfd,
                               const MemTableInfo& mem_table_info);
 

db/listener_test.cc

@@ -31,6 +31,7 @@
 #include "util/mutexlock.h"
 #include "util/rate_limiter_impl.h"
 #include "util/string_util.h"
+#include "utilities/fault_injection_fs.h"
 #include "utilities/merge_operators.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -316,6 +317,79 @@ TEST_F(EventListenerTest, OnCompactionPreCommitOrdering) {
   ROCKSDB_NAMESPACE::SyncPoint::GetInstance()->ClearAllCallBacks();
 }
 
+class TestDBShutdownBeginListener : public EventListener {
+ public:
+  void OnDBShutdownBegin(DB* db) override {
+    std::lock_guard<std::mutex> lock(mutex_);
+    ++shutdown_count_;
+    last_db_ = db;
+  }
+
+  int ShutdownCount() const {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return shutdown_count_;
+  }
+
+  DB* LastDB() const {
+    std::lock_guard<std::mutex> lock(mutex_);
+    return last_db_;
+  }
+
+ private:
+  mutable std::mutex mutex_;
+  int shutdown_count_ = 0;
+  DB* last_db_ = nullptr;
+};
+
+TEST_F(EventListenerTest, OnDBShutdownBeginOnceForCancelAndClose) {
+  Options options = CurrentOptions();
+  options.create_if_missing = true;
+  auto listener = std::make_shared<TestDBShutdownBeginListener>();
+  options.listeners.emplace_back(listener);
+  DestroyAndReopen(options);
+
+  DB* db = db_.get();
+  dbfull()->CancelAllBackgroundWork(false);
+  EXPECT_EQ(1, listener->ShutdownCount());
+  EXPECT_EQ(db, listener->LastDB());
+
+  Close();
+  EXPECT_EQ(1, listener->ShutdownCount());
+}
+
+TEST_F(EventListenerTest, OnDBShutdownBeginOnFailedOpen) {
+  Close();
+
+  std::shared_ptr<FaultInjectionTestFS> fs(
+      new FaultInjectionTestFS(env_->GetFileSystem()));
+  std::unique_ptr<Env> env(NewCompositeEnv(fs));
+
+  Options options = CurrentOptions();
+  options.env = env.get();
+  options.create_if_missing = true;
+  auto listener = std::make_shared<TestDBShutdownBeginListener>();
+  options.listeners.emplace_back(listener);
+
+  SyncPoint::GetInstance()->SetCallBack(
+      "PersistRocksDBOptions:create",
+      [&](void* /*arg*/) { fs->SetFilesystemActive(false); });
+  SyncPoint::GetInstance()->SetCallBack(
+      "PersistRocksDBOptions:written",
+      [&](void* /*arg*/) { fs->SetFilesystemActive(true); });
+  SyncPoint::GetInstance()->EnableProcessing();
+
+  std::unique_ptr<DB> db;
+  Status s = DB::Open(options, dbname_, &db);
+  ASSERT_TRUE(s.IsIOError()) << s.ToString();
+  EXPECT_EQ(1, listener->ShutdownCount());
+  EXPECT_NE(nullptr, listener->LastDB());
+  EXPECT_EQ(nullptr, db.get());
+
+  fs->SetFilesystemActive(true);
+  SyncPoint::GetInstance()->DisableProcessing();
+  SyncPoint::GetInstance()->ClearAllCallBacks();
+}
+
 // This simple Listener can only handle one flush at a time.
 class TestFlushListener : public EventListener {
  public:

db_stress_tool/db_stress_listener.h

@@ -63,12 +63,8 @@ class DbStressListener : public EventListener {
 
   ~DbStressListener() override { assert(num_pending_file_creations_ == 0); }
 
-  // Signal that the DB is about to shut down. Must be called before DB::Close()
-  // or CancelAllBackgroundWork(). Why? Listener notifications, especially for
-  // compaction, have different (mostly relaxed) contracts during shutdown, and
-  // not all callbacks take a DB object. Ideally, public API improvements would
-  // make this unnecessary in the future.
-  void NotifyShuttingDown() { shutting_down_.Store(true); }
+  void OnDBShutdownBegin(DB* /*db*/) override { shutting_down_.Store(true); }
+
   void OnFlushCompleted(DB* /*db*/, const FlushJobInfo& info) override {
     assert(IsValidColumnFamilyName(info.cf_name));
     VerifyFilePath(info.file_path);
@@ -486,7 +482,8 @@ class DbStressListener : public EventListener {
   // callback. Protected by compacting_files_mu_.
   std::mutex compacting_files_mu_;
   std::unordered_set<uint64_t> compacting_files_;
-  // Set before DB close to suppress false positives from stale tracking.
+  // Set when DBImpl begins shutdown to suppress false positives from stale
+  // tracking when compaction callbacks are skipped during shutdown.
   Atomic<bool> shutting_down_{false};
   // Jobs that have passed OnCompactionPreCommit but not yet
   // OnCompactionCompleted. Maps job_id -> input file numbers.

db_stress_tool/db_stress_test_base.cc

@@ -63,7 +63,74 @@ std::shared_ptr<const FilterPolicy> CreateFilterPolicy() {
   return std::shared_ptr<const FilterPolicy>(new_policy);
 }
 
-static bool NeedsFaultInjection() {
+struct OpenFaultInjectionConfig {
+  bool sync_fault = false;
+  int metadata_read_one_in = 0;
+  int metadata_write_one_in = 0;
+  int read_one_in = 0;
+  int write_one_in = 0;
+
+  bool Enabled() const {
+    return sync_fault || metadata_read_one_in > 0 ||
+           metadata_write_one_in > 0 || read_one_in > 0 || write_one_in > 0;
+  }
+
+  void Disable() {
+    sync_fault = false;
+    metadata_read_one_in = 0;
+    metadata_write_one_in = 0;
+    read_one_in = 0;
+    write_one_in = 0;
+  }
+};
+
+void EnableThreadLocalOpenFault(
+    const std::shared_ptr<FaultInjectionTestFS>& fault_injection_fs,
+    FaultInjectionIOType io_type, int one_in) {
+  assert(fault_injection_fs);
+  fault_injection_fs->SetThreadLocalErrorContext(
+      io_type, static_cast<uint32_t>(FLAGS_seed), one_in, false /* retryable */,
+      false /* has_data_loss */);
+  fault_injection_fs->EnableThreadLocalErrorInjection(io_type);
+}
+
+bool MaybeEnableOpenFaultInjection(
+    const std::shared_ptr<FaultInjectionTestFS>& fault_injection_fs,
+    const std::string& db_path,
+    const OpenFaultInjectionConfig& open_fault_injection) {
+  // Returns whether open-fault handling is active for this attempt. The caller
+  // uses this to run post-open cleanup/retry handling even if CURRENT did not
+  // exist and no fault was actually enabled.
+  if (!open_fault_injection.Enabled()) {
+    return false;
+  }
+  assert(fault_injection_fs);
+  if (!fault_injection_fs
+           ->FileExists(db_path + "/CURRENT", IOOptions(), nullptr)
+           .ok()) {
+    return true;
+  }
+
+  if (open_fault_injection.sync_fault ||
+      open_fault_injection.write_one_in > 0) {
+    fault_injection_fs->SetFilesystemDirectWritable(false);
+    fault_injection_fs->SetInjectUnsyncedDataLoss(
+        open_fault_injection.sync_fault);
+  }
+  EnableThreadLocalOpenFault(fault_injection_fs,
+                             FaultInjectionIOType::kMetadataRead,
+                             open_fault_injection.metadata_read_one_in);
+  EnableThreadLocalOpenFault(fault_injection_fs,
+                             FaultInjectionIOType::kMetadataWrite,
+                             open_fault_injection.metadata_write_one_in);
+  EnableThreadLocalOpenFault(fault_injection_fs, FaultInjectionIOType::kRead,
+                             open_fault_injection.read_one_in);
+  EnableThreadLocalOpenFault(fault_injection_fs, FaultInjectionIOType::kWrite,
+                             open_fault_injection.write_one_in);
+  return true;
+}
+
+bool NeedsFaultInjection() {
   return FLAGS_open_metadata_read_fault_one_in ||
          FLAGS_open_metadata_write_fault_one_in ||
          FLAGS_open_read_fault_one_in || FLAGS_open_write_fault_one_in ||
@@ -157,9 +224,6 @@ StressTest::StressTest()
 }
 
 void StressTest::CleanUp() {
-  // Notify listener before DB close so it can tolerate stale tracking
-  // from skipped notifications during shutdown.
-  NotifyListenerShuttingDown();
   CleanUpColumnFamilies();
   if (db_) {
     db_->Close();
@@ -170,15 +234,6 @@ void StressTest::CleanUp() {
   secondary_db_.reset();
 }
 
-void StressTest::NotifyListenerShuttingDown() {
-  for (auto& listener : options_.listeners) {
-    if (strcmp(listener->Name(), DbStressListener::kClassName()) == 0) {
-      static_cast_with_check<DbStressListener>(listener.get())
-          ->NotifyShuttingDown();
-    }
-  }
-}
-
 void StressTest::InitializeListenersForOpen(
     SharedState* shared,
     const std::vector<ColumnFamilyDescriptor>& cf_descriptors) {
@@ -4074,8 +4129,6 @@ void StressTest::Open(SharedState* shared, bool reopen) {
       column_family_names_.push_back(name);
     }
 
-    InitializeListenersForOpen(shared, cf_descriptors);
-
     // If this is for DB reopen,  error injection may have been enabled.
     // Disable it here in case there is no open fault injection.
     if (db_fault_injection_fs_) {
@@ -4083,54 +4136,20 @@ void StressTest::Open(SharedState* shared, bool reopen) {
     }
     // TODO; test transaction DB Open with fault injection
     if (!FLAGS_use_txn) {
-      bool inject_sync_fault = FLAGS_sync_fault_injection;
-      bool inject_open_meta_read_error =
-          FLAGS_open_metadata_read_fault_one_in > 0;
-      bool inject_open_meta_write_error =
-          FLAGS_open_metadata_write_fault_one_in > 0;
-      bool inject_open_read_error = FLAGS_open_read_fault_one_in > 0;
-      bool inject_open_write_error = FLAGS_open_write_fault_one_in > 0;
-      if ((inject_sync_fault || inject_open_meta_read_error ||
-           inject_open_meta_write_error || inject_open_read_error ||
-           inject_open_write_error) &&
-          db_fault_injection_fs_
-              ->FileExists(db_path + "/CURRENT", IOOptions(), nullptr)
-              .ok()) {
-        if (inject_sync_fault || inject_open_write_error) {
-          db_fault_injection_fs_->SetFilesystemDirectWritable(false);
-          db_fault_injection_fs_->SetInjectUnsyncedDataLoss(inject_sync_fault);
-        }
-        db_fault_injection_fs_->SetThreadLocalErrorContext(
-            FaultInjectionIOType::kMetadataRead,
-            static_cast<uint32_t>(FLAGS_seed),
-            FLAGS_open_metadata_read_fault_one_in, false /* retryable */,
-            false /* has_data_loss */);
-        db_fault_injection_fs_->EnableThreadLocalErrorInjection(
-            FaultInjectionIOType::kMetadataRead);
-
-        db_fault_injection_fs_->SetThreadLocalErrorContext(
-            FaultInjectionIOType::kMetadataWrite,
-            static_cast<uint32_t>(FLAGS_seed),
-            FLAGS_open_metadata_write_fault_one_in, false /* retryable */,
-            false /* has_data_loss */);
-        db_fault_injection_fs_->EnableThreadLocalErrorInjection(
-            FaultInjectionIOType::kMetadataWrite);
-
-        db_fault_injection_fs_->SetThreadLocalErrorContext(
-            FaultInjectionIOType::kRead, static_cast<uint32_t>(FLAGS_seed),
-            FLAGS_open_read_fault_one_in, false /* retryable */,
-            false /* has_data_loss */);
-        db_fault_injection_fs_->EnableThreadLocalErrorInjection(
-            FaultInjectionIOType::kRead);
-
-        db_fault_injection_fs_->SetThreadLocalErrorContext(
-            FaultInjectionIOType::kWrite, static_cast<uint32_t>(FLAGS_seed),
-            FLAGS_open_write_fault_one_in, false /* retryable */,
-            false /* has_data_loss */);
-        db_fault_injection_fs_->EnableThreadLocalErrorInjection(
-            FaultInjectionIOType::kWrite);
-      }
+      OpenFaultInjectionConfig open_fault_injection;
+      open_fault_injection.sync_fault = FLAGS_sync_fault_injection;
+      open_fault_injection.metadata_read_one_in =
+          FLAGS_open_metadata_read_fault_one_in;
+      open_fault_injection.metadata_write_one_in =
+          FLAGS_open_metadata_write_fault_one_in;
+      open_fault_injection.read_one_in = FLAGS_open_read_fault_one_in;
+      open_fault_injection.write_one_in = FLAGS_open_write_fault_one_in;
       while (true) {
+        InitializeListenersForOpen(shared, cf_descriptors);
+
+        const bool open_fault_injection_enabled = MaybeEnableOpenFaultInjection(
+            db_fault_injection_fs_, db_path, open_fault_injection);
+
         // StackableDB-based BlobDB
         if (FLAGS_use_blob_db) {
           blob_db::BlobDBOptions blob_db_options;
@@ -4159,9 +4178,7 @@ void StressTest::Open(SharedState* shared, bool reopen) {
           }
         }
 
-        if (inject_sync_fault || inject_open_meta_read_error ||
-            inject_open_meta_write_error || inject_open_read_error ||
-            inject_open_write_error) {
+        if (open_fault_injection_enabled) {
           db_fault_injection_fs_->DisableAllThreadLocalErrorInjection();
 
           if (s.ok()) {
@@ -4180,11 +4197,7 @@ void StressTest::Open(SharedState* shared, bool reopen) {
             // After failure to opening a DB due to IO error or unsynced data
             // loss, retry should successfully open the DB with correct data if
             // no IO error shows up.
-            inject_sync_fault = false;
-            inject_open_meta_read_error = false;
-            inject_open_meta_write_error = false;
-            inject_open_read_error = false;
-            inject_open_write_error = false;
+            open_fault_injection.Disable();
 
             // TODO: Unsynced data loss during DB reopen is not supported yet in
             //  stress test. Will need to recreate expected state if we decide
@@ -4201,16 +4214,14 @@ void StressTest::Open(SharedState* shared, bool reopen) {
                 db_fault_injection_fs_->DropRandomUnsyncedFileData(&rand);
               }
             }
-            // DB::Open() can fail after scheduling background compaction. Use
-            // fresh listeners on retry so per-DBImpl listener state is not
-            // carried across failed open attempts.
-            InitializeListenersForOpen(shared, cf_descriptors);
             continue;
           }
         }
         break;
       }
     } else {
+      InitializeListenersForOpen(shared, cf_descriptors);
+
       if (FLAGS_use_optimistic_txn) {
         OptimisticTransactionDBOptions optimistic_txn_db_options;
         optimistic_txn_db_options.validate_policy =
@@ -4589,10 +4600,6 @@ void StressTest::VerifyManifestNotRewritten() {
 }
 
 void StressTest::Reopen(ThreadState* thread) {
-  // Notify listener before DB close so it can tolerate stale tracking
-  // from skipped notifications during shutdown.
-  NotifyListenerShuttingDown();
-
   // BG jobs in WritePrepared must be canceled first because i) they can
   // access the db via a callbac ii) they hold on to a snapshot and the
   // upcoming

db_stress_tool/db_stress_test_base.h

@@ -86,7 +86,6 @@ class StressTest {
   void CleanUp();
 
  private:
-  void NotifyListenerShuttingDown();
   void InitializeListenersForOpen(
       SharedState* shared,
       const std::vector<ColumnFamilyDescriptor>& cf_descriptors);

db_stress_tool/multi_ops_txns_stress.h

@@ -9,6 +9,7 @@
 
 #ifdef GFLAGS
 #include "db_stress_tool/db_stress_common.h"
+#include "util/atomic.h"
 
 namespace ROCKSDB_NAMESPACE {
 
@@ -423,11 +424,16 @@ class MultiOpsTxnsStressListener : public EventListener {
 
   ~MultiOpsTxnsStressListener() override {}
 
+  void OnDBShutdownBegin(DB* /*db*/) override { shutting_down_.Store(true); }
+
   void OnFlushCompleted(DB* db, const FlushJobInfo& info) override {
     assert(db);
 #ifdef NDEBUG
     (void)db;
 #endif
+    if (shutting_down_.Load()) {
+      return;
+    }
     assert(info.cf_id == 0);
     const ReadOptions read_options(Env::IOActivity::kFlush);
     stress_test_->VerifyPkSkFast(read_options, info.job_id);
@@ -438,6 +444,9 @@ class MultiOpsTxnsStressListener : public EventListener {
 #ifdef NDEBUG
     (void)db;
 #endif
+    if (shutting_down_.Load()) {
+      return;
+    }
     assert(info.cf_id == 0);
     const ReadOptions read_options(Env::IOActivity::kCompaction);
     stress_test_->VerifyPkSkFast(read_options, info.job_id);
@@ -445,6 +454,7 @@ class MultiOpsTxnsStressListener : public EventListener {
 
  private:
   MultiOpsTxnsStressTest* const stress_test_ = nullptr;
+  Atomic<bool> shutting_down_{false};
 };
 
 }  // namespace ROCKSDB_NAMESPACE

include/rocksdb/listener.h

@@ -964,6 +964,17 @@ class EventListener : public Customizable {
   virtual void OnBackgroundJobPressureChanged(
       DB* /*db*/, const BackgroundJobPressure& /*pressure*/) {}
 
+  // A callback function for RocksDB which will be called once when a DB begins
+  // shutting down, before background work cancellation publishes the DB's
+  // shutdown state. This callback can also fire during cleanup of a failed
+  // DB::Open() attempt, in which case the DB pointer refers to the DB instance
+  // that failed to open and was never returned to the application.
+  //
+  // Background work may still be running when this callback fires, and other
+  // listener callbacks may still run concurrently or afterward. Implementations
+  // should not call blocking DB APIs or run for an extended period of time.
+  virtual void OnDBShutdownBegin(DB* /*db*/) {}
+
   ~EventListener() override {}
 };
 

unreleased_history/public_api_changes/on_db_shutdown_begin.md

@@ -0,0 +1 @@
+Added `EventListener::OnDBShutdownBegin`, a callback that fires once when a DB begins shutdown, including when RocksDB cleans up a failed `DB::Open()` attempt.