@nxtedition/rocksdb 13.1.4 → 13.2.0

package/deps/rocksdb/rocksdb/utilities/transactions/write_committed_transaction_ts_test.cc

@@ -418,6 +418,21 @@ TEST_P(WriteCommittedTxnWithTsTest, RecoverFromWal) {
   ASSERT_OK(txn3->Prepare());
   txn3.reset();
 
+  std::unique_ptr<Transaction> txn4(
+      NewTxn(WriteOptions(), TransactionOptions()));
+  assert(txn4);
+  ASSERT_OK(txn4->SetName("no_op_txn"));
+  txn4.reset();
+
+  std::unique_ptr<Transaction> rolled_back_txn(
+      NewTxn(write_opts, TransactionOptions()));
+  ASSERT_NE(nullptr, rolled_back_txn);
+  ASSERT_OK(rolled_back_txn->Put("non_exist0", "donotcare"));
+  ASSERT_OK(rolled_back_txn->Put(handles_[1], "non_exist1", "donotcare"));
+  ASSERT_OK(rolled_back_txn->SetName("rolled_back_txn"));
+  ASSERT_OK(rolled_back_txn->Rollback());
+  rolled_back_txn.reset();
+
   ASSERT_OK(ReOpenNoDelete(cf_descs, &handles_));
 
   {
@@ -452,9 +467,318 @@ TEST_P(WriteCommittedTxnWithTsTest, RecoverFromWal) {
 
     s = GetFromDb(ReadOptions(), handles_[1], "baz", /*ts=*/24, &value);
     ASSERT_TRUE(s.IsNotFound());
+
+    Transaction* no_op_txn = db->GetTransactionByName("no_op_txn");
+    ASSERT_EQ(nullptr, no_op_txn);
+
+    s = db->Get(ReadOptions(), handles_[0], "non_exist0", &value);
+    ASSERT_TRUE(s.IsNotFound());
+
+    s = GetFromDb(ReadOptions(), handles_[1], "non_exist1", /*ts=*/24, &value);
+    ASSERT_TRUE(s.IsNotFound());
+  }
+}
+
+TEST_P(WriteCommittedTxnWithTsTest, EnabledUDTDisabledRecoverFromWal) {
+  // This feature is not compatible with UDT in memtable only.
+  options.allow_concurrent_memtable_write = false;
+  ASSERT_OK(ReOpenNoDelete());
+
+  ColumnFamilyOptions cf_opts;
+  cf_opts.comparator = test::BytewiseComparatorWithU64TsWrapper();
+  cf_opts.persist_user_defined_timestamps = false;
+  const std::string test_cf_name = "test_cf";
+  ColumnFamilyHandle* cfh = nullptr;
+  assert(db);
+  ASSERT_OK(db->CreateColumnFamily(cf_opts, test_cf_name, &cfh));
+  delete cfh;
+  cfh = nullptr;
+
+  std::vector<ColumnFamilyDescriptor> cf_descs;
+  cf_descs.emplace_back(kDefaultColumnFamilyName, options);
+  cf_descs.emplace_back(test_cf_name, cf_opts);
+  options.avoid_flush_during_shutdown = true;
+  ASSERT_OK(ReOpenNoDelete(cf_descs, &handles_));
+
+  std::unique_ptr<Transaction> no_op_txn(
+      NewTxn(WriteOptions(), TransactionOptions()));
+  ASSERT_NE(nullptr, no_op_txn);
+  ASSERT_OK(no_op_txn->SetName("no_op_txn"));
+  no_op_txn.reset();
+
+  std::unique_ptr<Transaction> prepared_but_uncommitted_txn(
+      NewTxn(WriteOptions(), TransactionOptions()));
+  ASSERT_NE(nullptr, prepared_but_uncommitted_txn);
+  ASSERT_OK(prepared_but_uncommitted_txn->Put("foo0", "foo_value_0"));
+  ASSERT_OK(
+      prepared_but_uncommitted_txn->Put(handles_[1], "foo1", "foo_value_1"));
+  ASSERT_OK(
+      prepared_but_uncommitted_txn->SetName("prepared_but_uncommitted_txn"));
+  ASSERT_OK(prepared_but_uncommitted_txn->Prepare());
+
+  prepared_but_uncommitted_txn.reset();
+
+  WriteOptions write_opts;
+  write_opts.sync = true;
+  std::unique_ptr<Transaction> committed_txn(
+      NewTxn(write_opts, TransactionOptions()));
+  ASSERT_NE(nullptr, committed_txn);
+  ASSERT_OK(committed_txn->Put("bar0", "bar_value_0"));
+  ASSERT_OK(committed_txn->Put(handles_[1], "bar1", "bar_value_1"));
+  ASSERT_OK(committed_txn->SetName("committed_txn"));
+  ASSERT_OK(committed_txn->Prepare());
+  ASSERT_OK(committed_txn->SetCommitTimestamp(/*ts=*/23));
+  ASSERT_OK(committed_txn->Commit());
+  committed_txn.reset();
+
+  std::unique_ptr<Transaction> committed_without_prepare_txn(
+      NewTxn(write_opts, TransactionOptions()));
+  ASSERT_NE(nullptr, committed_without_prepare_txn);
+  ASSERT_OK(committed_without_prepare_txn->Put("baz0", "baz_value_0"));
+  ASSERT_OK(
+      committed_without_prepare_txn->Put(handles_[1], "baz1", "baz_value_1"));
+  ASSERT_OK(
+      committed_without_prepare_txn->SetName("committed_without_prepare_txn"));
+  ASSERT_OK(committed_without_prepare_txn->SetCommitTimestamp(/*ts=*/24));
+  ASSERT_OK(committed_without_prepare_txn->Commit());
+  committed_without_prepare_txn.reset();
+
+  std::unique_ptr<Transaction> rolled_back_txn(
+      NewTxn(write_opts, TransactionOptions()));
+  assert(rolled_back_txn);
+  ASSERT_OK(rolled_back_txn->Put("non_exist0", "donotcare"));
+  ASSERT_OK(rolled_back_txn->Put(handles_[1], "non_exist1", "donotcare"));
+  ASSERT_OK(rolled_back_txn->SetName("rolled_back_txn"));
+  ASSERT_OK(rolled_back_txn->Rollback());
+  rolled_back_txn.reset();
+
+  // Reopen and disable UDT to replay WAL entries.
+  cf_descs[1].options.comparator = BytewiseComparator();
+  ASSERT_OK(ReOpenNoDelete(cf_descs, &handles_));
+
+  {
+    Transaction* recovered_txn0 = db->GetTransactionByName("no_op_txn");
+    ASSERT_EQ(nullptr, recovered_txn0);
+
+    Transaction* recovered_txn1 =
+        db->GetTransactionByName("prepared_but_uncommitted_txn");
+    ASSERT_NE(nullptr, recovered_txn1);
+    std::string value;
+    ASSERT_OK(recovered_txn1->Commit());
+    Status s = db->Get(ReadOptions(), handles_[0], "foo0", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("foo_value_0", value);
+    s = db->Get(ReadOptions(), handles_[1], "foo1", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("foo_value_1", value);
+    delete recovered_txn1;
+
+    ASSERT_EQ(nullptr, db->GetTransactionByName("committed_txn"));
+    s = db->Get(ReadOptions(), handles_[0], "bar0", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("bar_value_0", value);
+    s = db->Get(ReadOptions(), handles_[1], "bar1", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("bar_value_1", value);
+
+    ASSERT_EQ(nullptr,
+              db->GetTransactionByName("committed_without_prepare_txn"));
+    s = db->Get(ReadOptions(), handles_[0], "baz0", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("baz_value_0", value);
+    s = db->Get(ReadOptions(), handles_[1], "baz1", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("baz_value_1", value);
+
+    ASSERT_EQ(nullptr, db->GetTransactionByName("rolled_back_txn"));
+    s = db->Get(ReadOptions(), handles_[0], "non_exist0", &value);
+    ASSERT_TRUE(s.IsNotFound());
+    s = db->Get(ReadOptions(), handles_[1], "non_exist1", &value);
+    ASSERT_TRUE(s.IsNotFound());
+  }
+}
+
+TEST_P(WriteCommittedTxnWithTsTest, UDTNewlyEnabledRecoverFromWal) {
+  ASSERT_OK(ReOpenNoDelete());
+
+  ColumnFamilyOptions cf_opts;
+  const std::string test_cf_name = "test_cf";
+  ColumnFamilyHandle* cfh = nullptr;
+  assert(db);
+  ASSERT_OK(db->CreateColumnFamily(cf_opts, test_cf_name, &cfh));
+  delete cfh;
+  cfh = nullptr;
+
+  std::vector<ColumnFamilyDescriptor> cf_descs;
+  cf_descs.emplace_back(kDefaultColumnFamilyName, options);
+  cf_descs.emplace_back(test_cf_name, cf_opts);
+  options.avoid_flush_during_shutdown = true;
+  ASSERT_OK(ReOpenNoDelete(cf_descs, &handles_));
+
+  std::unique_ptr<Transaction> no_op_txn(
+      NewTxn(WriteOptions(), TransactionOptions()));
+  ASSERT_NE(nullptr, no_op_txn);
+  ASSERT_OK(no_op_txn->SetName("no_op_txn"));
+  no_op_txn.reset();
+
+  std::unique_ptr<Transaction> prepared_but_uncommitted_txn(
+      NewTxn(WriteOptions(), TransactionOptions()));
+  ASSERT_NE(nullptr, prepared_but_uncommitted_txn);
+  ASSERT_OK(
+      prepared_but_uncommitted_txn->Put(handles_[0], "foo0", "foo_value_0"));
+  ASSERT_OK(
+      prepared_but_uncommitted_txn->Put(handles_[1], "foo1", "foo_value_1"));
+  ASSERT_OK(
+      prepared_but_uncommitted_txn->SetName("prepared_but_uncommitted_txn"));
+  ASSERT_OK(prepared_but_uncommitted_txn->Prepare());
+
+  prepared_but_uncommitted_txn.reset();
+
+  WriteOptions write_opts;
+  write_opts.sync = true;
+  std::unique_ptr<Transaction> committed_txn(
+      NewTxn(write_opts, TransactionOptions()));
+  ASSERT_NE(nullptr, committed_txn);
+  ASSERT_OK(committed_txn->Put("bar0", "bar_value_0"));
+  ASSERT_OK(committed_txn->Put(handles_[1], "bar1", "bar_value_1"));
+  ASSERT_OK(committed_txn->SetName("committed_txn"));
+  ASSERT_OK(committed_txn->Prepare());
+  ASSERT_OK(committed_txn->Commit());
+  committed_txn.reset();
+
+  std::unique_ptr<Transaction> committed_without_prepare_txn(
+      NewTxn(write_opts, TransactionOptions()));
+  assert(committed_without_prepare_txn);
+  ASSERT_OK(committed_without_prepare_txn->Put("baz0", "baz_value_0"));
+  ASSERT_OK(
+      committed_without_prepare_txn->Put(handles_[1], "baz1", "baz_value_1"));
+  ASSERT_OK(
+      committed_without_prepare_txn->SetName("committed_without_prepare_txn"));
+  ASSERT_OK(committed_without_prepare_txn->Commit());
+  committed_without_prepare_txn.reset();
+
+  std::unique_ptr<Transaction> rolled_back_txn(
+      NewTxn(write_opts, TransactionOptions()));
+  ASSERT_NE(nullptr, rolled_back_txn);
+  ASSERT_OK(rolled_back_txn->Put("non_exist0", "donotcare"));
+  ASSERT_OK(rolled_back_txn->Put(handles_[1], "non_exist1", "donotcare"));
+  ASSERT_OK(rolled_back_txn->SetName("rolled_back_txn"));
+  ASSERT_OK(rolled_back_txn->Rollback());
+  rolled_back_txn.reset();
+
+  // Reopen and enable UDT to replay WAL entries.
+  options.allow_concurrent_memtable_write = false;
+  cf_descs[1].options.comparator = test::BytewiseComparatorWithU64TsWrapper();
+  cf_descs[1].options.persist_user_defined_timestamps = false;
+  ASSERT_OK(ReOpenNoDelete(cf_descs, &handles_));
+
+  {
+    Transaction* recovered_txn1 =
+        db->GetTransactionByName("prepared_but_uncommitted_txn");
+    ASSERT_NE(nullptr, recovered_txn1);
+    std::string value;
+    ASSERT_OK(recovered_txn1->SetCommitTimestamp(23));
+    ASSERT_OK(recovered_txn1->Commit());
+    Status s = db->Get(ReadOptions(), handles_[0], "foo0", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("foo_value_0", value);
+    s = GetFromDb(ReadOptions(), handles_[1], "foo1", /*ts=*/23, &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("foo_value_1", value);
+    delete recovered_txn1;
+
+    ASSERT_EQ(nullptr, db->GetTransactionByName("committed_txn"));
+    s = db->Get(ReadOptions(), handles_[0], "bar0", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("bar_value_0", value);
+    s = GetFromDb(ReadOptions(), handles_[1], "bar1", /*ts=*/23, &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("bar_value_1", value);
+
+    ASSERT_EQ(nullptr,
+              db->GetTransactionByName("committed_without_prepare_txn"));
+    s = db->Get(ReadOptions(), handles_[0], "baz0", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("baz_value_0", value);
+    s = GetFromDb(ReadOptions(), handles_[1], "baz1", /*ts=*/23, &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("baz_value_1", value);
+
+    ASSERT_EQ(nullptr, db->GetTransactionByName("rolled_back_txn"));
+    s = db->Get(ReadOptions(), handles_[0], "non_exist0", &value);
+    ASSERT_TRUE(s.IsNotFound());
+    s = GetFromDb(ReadOptions(), handles_[1], "non_exist1", /*ts=*/23, &value);
+    ASSERT_TRUE(s.IsNotFound());
   }
 }
 
+TEST_P(WriteCommittedTxnWithTsTest, ChangeFromWriteCommittedAndDisableUDT) {
+  // This feature is not compatible with UDT in memtable only.
+  options.allow_concurrent_memtable_write = false;
+  ASSERT_OK(ReOpenNoDelete());
+
+  ColumnFamilyOptions cf_opts;
+  cf_opts.comparator = test::BytewiseComparatorWithU64TsWrapper();
+  cf_opts.persist_user_defined_timestamps = false;
+  const std::string test_cf_name = "test_cf";
+  ColumnFamilyHandle* cfh = nullptr;
+  assert(db);
+  ASSERT_OK(db->CreateColumnFamily(cf_opts, test_cf_name, &cfh));
+  delete cfh;
+  cfh = nullptr;
+
+  std::vector<ColumnFamilyDescriptor> cf_descs;
+  cf_descs.emplace_back(kDefaultColumnFamilyName, options);
+  cf_descs.emplace_back(test_cf_name, cf_opts);
+  options.avoid_flush_during_shutdown = true;
+  ASSERT_OK(ReOpenNoDelete(cf_descs, &handles_));
+
+  std::unique_ptr<Transaction> prepared_but_uncommitted_txn(
+      NewTxn(WriteOptions(), TransactionOptions()));
+  assert(prepared_but_uncommitted_txn);
+  ASSERT_OK(prepared_but_uncommitted_txn->Put("foo0", "foo_value_0"));
+  ASSERT_OK(
+      prepared_but_uncommitted_txn->Put(handles_[1], "foo1", "foo_value_1"));
+  ASSERT_OK(
+      prepared_but_uncommitted_txn->SetName("prepared_but_uncommitted_txn"));
+  ASSERT_OK(prepared_but_uncommitted_txn->Prepare());
+
+  prepared_but_uncommitted_txn.reset();
+
+  WriteOptions write_opts;
+  write_opts.sync = true;
+  std::unique_ptr<Transaction> committed_txn(
+      NewTxn(write_opts, TransactionOptions()));
+  assert(committed_txn);
+  ASSERT_OK(committed_txn->Put("bar0", "bar_value_0"));
+  ASSERT_OK(committed_txn->Put(handles_[1], "bar1", "bar_value_1"));
+  ASSERT_OK(committed_txn->SetName("committed_txn"));
+  ASSERT_OK(committed_txn->Prepare());
+  ASSERT_OK(committed_txn->SetCommitTimestamp(/*ts=*/23));
+  ASSERT_OK(committed_txn->Commit());
+  committed_txn.reset();
+
+  std::unique_ptr<Transaction> committed_without_prepare_txn(
+      NewTxn(write_opts, TransactionOptions()));
+  assert(committed_without_prepare_txn);
+  ASSERT_OK(committed_without_prepare_txn->Put("baz0", "baz_value_0"));
+  ASSERT_OK(
+      committed_without_prepare_txn->Put(handles_[1], "baz1", "baz_value_1"));
+  ASSERT_OK(
+      committed_without_prepare_txn->SetName("committed_without_prepare_txn"));
+  ASSERT_OK(committed_without_prepare_txn->SetCommitTimestamp(/*ts=*/24));
+  ASSERT_OK(committed_without_prepare_txn->Commit());
+  committed_without_prepare_txn.reset();
+
+  // Disable UDT and change write policy.
+  cf_descs[1].options.comparator = BytewiseComparator();
+  txn_db_options.write_policy = TxnDBWritePolicy::WRITE_PREPARED;
+  ASSERT_NOK(ReOpenNoDelete(cf_descs, &handles_));
+
+  txn_db_options.write_policy = TxnDBWritePolicy::WRITE_UNPREPARED;
+  ASSERT_NOK(ReOpenNoDelete(cf_descs, &handles_));
+}
+
 TEST_P(WriteCommittedTxnWithTsTest, TransactionDbLevelApi) {
   ASSERT_OK(ReOpenNoDelete());
 

package/deps/rocksdb/rocksdb/utilities/transactions/write_prepared_txn.cc

@@ -3,13 +3,13 @@
 //  COPYING file in the root directory) and Apache 2.0 License
 //  (found in the LICENSE.Apache file in the root directory).
 
-
 #include "utilities/transactions/write_prepared_txn.h"
 
 #include <cinttypes>
 #include <map>
 #include <set>
 
+#include "db/attribute_group_iterator_impl.h"
 #include "db/column_family.h"
 #include "db/db_impl/db_impl.h"
 #include "rocksdb/db.h"
@@ -135,6 +135,23 @@ Iterator* WritePreparedTxn::GetIterator(const ReadOptions& options,
   return write_batch_.NewIteratorWithBase(column_family, db_iter, &options);
 }
 
+std::unique_ptr<Iterator> WritePreparedTxn::GetCoalescingIterator(
+    const ReadOptions& /* read_options */,
+    const std::vector<ColumnFamilyHandle*>& /* column_families */) {
+  return std::unique_ptr<Iterator>(NewErrorIterator(
+      Status::NotSupported("GetCoalescingIterator not supported for "
+                           "write-prepared/write-unprepared transactions")));
+}
+
+std::unique_ptr<AttributeGroupIterator>
+WritePreparedTxn::GetAttributeGroupIterator(
+    const ReadOptions& /* read_options */,
+    const std::vector<ColumnFamilyHandle*>& /* column_families */) {
+  return NewAttributeGroupErrorIterator(
+      Status::NotSupported("GetAttributeGroupIterator not supported for "
+                           "write-prepared/write-unprepared transactions"));
+}
+
 Status WritePreparedTxn::PrepareInternal() {
   WriteOptions write_options = write_options_;
   write_options.disableWAL = false;

package/deps/rocksdb/rocksdb/utilities/transactions/write_prepared_txn.h

@@ -5,7 +5,6 @@
 
 #pragma once
 
-
 #include <algorithm>
 #include <atomic>
 #include <mutex>
@@ -71,6 +70,14 @@ class WritePreparedTxn : public PessimisticTransaction {
   Iterator* GetIterator(const ReadOptions& options,
                         ColumnFamilyHandle* column_family) override;
 
+  std::unique_ptr<Iterator> GetCoalescingIterator(
+      const ReadOptions& read_options,
+      const std::vector<ColumnFamilyHandle*>& column_families) override;
+
+  std::unique_ptr<AttributeGroupIterator> GetAttributeGroupIterator(
+      const ReadOptions& read_options,
+      const std::vector<ColumnFamilyHandle*>& column_families) override;
+
   void SetSnapshot() override;
 
  protected:

package/deps/rocksdb/rocksdb/utilities/write_batch_with_index/write_batch_with_index.cc

@@ -35,7 +35,8 @@ struct WriteBatchWithIndex::Rep {
         overwrite_key(_overwrite_key),
         last_entry_offset(0),
         last_sub_batch_offset(0),
-        sub_batch_cnt(1) {}
+        sub_batch_cnt(1),
+        track_cf_stat(false) {}
   ReadableWriteBatch write_batch;
   WriteBatchEntryComparator comparator;
   Arena arena;
@@ -49,6 +50,11 @@ struct WriteBatchWithIndex::Rep {
   // Total number of sub-batches in the write batch. Default is 1.
   size_t sub_batch_cnt;
 
+  bool track_cf_stat;
+  // Tracks ids of CFs that have updates in this WBWI, number of updates and
+  // number of overwritten single deletions per cf.
+  std::unordered_map<uint32_t, CFStat> cf_id_to_stat;
+
   // Remember current offset of internal write batch, which is used as
   // the starting offset of the next record.
   void SetLastEntryOffset() { last_entry_offset = write_batch.GetDataSize(); }
@@ -69,7 +75,7 @@ struct WriteBatchWithIndex::Rep {
 
   // Allocate an index entry pointing to the last entry in the write batch and
   // put it to skip list.
-  void AddNewEntry(uint32_t column_family_id);
+  void AddNewEntry(uint32_t column_family_id, WriteType type);
 
   // Clear all updates buffered in this batch.
   void Clear();
@@ -114,6 +120,16 @@ bool WriteBatchWithIndex::Rep::UpdateExistingEntryWithCfId(
     last_sub_batch_offset = last_entry_offset;
     sub_batch_cnt++;
   }
+  if (track_cf_stat) {
+    if (non_const_entry->has_single_del &&
+        !non_const_entry->has_overwritten_single_del) {
+      cf_id_to_stat[column_family_id].overwritten_sd_count++;
+      non_const_entry->has_overwritten_single_del = true;
+    }
+    if (type == kSingleDeleteRecord) {
+      non_const_entry->has_single_del = true;
+    }
+  }
   if (type == kMergeRecord) {
     return false;
   } else {
@@ -130,18 +146,19 @@ void WriteBatchWithIndex::Rep::AddOrUpdateIndex(
     if (cf_cmp != nullptr) {
       comparator.SetComparatorForCF(cf_id, cf_cmp);
     }
-    AddNewEntry(cf_id);
+    AddNewEntry(cf_id, type);
   }
 }
 
 void WriteBatchWithIndex::Rep::AddOrUpdateIndex(const Slice& key,
                                                 WriteType type) {
   if (!UpdateExistingEntryWithCfId(0, key, type)) {
-    AddNewEntry(0);
+    AddNewEntry(0, type);
   }
 }
 
-void WriteBatchWithIndex::Rep::AddNewEntry(uint32_t column_family_id) {
+void WriteBatchWithIndex::Rep::AddNewEntry(uint32_t column_family_id,
+                                           WriteType type) {
   const std::string& wb_data = write_batch.Data();
   Slice entry_ptr = Slice(wb_data.data() + last_entry_offset,
                           wb_data.size() - last_entry_offset);
@@ -166,10 +183,18 @@ void WriteBatchWithIndex::Rep::AddNewEntry(uint32_t column_family_id) {
       new (mem) WriteBatchIndexEntry(last_entry_offset, column_family_id,
                                      key.data() - wb_data.data(), key.size());
   skip_list.Insert(index_entry);
+
+  if (track_cf_stat) {
+    if (type == kSingleDeleteRecord) {
+      index_entry->has_single_del = true;
+    }
+    cf_id_to_stat[column_family_id].entry_count++;
+  }
 }
 
 void WriteBatchWithIndex::Rep::Clear() {
   write_batch.Clear();
+  cf_id_to_stat.clear();
   ClearIndex();
 }
 
@@ -220,7 +245,7 @@ Status WriteBatchWithIndex::Rep::ReBuildIndex() {
       case kTypeValue:
         found++;
         if (!UpdateExistingEntryWithCfId(column_family_id, key, kPutRecord)) {
-          AddNewEntry(column_family_id);
+          AddNewEntry(column_family_id, kPutRecord);
         }
         break;
       case kTypeColumnFamilyDeletion:
@@ -228,7 +253,7 @@ Status WriteBatchWithIndex::Rep::ReBuildIndex() {
         found++;
         if (!UpdateExistingEntryWithCfId(column_family_id, key,
                                          kDeleteRecord)) {
-          AddNewEntry(column_family_id);
+          AddNewEntry(column_family_id, kDeleteRecord);
         }
         break;
       case kTypeColumnFamilySingleDeletion:
@@ -236,14 +261,14 @@ Status WriteBatchWithIndex::Rep::ReBuildIndex() {
         found++;
         if (!UpdateExistingEntryWithCfId(column_family_id, key,
                                          kSingleDeleteRecord)) {
-          AddNewEntry(column_family_id);
+          AddNewEntry(column_family_id, kSingleDeleteRecord);
         }
         break;
       case kTypeColumnFamilyMerge:
       case kTypeMerge:
         found++;
         if (!UpdateExistingEntryWithCfId(column_family_id, key, kMergeRecord)) {
-          AddNewEntry(column_family_id);
+          AddNewEntry(column_family_id, kMergeRecord);
         }
         break;
       case kTypeLogData:
@@ -261,7 +286,7 @@ Status WriteBatchWithIndex::Rep::ReBuildIndex() {
         found++;
         if (!UpdateExistingEntryWithCfId(column_family_id, key,
                                          kPutEntityRecord)) {
-          AddNewEntry(column_family_id);
+          AddNewEntry(column_family_id, kPutEntityRecord);
         }
         break;
       case kTypeColumnFamilyValuePreferredSeqno:
@@ -313,6 +338,11 @@ WBWIIterator* WriteBatchWithIndex::NewIterator(
                               &(rep->comparator));
 }
 
+WBWIIterator* WriteBatchWithIndex::NewIterator(uint32_t cf_id) const {
+  return new WBWIIteratorImpl(cf_id, &(rep->skip_list), &rep->write_batch,
+                              &(rep->comparator));
+}
+
 Iterator* WriteBatchWithIndex::NewIteratorWithBase(
     ColumnFamilyHandle* column_family, Iterator* base_iterator,
     const ReadOptions* read_options) {
@@ -329,7 +359,8 @@ Iterator* WriteBatchWithIndex::NewIteratorWithBase(
   }
 
   return new BaseDeltaIterator(column_family, base_iterator, wbwiii,
-                               GetColumnFamilyUserComparator(column_family));
+                               GetColumnFamilyUserComparator(column_family),
+                               read_options);
 }
 
 Iterator* WriteBatchWithIndex::NewIteratorWithBase(Iterator* base_iterator) {
@@ -337,7 +368,8 @@ Iterator* WriteBatchWithIndex::NewIteratorWithBase(Iterator* base_iterator) {
   auto wbwiii = new WBWIIteratorImpl(0, &(rep->skip_list), &rep->write_batch,
                                      &rep->comparator);
   return new BaseDeltaIterator(nullptr, base_iterator, wbwiii,
-                               rep->comparator.default_comparator());
+                               rep->comparator.default_comparator(),
+                               /* read_options */ nullptr);
 }
 
 Status WriteBatchWithIndex::Put(ColumnFamilyHandle* column_family,
@@ -1131,4 +1163,17 @@ const Comparator* WriteBatchWithIndexInternal::GetUserComparator(
   return ucmps.GetComparator(cf_id);
 }
 
+void WriteBatchWithIndex::SetTrackPerCFStat(bool track) {
+  // Should be set when the wbwi contains no update.
+  assert(GetWriteBatch()->Count() == 0);
+  rep->track_cf_stat = track;
+}
+
+const std::unordered_map<uint32_t, WriteBatchWithIndex::CFStat>&
+WriteBatchWithIndex::GetCFStats() const {
+  assert(rep->track_cf_stat);
+  return rep->cf_id_to_stat;
+}
+
+bool WriteBatchWithIndex::GetOverwriteKey() const { return rep->overwrite_key; }
 }  // namespace ROCKSDB_NAMESPACE

package/deps/rocksdb/rocksdb/utilities/write_batch_with_index/write_batch_with_index_internal.cc

@@ -22,10 +22,13 @@ namespace ROCKSDB_NAMESPACE {
 BaseDeltaIterator::BaseDeltaIterator(ColumnFamilyHandle* column_family,
                                      Iterator* base_iterator,
                                      WBWIIteratorImpl* delta_iterator,
-                                     const Comparator* comparator)
+                                     const Comparator* comparator,
+                                     const ReadOptions* read_options)
     : forward_(true),
       current_at_base_(true),
       equal_keys_(false),
+      allow_unprepared_value_(
+          read_options ? read_options->allow_unprepared_value : false),
       status_(Status::OK()),
       column_family_(column_family),
       base_iterator_(base_iterator),
@@ -262,6 +265,17 @@ void BaseDeltaIterator::SetValueAndColumnsFromBase() {
   assert(value_.empty());
   assert(columns_.empty());
 
+  if (!base_iterator_->PrepareValue()) {
+    assert(!BaseValid());
+    assert(!base_iterator_->status().ok());
+
+    Invalidate(base_iterator_->status());
+
+    assert(!Valid());
+    assert(!status().ok());
+    return;
+  }
+
   value_ = base_iterator_->value();
   columns_ = base_iterator_->columns();
 }
@@ -314,6 +328,17 @@ void BaseDeltaIterator::SetValueAndColumnsFromDelta() {
         /* result_operand */ nullptr, &result_type);
   } else if (delta_entry.type == kMergeRecord) {
     if (equal_keys_) {
+      if (!base_iterator_->PrepareValue()) {
+        assert(!BaseValid());
+        assert(!base_iterator_->status().ok());
+
+        Invalidate(base_iterator_->status());
+
+        assert(!Valid());
+        assert(!status().ok());
+        return;
+      }
+
       if (WideColumnsHelper::HasDefaultColumnOnly(base_iterator_->columns())) {
         status_ = WriteBatchWithIndexInternal::MergeKeyWithBaseValue(
             column_family_, delta_entry.key, MergeHelper::kPlainBaseValue,
@@ -402,7 +427,9 @@ void BaseDeltaIterator::UpdateCurrent() {
     } else if (!DeltaValid()) {
       // Delta has finished.
       current_at_base_ = true;
-      SetValueAndColumnsFromBase();
+      if (!allow_unprepared_value_) {
+        SetValueAndColumnsFromBase();
+      }
       return;
     } else {
       int compare =
@@ -426,7 +453,9 @@ void BaseDeltaIterator::UpdateCurrent() {
         }
       } else {
         current_at_base_ = true;
-        SetValueAndColumnsFromBase();
+        if (!allow_unprepared_value_) {
+          SetValueAndColumnsFromBase();
+        }
         return;
       }
     }