Experimental API IngestWriteBatchWithIndex() (facebook#13550)

Summary:
add support for ingesting a WriteBatchWithIndex into the DB with the new API `IngestWriteBatchWithIndex()`. This ingestion works similarly as `TransactionOptions::commit_bypass_memtable` where the WBWI will be ingested as an immutable memtable. Since this skips memtable writes, it improves the write performance when writing a large write batch into the DB. Currently this API only supports `disableWAL=true`. Support for WAL write will be in a follow up if needed.

For a WBWI to be ingestable, we needed to call `SetTrackPerCFStat()` at WBWI creation. This PR removes this step for simpler usage and per CF stats will always be tracked in WBWI. `WBWIIteratorImpl::TestOutOfBound()` is optimized to offset the performance impact.

Pull Request resolved: facebook#13550

Test Plan:
- new unit test
- stress test option ingest_wbwi_one_in and ran a few runs of `python3 ./tools/db_crashtest.py blackbox --enable_pipelined_write=0 --use_timed_put_one_in=0 --use_put_entity_one_in=0 --ingest_wbwi_one_in=10 --test_batches_snapshots=0 --enable_blob_files=0 --preserve_unverified_changes=1 --avoid_flush_during_recovery=1 --disable_wal=1 --inplace_update_support=0 --interval=40`

Reviewed By: jowlyzhang

Differential Revision: D73152223

Pulled By: cbi42

fbshipit-source-id: 339f8ed26ac5a798238870df3ba857ba1add759b

Author: cbi42

db/db_impl/db_impl.h

@@ -256,6 +256,10 @@ class DBImpl : public DB {
   Status WriteWithCallback(const WriteOptions& options, WriteBatch* updates,
                            UserWriteCallback* user_write_cb) override;
 
+  Status IngestWriteBatchWithIndex(
+      const WriteOptions& options,
+      std::shared_ptr<WriteBatchWithIndex> wbwi) override;
+
   using DB::Get;
   Status Get(const ReadOptions& _read_options,
              ColumnFamilyHandle* column_family, const Slice& key,
@@ -1531,11 +1535,11 @@ class DBImpl : public DB {
   // ingests `wbwi` is done.
   // @param memtable_updated Whether the same write that ingests wbwi has
   // updated memtable. This is useful for determining whether to set bg
-  // error when IngestWBWI fails.
-  Status IngestWBWI(std::shared_ptr<WriteBatchWithIndex> wbwi,
-                    const WBWIMemTable::SeqnoRange& assigned_seqno,
-                    uint64_t min_prep_log, SequenceNumber last_seqno,
-                    bool memtable_updated, bool ignore_missing_cf);
+  // error when IngestWBWIAsMemtable fails.
+  Status IngestWBWIAsMemtable(std::shared_ptr<WriteBatchWithIndex> wbwi,
+                              const WBWIMemTable::SeqnoRange& assigned_seqno,
+                              uint64_t min_prep_log, SequenceNumber last_seqno,
+                              bool memtable_updated, bool ignore_missing_cf);
 
   // If disable_memtable is set the application logic must guarantee that the
   // batch will still be skipped from memtable during the recovery. An excption

db/db_impl/db_impl_write.cc

@@ -190,11 +190,38 @@ Status DBImpl::WriteWithCallback(const WriteOptions& write_options,
   return s;
 }
 
-Status DBImpl::IngestWBWI(std::shared_ptr<WriteBatchWithIndex> wbwi,
-                          const WBWIMemTable::SeqnoRange& assigned_seqno,
-                          uint64_t min_prep_log,
-                          SequenceNumber last_seqno_after_ingest,
-                          bool memtable_updated, bool ignore_missing_cf) {
+Status DBImpl::IngestWriteBatchWithIndex(
+    const WriteOptions& write_options,
+    std::shared_ptr<WriteBatchWithIndex> wbwi) {
+  if (!wbwi) {
+    return Status::InvalidArgument("Batch is nullptr!");
+  }
+  if (!write_options.disableWAL) {
+    return Status::NotSupported(
+        "IngestWriteBatchWithIndex does not support disableWAL=true");
+  }
+  Status s;
+  if (write_options.protection_bytes_per_key > 0) {
+    s = WriteBatchInternal::UpdateProtectionInfo(
+        wbwi->GetWriteBatch(), write_options.protection_bytes_per_key);
+  }
+  if (s.ok()) {
+    WriteBatch dummy_empty_batch;
+    s = WriteImpl(
+        write_options, /*updates=*/&dummy_empty_batch, /*callback=*/nullptr,
+        /*user_write_cb=*/nullptr, /*log_used=*/nullptr, /*log_ref=*/0,
+        /*disable_memtable=*/false, /*seq_used=*/nullptr,
+        /*batch_cnt=*/0, /*pre_release_callback=*/nullptr,
+        /*post_memtable_callback=*/nullptr, /*wbwi=*/wbwi);
+  }
+  return s;
+}
+
+Status DBImpl::IngestWBWIAsMemtable(
+    std::shared_ptr<WriteBatchWithIndex> wbwi,
+    const WBWIMemTable::SeqnoRange& assigned_seqno, uint64_t min_prep_log,
+    SequenceNumber last_seqno_after_ingest, bool memtable_updated,
+    bool ignore_missing_cf) {
   // Keys in new memtable have seqno > last_seqno_after_ingest >= keys in wbwi.
   assert(assigned_seqno.upper_bound <= last_seqno_after_ingest);
   // Keys in the current memtable have seqno <= LastSequence() < keys in wbwi.
@@ -436,9 +463,17 @@ Status DBImpl::WriteImpl(const WriteOptions& write_options,
     return Status::NotSupported(
         "DeleteRange is not compatible with row cache.");
   }
+  // Whether the WBWI is from transaction commit or a direct write
+  // (IngestWriteBatchWithIndex())
+  bool ingest_wbwi_for_commit = false;
   if (wbwi) {
-    assert(log_ref > 0);
-    // Used only in WriteCommittedTxn::CommitInternal() with no `callback`.
+    if (my_batch->HasCommit()) {
+      ingest_wbwi_for_commit = true;
+      assert(log_ref);
+    } else {
+      // Only supports disableWAL for directly ingesting WBWI for now.
+      assert(write_options.disableWAL);
+    }
     assert(!callback);
     if (immutable_db_options_.unordered_write) {
       return Status::NotSupported(
@@ -448,6 +483,10 @@ Status DBImpl::WriteImpl(const WriteOptions& write_options,
       return Status::NotSupported(
           "Ingesting WriteBatch does not support pipelined_write");
     }
+    if (!wbwi->GetOverwriteKey()) {
+      return Status::NotSupported(
+          "WriteBatchWithIndex ingestion requires overwrite_key=true");
+    }
   }
   // Otherwise IsLatestPersistentState optimization does not make sense
   assert(!WriteBatchInternal::IsLatestPersistentState(my_batch) ||
@@ -658,7 +697,13 @@ Status DBImpl::WriteImpl(const WriteOptions& write_options,
             continue;
           }
           // TODO: maybe handle the tracing status?
-          tracer_->Write(writer->batch).PermitUncheckedError();
+          if (wbwi && !ingest_wbwi_for_commit) {
+            // for transaction write, tracer only needs the commit marker which
+            // is in writer->batch
+            tracer_->Write(wbwi->GetWriteBatch()).PermitUncheckedError();
+          } else {
+            tracer_->Write(writer->batch).PermitUncheckedError();
+          }
         }
       }
     }
@@ -860,12 +905,13 @@ Status DBImpl::WriteImpl(const WriteOptions& write_options,
     // handle exit, false means somebody else did
     should_exit_batch_group = write_thread_.CompleteParallelMemTableWriter(&w);
   }
-  if (wbwi) {
-    if (status.ok() && w.status.ok()) {
+  if (wbwi && status.ok() && w.status.ok()) {
+    uint32_t wbwi_count = wbwi->GetWriteBatch()->Count();
+    // skip empty batch case
+    if (wbwi_count) {
       // w.batch contains (potentially empty) commit time batch updates,
       // only ingest wbwi if w.batch is applied to memtable successfully
       uint32_t memtable_update_count = w.batch->Count();
-      uint32_t wbwi_count = wbwi->GetWriteBatch()->Count();
       // Seqno assigned to this write are [last_seq + 1 - seq_inc, last_seq].
       // seq_inc includes w.batch (memtable updates) and wbwi
       // w.batch gets first `memtable_update_count` sequence numbers.
@@ -878,10 +924,11 @@ Status DBImpl::WriteImpl(const WriteOptions& write_options,
       if (two_write_queues_) {
         assert(ub <= versions_->LastAllocatedSequence());
       }
-      status = IngestWBWI(wbwi, {/*lower_bound=*/lb, /*upper_bound=*/ub},
-                          /*min_prep_log=*/log_ref, last_sequence,
-                          /*memtable_updated=*/memtable_update_count > 0,
-                          write_options.ignore_missing_column_families);
+      status =
+          IngestWBWIAsMemtable(wbwi, {/*lower_bound=*/lb, /*upper_bound=*/ub},
+                               /*min_prep_log=*/log_ref, last_sequence,
+                               /*memtable_updated=*/memtable_update_count > 0,
+                               write_options.ignore_missing_column_families);
     }
   }
 

db/db_write_test.cc

@@ -1000,6 +1000,80 @@ TEST_P(DBWriteTest, RecycleLogToggleTest) {
   ASSERT_EQ(Get(Key(1)), "val2");
 }
 
+TEST_P(DBWriteTest, IngestWriteBatchWithIndex) {
+  if (GetParam() == kPipelinedWrite) {
+    return;
+  }
+
+  Options options = GetOptions();
+  options.disable_auto_compactions = true;
+  Reopen(options);
+  Options cf_options = GetOptions();
+  cf_options.merge_operator = MergeOperators::CreateStringAppendOperator();
+  CreateColumnFamilies({"cf1", "cf2"}, cf_options);
+  ReopenWithColumnFamilies({"default", "cf1", "cf2"},
+                           {options, cf_options, cf_options});
+
+  // default cf
+  auto wbwi1 = std::make_shared<WriteBatchWithIndex>(options.comparator, 0,
+                                                     /*overwrite_key=*/true);
+  ASSERT_OK(wbwi1->Put("key1", "value1"));
+  ASSERT_OK(wbwi1->Put("key2", "value2"));
+  if (GetParam() == kPipelinedWrite) {
+    ASSERT_TRUE(db_->IngestWriteBatchWithIndex({}, wbwi1).IsNotSupported());
+    return;
+  }
+  // Test disableWAL=false
+  ASSERT_TRUE(db_->IngestWriteBatchWithIndex({}, wbwi1).IsNotSupported());
+
+  WriteOptions wo;
+  wo.disableWAL = true;
+  ASSERT_OK(db_->IngestWriteBatchWithIndex(wo, wbwi1));
+  ASSERT_EQ("value1", Get("key1"));
+  ASSERT_EQ("value2", Get("key2"));
+
+  // Test with overwrites
+  auto wbwi = std::make_shared<WriteBatchWithIndex>(options.comparator, 0,
+                                                    /*overwrite_key=*/true);
+  ASSERT_OK(wbwi->Put("key2", "value3"));
+  ASSERT_OK(wbwi->Delete("key1"));  // Delete an existing key
+  ASSERT_OK(db_->IngestWriteBatchWithIndex(wo, wbwi));
+  ASSERT_EQ("NOT_FOUND", Get("key1"));
+  ASSERT_EQ("value3", Get("key2"));
+
+  auto wbwi2 = std::make_shared<WriteBatchWithIndex>(options.comparator, 0,
+                                                     /*overwrite_key=*/true);
+  ASSERT_OK(wbwi2->Put(handles_[1], "cf1_key1", "cf1_value1"));
+  ASSERT_OK(wbwi2->Delete(handles_[1], "cf1_key2"));
+  // Test ingestion with column family
+  ASSERT_OK(db_->IngestWriteBatchWithIndex(wo, wbwi2));
+  ASSERT_EQ("cf1_value1", Get(1, "cf1_key1"));
+  ASSERT_EQ("NOT_FOUND", Get(1, "cf1_key2"));
+
+  auto wbwi3 = std::make_shared<WriteBatchWithIndex>(options.comparator, 0,
+                                                     /*overwrite_key=*/true);
+  ASSERT_OK(wbwi3->Merge(handles_[2], "cf2_key1", "cf2_value1"));
+  ASSERT_OK(wbwi3->Merge(handles_[2], "cf2_key1", "cf2_value2"));
+  // Test ingestion with merge operations
+  ASSERT_OK(db_->IngestWriteBatchWithIndex(wo, wbwi3));
+  ASSERT_EQ("cf2_value1,cf2_value2", Get(2, "cf2_key1"));
+
+  // Test with overwrite_key = false
+  auto wbwi_no_overwrite = std::make_shared<WriteBatchWithIndex>(
+      options.comparator, 0, /*overwrite_key=*/false);
+  ASSERT_OK(wbwi_no_overwrite->Put("key1", "value1"));
+  Status s = db_->IngestWriteBatchWithIndex(wo, wbwi_no_overwrite);
+  ASSERT_TRUE(s.IsNotSupported());
+
+  auto empty_wbwi = std::make_shared<WriteBatchWithIndex>(
+      options.comparator, 0, /*overwrite_key=*/true);
+  ASSERT_OK(db_->IngestWriteBatchWithIndex(wo, empty_wbwi));
+
+  DestroyAndReopen(options);
+  // Should fail when trying to ingest to non-existent column family
+  ASSERT_NOK(db_->IngestWriteBatchWithIndex(wo, wbwi2));
+}
+
 INSTANTIATE_TEST_CASE_P(DBWriteTestInstance, DBWriteTest,
                         testing::Values(DBTestBase::kDefault,
                                         DBTestBase::kConcurrentWALWrites,

db_stress_tool/db_stress_common.h

@@ -423,6 +423,7 @@ DECLARE_bool(track_and_verify_wals);
 DECLARE_bool(enable_remote_compaction);
 DECLARE_bool(auto_refresh_iterator_with_snapshot);
 DECLARE_uint32(memtable_op_scan_flush_trigger);
+DECLARE_uint32(ingest_wbwi_one_in);
 
 constexpr long KB = 1024;
 constexpr int kRandomValueMaxFactor = 3;

db_stress_tool/db_stress_gflags.cc

@@ -842,6 +842,11 @@ DEFINE_bool(track_and_verify_wals,
 DEFINE_bool(enable_remote_compaction, false,
             "Enable (simulated) Remote Compaction");
 
+DEFINE_uint32(ingest_wbwi_one_in, 0,
+              "If set, will call"
+              "IngestWriteBatchWithIndex() instead of regular write operations "
+              "once every N writes.");
+
 static bool ValidateInt32Percent(const char* flagname, int32_t value) {
   if (value < 0 || value > 100) {
     fprintf(stderr, "Invalid value for --%s: %d, 0<= pct <=100 \n", flagname,

db_stress_tool/no_batched_ops_stress.cc

@@ -1845,7 +1845,17 @@ class NonBatchedOpsStressTest : public StressTest {
       } else if (FLAGS_use_merge) {
         if (!FLAGS_use_txn) {
           if (FLAGS_user_timestamp_size == 0) {
-            s = db_->Merge(write_opts, cfh, k, v);
+            if (FLAGS_ingest_wbwi_one_in &&
+                thread->rand.OneIn(FLAGS_ingest_wbwi_one_in)) {
+              auto wbwi = std::make_shared<WriteBatchWithIndex>(
+                  options_.comparator, 0, /*overwrite_key=*/true);
+              s = wbwi->Merge(cfh, k, v);
+              if (s.ok()) {
+                s = db_->IngestWriteBatchWithIndex(write_opts, wbwi);
+              }
+            } else {
+              s = db_->Merge(write_opts, cfh, k, v);
+            }
           } else {
             s = db_->Merge(write_opts, cfh, k, write_ts, v);
           }
@@ -1857,7 +1867,17 @@ class NonBatchedOpsStressTest : public StressTest {
       } else {
         if (!FLAGS_use_txn) {
           if (FLAGS_user_timestamp_size == 0) {
-            s = db_->Put(write_opts, cfh, k, v);
+            if (FLAGS_ingest_wbwi_one_in &&
+                thread->rand.OneIn(FLAGS_ingest_wbwi_one_in)) {
+              auto wbwi = std::make_shared<WriteBatchWithIndex>(
+                  options_.comparator, 0, /*overwrite_key=*/true);
+              s = wbwi->Put(cfh, k, v);
+              if (s.ok()) {
+                s = db_->IngestWriteBatchWithIndex(write_opts, wbwi);
+              }
+            } else {
+              s = db_->Put(write_opts, cfh, k, v);
+            }
           } else {
             s = db_->Put(write_opts, cfh, k, write_ts, v);
           }
@@ -1949,7 +1969,17 @@ class NonBatchedOpsStressTest : public StressTest {
         }
         if (!FLAGS_use_txn) {
           if (FLAGS_user_timestamp_size == 0) {
-            s = db_->Delete(write_opts, cfh, key);
+            if (FLAGS_ingest_wbwi_one_in &&
+                thread->rand.OneIn(FLAGS_ingest_wbwi_one_in)) {
+              auto wbwi = std::make_shared<WriteBatchWithIndex>(
+                  options_.comparator, 0, /*overwrite_key=*/true);
+              s = wbwi->Delete(cfh, key);
+              if (s.ok()) {
+                s = db_->IngestWriteBatchWithIndex(write_opts, wbwi);
+              }
+            } else {
+              s = db_->Delete(write_opts, cfh, key);
+            }
           } else {
             s = db_->Delete(write_opts, cfh, key, write_ts);
           }
@@ -2006,7 +2036,17 @@ class NonBatchedOpsStressTest : public StressTest {
         }
         if (!FLAGS_use_txn) {
           if (FLAGS_user_timestamp_size == 0) {
-            s = db_->SingleDelete(write_opts, cfh, key);
+            if (FLAGS_ingest_wbwi_one_in &&
+                thread->rand.OneIn(FLAGS_ingest_wbwi_one_in)) {
+              auto wbwi = std::make_shared<WriteBatchWithIndex>(
+                  options_.comparator, 0, /*overwrite_key=*/true);
+              s = wbwi->SingleDelete(cfh, key);
+              if (s.ok()) {
+                s = db_->IngestWriteBatchWithIndex(write_opts, wbwi);
+              }
+            } else {
+              s = db_->SingleDelete(write_opts, cfh, key);
+            }
           } else {
             s = db_->SingleDelete(write_opts, cfh, key, write_ts);
           }

include/rocksdb/db.h

@@ -31,6 +31,7 @@
 #include "rocksdb/types.h"
 #include "rocksdb/user_write_callback.h"
 #include "rocksdb/utilities/table_properties_collectors.h"
+#include "rocksdb/utilities/write_batch_with_index.h"
 #include "rocksdb/version.h"
 #include "rocksdb/wide_columns.h"
 
@@ -633,6 +634,21 @@ class DB {
         "WriteWithCallback not implemented for this interface.");
   }
 
+  // EXPERIMENTAL, subject to change
+  // Ingest a WriteBatchWithIndex into DB, bypassing memtable writes for better
+  // write performance. Useful when there is a large number of updates
+  // in the write batch.
+  // The WriteBatchWithIndex must be created with overwrite_key=true.
+  // Currently this requires WriteOptions::disableWAL=true.
+  // The following options are currently not supported:
+  // - unordered_write
+  // - enable_pipelined_write
+  virtual Status IngestWriteBatchWithIndex(
+      const WriteOptions& /*options*/,
+      std::shared_ptr<WriteBatchWithIndex> /*wbwi*/) {
+    return Status::NotSupported("IngestWriteBatchWithIndex not implemented.");
+  }
+
   // If the column family specified by "column_family" contains an entry for
   // "key", return the corresponding value in "*value". If the entry is a plain
   // key-value, return the value as-is; if it is a wide-column entity, return

include/rocksdb/utilities/transaction_db.h

@@ -368,6 +368,8 @@ struct TransactionOptions {
   // Only supports write-committed policy. If set to true, the transaction will
   // skip memtable write and ingest into the DB directly during Commit(). This
   // makes Commit() much faster for transactions with many operations.
+  // Transaction neeeds to call Prepare() before Commit() for this option to
+  // take effect.
   // Transactions with Merge() or PutEntity() is not supported yet.
   //
   // Note that the transaction will be ingested as an immutable memtable for

include/rocksdb/utilities/write_batch_with_index.h

@@ -90,6 +90,8 @@ class WBWIIterator {
   // Returns n where the current entry is the n-th update to the current key.
   // The update count starts from 1.
   // Only valid if WBWI is created with overwrite_key = true.
+  // With overwrite_key=false, update count for each entry is not maintained,
+  // see UpdateExistingEntryWithCfId().
   virtual uint32_t GetUpdateCount() const { return 0; }
 };
 
@@ -374,9 +376,6 @@ class WriteBatchWithIndex : public WriteBatchBase {
     uint32_t entry_count = 0;
     uint32_t overwritten_sd_count = 0;
   };
-  // Will track CF ID, per CF entry count and overwritten sd count.
-  // Should be enabled when WBWI is empty for correct tracking.
-  void SetTrackPerCFStat(bool track);
   const std::unordered_map<uint32_t, CFStat>& GetCFStats() const;
 
   bool GetOverwriteKey() const;

memtable/wbwi_memtable.h

@@ -235,7 +235,7 @@ class WBWIMemTable final : public ReadOnlyMemTable {
   uint64_t num_entries_;
   // WBWI can contains updates to multiple CFs. `cf_id_` determines which CF
   // this memtable is for.
-  uint32_t cf_id_;
+  const uint32_t cf_id_;
 };
 
 class WBWIMemTableIterator final : public InternalIterator {