Publish external SST ingestion sequence reservations through the multi-batch commit queue

db/db_impl/db_impl.cc

@@ -5928,16 +5928,9 @@ 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();
+      const bool publish_seqno_through_commit_queue =
+          allow_write && immutable_db_options_.enable_multi_batch_write;
+      CommitRequest seqno_reservation(&w);
       SequenceNumber reserved_seqno_count = 0;
       if (allow_write) {
         // Each file consumes at most one sequence number. Jobs for different
@@ -5950,17 +5943,62 @@ Status DBImpl::IngestExternalFiles(
                            ingestion_jobs[i].files_to_ingest().size()));
         }
         assert(reserved_seqno_count > 0);
+
+        // Become the write queue leader before reserving sequence numbers.
+        if (publish_seqno_through_commit_queue) {
+          // Sequence reservation only needs write queue serialization. Keep
+          // the DB mutex unlocked so it cannot extend the write barrier.
+          mutex_.Unlock();
+          write_thread_.EnterUnbatched(&w);
+        } else {
+          write_thread_.EnterUnbatched(&w, &mutex_);
+        }
+        if (two_write_queues_) {
+          if (publish_seqno_through_commit_queue) {
+            nonmem_write_thread_.EnterUnbatched(&nonmem_w);
+          } else {
+            nonmem_write_thread_.EnterUnbatched(&nonmem_w, &mutex_);
+          }
+        }
+        if (!publish_seqno_through_commit_queue) {
+          WaitForPendingWrites();
+        }
+      }
+
+      SequenceNumber last_seqno =
+          publish_seqno_through_commit_queue
+              ? write_thread_.UpdateLastSequence(versions_->LastSequence())
+              : versions_->LastSequence();
+      if (allow_write) {
         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);
+        if (publish_seqno_through_commit_queue) {
+          // Publish the reservation after preceding multi-batch writers, but
+          // release the write queue immediately so later writers can proceed.
+          write_thread_.UpdateLastSequence(reserved_last_seqno);
+          seqno_reservation.commit_lsn = reserved_last_seqno;
+          write_thread_.EnterCommitQueue(&seqno_reservation);
+          // Keep global barriers from passing until the reservation is
+          // published.
+          ++pending_memtable_writes_;
+        } else {
+          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);
 
+        if (publish_seqno_through_commit_queue) {
+          TEST_SYNC_POINT(
+              "DBImpl::IngestExternalFiles:BeforeWaitForSeqnoReservation");
+          MultiBatchWriteCommit(&seqno_reservation);
+          mutex_.Lock();
+        }
+
         // 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");

db/db_impl/db_impl.h

@@ -2045,8 +2045,8 @@ class DBImpl : public DB {
       return;
     }
 
-    // Wait for writers that have allocated sequence numbers to finish their
-    // memtable writes and publish their sequences.
+    // Wait for allocated sequence numbers to be published, including after
+    // memtable writes and external SST ingestion reservations.
     if (pending_memtable_writes_.load() != 0) {
       TEST_SYNC_POINT("DBImpl::WaitForPendingWrites:PendingWrites");
       std::unique_lock<std::mutex> guard(switch_mutex_);
@@ -2715,8 +2715,8 @@ class DBImpl : public DB {
   // initialized with startup time.
   uint64_t delete_obsolete_files_last_run_;
 
-  // The thread that wants to switch memtable, can wait on this cv until the
-  // pending writes to memtable finishes.
+  // The thread that wants a global sequence barrier can wait on this cv until
+  // pending memtable writes and sequence reservations are published.
   std::condition_variable switch_cv_;
   // The mutex used by switch_cv_. mutex_ should be acquired beforehand.
   std::mutex switch_mutex_;

db/external_sst_file_test.cc

@@ -2529,20 +2529,35 @@ TEST_F(ExternalSSTFileTest, AllowWriteIngestWaitsForPendingMultiBatchWrite) {
   constexpr auto kStartIngest =
       "ExternalSSTFileTest::"
       "AllowWriteIngestWaitsForPendingMultiBatchWrite:StartIngest";
-  constexpr auto kReleaseWriter =
+  constexpr auto kStartSecondWriter =
+      "ExternalSSTFileTest::"
+      "AllowWriteIngestWaitsForPendingMultiBatchWrite:StartSecondWriter";
+  constexpr auto kSecondWriterPrepared =
+      "ExternalSSTFileTest::"
+      "AllowWriteIngestWaitsForPendingMultiBatchWrite:SecondWriterPrepared";
+  constexpr auto kSecondWriterObserved =
+      "ExternalSSTFileTest::"
+      "AllowWriteIngestWaitsForPendingMultiBatchWrite:SecondWriterObserved";
+  constexpr auto kReleaseFirstWriter =
       "ExternalSSTFileTest::"
-      "AllowWriteIngestWaitsForPendingMultiBatchWrite:ReleaseWriter";
+      "AllowWriteIngestWaitsForPendingMultiBatchWrite:ReleaseFirstWriter";
 
   auto* sync_point = SyncPoint::GetInstance();
   sync_point->DisableProcessing();
   sync_point->ClearAllCallBacks();
   sync_point->LoadDependency(
       {{"DBImpl::WriteImpl:CommitAfterWriteWAL", kStartIngest},
-       {kReleaseWriter, "DBImpl::WriteImpl:BeforePipelineWriteMemtable"}});
+       {"DBImpl::IngestExternalFiles:BeforeWaitForSeqnoReservation",
+        kStartSecondWriter},
+       {kSecondWriterPrepared, kSecondWriterObserved},
+       {kReleaseFirstWriter, "DBImpl::WriteImpl:BeforePipelineWriteMemtable"}});
 
-  const auto release_writer = [&](void*) { TEST_SYNC_POINT(kReleaseWriter); };
-  sync_point->SetCallBack("DBImpl::WaitForPendingWrites:BeforeBlock",
-                          release_writer);
+  std::atomic<int> writer_count{0};
+  sync_point->SetCallBack("DBImpl::WriteImpl:CommitAfterWriteWAL", [&](void*) {
+    if (writer_count.fetch_add(1) == 1) {
+      TEST_SYNC_POINT(kSecondWriterPrepared);
+    }
+  });
 
   SequenceNumber assigned_seqno = 0;
   sync_point->SetCallBack("ExternalSstFileIngestionJob::Run", [&](void* arg) {
@@ -2558,35 +2573,141 @@ TEST_F(ExternalSSTFileTest, AllowWriteIngestWaitsForPendingMultiBatchWrite) {
   const SequenceNumber last_seqno = db_->GetLatestSequenceNumber();
   // Force the ingested file to consume a global sequence number.
   const Snapshot* snapshot = db_->GetSnapshot();
-  Status write_status;
+  Status first_write_status;
+  Status second_write_status;
+  Status ingest_status;
 
   sync_point->EnableProcessing();
-  port::Thread writer([&]() { write_status = Put("bar", "v1"); });
+  port::Thread first_writer([&]() { first_write_status = Put("bar1", "v1"); });
 
   // The writer has allocated a sequence and released write_thread_, but has
   // not inserted into the memtable or published the sequence yet.
   TEST_SYNC_POINT(kStartIngest);
-  Status ingest_status =
-      GenerateAndAddExternalFile(options, {{"foo", "v"}}, -1, true, false, true,
-                                 false, false, true /* allow_write */);
+  port::Thread ingest_thread([&]() {
+    ingest_status =
+        GenerateAndAddExternalFile(options, {{"foo", "v"}}, -1, true, false,
+                                   true, false, false, true /* allow_write */);
+  });
 
-  // If ingestion does not wait for the pending writer, release it here so the
-  // sequence assertions fail instead of hanging in writer.join().
-  release_writer(nullptr);
-  writer.join();
+  // Ingestion has reserved its sequence and released write_thread_, so a new
+  // writer can proceed while ingestion waits for the first writer to publish.
+  TEST_SYNC_POINT(kStartSecondWriter);
+  port::Thread second_writer(
+      [&]() { second_write_status = Put("bar2", "v2"); });
+  TEST_SYNC_POINT(kSecondWriterObserved);
+  TEST_SYNC_POINT(kReleaseFirstWriter);
+
+  first_writer.join();
+  second_writer.join();
+  ingest_thread.join();
 
   ASSERT_OK(ingest_status);
-  ASSERT_OK(write_status);
+  ASSERT_OK(first_write_status);
+  ASSERT_OK(second_write_status);
   ASSERT_EQ(last_seqno + 2, assigned_seqno);
-  ASSERT_EQ(last_seqno + 2, db_->GetLatestSequenceNumber());
-  ASSERT_EQ("v1", Get("bar"));
+  ASSERT_EQ(last_seqno + 3, db_->GetLatestSequenceNumber());
+  ASSERT_EQ("v1", Get("bar1"));
+  ASSERT_EQ("v2", Get("bar2"));
   ASSERT_EQ("v", Get("foo"));
   db_->ReleaseSnapshot(snapshot);
 
   sync_point->DisableProcessing();
   sync_point->ClearAllCallBacks();
 }
 
+TEST_F(ExternalSSTFileTest, GlobalBarrierWaitsForSeqnoReservation) {
+  constexpr auto kReservationReady =
+      "ExternalSSTFileTest::GlobalBarrierWaitsForSeqnoReservation:"
+      "ReservationReady";
+  constexpr auto kStartBarrier =
+      "ExternalSSTFileTest::GlobalBarrierWaitsForSeqnoReservation:"
+      "StartBarrier";
+  constexpr auto kBarrierReached =
+      "ExternalSSTFileTest::GlobalBarrierWaitsForSeqnoReservation:"
+      "BarrierReached";
+  constexpr auto kBarrierObserved =
+      "ExternalSSTFileTest::GlobalBarrierWaitsForSeqnoReservation:"
+      "BarrierObserved";
+  constexpr auto kReleaseReservation =
+      "ExternalSSTFileTest::GlobalBarrierWaitsForSeqnoReservation:"
+      "ReleaseReservation";
+  constexpr auto kContinueReservation =
+      "ExternalSSTFileTest::GlobalBarrierWaitsForSeqnoReservation:"
+      "ContinueReservation";
+
+  auto* sync_point = SyncPoint::GetInstance();
+  sync_point->DisableProcessing();
+  sync_point->ClearAllCallBacks();
+  sync_point->LoadDependency({{kReservationReady, kStartBarrier},
+                              {kBarrierReached, kBarrierObserved},
+                              {kReleaseReservation, kContinueReservation}});
+
+  sync_point->SetCallBack(
+      "DBImpl::IngestExternalFiles:BeforeWaitForSeqnoReservation", [&](void*) {
+        TEST_SYNC_POINT(kReservationReady);
+        TEST_SYNC_POINT(kContinueReservation);
+      });
+
+  std::atomic<bool> barrier_waited{false};
+  std::atomic<bool> barrier_reached{false};
+  const auto signal_barrier_reached = [&] {
+    if (!barrier_reached.exchange(true)) {
+      TEST_SYNC_POINT(kBarrierReached);
+    }
+  };
+  sync_point->SetCallBack("DBImpl::WaitForPendingWrites:PendingWrites",
+                          [&](void*) {
+                            barrier_waited.store(true);
+                            signal_barrier_reached();
+                          });
+
+  sync_point->SetCallBack("ExternalSstFileIngestionJob::Run", [&](void*) {
+    // If the barrier fails to wait, prevent the test from hanging.
+    signal_barrier_reached();
+  });
+
+  Options options = CurrentOptions();
+  options.enable_pipelined_write = false;
+  options.unordered_write = false;
+  options.enable_multi_batch_write = true;
+  DestroyAndReopen(options);
+
+  const SequenceNumber last_seqno = db_->GetLatestSequenceNumber();
+  const Snapshot* snapshot = db_->GetSnapshot();
+  Status allow_write_status;
+  Status barrier_status;
+
+  sync_point->EnableProcessing();
+  port::Thread allow_write_ingest([&]() {
+    allow_write_status =
+        GenerateAndAddExternalFile(options, {{"foo1", "v1"}}, 1, true, false,
+                                   true, false, false, true /* allow_write */);
+  });
+
+  TEST_SYNC_POINT(kStartBarrier);
+  port::Thread barrier_ingest([&]() {
+    barrier_status =
+        GenerateAndAddExternalFile(options, {{"foo2", "v2"}}, 2, true, false,
+                                   true, false, false, false /* allow_write */);
+  });
+  TEST_SYNC_POINT(kBarrierObserved);
+  TEST_SYNC_POINT(kReleaseReservation);
+
+  allow_write_ingest.join();
+  barrier_ingest.join();
+
+  ASSERT_TRUE(barrier_waited.load());
+  ASSERT_OK(allow_write_status);
+  ASSERT_OK(barrier_status);
+  ASSERT_EQ(last_seqno + 2, db_->GetLatestSequenceNumber());
+  ASSERT_EQ("v1", Get("foo1"));
+  ASSERT_EQ("v2", Get("foo2"));
+  db_->ReleaseSnapshot(snapshot);
+
+  sync_point->DisableProcessing();
+  sync_point->ClearAllCallBacks();
+}
+
 TEST_F(ExternalSSTFileTest, AllowWriteIngestWaitsForFailedMultiBatchWrite) {
   constexpr auto kWriterPrepared =
       "ExternalSSTFileTest::"
@@ -2612,15 +2733,17 @@ TEST_F(ExternalSSTFileTest, AllowWriteIngestWaitsForFailedMultiBatchWrite) {
     TEST_SYNC_POINT(kContinueWriter);
   });
 
-  const auto release_writer = [&](void*) { TEST_SYNC_POINT(kReleaseWriter); };
   sync_point->SetCallBack("DBImpl::WaitForPendingWrites:PendingWrites",
-                          release_writer);
+                          [&](void*) { TEST_SYNC_POINT(kReleaseWriter); });
+  sync_point->SetCallBack(
+      "DBImpl::IngestExternalFiles:BeforeWaitForSeqnoReservation",
+      [&](void*) { TEST_SYNC_POINT(kReleaseWriter); });
 
   SequenceNumber assigned_seqno = 0;
   sync_point->SetCallBack("ExternalSstFileIngestionJob::Run", [&](void* arg) {
     // Release the writer here as a fallback so that a failure to wait for it
     // causes the sequence assertions to fail instead of a deadlock.
-    release_writer(nullptr);
+    TEST_SYNC_POINT(kReleaseWriter);
     assigned_seqno = *static_cast<SequenceNumber*>(arg);
   });
 

db/write_thread.cc

@@ -804,9 +804,8 @@ void WriteThread::ExitAsBatchGroupLeader(WriteGroup& write_group,
 }
 
 static WriteThread::AdaptationContext eu_ctx("EnterUnbatched");
-void WriteThread::EnterUnbatched(Writer* w, InstrumentedMutex* mu) {
+void WriteThread::EnterUnbatched(Writer* w) {
   assert(w != nullptr && w->multi_batch.batches.empty());
-  mu->Unlock();
   bool linked_as_leader = LinkOne(w, &newest_writer_);
   if (!linked_as_leader) {
     TEST_SYNC_POINT("WriteThread::EnterUnbatched:Wait");
@@ -816,6 +815,11 @@ void WriteThread::EnterUnbatched(Writer* w, InstrumentedMutex* mu) {
   if (enable_pipelined_write_) {
     WaitForMemTableWriters();
   }
+}
+
+void WriteThread::EnterUnbatched(Writer* w, InstrumentedMutex* mu) {
+  mu->Unlock();
+  EnterUnbatched(w);
   mu->Lock();
 }
 

db/write_thread.h

@@ -397,10 +397,8 @@ class WriteThread {
   virtual ~WriteThread() = default;
 
   // IMPORTANT: None of the methods in this class rely on the db mutex
-  // for correctness. All of the methods except JoinBatchGroup and
-  // EnterUnbatched may be called either with or without the db mutex held.
-  // Correctness is maintained by ensuring that only a single thread is
-  // a leader at a time.
+  // for correctness. Correctness is maintained by ensuring that only a
+  // single thread is a leader at a time.
 
   // Registers w as ready to become part of a batch group, waits until the
   // caller should perform some work, and returns the current state of the
@@ -459,8 +457,12 @@ class WriteThread {
   // someone else has already taken responsibility for that.
   bool CompleteParallelMemTableWriter(Writer* w);
 
-  // Waits for all preceding writers (unlocking mu while waiting), then
-  // registers w as the currently proceeding writer.
+  // Waits for all preceding writers, then registers w as the currently
+  // proceeding writer. The db mutex must not be held.
+  void EnterUnbatched(Writer* w);
+
+  // Same as above, but unlocks the db mutex while waiting and reacquires it
+  // before returning.
   //
   // Writer* w:              A Writer not eligible for batching
   // InstrumentedMutex* mu:  The db mutex, to unlock while waiting