From 36838bbf51bbc8bfd0e9fe94cfb06503c51be096 Mon Sep 17 00:00:00 2001 From: Changyu Bi Date: Wed, 19 Feb 2025 12:35:56 -0800 Subject: [PATCH] Update sequence number assignment method in `WBWIMemTable` (#13400) Summary: This is a preparation for supporting merge in `WBWIMemTable`. This PR updates the sequence number assignment method so that it allows efficient and simple assignment when there are multiple entries with the same user key. This can happen when the WBWI contains Merge operations. This assignment relies on tracking the number of updates issued for each key in each WBWI entry (`WriteBatchIndexEntry::update_count`). Some refactoring is done in WBWI to remove `last_entry_offset` as part of the WBWI state which I find it harder to use correctly. Pull Request resolved: https://github.com/facebook/rocksdb/pull/13400 Test Plan: updated unit tests to check that update count is tracked correctly and WBWIMemTable is assigning sequence number as expected. Reviewed By: pdillinger Differential Revision: D69666462 Pulled By: cbi42 fbshipit-source-id: 9b18291825017a67c4da3318e8a556aa2971326b --- .../utilities/write_batch_with_index.h | 7 +- memtable/wbwi_memtable.cc | 4 + memtable/wbwi_memtable.h | 91 +++++---- .../write_batch_with_index.cc | 180 +++++++++--------- .../write_batch_with_index_internal.h | 13 +- .../write_batch_with_index_test.cc | 133 ++++++++++--- 6 files changed, 267 insertions(+), 161 deletions(-) diff --git a/include/rocksdb/utilities/write_batch_with_index.h b/include/rocksdb/utilities/write_batch_with_index.h index 9cb8632d18..6ff8b58709 100644 --- a/include/rocksdb/utilities/write_batch_with_index.h +++ b/include/rocksdb/utilities/write_batch_with_index.h @@ -77,6 +77,8 @@ class WBWIIterator { virtual void Prev() = 0; + virtual Status status() const = 0; + // The returned WriteEntry is only valid until the next mutation of // WriteBatchWithIndex. virtual WriteEntry Entry() const = 0; @@ -85,7 +87,10 @@ class WBWIIterator { // and it was overwritten by another update. virtual bool HasOverWrittenSingleDel() const { return false; } - virtual Status status() const = 0; + // 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. + virtual uint32_t GetUpdateCount() const { return 0; } }; // A WriteBatchWithIndex with a binary searchable index built for all the keys diff --git a/memtable/wbwi_memtable.cc b/memtable/wbwi_memtable.cc index 9ebd1b8d78..72f2f05bb7 100644 --- a/memtable/wbwi_memtable.cc +++ b/memtable/wbwi_memtable.cc @@ -62,6 +62,10 @@ bool WBWIMemTable::Get(const LookupKey& key, std::string* value, [[maybe_unused]] SequenceNumber read_seq = GetInternalKeySeqno(key.internal_key()); + // This is memtable is a single write batch, no snapshot can be taken within + // assigned seqnos for this memtable. + assert(read_seq >= assigned_seqno_.upper_bound || + read_seq < assigned_seqno_.lower_bound); std::unique_ptr iter{NewIterator()}; iter->Seek(key.internal_key()); const Slice lookup_user_key = key.user_key(); diff --git a/memtable/wbwi_memtable.h b/memtable/wbwi_memtable.h index 98b7c32418..c8e16910e4 100644 --- a/memtable/wbwi_memtable.h +++ b/memtable/wbwi_memtable.h @@ -13,12 +13,21 @@ namespace ROCKSDB_NAMESPACE { // a transaction (which is based on WBWI) into the DB as an immutable memtable. // // REQUIRE overwrite_key to be true for the WBWI -// Since the keys in WBWI do not have sequence number, the memtable needs to be +// Since the keys in WBWI do not have sequence number, this class is responsible +// for assigning sequence numbers to the keys. This memtable needs to be // assigned a range of sequence numbers through AssignSequenceNumbers(seqno) // before being available for reads. -// With overwrite_key = true, WBWI keeps track of the most recent update for -// each key, and each such key will be assigned seqno.upper_bound during reads. -// One exception is during flush, consider the following scenario: +// +// The sequence number assignment uses the update count for each key +// tracked in WBWI (see WBWIIterator::GetUpdateCount()). For each key, the +// sequence number assigned is seqno.lower_bound + update_count - 1. So more +// recent updates will have higher sequence number. +// +// WBWI with overwrite mode keeps track of the most recent update for each key, +// so this memtable contains one update per key usually. However, there is a +// special case where this memtable needs to emit an extra SingleDelete even +// when the SD is overwritten by another update. +// Consider the following scenario: // - WBWI has SD(k) then PUT(k, v1) // - DB has PUT(k, v2) in L1 // - flush WBWI adds PUT(k, v1) into L0 @@ -26,8 +35,7 @@ namespace ROCKSDB_NAMESPACE { // - flush live memtable and compact it with L0 will drop SD(k) and PUT(k, v1) // - the PUT(k, v2) in L1 incorrectly becomes visible // So during flush, iterator from this memtable will need emit overwritten -// single deletion. These single deletion entries will be -// assigned seqno.upper_bound - 1. +// single deletion. This SD will be assigned seqno.lower_bound. class WBWIMemTable final : public ReadOnlyMemTable { public: struct SeqnoRange { @@ -258,51 +266,40 @@ class WBWIMemTableIterator final : public InternalIterator { } void Seek(const Slice& target) override { + // `emit_overwritten_single_del_` is only used for flush, which does + // sequential forward scan from the beginning. + assert(!emit_overwritten_single_del_); Slice target_user_key = ExtractUserKey(target); + // Moves to first update >= target_user_key it_->Seek(target_user_key); - if (it_->Valid()) { - // compare seqno - SequenceNumber seqno = GetInternalKeySeqno(target); - assert(!emit_overwritten_single_del_); - // For now all keys are assigned seqno_ub_, this may change after merge - // is supported. - assert(seqno <= assigned_seqno_.lower_bound || - seqno >= assigned_seqno_.upper_bound); - if (seqno < assigned_seqno_.upper_bound && - comparator_->Compare(it_->Entry().key, target_user_key) == 0) { - it_->Next(); - // TODO: cannot assume distinct keys once Merge is supported - if (it_->Valid()) { - assert(comparator_->Compare(it_->Entry().key, target_user_key) > 0); - } - } + SequenceNumber target_seqno = GetInternalKeySeqno(target); + // Move to the first entry with seqno <= target_seqno for the same + // user key or a different user key. + while (it_->Valid() && + comparator_->Compare(it_->Entry().key, target_user_key) == 0 && + target_seqno < CurrentKeySeqno()) { + it_->Next(); } UpdateKey(); } void SeekForPrev(const Slice& target) override { + assert(!emit_overwritten_single_del_); Slice target_user_key = ExtractUserKey(target); + // Moves to last update <= target_user_key it_->SeekForPrev(target_user_key); - if (it_->Valid()) { - SequenceNumber seqno = GetInternalKeySeqno(target); - assert(seqno <= assigned_seqno_.lower_bound || - seqno >= assigned_seqno_.upper_bound); - if (seqno > assigned_seqno_.upper_bound && - comparator_->Compare(it_->Entry().key, target_user_key) == 0) { - it_->Prev(); - if (it_->Valid()) { - // TODO: cannot assume distinct keys once Merge is supported - assert(comparator_->Compare(it_->Entry().key, target_user_key) < 0); - } - } + SequenceNumber target_seqno = GetInternalKeySeqno(target); + // Move to the first entry with seqno >= target_seqno for the same + // user key or a different user key. + while (it_->Valid() && + comparator_->Compare(it_->Entry().key, target_user_key) == 0 && + CurrentKeySeqno() < target_seqno) { + it_->Prev(); } UpdateKey(); } void Next() override { - // Only need to emit single deletion during flush. Since Flush does - // sequential forward scan, we only need to emit single deletion in Next(), - // and do not need to consider iterator direction change. assert(Valid()); if (emit_overwritten_single_del_) { if (it_->HasOverWrittenSingleDel() && !at_overwritten_single_del_) { @@ -329,6 +326,7 @@ class WBWIMemTableIterator final : public InternalIterator { } void Prev() override { + assert(!emit_overwritten_single_del_); assert(Valid()); it_->Prev(); UpdateKey(); @@ -341,7 +339,6 @@ class WBWIMemTableIterator final : public InternalIterator { Slice value() const override { assert(Valid()); - // TODO: it_->Entry() is not trivial, cache it return it_->Entry().value; } @@ -355,6 +352,16 @@ class WBWIMemTableIterator final : public InternalIterator { private: static const std::unordered_map WriteTypeToValueTypeMap; + SequenceNumber CurrentKeySeqno() { + assert(it_->Valid()); + assert(it_->GetUpdateCount() >= 1); + auto seq = assigned_seqno_.lower_bound + it_->GetUpdateCount() - 1; + assert(seq <= assigned_seqno_.upper_bound); + return seq; + } + + // If it_ is valid, udate key_ to an internal key containing it_ current + // key, CurrentKeySeqno() and a type corresponding to it_ current entry type. void UpdateKey() { valid_ = it_->Valid(); if (!Valid()) { @@ -370,16 +377,16 @@ class WBWIMemTableIterator final : public InternalIterator { std::to_string(it_->Entry().type)); return; } - key_buf_.SetInternalKey(it_->Entry().key, assigned_seqno_.upper_bound, - t->second); + key_buf_.SetInternalKey(it_->Entry().key, CurrentKeySeqno(), t->second); key_ = key_buf_.GetInternalKey(); } void UpdateSingleDeleteKey() { assert(it_->Valid()); assert(Valid()); - assert(assigned_seqno_.lower_bound < assigned_seqno_.upper_bound); - key_buf_.SetInternalKey(it_->Entry().key, assigned_seqno_.upper_bound - 1, + // The key that overwrites this SingleDelete will be assigned at least + // seqno lower_bound + 1 (see CurrentKeySeqno()). + key_buf_.SetInternalKey(it_->Entry().key, assigned_seqno_.lower_bound, kTypeSingleDeletion); key_ = key_buf_.GetInternalKey(); at_overwritten_single_del_ = true; diff --git a/utilities/write_batch_with_index/write_batch_with_index.cc b/utilities/write_batch_with_index/write_batch_with_index.cc index e2cbb7d2b7..7a6540b5ac 100644 --- a/utilities/write_batch_with_index/write_batch_with_index.cc +++ b/utilities/write_batch_with_index/write_batch_with_index.cc @@ -15,8 +15,6 @@ #include "db/merge_helper.h" #include "db/wide/wide_columns_helper.h" #include "memory/arena.h" -#include "memtable/skiplist.h" -#include "options/db_options.h" #include "rocksdb/comparator.h" #include "rocksdb/iterator.h" #include "util/cast_util.h" @@ -32,17 +30,14 @@ struct WriteBatchWithIndex::Rep { index_comparator ? index_comparator->timestamp_size() : 0), comparator(index_comparator, &write_batch), skip_list(comparator, &arena), - overwrite_key(_overwrite_key), - last_entry_offset(0), last_sub_batch_offset(0), sub_batch_cnt(1), + overwrite_key(_overwrite_key), track_cf_stat(false) {} ReadableWriteBatch write_batch; WriteBatchEntryComparator comparator; Arena arena; WriteBatchEntrySkipList skip_list; - bool overwrite_key; - size_t last_entry_offset; // The starting offset of the last sub-batch. A sub-batch starts right before // inserting a key that is a duplicate of a key in the last sub-batch. Zero, // the default, means that no duplicate key is detected so far. @@ -50,32 +45,51 @@ struct WriteBatchWithIndex::Rep { // Total number of sub-batches in the write batch. Default is 1. size_t sub_batch_cnt; + bool overwrite_key; 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 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(); } - // In overwrite mode, find the existing entry for the same key and update it - // to point to the current entry. + // to point to the current entry if this is not a Merge operation. // Return true if the key is found and updated. - bool UpdateExistingEntry(ColumnFamilyHandle* column_family, const Slice& key, - WriteType type); + // most_recent_entry_update_count will be set to the update count of the + // most recent index entry for `key` if exists. bool UpdateExistingEntryWithCfId(uint32_t column_family_id, const Slice& key, - WriteType type); + WriteType type, size_t last_entry_offset, + uint32_t* most_recent_entry_update_count); - // Add the recent entry to the update. - // In overwrite mode, if key already exists in the index, update it. - void AddOrUpdateIndex(ColumnFamilyHandle* column_family, const Slice& key, - WriteType type); - void AddOrUpdateIndex(const Slice& key, WriteType type); + // Add or update the index for the given `key` of `type` + // at `last_entry_offset` in the write batch. + // + // The index is updated in place if this WBWI is in overwrite mode + // and this is not a Merge operation. + // + // If `cf_cmp` is provided and there is no compartor stored + // for this cf, it will be stored as the comparator to use for this cf. + void AddOrUpdateIndexWithCfId(uint32_t cf_id, const Slice& key, + WriteType type, size_t last_entry_offset, + const Comparator* cf_cmp = nullptr); - // 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, WriteType type); + void AddOrUpdateIndex(ColumnFamilyHandle* cfh, const Slice& key, + WriteType type, size_t last_entry_offset) { + AddOrUpdateIndexWithCfId(GetColumnFamilyID(cfh), key, type, + last_entry_offset, + GetColumnFamilyUserComparator(cfh)); + } + + void AddOrUpdateIndex(const Slice& key, WriteType type, + size_t last_entry_offset) { + AddOrUpdateIndexWithCfId(0, key, type, last_entry_offset); + } + + // Add a new index entry pointing to the the entry at `last_entry_offset` + // in the write batch. `most_recent_entry_update_count` will be used to + // initialize the update count of the new index entry. + void AddNewEntry(uint32_t column_family_id, WriteType type, + size_t last_entry_offset, + uint32_t most_recent_entry_update_count); // Clear all updates buffered in this batch. void Clear(); @@ -86,14 +100,10 @@ struct WriteBatchWithIndex::Rep { Status ReBuildIndex(); }; -bool WriteBatchWithIndex::Rep::UpdateExistingEntry( - ColumnFamilyHandle* column_family, const Slice& key, WriteType type) { - uint32_t cf_id = GetColumnFamilyID(column_family); - return UpdateExistingEntryWithCfId(cf_id, key, type); -} - bool WriteBatchWithIndex::Rep::UpdateExistingEntryWithCfId( - uint32_t column_family_id, const Slice& key, WriteType type) { + uint32_t column_family_id, const Slice& key, WriteType type, + size_t last_entry_offset, uint32_t* most_recent_entry_update_count) { + assert(most_recent_entry_update_count); if (!overwrite_key) { return false; } @@ -110,6 +120,8 @@ bool WriteBatchWithIndex::Rep::UpdateExistingEntryWithCfId( // key == `key`, which is the most recently update to `key`. WriteBatchIndexEntry* most_recent_entry = const_cast(iter.GetRawEntry()); + *most_recent_entry_update_count = most_recent_entry->update_count; + if (LIKELY(last_sub_batch_offset <= most_recent_entry->offset)) { last_sub_batch_offset = last_entry_offset; sub_batch_cnt++; @@ -127,32 +139,33 @@ bool WriteBatchWithIndex::Rep::UpdateExistingEntryWithCfId( if (type == kMergeRecord) { return false; } else { + // We still increment the update count when updating in-place. This is + // useful for WBWIMemTable when it needs to emit overwritten SingleDeletes. + // The overwritten SingleDelete will be assigned sequence number as if + // update_count is 0. So any later update should have update_count > 0. + most_recent_entry->update_count++; most_recent_entry->offset = last_entry_offset; return true; } } -void WriteBatchWithIndex::Rep::AddOrUpdateIndex( - ColumnFamilyHandle* column_family, const Slice& key, WriteType type) { - if (!UpdateExistingEntry(column_family, key, type)) { - uint32_t cf_id = GetColumnFamilyID(column_family); - const auto* cf_cmp = GetColumnFamilyUserComparator(column_family); - if (cf_cmp != nullptr) { +void WriteBatchWithIndex::Rep::AddOrUpdateIndexWithCfId( + uint32_t cf_id, const Slice& key, WriteType type, size_t last_entry_offset, + const Comparator* cf_cmp) { + uint32_t update_count = 0; + if (!UpdateExistingEntryWithCfId(cf_id, key, type, last_entry_offset, + &update_count)) { + if (cf_cmp) { comparator.SetComparatorForCF(cf_id, cf_cmp); } - AddNewEntry(cf_id, type); - } -} - -void WriteBatchWithIndex::Rep::AddOrUpdateIndex(const Slice& key, - WriteType type) { - if (!UpdateExistingEntryWithCfId(0, key, type)) { - AddNewEntry(0, type); + AddNewEntry(cf_id, type, last_entry_offset, update_count + 1); } } void WriteBatchWithIndex::Rep::AddNewEntry(uint32_t column_family_id, - WriteType type) { + WriteType type, + size_t last_entry_offset, + uint32_t update_count) { const std::string& wb_data = write_batch.Data(); Slice entry_ptr = Slice(wb_data.data() + last_entry_offset, wb_data.size() - last_entry_offset); @@ -173,9 +186,9 @@ void WriteBatchWithIndex::Rep::AddNewEntry(uint32_t column_family_id, } auto* mem = arena.Allocate(sizeof(WriteBatchIndexEntry)); - auto* index_entry = - new (mem) WriteBatchIndexEntry(last_entry_offset, column_family_id, - key.data() - wb_data.data(), key.size()); + auto* index_entry = new (mem) WriteBatchIndexEntry( + last_entry_offset, column_family_id, key.data() - wb_data.data(), + key.size(), update_count); skip_list.Insert(index_entry); if (track_cf_stat) { @@ -197,7 +210,6 @@ void WriteBatchWithIndex::Rep::ClearIndex() { arena.~Arena(); new (&arena) Arena(); new (&skip_list) WriteBatchEntrySkipList(comparator, &arena); - last_entry_offset = 0; last_sub_batch_offset = 0; sub_batch_cnt = 1; } @@ -226,7 +238,7 @@ Status WriteBatchWithIndex::Rep::ReBuildIndex() { char tag = 0; // set offset of current entry for call to AddNewEntry() - last_entry_offset = input.data() - write_batch.Data().data(); + size_t last_entry_offset = input.data() - write_batch.Data().data(); s = ReadRecordFromWriteBatch(&input, &tag, &column_family_id, &key, &value, &blob, &xid, &unix_write_time); @@ -238,32 +250,26 @@ Status WriteBatchWithIndex::Rep::ReBuildIndex() { case kTypeColumnFamilyValue: case kTypeValue: found++; - if (!UpdateExistingEntryWithCfId(column_family_id, key, kPutRecord)) { - AddNewEntry(column_family_id, kPutRecord); - } + AddOrUpdateIndexWithCfId(column_family_id, key, kPutRecord, + last_entry_offset); break; case kTypeColumnFamilyDeletion: case kTypeDeletion: found++; - if (!UpdateExistingEntryWithCfId(column_family_id, key, - kDeleteRecord)) { - AddNewEntry(column_family_id, kDeleteRecord); - } + AddOrUpdateIndexWithCfId(column_family_id, key, kDeleteRecord, + last_entry_offset); break; case kTypeColumnFamilySingleDeletion: case kTypeSingleDeletion: found++; - if (!UpdateExistingEntryWithCfId(column_family_id, key, - kSingleDeleteRecord)) { - AddNewEntry(column_family_id, kSingleDeleteRecord); - } + AddOrUpdateIndexWithCfId(column_family_id, key, kSingleDeleteRecord, + last_entry_offset); break; case kTypeColumnFamilyMerge: case kTypeMerge: found++; - if (!UpdateExistingEntryWithCfId(column_family_id, key, kMergeRecord)) { - AddNewEntry(column_family_id, kMergeRecord); - } + AddOrUpdateIndexWithCfId(column_family_id, key, kMergeRecord, + last_entry_offset); break; case kTypeLogData: case kTypeBeginPrepareXID: @@ -278,10 +284,8 @@ Status WriteBatchWithIndex::Rep::ReBuildIndex() { case kTypeColumnFamilyWideColumnEntity: case kTypeWideColumnEntity: found++; - if (!UpdateExistingEntryWithCfId(column_family_id, key, - kPutEntityRecord)) { - AddNewEntry(column_family_id, kPutEntityRecord); - } + AddOrUpdateIndexWithCfId(column_family_id, key, kPutEntityRecord, + last_entry_offset); break; case kTypeColumnFamilyValuePreferredSeqno: case kTypeValuePreferredSeqno: @@ -368,19 +372,19 @@ Iterator* WriteBatchWithIndex::NewIteratorWithBase(Iterator* base_iterator) { Status WriteBatchWithIndex::Put(ColumnFamilyHandle* column_family, const Slice& key, const Slice& value) { - rep->SetLastEntryOffset(); + size_t last_entry_offset = rep->write_batch.GetDataSize(); auto s = rep->write_batch.Put(column_family, key, value); if (s.ok()) { - rep->AddOrUpdateIndex(column_family, key, kPutRecord); + rep->AddOrUpdateIndex(column_family, key, kPutRecord, last_entry_offset); } return s; } Status WriteBatchWithIndex::Put(const Slice& key, const Slice& value) { - rep->SetLastEntryOffset(); + size_t last_entry_offset = rep->write_batch.GetDataSize(); auto s = rep->write_batch.Put(key, value); if (s.ok()) { - rep->AddOrUpdateIndex(key, kPutRecord); + rep->AddOrUpdateIndex(key, kPutRecord, last_entry_offset); } return s; } @@ -399,33 +403,30 @@ Status WriteBatchWithIndex::PutEntity(ColumnFamilyHandle* column_family, const Slice& key, const WideColumns& columns) { assert(rep); - - rep->SetLastEntryOffset(); - + size_t last_entry_offset = rep->write_batch.GetDataSize(); const Status s = rep->write_batch.PutEntity(column_family, key, columns); - if (s.ok()) { - rep->AddOrUpdateIndex(column_family, key, kPutEntityRecord); + rep->AddOrUpdateIndex(column_family, key, kPutEntityRecord, + last_entry_offset); } - return s; } Status WriteBatchWithIndex::Delete(ColumnFamilyHandle* column_family, const Slice& key) { - rep->SetLastEntryOffset(); + size_t last_entry_offset = rep->write_batch.GetDataSize(); auto s = rep->write_batch.Delete(column_family, key); if (s.ok()) { - rep->AddOrUpdateIndex(column_family, key, kDeleteRecord); + rep->AddOrUpdateIndex(column_family, key, kDeleteRecord, last_entry_offset); } return s; } Status WriteBatchWithIndex::Delete(const Slice& key) { - rep->SetLastEntryOffset(); + size_t last_entry_offset = rep->write_batch.GetDataSize(); auto s = rep->write_batch.Delete(key); if (s.ok()) { - rep->AddOrUpdateIndex(key, kDeleteRecord); + rep->AddOrUpdateIndex(key, kDeleteRecord, last_entry_offset); } return s; } @@ -441,19 +442,20 @@ Status WriteBatchWithIndex::Delete(ColumnFamilyHandle* column_family, Status WriteBatchWithIndex::SingleDelete(ColumnFamilyHandle* column_family, const Slice& key) { - rep->SetLastEntryOffset(); + size_t last_entry_offset = rep->write_batch.GetDataSize(); auto s = rep->write_batch.SingleDelete(column_family, key); if (s.ok()) { - rep->AddOrUpdateIndex(column_family, key, kSingleDeleteRecord); + rep->AddOrUpdateIndex(column_family, key, kSingleDeleteRecord, + last_entry_offset); } return s; } Status WriteBatchWithIndex::SingleDelete(const Slice& key) { - rep->SetLastEntryOffset(); + size_t last_entry_offset = rep->write_batch.GetDataSize(); auto s = rep->write_batch.SingleDelete(key); if (s.ok()) { - rep->AddOrUpdateIndex(key, kSingleDeleteRecord); + rep->AddOrUpdateIndex(key, kSingleDeleteRecord, last_entry_offset); } return s; } @@ -470,19 +472,19 @@ Status WriteBatchWithIndex::SingleDelete(ColumnFamilyHandle* column_family, Status WriteBatchWithIndex::Merge(ColumnFamilyHandle* column_family, const Slice& key, const Slice& value) { - rep->SetLastEntryOffset(); + size_t last_entry_offset = rep->write_batch.GetDataSize(); auto s = rep->write_batch.Merge(column_family, key, value); if (s.ok()) { - rep->AddOrUpdateIndex(column_family, key, kMergeRecord); + rep->AddOrUpdateIndex(column_family, key, kMergeRecord, last_entry_offset); } return s; } Status WriteBatchWithIndex::Merge(const Slice& key, const Slice& value) { - rep->SetLastEntryOffset(); + size_t last_entry_offset = rep->write_batch.GetDataSize(); auto s = rep->write_batch.Merge(key, value); if (s.ok()) { - rep->AddOrUpdateIndex(key, kMergeRecord); + rep->AddOrUpdateIndex(key, kMergeRecord, last_entry_offset); } return s; } diff --git a/utilities/write_batch_with_index/write_batch_with_index_internal.h b/utilities/write_batch_with_index/write_batch_with_index_internal.h index da6f53e572..79134217e2 100644 --- a/utilities/write_batch_with_index/write_batch_with_index_internal.h +++ b/utilities/write_batch_with_index/write_batch_with_index_internal.h @@ -110,9 +110,10 @@ class BaseDeltaIterator : public Iterator { // Key used by skip list, as the binary searchable index of WriteBatchWithIndex. struct WriteBatchIndexEntry { - WriteBatchIndexEntry(size_t o, uint32_t c, size_t ko, size_t ksz) + WriteBatchIndexEntry(size_t o, uint32_t c, size_t ko, size_t ksz, uint32_t uc) : offset(o), column_family(c), + update_count(uc), has_single_del(false), has_overwritten_single_del(false), key_offset(ko), @@ -133,6 +134,7 @@ struct WriteBatchIndexEntry { // Keys are ordered by descending offset. : offset(is_forward_direction ? std::numeric_limits::max() : 0), column_family(_column_family), + update_count(1), has_single_del(false), has_overwritten_single_del(false), key_offset(0), @@ -161,6 +163,10 @@ struct WriteBatchIndexEntry { // SeekForPrev() will see all the keys with the same key. size_t offset; uint32_t column_family; // column family of the entry. + // The following three fields are only maintained when the WBWI is created + // with overwrite_key = true. + uint32_t update_count; // The number of updates (1-based) for this key up to + // this entry. bool has_single_del; // whether single del was issued for this key bool has_overwritten_single_del; // whether a single del for this key was // overwritten by another key @@ -348,6 +354,11 @@ class WBWIIteratorImpl final : public WBWIIterator { return skip_list_iter_.key()->has_overwritten_single_del; } + uint32_t GetUpdateCount() const override { + assert(Valid()); + return skip_list_iter_.key()->update_count; + } + Status status() const override { // this is in-memory data structure, so the only way status can be non-ok is // through memory corruption diff --git a/utilities/write_batch_with_index/write_batch_with_index_test.cc b/utilities/write_batch_with_index/write_batch_with_index_test.cc index 7189b0ac3c..70528657ff 100644 --- a/utilities/write_batch_with_index/write_batch_with_index_test.cc +++ b/utilities/write_batch_with_index/write_batch_with_index_test.cc @@ -51,6 +51,7 @@ struct Entry { struct TestHandler : public WriteBatch::Handler { std::map> seen; + std::map> cf_key_count; Status PutCF(uint32_t column_family_id, const Slice& key, const Slice& value) override { Entry e; @@ -58,6 +59,7 @@ struct TestHandler : public WriteBatch::Handler { e.value = value.ToString(); e.type = kPutRecord; seen[column_family_id].push_back(e); + cf_key_count[column_family_id][e.key]++; return Status::OK(); } Status MergeCF(uint32_t column_family_id, const Slice& key, @@ -67,6 +69,7 @@ struct TestHandler : public WriteBatch::Handler { e.value = value.ToString(); e.type = kMergeRecord; seen[column_family_id].push_back(e); + cf_key_count[column_family_id][e.key]++; return Status::OK(); } void LogData(const Slice& /*blob*/) override {} @@ -76,6 +79,16 @@ struct TestHandler : public WriteBatch::Handler { e.value = ""; e.type = kDeleteRecord; seen[column_family_id].push_back(e); + cf_key_count[column_family_id][e.key]++; + return Status::OK(); + } + Status SingleDeleteCF(uint32_t column_family_id, const Slice& key) override { + Entry e; + e.key = key.ToString(); + e.value = ""; + e.type = kDeleteRecord; + seen[column_family_id].push_back(e); + cf_key_count[column_family_id][e.key]++; return Status::OK(); } }; @@ -390,6 +403,33 @@ class WBWIBaseTest : public testing::Test { } } + void VerifyWBWIIterUpdateCount(ColumnFamilyHandle* cfh = nullptr) { + TestHandler handler; + ASSERT_OK(batch_->GetWriteBatch()->Iterate(&handler)); + auto& count_map = handler.cf_key_count[GetColumnFamilyID(cfh)]; + // handler for going through write batch + std::unique_ptr iter(batch_->NewIterator(cfh)); + iter->SeekToFirst(); + size_t key_seen = 0; + while (iter->Valid()) { + ASSERT_OK(iter->status()); + std::string key = iter->Entry().key.ToString(); + auto count = count_map[key]; + ASSERT_EQ(iter->GetUpdateCount(), count); + iter->Next(); + // Merge is not overwritten when WBWI is created with overrwite_key + // We can only check for decreasing update count since the update before + // Merge can be overwritten multiple times. + while (iter->Valid() && key == iter->Entry().key.ToString()) { + ASSERT_GT(count, iter->GetUpdateCount()); + count = iter->GetUpdateCount(); + iter->Next(); + } + ++key_seen; + } + ASSERT_EQ(key_seen, count_map.size()); + } + public: DB* db_; std::string dbname_; @@ -707,6 +747,9 @@ TEST_P(WriteBatchWithIndexTest, WBWIIteratorImpl) { // put then merge ASSERT_OK(batch_->Put("k3", "k3p0")); ASSERT_OK(batch_->Merge("k3", "k3m1")); + if (GetParam()) { + VerifyWBWIIterUpdateCount(); + } std::unique_ptr iter( static_cast(batch_->NewIterator())); @@ -1021,6 +1064,9 @@ TEST_P(WriteBatchWithIndexTest, TestWBWIIterator) { ASSERT_OK(batch_->Delete(&cf2, "f")); AssertIterEqual(iter1.get(), {"a", "c", "e"}); AssertIterEqual(iter2.get(), {"a", "b", "d", "f"}); + if (GetParam()) { + VerifyWBWIIterUpdateCount(); + } } TEST_P(WriteBatchWithIndexTest, TestRandomIteraratorWithBase) { @@ -1142,6 +1188,11 @@ TEST_P(WriteBatchWithIndexTest, TestRandomIteraratorWithBase) { } ASSERT_OK(iter->status()); + if (GetParam()) { + VerifyWBWIIterUpdateCount(&cf1); + VerifyWBWIIterUpdateCount(&cf2); + VerifyWBWIIterUpdateCount(&cf3); + } } } @@ -1390,6 +1441,12 @@ TEST_P(WriteBatchWithIndexTest, TestIteraratorWithBaseReverseCmp) { iter->Seek("0"); AssertIter(iter.get(), "a", "b"); } + + if (GetParam()) { + VerifyWBWIIterUpdateCount(); + VerifyWBWIIterUpdateCount(&cf1); + VerifyWBWIIterUpdateCount(&cf2); + } } TEST_P(WriteBatchWithIndexTest, TestGetFromBatch) { @@ -1896,6 +1953,7 @@ TEST_F(WBWIOverwriteTest, MutateWhileIteratingBaseStressTest) { } } ASSERT_OK(iter->status()); + VerifyWBWIIterUpdateCount(); } TEST_P(WriteBatchWithIndexTest, TestNewIteratorWithBaseFromWbwi) { @@ -3675,16 +3733,16 @@ TEST_F(WBWIMemTableTest, ReadFromWBWIMemtable) { cmp, 0, /*overwrite_key=*/true, 0, 0); wbwi->SetTrackPerCFStat(true); std::vector> expected; - expected.resize(10000); - for (int i = 0; i < 10000; ++i) { + const int kNumUpdate = 10000; + expected.resize(kNumUpdate); + for (int i = 0; i < kNumUpdate; ++i) { // Leave a non-existing key 9999 in between existing keys to test read. std::string key = i < 9999 ? DBTestBase::Key(i) : DBTestBase::Key(i + 1); bool del = rnd.OneIn(2); std::string val = del ? "NOT_FOUND" : rnd.RandomString(50); expected[i] = std::make_pair(key, val); } - // Random insertion order - RandomShuffle(expected.begin(), expected.end()); + std::unique_ptr wbwi_mem{ new WBWIMemTable(wbwi, cmp, /*cf_id=*/0, &immutable_opts, &mutable_cf_options, @@ -3693,9 +3751,13 @@ TEST_F(WBWIMemTableTest, ReadFromWBWIMemtable) { ASSERT_TRUE(wbwi_mem->IsEmpty()); constexpr SequenceNumber visible_seq = 10002; constexpr SequenceNumber non_visible_seq = 1; - constexpr WBWIMemTable::SeqnoRange assigned_seq = {2, 10001}; + const SequenceNumber seqno_lb = 2; + constexpr WBWIMemTable::SeqnoRange assigned_seq = {seqno_lb, + seqno_lb + kNumUpdate - 1}; wbwi_mem->AssignSequenceNumbers(assigned_seq); + // Random insertion order + RandomShuffle(expected.begin(), expected.end()); bool found_final_value = false; for (const auto& [key, val] : expected) { if (val == "NOT_FOUND") { @@ -3719,28 +3781,38 @@ TEST_F(WBWIMemTableTest, ReadFromWBWIMemtable) { ColumnFamilyHandleImplDummy meta_cf(/*id=*/1, BytewiseComparator()); ASSERT_OK(wbwi->Put(&meta_cf, DBTestBase::Key(0), "foo")); - RandomShuffle(expected.begin(), expected.end()); + // Sort keys to compare with iterator + std::sort(expected.begin(), expected.end(), + [](const std::pair& a, + const std::pair& b) { + return a.first < b.first; + }); + // overwrites + std::vector expected_seqno(kNumUpdate, seqno_lb); for (size_t i = 0; i < 2000; ++i) { // We don't expect to mix SD and DEL, or issue multiple SD consecutively in // a DB. Read from WBWI should still work so we do it here to keep the test // simple. + size_t idx = rnd.Uniform(kNumUpdate); if (rnd.OneIn(2)) { std::string val = rnd.RandomString(100); - expected[i].second = val; - ASSERT_OK(wbwi->Put(expected[i].first, val)); + expected[idx].second = val; + ASSERT_OK(wbwi->Put(expected[idx].first, val)); } else { - expected[i].second = "NOT_FOUND"; + expected[idx].second = "NOT_FOUND"; if (rnd.OneIn(2)) { - ASSERT_OK(wbwi->SingleDelete(expected[i].first)); + ASSERT_OK(wbwi->SingleDelete(expected[idx].first)); } else { - ASSERT_OK(wbwi->Delete(expected[i].first)); + ASSERT_OK(wbwi->Delete(expected[idx].first)); } } found_final_value = false; - ASSERT_TRUE(expected[i].second == Get(expected[i].first, wbwi_mem, - visible_seq, &found_final_value)); + ASSERT_TRUE(expected[idx].second == Get(expected[idx].first, wbwi_mem, + visible_seq, &found_final_value)); ASSERT_TRUE(found_final_value); + // See comment for WBWIMemTable for sequence number assignment method. + expected_seqno[idx]++; } // Get a non-existing key found_final_value = false; @@ -3755,7 +3827,6 @@ TEST_F(WBWIMemTableTest, ReadFromWBWIMemtable) { &found_final_value)); ASSERT_FALSE(found_final_value); // Get existing keys - RandomShuffle(expected.begin(), expected.end()); for (const auto& [key, val] : expected) { found_final_value = false; ASSERT_TRUE(val == Get(key, wbwi_mem, visible_seq, &found_final_value)); @@ -3814,12 +3885,6 @@ TEST_F(WBWIMemTableTest, ReadFromWBWIMemtable) { } } - // Sort keys to compare with iterator - std::sort(expected.begin(), expected.end(), - [](const std::pair& a, - const std::pair& b) { - return a.first < b.first; - }); Arena arena; InternalIterator* iter = wbwi_mem->NewIterator( ReadOptions(), /*seqno_to_time_mapping=*/nullptr, &arena, @@ -3835,7 +3900,7 @@ TEST_F(WBWIMemTableTest, ReadFromWBWIMemtable) { ValueType val_type; UnPackSequenceAndType(ExtractInternalKeyFooter(iter->key()), &seq, &val_type); - ASSERT_EQ(seq, assigned_seq.upper_bound); + ASSERT_EQ(seq, expected_seqno[idx]); if (expected[idx].second == "NOT_FOUND") { ASSERT_TRUE(val_type == kTypeDeletion || val_type == kTypeSingleDeletion); } else { @@ -3936,9 +4001,14 @@ TEST_F(WBWIMemTableTest, IterEmitSingleDelete) { ASSERT_OK(wbwi->SingleDelete(DBTestBase::Key(0))); ASSERT_OK(wbwi->SingleDelete(DBTestBase::Key(1))); ASSERT_OK(wbwi->SingleDelete(DBTestBase::Key(2))); - ASSERT_OK(wbwi->Put(DBTestBase::Key(3), "val3")); // SD at key1 overwritten ASSERT_OK(wbwi->Put(DBTestBase::Key(1), "val1")); + // For key3, most recent update is SD, and a SD is overwritten, + // here we test that these two SD get assigned different seqnos. + // FIXME: ideally we should only emit at most one SD per key. + ASSERT_OK(wbwi->SingleDelete(DBTestBase::Key(3))); + ASSERT_OK(wbwi->Put(DBTestBase::Key(3), "val3")); + ASSERT_OK(wbwi->SingleDelete(DBTestBase::Key(3))); std::unique_ptr wbwi_mem{ new WBWIMemTable(wbwi, cmp, @@ -3961,16 +4031,23 @@ TEST_F(WBWIMemTableTest, IterEmitSingleDelete) { ASSERT_TRUE(iter_for_flush->Valid()); ASSERT_EQ(iter->key(), iter_for_flush->key()); - iter->Next(); - iter_for_flush->Next(); - if (i == 1) { - // overwritten SD at key1 + if (i == 1 || i == 3) { + // Most recent update should have a higher seqno than the overwritten SD + SequenceNumber current_seq; + ValueType _; + UnPackSequenceAndType(ExtractInternalKeyFooter(iter->key()), ¤t_seq, + &_); + ASSERT_GT(current_seq, assigned_seqno.lower_bound); + + iter_for_flush->Next(); + // overwritten SD at key1 and key3 // See WBWIMemTableIterator::UpdateSingleDeleteKey() for seqno assignment - InternalKey ikey(DBTestBase::Key(1), assigned_seqno.upper_bound - 1, + InternalKey ikey(DBTestBase::Key(i), assigned_seqno.lower_bound, kTypeSingleDeletion); ASSERT_EQ(ikey.Encode(), iter_for_flush->key()); - iter_for_flush->Next(); } + iter->Next(); + iter_for_flush->Next(); } ASSERT_FALSE(iter->Valid()); ASSERT_FALSE(iter_for_flush->Valid());