Split ingest external files into prepare and commit APIs (#14849)

Summary:

This change splits ingestion into two public calls: `PrepareFileIngestion()` performs all of the off-mutex work and returns an opaque `FileIngestionHandle`, and `CommitFileIngestionHandle()` makes the prepared files visible under the mutex.

Even though the "Prepare" phase was off the DB mutex, the application layer may still have application logic that depends on the completion of `IngestExternalFile`.

`IngestExternalFiles(args)` is also just `PrepareFileIngestion(args)` followed by `CommitFileIngestionHandle()`, so its behavior is unchanged.

`CommitFileIngestionHandles()` also supports committing multiple handles atomically. However, it is possible that a CF may be present in multiple handles. In this case, we merge the ingestion jobs together. This allows applications to prepare file ingestions at different times, but still provide a single atomic commit.

Differential Revision: D108225105

Author: joshkang97

db/db_impl/db_impl.cc

@@ -76,6 +76,7 @@ namespace ROCKSDB_NAMESPACE {
 
 class Arena;
 class ArenaWrappedDBIter;
+class FileIngestionHandleImpl;
 class InMemoryStatsHistoryIterator;
 class MemTable;
 class PersistentStatsHistoryIterator;
@@ -604,6 +605,14 @@ class DBImpl : public DB {
   Status IngestExternalFiles(
       const std::vector<IngestExternalFileArg>& args) override;
 
+  using DB::PrepareFileIngestion;
+  Status PrepareFileIngestion(
+      const std::vector<IngestExternalFileArg>& args,
+      std::unique_ptr<FileIngestionHandle>* handle) override;
+
+  Status CommitFileIngestionHandles(
+      std::vector<std::unique_ptr<FileIngestionHandle>> handles) override;
+
   using DB::CreateColumnFamilyWithImport;
   Status CreateColumnFamilyWithImport(
       const ColumnFamilyOptions& options, const std::string& column_family_name,
@@ -1916,6 +1925,7 @@ class DBImpl : public DB {
   friend class WriteBatchWithIndex;
   friend class WriteUnpreparedTxnDB;
   friend class WriteUnpreparedTxn;
+  friend class FileIngestionHandleImpl;
 
   friend class ForwardIterator;
   friend struct SuperVersion;
@@ -2253,6 +2263,10 @@ class DBImpl : public DB {
   void ReleaseFileNumberFromPendingOutputs(
       std::unique_ptr<std::list<uint64_t>::iterator>& v);
 
+  // Rolls back one prepared file ingestion (delete its staged files, release
+  // the reserved file numbers)
+  void RollbackPreparedFileIngestion(FileIngestionHandleImpl* const h);
+
   // Similar to pending_outputs, preserve OPTIONS file. Used for remote
   // compaction.
   std::list<uint64_t>::iterator CaptureOptionsFileNumber();
@@ -3472,6 +3486,10 @@ class DBImpl : public DB {
   // REQUIRES: mutex held
   int num_running_ingest_file_ = 0;
 
+  // Number of FileIngestionHandle objects produced by PrepareFileIngestion()
+  // that have not been committed or destroyed yet.
+  std::atomic<uint32_t> num_outstanding_prepared_ingestions_{0};
+
   WalManager wal_manager_;
 
   // A value of > 0 temporarily disables scheduling of background work

db/db_impl/db_impl.h

@@ -134,6 +134,19 @@ class DBImplReadOnly : public DBImpl {
     return Status::NotSupported("Not supported operation in read only mode.");
   }
 
+  using DB::PrepareFileIngestion;
+  Status PrepareFileIngestion(
+      const std::vector<IngestExternalFileArg>& /*args*/,
+      std::unique_ptr<FileIngestionHandle>* /*handle*/) override {
+    return Status::NotSupported("Not supported operation in read only mode.");
+  }
+
+  using DB::CommitFileIngestionHandles;
+  Status CommitFileIngestionHandles(
+      std::vector<std::unique_ptr<FileIngestionHandle>> /*handles*/) override {
+    return Status::NotSupported("Not supported operation in read only mode.");
+  }
+
   using DB::CreateColumnFamilyWithImport;
   Status CreateColumnFamilyWithImport(
       const ColumnFamilyOptions& /*options*/,

db/db_impl/db_impl_readonly.h

@@ -239,6 +239,19 @@ class DBImplSecondary : public DBImpl {
     return Status::NotSupported("Not supported operation in secondary mode.");
   }
 
+  using DB::PrepareFileIngestion;
+  Status PrepareFileIngestion(
+      const std::vector<IngestExternalFileArg>& /*args*/,
+      std::unique_ptr<FileIngestionHandle>* /*handle*/) override {
+    return Status::NotSupported("Not supported operation in secondary mode.");
+  }
+
+  using DB::CommitFileIngestionHandles;
+  Status CommitFileIngestionHandles(
+      std::vector<std::unique_ptr<FileIngestionHandle>> /*handles*/) override {
+    return Status::NotSupported("Not supported operation in secondary mode.");
+  }
+
   // Try to catch up with the primary by reading as much as possible from the
   // log files until there is nothing more to read or encounters an error. If
   // the amount of information in the log files to process is huge, this

db/db_impl/db_impl_secondary.h

@@ -3516,6 +3516,18 @@ class ModelDB : public DB {
     return Status::NotSupported("Not implemented");
   }
 
+  using DB::PrepareFileIngestion;
+  Status PrepareFileIngestion(
+      const std::vector<IngestExternalFileArg>& /*args*/,
+      std::unique_ptr<FileIngestionHandle>* /*handle*/) override {
+    return Status::NotSupported("Not implemented.");
+  }
+
+  Status CommitFileIngestionHandles(
+      std::vector<std::unique_ptr<FileIngestionHandle>> /*handles*/) override {
+    return Status::NotSupported("Not implemented.");
+  }
+
   using DB::CreateColumnFamilyWithImport;
   Status CreateColumnFamilyWithImport(
       const ColumnFamilyOptions& /*options*/,

db/db_test.cc

@@ -76,23 +76,10 @@ Status ExternalSstFileIngestionJob::Prepare(
   auto num_files = files_to_ingest_.size();
   if (num_files == 0) {
     return Status::InvalidArgument("The list of files is empty");
-  } else if (num_files > 1) {
-    // Verify that passed files don't have overlapping ranges
-    autovector<const IngestedFileInfo*> sorted_files;
-    for (size_t i = 0; i < num_files; i++) {
-      sorted_files.push_back(&files_to_ingest_[i]);
-    }
-
-    std::sort(sorted_files.begin(), sorted_files.end(), file_range_checker_);
-
-    for (size_t i = 0; i + 1 < num_files; i++) {
-      if (file_range_checker_.Overlaps(*sorted_files[i], *sorted_files[i + 1],
-                                       /* known_sorted= */ true)) {
-        files_overlap_ = true;
-        break;
-      }
-    }
   }
+  // Detect whether the input files overlap one another; this drives how they
+  // are divided into batches below.
+  files_overlap_ = ComputeFilesOverlap(files_to_ingest_);
 
   if (atomic_replace_range.has_value()) {
     atomic_replace_range_.emplace();
@@ -433,6 +420,60 @@ void ExternalSstFileIngestionJob::DivideInputFilesIntoBatches() {
   }
 }
 
+bool ExternalSstFileIngestionJob::ComputeFilesOverlap(
+    const autovector<IngestedFileInfo>& files) const {
+  const size_t num_files = files.size();
+  if (num_files <= 1) {
+    return false;
+  }
+  // Verify whether the files have overlapping ranges by sorting copies of the
+  // file ranges and checking adjacent pairs.
+  autovector<const IngestedFileInfo*> sorted_files;
+  for (size_t i = 0; i < num_files; i++) {
+    sorted_files.push_back(&files[i]);
+  }
+  std::sort(sorted_files.begin(), sorted_files.end(), file_range_checker_);
+  for (size_t i = 0; i + 1 < num_files; i++) {
+    if (file_range_checker_.Overlaps(*sorted_files[i], *sorted_files[i + 1],
+                                     /* known_sorted= */ true)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+Status ExternalSstFileIngestionJob::MergeForSameColumnFamily(
+    ExternalSstFileIngestionJob* other) {
+  assert(other != nullptr);
+  assert(other != this);
+  assert(cfd_ == other->cfd_);
+  if (atomic_replace_range_.has_value() ||
+      other->atomic_replace_range_.has_value()) {
+    return Status::NotSupported(
+        "cannot merge file ingestion handles for the same column family when "
+        "atomic_replace_range is used");
+  }
+  if (!(ingestion_options_ == other->ingestion_options_)) {
+    return Status::InvalidArgument(
+        "file ingestion handles for the same column family must be prepared "
+        "with the same IngestExternalFileOptions");
+  }
+  // Append the other job's prepared files after this job's so that, for any
+  // overlapping keys, the other job's data wins via a higher assigned sequence
+  // number -- the same semantics as passing all the files to a single ingestion
+  // call in this order. Recompute overlap and rebuild the batches over the
+  // union.
+  for (IngestedFileInfo& file : other->files_to_ingest_) {
+    files_to_ingest_.push_back(std::move(file));
+  }
+  other->files_to_ingest_.clear();
+  other->file_batches_to_ingest_.clear();
+  files_overlap_ = ComputeFilesOverlap(files_to_ingest_);
+  file_batches_to_ingest_.clear();
+  DivideInputFilesIntoBatches();
+  return Status::OK();
+}
+
 Status ExternalSstFileIngestionJob::NeedsFlush(bool* flush_needed,
                                                SuperVersion* super_version) {
   Status status;

db/external_sst_file_ingestion_job.cc

@@ -301,6 +301,14 @@ class ExternalSstFileIngestionJob {
     return max_assigned_seqno_;
   }
 
+  // Merge another already-Prepare()d job for the SAME column family into this
+  // one so both sets of files are committed by a single Run(). The other job's
+  // files are appended after this job's, so for any overlapping keys the other
+  // job's data wins via a higher assigned sequence number -- the same semantics
+  // as passing all the files to a single ingestion call in this order. The
+  // other job is left empty.
+  Status MergeForSameColumnFamily(ExternalSstFileIngestionJob* other);
+
  private:
   Status ResetTableReader(const std::string& external_file,
                           uint64_t new_file_number,
@@ -334,6 +342,10 @@ class ExternalSstFileIngestionJob {
   // make one batch.
   void DivideInputFilesIntoBatches();
 
+  // Returns whether any two files in `files` have overlapping key ranges, by
+  // sorting the file ranges and checking adjacent pairs.
+  bool ComputeFilesOverlap(const autovector<IngestedFileInfo>& files) const;
+
   // Assign level for the files in one batch. The files within one batch are not
   // overlapping, and we assign level to each file one after another.
   // If `prev_batch_uppermost_level` is specified, all files in this batch will
@@ -407,7 +419,7 @@ class ExternalSstFileIngestionJob {
   SnapshotList* db_snapshots_;
   autovector<IngestedFileInfo> files_to_ingest_;
   std::vector<FileBatchInfo> file_batches_to_ingest_;
-  const IngestExternalFileOptions& ingestion_options_;
+  const IngestExternalFileOptions ingestion_options_;
   std::optional<KeyRangeInfo> atomic_replace_range_;
   Directories* directories_;
   EventLogger* event_logger_;