Fix sequence publication race in allow-write external SST ingestion

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<ColumnFamilyHandleImpl*>(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<bool> final_need_flushes(num_cfs, false);
+      for (size_t i = 0; i != num_cfs; ++i) {
+        auto* cfd =
+            static_cast<ColumnFamilyHandleImpl*>(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<ColumnFamilyHandleImpl*>(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);
       }

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);

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);