diff --git a/db/db_basic_test.cc b/db/db_basic_test.cc index edb10693aff..db04a355c84 100644 --- a/db/db_basic_test.cc +++ b/db/db_basic_test.cc @@ -1336,8 +1336,15 @@ TEST_F(DBBasicTest, DBClose) { ASSERT_OK(s); ASSERT_TRUE(db != nullptr); - s = db->Close(); + bool prepare_close_fn_called = false; + CloseOptions close_options; + close_options.prepare_close_fn = [&prepare_close_fn_called]() { + prepare_close_fn_called = true; + }; + + s = db->Close(close_options); ASSERT_EQ(s, Status::OK()); + ASSERT_TRUE(prepare_close_fn_called); delete db; ASSERT_EQ(env->GetCloseCount(), 3); options.info_log.reset(); diff --git a/db/db_compaction_test.cc b/db/db_compaction_test.cc index c71cf2e0c13..2351f59a384 100644 --- a/db/db_compaction_test.cc +++ b/db/db_compaction_test.cc @@ -3622,7 +3622,18 @@ TEST_P(DBCompactionWaitForCompactTest, ASSERT_EQ(0, flush_finished); ASSERT_EQ("2", FilesPerLevel()); + bool prepare_close_fn_called = false; + if (close_db_) { + CloseOptions close_options; + close_options.prepare_close_fn = [&prepare_close_fn_called]() { + prepare_close_fn_called = true; + }; + wait_for_compact_options_.close_options = std::move(close_options); + } ASSERT_OK(dbfull()->WaitForCompact(wait_for_compact_options_)); + if (close_db_) { + ASSERT_TRUE(prepare_close_fn_called); + } int expected_flush_count = flush_ || close_db_; ASSERT_EQ(expected_flush_count, flush_finished); diff --git a/db/db_impl/db_impl.cc b/db/db_impl/db_impl.cc index 3bc2be1f047..22112f8eb02 100644 --- a/db/db_impl/db_impl.cc +++ b/db/db_impl/db_impl.cc @@ -5318,7 +5318,7 @@ Status DB::DestroyColumnFamilyHandle(ColumnFamilyHandle* column_family) { DB::~DB() = default; -Status DBImpl::Close() { +Status DBImpl::Close(const CloseOptions& close_options) { InstrumentedMutexLock closing_lock_guard(&closing_mutex_); if (closed_) { return closing_status_; @@ -5330,7 +5330,9 @@ Status DBImpl::Close() { return s; } } - + if (close_options.prepare_close_fn) { + close_options.prepare_close_fn(); + } closing_status_ = CloseImpl(); closed_ = true; return closing_status_; diff --git a/db/db_impl/db_impl.h b/db/db_impl/db_impl.h index 6621b20f64a..d9fa8c9d7f0 100644 --- a/db/db_impl/db_impl.h +++ b/db/db_impl/db_impl.h @@ -502,7 +502,8 @@ class DBImpl : public DB { ColumnFamilyHandle* PersistentStatsColumnFamily() const; - Status Close() override; + using DB::Close; + Status Close(const CloseOptions& close_options) override; Status DisableFileDeletions() override; diff --git a/db/db_impl/db_impl_compaction_flush.cc b/db/db_impl/db_impl_compaction_flush.cc index 0099af88cec..a0fb57d1b62 100644 --- a/db/db_impl/db_impl_compaction_flush.cc +++ b/db/db_impl/db_impl_compaction_flush.cc @@ -4429,7 +4429,7 @@ Status DBImpl::WaitForCompact( } else if (wait_for_compact_options.close_db) { reject_new_background_jobs_ = true; mutex_.Unlock(); - Status s = Close(); + Status s = Close(wait_for_compact_options.close_options); mutex_.Lock(); if (!s.ok()) { reject_new_background_jobs_ = false; diff --git a/db/db_impl/db_impl_follower.cc b/db/db_impl/db_impl_follower.cc index a104a83d6bf..f23952322a1 100644 --- a/db/db_impl/db_impl_follower.cc +++ b/db/db_impl/db_impl_follower.cc @@ -224,7 +224,7 @@ void DBImplFollower::PeriodicRefresh() { } } -Status DBImplFollower::Close() { +Status DBImplFollower::Close(const CloseOptions& close_options) { if (catch_up_thread_) { stop_requested_.store(true); { @@ -237,7 +237,7 @@ Status DBImplFollower::Close() { ReleaseFileNumberFromPendingOutputs(pending_outputs_inserted_elem_); - return DBImpl::Close(); + return DBImpl::Close(close_options); } Status DB::OpenAsFollower(const Options& options, const std::string& dbname, diff --git a/db/db_impl/db_impl_follower.h b/db/db_impl/db_impl_follower.h index 374c60d5c64..f5282066ea5 100644 --- a/db/db_impl/db_impl_follower.h +++ b/db/db_impl/db_impl_follower.h @@ -21,7 +21,8 @@ class DBImplFollower : public DBImplSecondary { const std::string& dbname, std::string src_path); ~DBImplFollower(); - Status Close() override; + using DBImpl::Close; + Status Close(const CloseOptions& close_options) override; protected: bool OwnTablesAndLogs() const override { diff --git a/db/db_test.cc b/db/db_test.cc index 0975407c156..36d1db53e6a 100644 --- a/db/db_test.cc +++ b/db/db_test.cc @@ -3061,7 +3061,7 @@ class ModelDB : public DB { } using DB::Close; - Status Close() override { return Status::OK(); } + Status Close(const CloseOptions&) override { return Status::OK(); } using DB::Delete; Status Delete(const WriteOptions& o, ColumnFamilyHandle* cf, const Slice& key) override { diff --git a/include/rocksdb/db.h b/include/rocksdb/db.h index 9beb4e74f6d..b3c1e2b30fb 100644 --- a/include/rocksdb/db.h +++ b/include/rocksdb/db.h @@ -383,7 +383,10 @@ class DB { // WaitForCompact() with WaitForCompactOptions.close_db=true will be a good // choice for users who want to wait for background work before closing // (rather than aborting and potentially redoing some work on re-open) - virtual Status Close() { return Status::NotSupported(); } + Status Close() { return Close(CloseOptions()); } + virtual Status Close(const CloseOptions& /*close_options*/) { + return Status::NotSupported(); + } // ListColumnFamilies will open the DB specified by argument name // and return the list of all column families in that DB diff --git a/include/rocksdb/options.h b/include/rocksdb/options.h index b27f53b4a84..34ea16a2c5f 100644 --- a/include/rocksdb/options.h +++ b/include/rocksdb/options.h @@ -2391,6 +2391,10 @@ struct LiveFilesStorageInfoOptions { uint64_t wal_size_for_flush = 0; }; +struct CloseOptions { + std::function prepare_close_fn; +}; + struct WaitForCompactOptions { // A boolean to abort waiting in case of a pause (PauseBackgroundWork() // called) If true, Status::Aborted will be returned immediately. If false, @@ -2412,6 +2416,8 @@ struct WaitForCompactOptions { // returned Aborted status due to unreleased snapshots in the system. See // comments in DB::Close() for details. bool close_db = false; + // Options to be used for closing db. If close_db is false, this is ignored + CloseOptions close_options; // Timeout in microseconds for waiting for compaction to complete. // Status::TimedOut will be returned if timeout expires. diff --git a/include/rocksdb/utilities/stackable_db.h b/include/rocksdb/utilities/stackable_db.h index 414a28dfbfe..4af501a5358 100644 --- a/include/rocksdb/utilities/stackable_db.h +++ b/include/rocksdb/utilities/stackable_db.h @@ -36,7 +36,10 @@ class StackableDB : public DB { db_ = nullptr; } - Status Close() override { return db_->Close(); } + using DB::Close; + Status Close(const CloseOptions& close_options) override { + return db_->Close(close_options); + } virtual DB* GetBaseDB() { return db_; } diff --git a/unreleased_history/public_api_changes/close_options.md b/unreleased_history/public_api_changes/close_options.md new file mode 100644 index 00000000000..00affcfd8df --- /dev/null +++ b/unreleased_history/public_api_changes/close_options.md @@ -0,0 +1 @@ +Introduce `CloseOptions` for `Close()` API. `close_options.prepare_close_fn` allows users to inject a necessary external functions to be called before closing the DB. This option is also integrated with `WaitForCompact()` API when `WaitForCompactOptions::close_db=true`. `WaitForCompactOptions::close_options` can be specified, but is ignored if close_db is false. diff --git a/utilities/blob_db/blob_db.h b/utilities/blob_db/blob_db.h index 503d476fa51..6fb824a0b16 100644 --- a/utilities/blob_db/blob_db.h +++ b/utilities/blob_db/blob_db.h @@ -199,7 +199,7 @@ class BlobDB : public StackableDB { } using ROCKSDB_NAMESPACE::StackableDB::Close; - Status Close() override = 0; + Status Close(const CloseOptions& close_options) override = 0; // Opening blob db. static Status Open(const Options& options, const BlobDBOptions& bdb_options, diff --git a/utilities/blob_db/blob_db_impl.cc b/utilities/blob_db/blob_db_impl.cc index c88c3e8a709..e31ea0b1937 100644 --- a/utilities/blob_db/blob_db_impl.cc +++ b/utilities/blob_db/blob_db_impl.cc @@ -105,16 +105,16 @@ BlobDBImpl::~BlobDBImpl() { assert(s.ok()); } -Status BlobDBImpl::Close() { +Status BlobDBImpl::Close(const CloseOptions& close_options) { ThreadStatus::OperationType cur_op_type = ThreadStatusUtil::GetThreadOperation(); ThreadStatusUtil::SetThreadOperation(ThreadStatus::OperationType::OP_UNKNOWN); - Status s = CloseImpl(); + Status s = CloseImpl(close_options); ThreadStatusUtil::SetThreadOperation(cur_op_type); return s; } -Status BlobDBImpl::CloseImpl() { +Status BlobDBImpl::CloseImpl(const CloseOptions& close_options) { if (closed_) { return Status::OK(); } @@ -122,7 +122,7 @@ Status BlobDBImpl::CloseImpl() { // Close base DB before BlobDBImpl destructs to stop event listener and // compaction filter call. - Status s = db_->Close(); + Status s = db_->Close(close_options); // delete db_ anyway even if close failed. delete db_; // Reset pointers to avoid StackableDB delete the pointer again. diff --git a/utilities/blob_db/blob_db_impl.h b/utilities/blob_db/blob_db_impl.h index 75776e6a8a7..e706cd6a754 100644 --- a/utilities/blob_db/blob_db_impl.h +++ b/utilities/blob_db/blob_db_impl.h @@ -130,7 +130,8 @@ class BlobDBImpl : public BlobDB { using BlobDB::Write; Status Write(const WriteOptions& opts, WriteBatch* updates) override; - Status Close() override; + using BlobDB::Close; + Status Close(const CloseOptions& close_options) override; using BlobDB::PutWithTTL; Status PutWithTTL(const WriteOptions& options, const Slice& key, @@ -406,7 +407,7 @@ class BlobDBImpl : public BlobDB { uint64_t blob_size, bool force_evict = false); - Status CloseImpl(); + Status CloseImpl(const CloseOptions& close_options); // name of the database directory std::string dbname_; diff --git a/utilities/ttl/db_ttl_impl.cc b/utilities/ttl/db_ttl_impl.cc index 55354c6cbce..acfaaa3b351 100644 --- a/utilities/ttl/db_ttl_impl.cc +++ b/utilities/ttl/db_ttl_impl.cc @@ -313,13 +313,13 @@ DBWithTTLImpl::~DBWithTTLImpl() { } } -Status DBWithTTLImpl::Close() { +Status DBWithTTLImpl::Close(const CloseOptions& close_options) { Status ret = Status::OK(); if (!closed_) { Options default_options = GetOptions(); // Need to stop background compaction before getting rid of the filter CancelAllBackgroundWork(db_, /* wait = */ true); - ret = db_->Close(); + ret = db_->Close(close_options); delete default_options.compaction_filter; closed_ = true; } diff --git a/utilities/ttl/db_ttl_impl.h b/utilities/ttl/db_ttl_impl.h index 731cd3955fe..c6bc7fcfc34 100644 --- a/utilities/ttl/db_ttl_impl.h +++ b/utilities/ttl/db_ttl_impl.h @@ -36,7 +36,8 @@ class DBWithTTLImpl : public DBWithTTL { virtual ~DBWithTTLImpl(); - Status Close() override; + using StackableDB::Close; + Status Close(const CloseOptions& close_options) override; Status CreateColumnFamilyWithTtl(const ColumnFamilyOptions& options, const std::string& column_family_name,