external_sst_file_test.cc

//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under both the GPLv2 (found in the
//  COPYING file in the root directory) and Apache 2.0 License
//  (found in the LICENSE.Apache file in the root directory).

#include <table/block_based/block_based_table_factory.h>

#include <atomic>
#include <functional>
#include <memory>
#include <sstream>

#include "db/db_test_util.h"
#include "db/dbformat.h"
#include "file/filename.h"
#include "options/options_helper.h"
#include "port/port.h"
#include "port/stack_trace.h"
#include "rocksdb/sst_file_reader.h"
#include "rocksdb/sst_file_writer.h"
#include "table/prepared_file_info.h"
#include "test_util/testutil.h"
#include "util/random.h"
#include "util/thread_guard.h"
#include "utilities/fault_injection_env.h"

namespace ROCKSDB_NAMESPACE {

// A test environment that can be configured to fail the Link operation.
class ExternalSSTTestFS : public FileSystemWrapper {
 public:
  ExternalSSTTestFS(const std::shared_ptr<FileSystem>& t, bool fail_link)
      : FileSystemWrapper(t), fail_link_(fail_link) {}
  static const char* kClassName() { return "ExternalSSTTestFS"; }
  const char* Name() const override { return kClassName(); }

  IOStatus LinkFile(const std::string& s, const std::string& t,
                    const IOOptions& options, IODebugContext* dbg) override {
    if (fail_link_) {
      return IOStatus::NotSupported("Link failed");
    }
    return target()->LinkFile(s, t, options, dbg);
  }

  void set_fail_link(bool fail_link) { fail_link_ = fail_link; }

 private:
  bool fail_link_;
};

class ExternalSSTFileTestBase : public DBTestBase {
 public:
  ExternalSSTFileTestBase()
      : DBTestBase("external_sst_file_test", /*env_do_fsync=*/true) {
    sst_files_dir_ = dbname_ + "/sst_files/";
    DestroyAndRecreateExternalSSTFilesDir();
  }

  void DestroyAndRecreateExternalSSTFilesDir() {
    ASSERT_OK(DestroyDir(env_, sst_files_dir_));
    ASSERT_OK(env_->CreateDir(sst_files_dir_));
  }

  ~ExternalSSTFileTestBase() override {
    DestroyDir(env_, sst_files_dir_).PermitUncheckedError();
  }

 protected:
  std::string sst_files_dir_;
};

class ExternSSTFileLinkFailFallbackTest
    : public ExternalSSTFileTestBase,
      public ::testing::WithParamInterface<std::tuple<bool, bool, bool>> {
 public:
  ExternSSTFileLinkFailFallbackTest() {
    fs_ = std::make_shared<ExternalSSTTestFS>(env_->GetFileSystem(), true);
    test_env_.reset(new CompositeEnvWrapper(env_, fs_));
    options_ = CurrentOptions();
    options_.disable_auto_compactions = true;
    options_.env = test_env_.get();
  }

  void TearDown() override {
    db_.reset();
    ASSERT_OK(DestroyDB(dbname_, options_));
  }

 protected:
  Options options_;
  std::shared_ptr<ExternalSSTTestFS> fs_;
  std::unique_ptr<Env> test_env_;
};

class ExternalSSTFileTest
    : public ExternalSSTFileTestBase,
      public ::testing::WithParamInterface<std::tuple<bool, bool, bool>> {
 public:
  ExternalSSTFileTest() = default;

  void SetUp() override {
    const auto* info = ::testing::UnitTest::GetInstance()->current_test_info();
    if (info != nullptr && info->value_param() != nullptr) {
      two_phase_ingest_ = std::get<2>(GetParam());
    }
  }

  Status GenerateOneExternalFile(
      const Options& options, ColumnFamilyHandle* cfh,
      std::vector<std::pair<std::string, std::string>>& data, int file_id,
      bool sort_data, std::string* external_file_path,
      std::map<std::string, std::string>* true_data) {
    // Generate a file id if not provided
    if (-1 == file_id) {
      file_id = (++last_file_id_);
    }
    // Sort data if asked to do so
    if (sort_data) {
      std::sort(data.begin(), data.end(),
                [&](const std::pair<std::string, std::string>& e1,
                    const std::pair<std::string, std::string>& e2) {
                  return options.comparator->Compare(e1.first, e2.first) < 0;
                });
      auto uniq_iter = std::unique(
          data.begin(), data.end(),
          [&](const std::pair<std::string, std::string>& e1,
              const std::pair<std::string, std::string>& e2) {
            return options.comparator->Compare(e1.first, e2.first) == 0;
          });
      data.resize(uniq_iter - data.begin());
    }
    std::string file_path = sst_files_dir_ + std::to_string(file_id);
    SstFileWriter sst_file_writer(EnvOptions(), options, cfh);
    Status s = sst_file_writer.Open(file_path);
    if (!s.ok()) {
      return s;
    }
    for (const auto& entry : data) {
      s = sst_file_writer.Put(entry.first, entry.second);
      if (!s.ok()) {
        sst_file_writer.Finish().PermitUncheckedError();
        return s;
      }
    }
    s = sst_file_writer.Finish();
    if (s.ok() && external_file_path != nullptr) {
      *external_file_path = file_path;
    }
    if (s.ok() && nullptr != true_data) {
      for (const auto& entry : data) {
        true_data->insert({entry.first, entry.second});
      }
    }
    return s;
  }

  Status GenerateAndAddExternalFile(
      const Options options,
      std::vector<std::pair<std::string, std::string>> data, int file_id = -1,
      bool allow_global_seqno = false, bool write_global_seqno = false,
      bool verify_checksums_before_ingest = true, bool ingest_behind = false,
      bool sort_data = false,
      std::map<std::string, std::string>* true_data = nullptr,
      ColumnFamilyHandle* cfh = nullptr, bool fill_cache = false,
      bool ingest_with_file_info = false) {
    // Generate a file id if not provided
    if (file_id == -1) {
      file_id = last_file_id_ + 1;
      last_file_id_++;
    }

    // Sort data if asked to do so
    if (sort_data) {
      std::sort(data.begin(), data.end(),
                [&](const std::pair<std::string, std::string>& e1,
                    const std::pair<std::string, std::string>& e2) {
                  return options.comparator->Compare(e1.first, e2.first) < 0;
                });
      auto uniq_iter = std::unique(
          data.begin(), data.end(),
          [&](const std::pair<std::string, std::string>& e1,
              const std::pair<std::string, std::string>& e2) {
            return options.comparator->Compare(e1.first, e2.first) == 0;
          });
      data.resize(uniq_iter - data.begin());
    }
    std::string file_path = sst_files_dir_ + std::to_string(file_id);
    SstFileWriter sst_file_writer(EnvOptions(), options, cfh);

    Status s = sst_file_writer.Open(file_path);
    if (!s.ok()) {
      return s;
    }
    for (auto& entry : data) {
      s = sst_file_writer.Put(entry.first, entry.second);
      if (!s.ok()) {
        sst_file_writer.Finish().PermitUncheckedError();
        return s;
      }
    }
    ExternalSstFileInfo file_info;
    if (ingest_with_file_info) {
      s = sst_file_writer.Finish(&file_info);
    } else {
      s = sst_file_writer.Finish();
    }

    if (s.ok()) {
      IngestExternalFileOptions ifo;
      ifo.allow_global_seqno = allow_global_seqno;
      ifo.write_global_seqno = allow_global_seqno ? write_global_seqno : false;
      ifo.verify_checksums_before_ingest = verify_checksums_before_ingest;
      ifo.ingest_behind = ingest_behind;
      ifo.fill_cache = fill_cache;
      if (two_phase_ingest_) {
        std::unique_ptr<FileIngestionHandle> handle;
        s = db_->PrepareFileIngestion(cfh ? cfh : db_->DefaultColumnFamily(),
                                      {file_path}, ifo, &handle);
        if (s.ok()) {
          s = db_->CommitFileIngestionHandle(std::move(handle));
        }
      } else if (ingest_with_file_info) {
        IngestExternalFileArg arg;
        arg.column_family = cfh ? cfh : db_->DefaultColumnFamily();
        arg.external_files = {file_path};
        arg.file_infos = {file_info.prepared_file_info.get()};
        arg.options = ifo;
        s = db_->IngestExternalFiles({arg});
      } else if (cfh) {
        s = db_->IngestExternalFile(cfh, {file_path}, ifo);
      } else {
        s = db_->IngestExternalFile({file_path}, ifo);
      }
    }

    if (s.ok() && true_data) {
      for (auto& entry : data) {
        (*true_data)[entry.first] = entry.second;
      }
    }

    return s;
  }

  Status GenerateAndAddExternalFiles(
      const Options& options,
      const std::vector<ColumnFamilyHandle*>& column_families,
      const std::vector<IngestExternalFileOptions>& ifos,
      std::vector<std::vector<std::pair<std::string, std::string>>>& data,
      int file_id, bool sort_data,
      std::vector<std::map<std::string, std::string>>& true_data) {
    if (-1 == file_id) {
      file_id = (++last_file_id_);
    }
    // Generate external SST files, one for each column family
    size_t num_cfs = column_families.size();
    assert(ifos.size() == num_cfs);
    assert(data.size() == num_cfs);
    std::vector<IngestExternalFileArg> args(num_cfs);
    for (size_t i = 0; i != num_cfs; ++i) {
      std::string external_file_path;
      Status s = GenerateOneExternalFile(
          options, column_families[i], data[i], file_id, sort_data,
          &external_file_path,
          true_data.size() == num_cfs ? &true_data[i] : nullptr);
      if (!s.ok()) {
        return s;
      }
      ++file_id;

      args[i].column_family = column_families[i];
      args[i].external_files.push_back(external_file_path);
      args[i].options = ifos[i];
    }
    if (two_phase_ingest_) {
      std::unique_ptr<FileIngestionHandle> handle;
      Status s = db_->PrepareFileIngestion(args, &handle);
      if (!s.ok()) {
        return s;
      }
      return db_->CommitFileIngestionHandle(std::move(handle));
    }
    return db_->IngestExternalFiles(args);
  }

  Status GenerateAndAddExternalFile(
      const Options options, std::vector<std::pair<int, std::string>> data,
      int file_id = -1, bool allow_global_seqno = false,
      bool write_global_seqno = false,
      bool verify_checksums_before_ingest = true, bool ingest_behind = false,
      bool sort_data = false,
      std::map<std::string, std::string>* true_data = nullptr,
      ColumnFamilyHandle* cfh = nullptr) {
    std::vector<std::pair<std::string, std::string>> file_data;
    for (auto& entry : data) {
      file_data.emplace_back(Key(entry.first), entry.second);
    }
    return GenerateAndAddExternalFile(options, file_data, file_id,
                                      allow_global_seqno, write_global_seqno,
                                      verify_checksums_before_ingest,
                                      ingest_behind, sort_data, true_data, cfh);
  }

  Status GenerateAndAddExternalFile(
      const Options options, std::vector<int> keys, int file_id = -1,
      bool allow_global_seqno = false, bool write_global_seqno = false,
      bool verify_checksums_before_ingest = true, bool ingest_behind = false,
      bool sort_data = false,
      std::map<std::string, std::string>* true_data = nullptr,
      ColumnFamilyHandle* cfh = nullptr, bool fill_cache = false) {
    std::vector<std::pair<std::string, std::string>> file_data;
    for (auto& k : keys) {
      file_data.emplace_back(Key(k), Key(k) + std::to_string(file_id));
    }
    return GenerateAndAddExternalFile(
        options, file_data, file_id, allow_global_seqno, write_global_seqno,
        verify_checksums_before_ingest, ingest_behind, sort_data, true_data,
        cfh, fill_cache);
  }

  Status DeprecatedAddFile(const std::vector<std::string>& files,
                           bool move_files = false,
                           bool skip_snapshot_check = false,
                           bool skip_write_global_seqno = false) {
    IngestExternalFileOptions opts;
    opts.move_files = move_files;
    opts.snapshot_consistency = !skip_snapshot_check;
    opts.allow_global_seqno = false;
    opts.allow_blocking_flush = false;
    opts.write_global_seqno = !skip_write_global_seqno;
    return db_->IngestExternalFile(files, opts);
  }

  // Writes `data` (sorted by user key) into a new external SST and returns its
  // path in *file_path WITHOUT ingesting it, so a test can drive the two-phase
  // PrepareFileIngestion()/CommitFileIngestion() API directly.
  Status GenerateExternalFileOnly(
      const Options& options,
      std::vector<std::pair<std::string, std::string>> data,
      std::string* file_path) {
    std::sort(data.begin(), data.end(),
              [&](const std::pair<std::string, std::string>& e1,
                  const std::pair<std::string, std::string>& e2) {
                return options.comparator->Compare(e1.first, e2.first) < 0;
              });
    *file_path = sst_files_dir_ + "only_" + std::to_string(++last_file_id_);
    SstFileWriter sst_file_writer(EnvOptions(), options);
    Status s = sst_file_writer.Open(*file_path);
    if (!s.ok()) {
      return s;
    }
    for (const auto& entry : data) {
      s = sst_file_writer.Put(entry.first, entry.second);
      if (!s.ok()) {
        sst_file_writer.Finish().PermitUncheckedError();
        return s;
      }
    }
    return sst_file_writer.Finish();
  }

 protected:
  int last_file_id_ = 0;
  bool two_phase_ingest_ = false;
};

TEST_F(ExternalSSTFileTest, IngestionTimingHistogram) {
  Options options = CurrentOptions();
  options.statistics = CreateDBStatistics();
  options.statistics->set_stats_level(StatsLevel::kAll);  // enable timers
  DestroyAndReopen(options);

  ASSERT_OK(GenerateAndAddExternalFile(
      options, {{"k1", "v1"}, {"k2", "v2"}, {"k3", "v3"}}, /*file_id=*/-1,
      /*allow_global_seqno=*/true, /*write_global_seqno=*/false,
      /*verify_checksums_before_ingest=*/true, /*ingest_behind=*/false,
      /*sort_data=*/true));

  auto hist = [&](Histograms h) {
    HistogramData hd;
    options.statistics->histogramData(h, &hd);
    return hd;
  };

  // One sample per call (not per CF) in each phase histogram.
  HistogramData prepare_hd = hist(INGEST_EXTERNAL_FILE_PREPARE_TIME);
  HistogramData run_hd = hist(INGEST_EXTERNAL_FILE_RUN_TIME);
  ASSERT_EQ(1, prepare_hd.count);
  ASSERT_EQ(1, run_hd.count);
  ASSERT_GT(prepare_hd.max, 0.0);
  ASSERT_GT(run_hd.max, 0.0);

  // A second call adds exactly one more sample to each histogram.
  ASSERT_OK(GenerateAndAddExternalFile(
      options, {{"k4", "v4"}, {"k5", "v5"}}, /*file_id=*/-1,
      /*allow_global_seqno=*/true, /*write_global_seqno=*/false,
      /*verify_checksums_before_ingest=*/true, /*ingest_behind=*/false,
      /*sort_data=*/true));
  ASSERT_EQ(2, hist(INGEST_EXTERNAL_FILE_PREPARE_TIME).count);
  ASSERT_EQ(2, hist(INGEST_EXTERNAL_FILE_RUN_TIME).count);

  // A failed ingestion records nothing (latency is logged only on success).
  ASSERT_NOK(db_->IngestExternalFile({"/path/does/not/exist.sst"},
                                     IngestExternalFileOptions()));
  ASSERT_EQ(2, hist(INGEST_EXTERNAL_FILE_PREPARE_TIME).count);
  ASSERT_EQ(2, hist(INGEST_EXTERNAL_FILE_RUN_TIME).count);
}

TEST_F(ExternalSSTFileTest, PrepareThenCommit) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);

  std::string file_path;
  ASSERT_OK(GenerateExternalFileOnly(
      options, {{"k1", "v1"}, {"k2", "v2"}, {"k3", "v3"}}, &file_path));

  IngestExternalFileArg arg;
  arg.column_family = db_->DefaultColumnFamily();
  arg.external_files = {file_path};

  std::unique_ptr<FileIngestionHandle> handle;
  ASSERT_OK(db_->PrepareFileIngestion({arg}, &handle));
  ASSERT_TRUE(handle != nullptr);

  // The prepared file is staged in the DB but not yet visible.
  ASSERT_EQ("NOT_FOUND", Get("k1"));

  ASSERT_OK(db_->CommitFileIngestionHandle(std::move(handle)));
  ASSERT_EQ("v1", Get("k1"));
  ASSERT_EQ("v2", Get("k2"));
  ASSERT_EQ("v3", Get("k3"));
}

TEST_F(ExternalSSTFileTest, ParallelFileOpenWithFileOpeningThreads) {
  // Ingest many non-overlapping files with file_opening_threads > 1 so commit
  // opens table readers with multiple threads. max_open_files == -1 makes
  // commit open every new file.
  Options options = CurrentOptions();
  options.max_open_files = -1;
  DestroyAndReopen(options);

  constexpr int kNumFiles = 8;
  constexpr int kKeysPerFile = 50;
  std::vector<std::string> files;
  std::map<std::string, std::string> true_data;
  for (int f = 0; f < kNumFiles; f++) {
    std::vector<std::pair<std::string, std::string>> data;
    for (int k = 0; k < kKeysPerFile; k++) {
      const int key = f * kKeysPerFile + k;
      const std::string value = "val" + std::to_string(key);
      data.emplace_back(Key(key), value);
      true_data[Key(key)] = value;
    }
    std::string file_path;
    ASSERT_OK(GenerateExternalFileOnly(options, data, &file_path));
    files.push_back(file_path);
  }

  IngestExternalFileOptions ifo;
  ifo.file_opening_threads = 4;
  ASSERT_OK(db_->IngestExternalFile(files, ifo));

  for (const auto& [key, value] : true_data) {
    ASSERT_EQ(value, Get(key));
  }
}

TEST_F(ExternalSSTFileTest, LmaxPrefetchSkipDoesNotDisableL0Prefetch) {
  LRUCacheOptions co;
  co.capacity = 32 << 20;
  std::shared_ptr<Cache> cache = NewLRUCache(co);
  BlockBasedTableOptions table_options;
  table_options.block_cache = cache;
  table_options.cache_index_and_filter_blocks = true;
  table_options.filter_policy.reset(NewBloomFilterPolicy(10));

  Options options = CurrentOptions();
  options.disable_auto_compactions = true;
  options.max_open_files = -1;
  options.num_levels = 2;
  options.optimize_filters_for_hits = false;
  options.statistics = CreateDBStatistics();
  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
  DestroyAndReopen(options);

  const auto write_file = [&](const std::string& file_path,
                              const std::string& value,
                              ExternalSstFileInfo* file_info) {
    SstFileWriter writer(EnvOptions(), options);
    ASSERT_OK(writer.Open(file_path));
    ASSERT_OK(writer.Put(Key(10), value));
    ASSERT_OK(writer.Finish(file_info));
  };

  const auto ingest_file = [&](const std::string& file_path,
                               const ExternalSstFileInfo& file_info,
                               bool fail_if_not_bottommost_level) {
    IngestExternalFileArg arg;
    arg.column_family = db_->DefaultColumnFamily();
    arg.external_files = {file_path};
    arg.file_infos = {file_info.prepared_file_info.get()};
    arg.options.fail_if_not_bottommost_level = fail_if_not_bottommost_level;
    arg.options.verify_checksums_before_ingest = false;
    arg.options.prefetch_lmax_index_and_filter_blocks = false;
    ASSERT_OK(db_->IngestExternalFiles({arg}));
  };

  const std::string lmax_file_path = sst_files_dir_ + "lazy_lmax.sst";
  ExternalSstFileInfo lmax_file_info;
  write_file(lmax_file_path, "v10", &lmax_file_info);

  ASSERT_EQ(0, options.statistics->getAndResetTickerCount(
                   Tickers::BLOCK_CACHE_INDEX_ADD));
  ASSERT_EQ(0, options.statistics->getAndResetTickerCount(
                   Tickers::BLOCK_CACHE_FILTER_ADD));

  ingest_file(lmax_file_path, lmax_file_info,
              true /* fail_if_not_bottommost_level */);
  ASSERT_EQ("0,1", FilesPerLevel());

  ASSERT_EQ(0, options.statistics->getAndResetTickerCount(
                   Tickers::BLOCK_CACHE_INDEX_ADD));
  ASSERT_EQ(0, options.statistics->getAndResetTickerCount(
                   Tickers::BLOCK_CACHE_FILTER_ADD));

  const std::string l0_file_path = sst_files_dir_ + "l0.sst";
  ExternalSstFileInfo l0_file_info;
  write_file(l0_file_path, "v11", &l0_file_info);

  ASSERT_EQ(0, options.statistics->getAndResetTickerCount(
                   Tickers::BLOCK_CACHE_INDEX_ADD));
  ASSERT_EQ(0, options.statistics->getAndResetTickerCount(
                   Tickers::BLOCK_CACHE_FILTER_ADD));

  ingest_file(l0_file_path, l0_file_info,
              false /* fail_if_not_bottommost_level */);
  ASSERT_EQ("1,1", FilesPerLevel());

  EXPECT_GT(options.statistics->getAndResetTickerCount(
                Tickers::BLOCK_CACHE_INDEX_ADD),
            0);
  EXPECT_GT(options.statistics->getAndResetTickerCount(
                Tickers::BLOCK_CACHE_FILTER_ADD),
            0);
}

TEST_F(ExternalSSTFileTest, AbortPreparedIngestion) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);

  std::string file_path;
  ASSERT_OK(GenerateExternalFileOnly(options, {{"x", "1"}}, &file_path));
  IngestExternalFileArg arg;
  arg.column_family = db_->DefaultColumnFamily();
  arg.external_files = {file_path};

  std::unique_ptr<FileIngestionHandle> handle;
  ASSERT_OK(db_->PrepareFileIngestion({arg}, &handle));
  ASSERT_OK(handle->Abort());
  ASSERT_EQ("NOT_FOUND", Get("x"));

  // Committing an already-aborted handle returns an error.
  ASSERT_TRUE(
      db_->CommitFileIngestionHandle(std::move(handle)).IsInvalidArgument());

  // The DB remains healthy after an abort: a fresh prepare + commit works.
  std::unique_ptr<FileIngestionHandle> handle2;
  ASSERT_OK(db_->PrepareFileIngestion({arg}, &handle2));
  ASSERT_OK(db_->CommitFileIngestionHandle(std::move(handle2)));
  ASSERT_EQ("1", Get("x"));
}

TEST_F(ExternalSSTFileTest, MultiHandleAtomicCommit) {
  Options options = CurrentOptions();
  CreateAndReopenWithCF({"cf1", "cf2", "cf3"}, options);

  // Two handles whose column families overlap on cf2:
  //   handle 1 -> cf1, cf2     handle 2 -> cf2, cf3
  // Committing them together installs cf1 and cf3 each from one handle, while
  // cf2 receives files from BOTH handles, merged into a single per-CF commit.
  std::string f_cf1;
  std::string f_cf2_h1;
  std::string f_cf2_h2;
  std::string f_cf3;
  ASSERT_OK(GenerateExternalFileOnly(options, {{"a", "1"}}, &f_cf1));
  // The cf2 files share key "b": handle 2 is committed later, so its value must
  // win via a higher assigned sequence number.
  ASSERT_OK(GenerateExternalFileOnly(options, {{"b", "h1"}, {"b1", "10"}},
                                     &f_cf2_h1));
  ASSERT_OK(GenerateExternalFileOnly(options, {{"b", "h2"}, {"b2", "20"}},
                                     &f_cf2_h2));
  ASSERT_OK(GenerateExternalFileOnly(options, {{"c", "3"}}, &f_cf3));

  std::vector<IngestExternalFileArg> args1(2);
  args1[0].column_family = handles_[1];  // cf1
  args1[0].external_files = {f_cf1};
  args1[1].column_family = handles_[2];  // cf2
  args1[1].external_files = {f_cf2_h1};

  std::vector<IngestExternalFileArg> args2(2);
  args2[0].column_family =
      handles_[2];  // cf2 -> merges with handle 1's cf2 job
  args2[0].external_files = {f_cf2_h2};
  args2[1].column_family = handles_[3];  // cf3
  args2[1].external_files = {f_cf3};

  std::unique_ptr<FileIngestionHandle> h1;
  std::unique_ptr<FileIngestionHandle> h2;
  ASSERT_OK(db_->PrepareFileIngestion(args1, &h1));
  ASSERT_OK(db_->PrepareFileIngestion(args2, &h2));
  ASSERT_EQ("NOT_FOUND", Get(1, "a"));
  ASSERT_EQ("NOT_FOUND", Get(2, "b1"));
  ASSERT_EQ("NOT_FOUND", Get(3, "c"));

  // Commit both handles atomically.
  std::vector<std::unique_ptr<FileIngestionHandle>> batch;
  batch.push_back(std::move(h1));
  batch.push_back(std::move(h2));
  ASSERT_OK(db_->CommitFileIngestionHandles(std::move(batch)));

  ASSERT_EQ("1", Get(1, "a"));
  ASSERT_EQ("h2", Get(2, "b"));   // overlapping key: later handle (h2) wins
  ASSERT_EQ("10", Get(2, "b1"));  // cf2 from handle 1
  ASSERT_EQ("20", Get(2, "b2"));  // cf2 from handle 2 (merged)
  ASSERT_EQ("3", Get(3, "c"));
}

TEST_F(ExternalSSTFileTest, MultiHandleSameColumnFamilyIncompatibleOptions) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);

  std::string f1;
  std::string f2;
  ASSERT_OK(GenerateExternalFileOnly(options, {{"a", "1"}}, &f1));
  ASSERT_OK(GenerateExternalFileOnly(options, {{"b", "2"}}, &f2));

  // Same column family but with options the merged Run() would consult set
  // differently across the two handles.
  IngestExternalFileOptions opts_fill;
  IngestExternalFileOptions opts_no_fill;
  opts_no_fill.fill_cache = false;
  std::unique_ptr<FileIngestionHandle> h1;
  std::unique_ptr<FileIngestionHandle> h2;
  ASSERT_OK(db_->PrepareFileIngestion(db_->DefaultColumnFamily(), {f1},
                                      opts_fill, &h1));
  ASSERT_OK(db_->PrepareFileIngestion(db_->DefaultColumnFamily(), {f2},
                                      opts_no_fill, &h2));

  std::vector<std::unique_ptr<FileIngestionHandle>> batch;
  batch.push_back(std::move(h1));
  batch.push_back(std::move(h2));
  ASSERT_TRUE(
      db_->CommitFileIngestionHandles(std::move(batch)).IsInvalidArgument());

  // The rejected commit rolled both handles back; the CF is still usable.
  ASSERT_EQ("NOT_FOUND", Get("a"));
  ASSERT_EQ("NOT_FOUND", Get("b"));
  std::unique_ptr<FileIngestionHandle> h3;
  ASSERT_OK(db_->PrepareFileIngestion(db_->DefaultColumnFamily(), {f1},
                                      IngestExternalFileOptions(), &h3));
  ASSERT_OK(db_->CommitFileIngestionHandle(std::move(h3)));
  ASSERT_EQ("1", Get("a"));
}

TEST_F(ExternalSSTFileTest, PrepareFileIngestionValidationErrors) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);

  std::unique_ptr<FileIngestionHandle> handle;

  ASSERT_TRUE(
      db_->PrepareFileIngestion(std::vector<IngestExternalFileArg>{}, nullptr)
          .IsInvalidArgument());

  // Empty arg list.
  ASSERT_TRUE(
      db_->PrepareFileIngestion(std::vector<IngestExternalFileArg>{}, &handle)
          .IsInvalidArgument());
  ASSERT_TRUE(handle == nullptr);

  // move_files and link_files are mutually exclusive.
  std::string file_path;
  ASSERT_OK(GenerateExternalFileOnly(options, {{"k", "v"}}, &file_path));
  IngestExternalFileArg arg;
  arg.column_family = db_->DefaultColumnFamily();
  arg.external_files = {file_path};
  arg.options.move_files = true;
  arg.options.link_files = true;
  ASSERT_TRUE(db_->PrepareFileIngestion({arg}, &handle).IsInvalidArgument());
  ASSERT_TRUE(handle == nullptr);
}

// Ingestion that reuses the SstFileWriter's metadata via
// IngestExternalFileArg::file_infos skips the open-and-scan path and still
// produces correct data, including overlapping data that is reassigned a new
// global sequence number.
TEST_F(ExternalSSTFileTest, IngestWithFileInfo) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);
  std::map<std::string, std::string> true_data;

  // Count entries into the open-and-scan path so we can assert it is skipped.
  std::atomic<int> read_path_count{0};
  SyncPoint::GetInstance()->SetCallBack(
      "ExternalSstFileIngestionJob::GetIngestedFileInfo:ReadPath",
      [&](void*) { read_path_count.fetch_add(1); });
  SyncPoint::GetInstance()->EnableProcessing();

  // Reuse the writer's metadata, with the normal ingestion options: a global
  // sequence number is assigned (allow_db_generated_files is NOT set) and the
  // checksum is verified (the fast-path opens the file only to verify it, which
  // is not a metadata scan).
  auto ingest_fi = [&](std::vector<std::pair<std::string, std::string>> data) {
    return GenerateAndAddExternalFile(
        options, std::move(data), /*file_id=*/-1, /*allow_global_seqno=*/true,
        /*write_global_seqno=*/false, /*verify_checksums_before_ingest=*/true,
        /*ingest_behind=*/false, /*sort_data=*/true, &true_data,
        /*cfh=*/nullptr, /*fill_cache=*/false, /*ingest_with_file_info=*/true);
  };

  // Non-overlapping ingest: data reads back and the file is not scanned.
  ASSERT_OK(ingest_fi({{Key(1), "a1"}, {Key(2), "a2"}, {Key(3), "a3"}}));
  ASSERT_EQ(read_path_count.load(), 0);

  // Overlap existing (flushed) data so a non-zero global seqno is assigned and
  // the newer values win -- the normal seqno-reassignment path through
  // file_infos.
  ASSERT_OK(db_->Put(WriteOptions(), Key(2), "old2"));
  true_data[Key(2)] = "old2";
  ASSERT_OK(Flush());
  ASSERT_OK(ingest_fi({{Key(2), "new2"}, {Key(4), "new4"}}));
  ASSERT_EQ(read_path_count.load(), 0);

  // Control: ingesting the same way but WITHOUT file_infos does enter the scan.
  ASSERT_OK(GenerateAndAddExternalFile(
      options, {{Key(8), "c8"}}, /*file_id=*/-1, /*allow_global_seqno=*/true,
      /*write_global_seqno=*/false, /*verify_checksums_before_ingest=*/true,
      /*ingest_behind=*/false, /*sort_data=*/true, &true_data));
  ASSERT_GT(read_path_count.load(), 0);

  SyncPoint::GetInstance()->DisableProcessing();
  SyncPoint::GetInstance()->ClearAllCallBacks();
  size_t kcnt = 0;
  VerifyDBFromMap(true_data, &kcnt, false);
}

// Range deletions are honored on the file_infos fast-path: the writer's
// range-del bounds are carried in the metadata, so a DeleteRange covering
// existing keys still shadows them.
TEST_F(ExternalSSTFileTest, IngestWithFileInfoRangeDeletion) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);
  ASSERT_OK(db_->Put(WriteOptions(), Key(20), "live20"));
  ASSERT_OK(db_->Put(WriteOptions(), Key(25), "live25"));
  ASSERT_OK(Flush());

  const std::string f = sst_files_dir_ + "rangedel.sst";
  SstFileWriter w(EnvOptions(), options);
  ASSERT_OK(w.Open(f));
  ASSERT_OK(w.Put(Key(30), "v30"));
  ASSERT_OK(w.DeleteRange(Key(20), Key(26)));  // covers keys 20..25
  ExternalSstFileInfo file_info;
  ASSERT_OK(w.Finish(&file_info));

  IngestExternalFileArg arg;
  arg.column_family = db_->DefaultColumnFamily();
  arg.external_files = {f};
  arg.file_infos = {file_info.prepared_file_info.get()};
  arg.options.allow_global_seqno = true;
  ASSERT_OK(db_->IngestExternalFiles({arg}));

  ReadOptions ro;
  std::string val;
  ASSERT_TRUE(db_->Get(ro, Key(20), &val).IsNotFound());
  ASSERT_TRUE(db_->Get(ro, Key(25), &val).IsNotFound());
  ASSERT_OK(db_->Get(ro, Key(30), &val));
  ASSERT_EQ(val, "v30");
}

// verify_checksums_before_ingest is honored on the file_infos fast-path: it
// opens the file just to verify, so corruption is still detected.
TEST_F(ExternalSSTFileTest, IngestWithFileInfoVerifiesChecksum) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);

  const std::string f = sst_files_dir_ + "corrupt.sst";
  SstFileWriter w(EnvOptions(), options);
  ASSERT_OK(w.Open(f));
  for (int k = 0; k < 100; k++) {
    ASSERT_OK(w.Put(Key(k), "v" + Key(k)));
  }
  ExternalSstFileInfo file_info;
  ASSERT_OK(w.Finish(&file_info));
  ASSERT_OK(test::CorruptFile(options.env, f, /*offset=*/64,
                              /*bytes_to_corrupt=*/8));

  IngestExternalFileArg arg;
  arg.column_family = db_->DefaultColumnFamily();
  arg.external_files = {f};
  arg.file_infos = {file_info.prepared_file_info.get()};
  arg.options.verify_checksums_before_ingest = true;
  ASSERT_TRUE(db_->IngestExternalFiles({arg}).IsCorruption());
}

// write_global_seqno is incompatible with file_infos: the file is not opened,
// so the seqno cannot be written back into it.
TEST_F(ExternalSSTFileTest, IngestWithFileInfoRejectsWriteGlobalSeqno) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);

  const std::string f = sst_files_dir_ + "wgs.sst";
  SstFileWriter w(EnvOptions(), options);
  ASSERT_OK(w.Open(f));
  ASSERT_OK(w.Put(Key(1), "v1"));
  ExternalSstFileInfo file_info;
  ASSERT_OK(w.Finish(&file_info));

  IngestExternalFileArg arg;
  arg.column_family = db_->DefaultColumnFamily();
  arg.external_files = {f};
  arg.file_infos = {file_info.prepared_file_info.get()};
  arg.options.write_global_seqno = true;
  ASSERT_TRUE(db_->IngestExternalFiles({arg}).IsInvalidArgument());
}

TEST_F(ExternalSSTFileTest, GetPreparedFileInfoForExternalSstIngestion) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);

  const std::string source_dbname = dbname_ + "_source";
  ASSERT_OK(DestroyDB(source_dbname, options));
  std::unique_ptr<DB> source_db;
  ASSERT_OK(DB::Open(options, source_dbname, &source_db));

  ASSERT_OK(source_db->Put(WriteOptions(), Key(10), "v10"));
  ASSERT_OK(source_db->Flush(FlushOptions()));

  std::vector<LiveFileMetaData> live_meta;
  source_db->GetLiveFilesMetaData(&live_meta);
  ASSERT_EQ(live_meta.size(), 1);
  ASSERT_GT(live_meta[0].largest_seqno, 0);
  const std::string file_path =
      live_meta[0].directory + "/" + live_meta[0].relative_filename;

  std::shared_ptr<const PreparedFileInfo> prepared_file_info;
  ASSERT_TRUE(source_db
                  ->GetPreparedFileInfoForExternalSstIngestion(
                      live_meta[0].relative_filename, &prepared_file_info)
                  .IsInvalidArgument());
  ASSERT_EQ(prepared_file_info, nullptr);
  prepared_file_info = std::make_shared<PreparedFileInfo>();
  ASSERT_TRUE(source_db
                  ->GetPreparedFileInfoForExternalSstIngestion(
                      dbname_ + "_wrong/" + live_meta[0].relative_filename,
                      &prepared_file_info)
                  .IsInvalidArgument());
  ASSERT_EQ(prepared_file_info, nullptr);

  ASSERT_OK(source_db->GetPreparedFileInfoForExternalSstIngestion(
      file_path, &prepared_file_info));
  ASSERT_TRUE(prepared_file_info != nullptr);
  ASSERT_EQ(prepared_file_info->file_size, live_meta[0].size);
  ASSERT_EQ(prepared_file_info->table_properties.key_largest_seqno,
            live_meta[0].largest_seqno);

  IngestExternalFileArg arg;
  arg.column_family = db_->DefaultColumnFamily();
  arg.external_files = {file_path};
  arg.file_infos = {prepared_file_info.get()};
  arg.options.allow_db_generated_files = true;
  arg.options.snapshot_consistency = false;

  SyncPoint::GetInstance()->SetCallBack(
      "ExternalSstFileIngestionJob::GetIngestedFileInfo:ReadPath",
      [](void*) { FAIL() << "Prepared file info should skip the read path"; });
  SyncPoint::GetInstance()->EnableProcessing();

  ASSERT_OK(db_->IngestExternalFiles({arg}));

  ASSERT_EQ(Get(Key(10)), "v10");

  ASSERT_OK(source_db->Close());
  source_db.reset();
  ASSERT_OK(DestroyDB(source_dbname, options));

  SyncPoint::GetInstance()->DisableProcessing();
  SyncPoint::GetInstance()->ClearAllCallBacks();
}

TEST_F(ExternalSSTFileTest,
       IngestWithFileInfoRejectsMissingDbGeneratedSeqnoBounds) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);

  const std::string f = sst_files_dir_ + "missing_seqno_bounds.sst";
  SstFileWriter w(EnvOptions(), options);
  ASSERT_OK(w.Open(f));
  ASSERT_OK(w.Put(Key(1), "v1"));
  ExternalSstFileInfo file_info;
  ASSERT_OK(w.Finish(&file_info));
  auto prepared_file_info = file_info.prepared_file_info;

  auto missing_seqno_bounds =
      std::make_shared<PreparedFileInfo>(*prepared_file_info);
  missing_seqno_bounds->table_properties.key_smallest_seqno = UINT64_MAX;
  missing_seqno_bounds->table_properties.key_largest_seqno = UINT64_MAX;

  IngestExternalFileArg arg;
  arg.column_family = db_->DefaultColumnFamily();
  arg.external_files = {f};
  arg.file_infos = {missing_seqno_bounds.get()};
  arg.options.allow_db_generated_files = true;
  arg.options.snapshot_consistency = false;
  Status s = db_->IngestExternalFiles({arg});
  ASSERT_TRUE(s.IsCorruption());
  ASSERT_NE(s.ToString().find("Unknown largest seqno"), std::string::npos);
}

TEST_F(ExternalSSTFileTest, IngestWithFileInfoPreparedKeyBoundTypes) {
  Options options = CurrentOptions();
  DestroyAndReopen(options);

  const std::string f = sst_files_dir_ + "prepared_key_bound_types.sst";
  SstFileWriter w(EnvOptions(), options);
  ASSERT_OK(w.Open(f));
  ASSERT_OK(w.Put(Key(10), "v10"));
  ASSERT_OK(w.DeleteRange(Key(20), Key(30)));
  ExternalSstFileInfo file_info;
  ASSERT_OK(w.Finish(&file_info));
  auto prepared_file_info = file_info.prepared_file_info;

  ParsedInternalKey smallest;
  ASSERT_OK(ParseInternalKey(prepared_file_info->smallest.Encode(), &smallest,
                             false /* log_err_key */));
  ASSERT_EQ(0, smallest.sequence);
  ASSERT_EQ(kTypeValue, smallest.type);
  ASSERT_EQ(Key(10), smallest.user_key.ToString());

  ParsedInternalKey largest;
  ASSERT_OK(ParseInternalKey(prepared_file_info->largest.Encode(), &largest,
                             false /* log_err_key */));
  ASSERT_EQ(kMaxSequenceNumber, largest.sequence);
  ASSERT_EQ(kTypeRangeDeletion, largest.type);
  ASSERT_EQ(Key(30), largest.user_key.ToString());
}

TEST_F(ExternalSSTFileTest, ComparatorMismatch) {
  Options options = CurrentOptions();
  Options options_diff_ucmp = options;

  options.comparator = BytewiseComparator();
  options_diff_ucmp.comparator = ReverseBytewiseComparator();

  SstFileWriter sst_file_writer(EnvOptions(), options_diff_ucmp);

  std::string file = sst_files_dir_ + "file.sst";
  ASSERT_OK(sst_file_writer.Open(file));
  ASSERT_OK(sst_file_writer.Put("foo", "val"));
  ASSERT_OK(sst_file_writer.Put("bar", "val1"));
  ASSERT_OK(sst_file_writer.Finish());

  DestroyAndReopen(options);
  ASSERT_NOK(DeprecatedAddFile({file}));
}

TEST_F(ExternalSSTFileTest, NoBlockCache) {
  LRUCacheOptions co;
  co.capacity = 32 << 20;
  std::shared_ptr<Cache> cache = NewLRUCache(co);
  BlockBasedTableOptions table_options;
  table_options.block_cache = cache;
  table_options.filter_policy.reset(NewBloomFilterPolicy(10));
  table_options.cache_index_and_filter_blocks = true;
  Options options = CurrentOptions();
  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
  Reopen(options);

  size_t usage_before_ingestion = cache->GetUsage();
  std::map<std::string, std::string> true_data;
  // Ingest with fill_cache = true
  ASSERT_OK(GenerateAndAddExternalFile(options, {1, 2}, -1, false, false, true,
                                       false, false, &true_data, nullptr,
                                       /*fill_cache=*/true));
  ASSERT_EQ(FilesPerLevel(), "0,0,0,0,0,0,1");
  EXPECT_GT(cache->GetUsage(), usage_before_ingestion);

  TablePropertiesCollection tp;
  ASSERT_OK(db_->GetPropertiesOfAllTables(&tp));
  for (const auto& entry : tp) {
    EXPECT_GT(entry.second->index_size, 0);
    EXPECT_GT(entry.second->filter_size, 0);
  }

  usage_before_ingestion = cache->GetUsage();
  // Ingest with fill_cache = false
  ASSERT_OK(GenerateAndAddExternalFile(options, {3, 4}, -1, false, false, true,
                                       false, false, &true_data, nullptr,
                                       /*fill_cache=*/false));
  EXPECT_EQ(usage_before_ingestion, cache->GetUsage());