Fix sequence number race when allowing writes during external SST ingestion

db/db_impl/db_impl.cc

@@ -5928,12 +5928,54 @@ Status DBImpl::IngestExternalFiles(
     }
     // Run ingestion jobs.
     if (status.ok()) {
+      if (allow_write) {
+        // Briefly stop writes while reserving sequence numbers for ingestion.
+        write_thread_.EnterUnbatched(&w, &mutex_);
+        if (two_write_queues_) {
+          nonmem_write_thread_.EnterUnbatched(&nonmem_w, &mutex_);
+        }
+        WaitForPendingWrites();
+      }
+
+      SequenceNumber last_seqno = versions_->LastSequence();
+      SequenceNumber reserved_seqno_count = 0;
+      if (allow_write) {
+        // Each file consumes at most one sequence number. Jobs for different
+        // column families share the same sequence range, so reserve the maximum
+        // file count. Unused sequence numbers are harmless gaps.
+        for (size_t i = 0; i != num_cfs; ++i) {
+          reserved_seqno_count =
+              std::max(reserved_seqno_count,
+                       static_cast<SequenceNumber>(
+                           ingestion_jobs[i].files_to_ingest().size()));
+        }
+        assert(reserved_seqno_count > 0);
+        const SequenceNumber reserved_last_seqno =
+            last_seqno + reserved_seqno_count;
+        versions_->SetLastAllocatedSequence(reserved_last_seqno);
+        versions_->SetLastPublishedSequence(reserved_last_seqno);
+        versions_->SetLastSequence(reserved_last_seqno);
+        // Resume writes
+        if (two_write_queues_) {
+          nonmem_write_thread_.ExitUnbatched(&nonmem_w);
+        }
+        write_thread_.ExitUnbatched(&w);
+
+        // The reservation cannot be rolled back if ingestion fails because a
+        // foreground write may have already consumed a later sequence number.
+        TEST_SYNC_POINT("DBImpl::IngestExternalFiles:AfterReserveSeqno");
+      }
+
       for (size_t i = 0; i != num_cfs; ++i) {
         mutex_.AssertHeld();
-        status = ingestion_jobs[i].Run();
+        status = ingestion_jobs[i].Run(last_seqno);
         if (!status.ok()) {
           break;
         }
+        assert(!allow_write ||
+               static_cast<SequenceNumber>(
+                   ingestion_jobs[i].ConsumedSequenceNumbersCount()) <=
+                   reserved_seqno_count);
         ingestion_jobs[i].RegisterRange();
       }
     }
@@ -5964,29 +6006,37 @@ Status DBImpl::IngestExternalFiles(
         }
         assert(0 == num_entries);
       }
+      // With allow_write, a concurrent flush may persist a higher last sequence
+      // before this ingestion edit is applied. LogAndApply updates the edit as
+      // needed to keep VersionEdit::last_sequence non-decreasing in the
+      // MANIFEST.
       status = versions_->LogAndApply(cfds_to_commit, mutable_cf_options_list,
                                       read_options, edit_lists, &mutex_,
                                       directories_.GetDbDir());
-      // It is safe to update VersionSet last seqno here after LogAndApply since
-      // LogAndApply persists last sequence number from VersionEdits,
-      // which are from file's largest seqno and not from VersionSet.
-      //
-      // It is necessary to update last seqno here since LogAndApply releases
-      // mutex when persisting MANIFEST file, and the snapshots taken during
-      // that period will not be stable if VersionSet last seqno is updated
-      // before LogAndApply.
-      int consumed_seqno_count =
-          ingestion_jobs[0].ConsumedSequenceNumbersCount();
-      for (size_t i = 1; i != num_cfs; ++i) {
-        consumed_seqno_count =
-            std::max(consumed_seqno_count,
-                     ingestion_jobs[i].ConsumedSequenceNumbersCount());
-      }
-      if (consumed_seqno_count > 0) {
-        const SequenceNumber last_seqno = versions_->LastSequence();
-        versions_->SetLastAllocatedSequence(last_seqno + consumed_seqno_count);
-        versions_->SetLastPublishedSequence(last_seqno + consumed_seqno_count);
-        versions_->SetLastSequence(last_seqno + consumed_seqno_count);
+      if (!allow_write) {
+        // It is safe to update VersionSet last seqno here after LogAndApply
+        // since LogAndApply persists last sequence number from VersionEdits,
+        // which are from file's largest seqno and not from VersionSet.
+        //
+        // It is necessary to update last seqno here since LogAndApply releases
+        // mutex when persisting MANIFEST file, and the snapshots taken during
+        // that period will not be stable if VersionSet last seqno is updated
+        // before LogAndApply.
+        int consumed_seqno_count =
+            ingestion_jobs[0].ConsumedSequenceNumbersCount();
+        for (size_t i = 1; i != num_cfs; ++i) {
+          consumed_seqno_count =
+              std::max(consumed_seqno_count,
+                       ingestion_jobs[i].ConsumedSequenceNumbersCount());
+        }
+        if (consumed_seqno_count > 0) {
+          const SequenceNumber last_seqno = versions_->LastSequence();
+          versions_->SetLastAllocatedSequence(last_seqno +
+                                              consumed_seqno_count);
+          versions_->SetLastPublishedSequence(last_seqno +
+                                              consumed_seqno_count);
+          versions_->SetLastSequence(last_seqno + consumed_seqno_count);
+        }
       }
     }
 

db/db_impl/db_impl.h

@@ -2045,9 +2045,10 @@ class DBImpl : public DB {
       return;
     }
 
-    // Wait for the ones who already wrote to the WAL to finish their
-    // memtable write.
+    // Wait for writers that have allocated sequence numbers to finish their
+    // memtable writes and publish their sequences.
     if (pending_memtable_writes_.load() != 0) {
+      TEST_SYNC_POINT("DBImpl::WaitForPendingWrites:PendingWrites");
       std::unique_lock<std::mutex> guard(switch_mutex_);
       switch_cv_.wait(guard,
                       [&] { return pending_memtable_writes_.load() == 0; });
@@ -2719,7 +2720,7 @@ class DBImpl : public DB {
   std::condition_variable switch_cv_;
   // The mutex used by switch_cv_. mutex_ should be acquired beforehand.
   std::mutex switch_mutex_;
-  // Number of threads intending to write to memtable
+  // Number of writers with pending memtable writes or sequence publication.
   std::atomic<size_t> pending_memtable_writes_ = {};
 
   // A flag indicating whether the current rocksdb database has any

db/db_impl/db_impl_write.cc

@@ -366,13 +366,14 @@ Status DBImpl::MultiBatchWriteImpl(const WriteOptions& write_options,
       const ReadOptions read_options;
       writer.status = ApplyWALToManifest(read_options, &synced_wals);
     }
-    if (writer.status.ok()) {
-      pending_memtable_writes_ += memtable_write_cnt;
-    } else {
+    if (!writer.status.ok()) {
       // The `pending_wb_cnt` must be reset to avoid other writers helping
       // the front writer write its WBs after it failed to write the WAL.
       writer.ResetPendingWBCnt();
     }
+    // Every writer in the commit queue calls MultiBatchWriteCommit, including
+    // when the write fails before its memtable write.
+    pending_memtable_writes_ += memtable_write_cnt;
     write_thread_.ExitAsBatchGroupLeader(wal_write_group, writer.status);
   }
 

db/external_sst_file_ingestion_job.cc

@@ -372,9 +372,7 @@ Status ExternalSstFileIngestionJob::NeedsFlush(bool* flush_needed,
   return status;
 }
 
-// REQUIRES: we have become the only writer by entering both write_thread_ and
-// nonmem_write_thread_
-Status ExternalSstFileIngestionJob::Run() {
+Status ExternalSstFileIngestionJob::Run(SequenceNumber last_seqno) {
   Status status;
   SuperVersion* super_version = cfd_->GetSuperVersion();
 #ifndef NDEBUG
@@ -398,9 +396,6 @@ Status ExternalSstFileIngestionJob::Run() {
     // if the don't overlap with any ranges since we have snapshots
     force_global_seqno = true;
   }
-  // It is safe to use this instead of LastAllocatedSequence since we are
-  // the only active writer, and hence they are equal
-  SequenceNumber last_seqno = versions_->LastSequence();
   edit_.SetColumnFamily(cfd_->GetID());
   // The levels that the files will be ingested into
 

db/external_sst_file_ingestion_job.h

@@ -132,7 +132,7 @@ class ExternalSstFileIngestionJob {
 
   // Will execute the ingestion job and prepare edit() to be applied.
   // REQUIRES: Mutex held
-  Status Run();
+  Status Run(SequenceNumber last_seqno);
 
   // Register key range involved in this ingestion job
   // to prevent key range conflict with other ongoing compaction/file ingestion