From 757860fe3fc57eea0fe61fe6ee3abfc6da2dc272 Mon Sep 17 00:00:00 2001 From: lucasliang Date: Thu, 30 Jul 2026 19:30:29 +0800 Subject: [PATCH] Fix sequence publication race in allow-write external SST ingestion Signed-off-by: lucasliang --- db/db_impl/db_impl.cc | 89 ++++++++++++++++++++++++++++++++---- db/db_impl/db_impl_write.cc | 1 + db/external_sst_file_test.cc | 75 +++++++++++++++++++++++++++++- 3 files changed, 153 insertions(+), 12 deletions(-) diff --git a/db/db_impl/db_impl.cc b/db/db_impl/db_impl.cc index d291021c82ee..2579e3aac9c9 100644 --- a/db/db_impl/db_impl.cc +++ b/db/db_impl/db_impl.cc @@ -5855,8 +5855,8 @@ Status DBImpl::IngestExternalFiles( WriteThread::Writer w; WriteThread::Writer nonmem_w; - if (!allow_write) { - // Stop writes to the DB by entering both write threads. + bool write_threads_entered = false; + const auto enter_write_threads = [&]() { write_thread_.EnterUnbatched(&w, &mutex_); if (two_write_queues_) { nonmem_write_thread_.EnterUnbatched(&nonmem_w, &mutex_); @@ -5868,6 +5868,11 @@ Status DBImpl::IngestExternalFiles( // memtable flush. // So wait here to ensure there is no pending write to memtable. WaitForPendingWrites(); + write_threads_entered = true; + }; + if (!allow_write) { + // Stop writes to the DB by entering both write threads. + enter_write_threads(); } TEST_SYNC_POINT_CALLBACK("DBImpl::IngestExternalFile:AfterAllowWriteCheck", @@ -5905,10 +5910,10 @@ Status DBImpl::IngestExternalFiles( flush_opts.check_if_compaction_disabled = true; if (immutable_db_options_.atomic_flush) { mutex_.Unlock(); - status = AtomicFlushMemTables(flush_opts, - FlushReason::kExternalFileIngestion, - {} /* provided_candidate_cfds */, - !allow_write /* entered_write_thread */); + status = AtomicFlushMemTables( + flush_opts, FlushReason::kExternalFileIngestion, + {} /* provided_candidate_cfds */, + write_threads_entered /* entered_write_thread */); mutex_.Lock(); } else { for (size_t i = 0; i != num_cfs; ++i) { @@ -5917,9 +5922,9 @@ Status DBImpl::IngestExternalFiles( auto* cfd = static_cast(args[i].column_family) ->cfd(); - status = FlushMemTable(cfd, flush_opts, - FlushReason::kExternalFileIngestion, - !allow_write /* entered_write_thread */); + status = FlushMemTable( + cfd, flush_opts, FlushReason::kExternalFileIngestion, + write_threads_entered /* entered_write_thread */); mutex_.Lock(); if (!status.ok()) { break; @@ -5928,6 +5933,68 @@ Status DBImpl::IngestExternalFiles( } } } + + if (status.ok() && allow_write) { + // Run() assigns global sequence numbers and the sequence numbers are + // published after LogAndApply below. It requires there to be no active + // writers during both operations. Keep writes allowed while preparing + // and flushing the ingestion, but serialize this final phase with normal + // writes. + enter_write_threads(); + + // Writes were allowed during the earlier NeedsFlush check and flush, so + // the result can have become stale. Recheck after writer exclusion. If a + // new overlapping memtable appeared, flush it before calling Run(). + bool final_need_flush = false; + std::vector final_need_flushes(num_cfs, false); + for (size_t i = 0; i != num_cfs; ++i) { + auto* cfd = + static_cast(args[i].column_family)->cfd(); + if (cfd->IsDropped()) { + status = Status::InvalidArgument( + "cannot ingest an external file into a dropped CF"); + break; + } + bool tmp = false; + status = ingestion_jobs[i].NeedsFlush(&tmp, cfd->GetSuperVersion()); + final_need_flushes[i] = tmp; + final_need_flush = final_need_flush || tmp; + if (!status.ok()) { + break; + } + } + + if (status.ok() && final_need_flush) { + FlushOptions flush_opts; + flush_opts.allow_write_stall = true; + flush_opts.check_if_compaction_disabled = true; + if (immutable_db_options_.atomic_flush) { + mutex_.Unlock(); + status = AtomicFlushMemTables(flush_opts, + FlushReason::kExternalFileIngestion, + {} /* provided_candidate_cfds */, + true /* entered_write_thread */); + mutex_.Lock(); + } else { + for (size_t i = 0; i != num_cfs; ++i) { + if (final_need_flushes[i]) { + mutex_.Unlock(); + auto* cfd = + static_cast(args[i].column_family) + ->cfd(); + status = FlushMemTable(cfd, flush_opts, + FlushReason::kExternalFileIngestion, + true /* entered_write_thread */); + mutex_.Lock(); + if (!status.ok()) { + break; + } + } + } + } + } + } + // Run ingestion jobs. if (status.ok()) { for (size_t i = 0; i != num_cfs; ++i) { @@ -5986,6 +6053,8 @@ Status DBImpl::IngestExternalFiles( } if (consumed_seqno_count > 0) { const SequenceNumber last_seqno = versions_->LastSequence(); + TEST_SYNC_POINT( + "DBImpl::IngestExternalFiles:AfterReadLastSequenceBeforePublish"); versions_->SetLastAllocatedSequence(last_seqno + consumed_seqno_count); versions_->SetLastPublishedSequence(last_seqno + consumed_seqno_count); versions_->SetLastSequence(last_seqno + consumed_seqno_count); @@ -6024,7 +6093,7 @@ Status DBImpl::IngestExternalFiles( error_handler_.SetBGError(io_s, BackgroundErrorReason::kManifestWrite); } - if (!allow_write) { + if (write_threads_entered) { if (two_write_queues_) { nonmem_write_thread_.ExitUnbatched(&nonmem_w); } diff --git a/db/db_impl/db_impl_write.cc b/db/db_impl/db_impl_write.cc index c74ec8dba933..790e84e10cf7 100644 --- a/db/db_impl/db_impl_write.cc +++ b/db/db_impl/db_impl_write.cc @@ -664,6 +664,7 @@ Status DBImpl::WriteImpl(const WriteOptions& write_options, // and protects against concurrent loggers and concurrent writes // into memtables + TEST_SYNC_POINT("DBImpl::WriteImpl:BeforeLeaderEnters:0"); TEST_SYNC_POINT("DBImpl::WriteImpl:BeforeLeaderEnters"); last_batch_group_size_ = write_thread_.EnterAsBatchGroupLeader(&w, &write_group); diff --git a/db/external_sst_file_test.cc b/db/external_sst_file_test.cc index f86af51e2a41..7edc6db1b2d1 100644 --- a/db/external_sst_file_test.cc +++ b/db/external_sst_file_test.cc @@ -14,6 +14,7 @@ #include "port/stack_trace.h" #include "rocksdb/sst_file_reader.h" #include "rocksdb/sst_file_writer.h" +#include "rocksdb/write_batch.h" #include "test_util/testutil.h" #include "util/random.h" #include "util/thread_guard.h" @@ -2394,14 +2395,16 @@ TEST_P(ExternalSSTFileTest, IngestBehind) { TEST_P(ExternalSSTFileTest, WriteDuringIngest) { SyncPoint::GetInstance()->DisableProcessing(); - // Set callback to simulate concurrent write during ingestion + // Set callback to simulate a write before the initial NeedsFlush check. SyncPoint::GetInstance()->SetCallBack( "DBImpl::IngestExternalFile:AfterAllowWriteCheck", [&](void*) { - // Write a non-overlapping key + // Write the same key as the external file. ASSERT_OK(Put("foo", "v1")); }); Options options = CurrentOptions(); + options.two_write_queues = std::get<0>(GetParam()); + options.atomic_flush = std::get<1>(GetParam()); DestroyAndReopen(options); SyncPoint::GetInstance()->EnableProcessing(); @@ -2416,6 +2419,74 @@ TEST_P(ExternalSSTFileTest, WriteDuringIngest) { SyncPoint::GetInstance()->ClearAllCallBacks(); } +TEST_P(ExternalSSTFileTest, AllowWriteIngestSequencePublication) { + constexpr auto kReleaseWriter = + "ExternalSSTFileTest::AllowWriteIngestSequencePublication:" + "ReleaseWriter"; + constexpr auto kWriterPrepared = + "ExternalSSTFileTest::AllowWriteIngestSequencePublication:" + "WriterPrepared"; + + auto* sync_point = SyncPoint::GetInstance(); + sync_point->DisableProcessing(); + sync_point->ClearAllCallBacks(); + sync_point->LoadDependency( + {{"DBImpl::WriteImpl:BeforeLeaderEnters:0", kWriterPrepared}, + {kReleaseWriter, "DBImpl::WriteImpl:BeforeLeaderEnters"}}); + const auto release_writer = [&](void*) { TEST_SYNC_POINT(kReleaseWriter); }; + sync_point->SetCallBack("WriteThread::EnterUnbatched:Wait", release_writer); + sync_point->SetCallBack( + "DBImpl::IngestExternalFiles:AfterReadLastSequenceBeforePublish", + release_writer); + + Options options = CurrentOptions(); + options.two_write_queues = std::get<0>(GetParam()); + options.atomic_flush = std::get<1>(GetParam()); + DestroyAndReopen(options); + const Snapshot* old_snapshot = db_->GetSnapshot(); + ASSERT_NE(nullptr, old_snapshot); + + sync_point->EnableProcessing(); + port::Thread writer([&]() { + WriteBatch batch; + ASSERT_OK(batch.Put("write-1", "v1")); + ASSERT_OK(batch.Put("write-2", "v2")); + ASSERT_OK(db_->Write(WriteOptions(), &batch)); + }); + + // The writer has read its sequence base but is paused before entering the + // batch group. With the fix, the late ingest writer gate waits behind it and + // releases it. Without the gate, it is released only after ingestion reads + // its stale sequence base before publication. + TEST_SYNC_POINT(kWriterPrepared); + port::Thread ingest([&]() { + ASSERT_OK(GenerateAndAddExternalFile( + options, {{"ingest", "v"}}, -1, true /* allow_global_seqno */, + false /* write_global_seqno */, + true /* verify_checksums_before_ingest */, false /* ingest_behind */, + false /* sort_data */, true /* allow_write */)); + }); + ingest.join(); + writer.join(); + + ASSERT_GE(db_->GetLatestSequenceNumber(), 3U); + ASSERT_EQ("v1", Get("write-1")); + ASSERT_EQ("v2", Get("write-2")); + ASSERT_EQ("v", Get("ingest")); + const Snapshot* new_snapshot = db_->GetSnapshot(); + ASSERT_NE(nullptr, new_snapshot); + ASSERT_EQ("v1", Get("write-1", new_snapshot)); + ASSERT_EQ("v2", Get("write-2", new_snapshot)); + ASSERT_EQ("v", Get("ingest", new_snapshot)); + db_->ReleaseSnapshot(new_snapshot); + ASSERT_EQ("NOT_FOUND", Get("write-2", old_snapshot)); + ASSERT_EQ("NOT_FOUND", Get("ingest", old_snapshot)); + db_->ReleaseSnapshot(old_snapshot); + + sync_point->DisableProcessing(); + sync_point->ClearAllCallBacks(); +} + TEST_P(ExternalSSTFileTest, InconsistentAllowWriteArguments) { Options options = CurrentOptions(); CreateAndReopenWithCF({"koko"}, options);