diff --git a/db/compaction/compaction_outputs.cc b/db/compaction/compaction_outputs.cc index 8c1d9b9f3f..e2a67a538c 100644 --- a/db/compaction/compaction_outputs.cc +++ b/db/compaction/compaction_outputs.cc @@ -10,7 +10,10 @@ #include "db/compaction/compaction_outputs.h" +#include + #include "db/builder.h" +#include "logging/logging.h" namespace ROCKSDB_NAMESPACE { @@ -373,6 +376,24 @@ bool CompactionOutputs::ShouldStopBefore(const CompactionIterator& c_iter) { return false; } +void CompactionOutputs::PrepareNewOutputForFirstKey( + const CompactionIterator& c_iter) { + // A new output file starts a fresh partition-tracking window. Establish the + // partitioner fast path from this file's first key: if the partitioner can + // express its next boundary as "first key not sharing this prefix", cache + // that prefix; otherwise fall back to per-key ShouldPartition(). This is the + // only place the prefix needs to be (re)derived, since a partition boundary + // always forces a file cut, and re-deriving on non-partition cuts just + // recomputes the same still-valid prefix. + if (partitioner_) { + OptSlice prefix = partitioner_->ShouldPartitionByPrefix(c_iter.user_key()); + partitioner_prefix_valid_ = prefix.has_value(); + if (partitioner_prefix_valid_) { + partitioner_prefix_.assign(prefix->data(), prefix->size()); + } + } +} + Status CompactionOutputs::AddToOutput( const CompactionIterator& c_iter, const CompactionFileOpenFunc& open_file_func, @@ -412,42 +433,117 @@ Status CompactionOutputs::AddToOutput( if (!s.ok()) { return s; } - // A new output file starts a fresh partition-tracking window. Establish the - // partitioner fast path from this file's first key: if the partitioner can - // express its next boundary as "first key not sharing this prefix", cache - // that prefix; otherwise fall back to per-key ShouldPartition(). This is - // the only place the prefix needs to be (re)derived, since a partition - // boundary always forces a file cut, and re-deriving on non-partition cuts - // just recomputes the same still-valid prefix. - if (partitioner_) { - OptSlice prefix = - partitioner_->ShouldPartitionByPrefix(c_iter.user_key()); - partitioner_prefix_valid_ = prefix.has_value(); - if (partitioner_prefix_valid_) { - partitioner_prefix_.assign(prefix->data(), prefix->size()); - } - } - } - - // c_iter may emit range deletion keys, so update `last_key_for_partitioner_` - // here before returning below when `is_range_del` is true. Only needed in the - // fallback mode; the prefix fast path does not use it. - if (partitioner_ && !partitioner_prefix_valid_) { - last_key_for_partitioner_.assign(c_iter.user_key().data_, - c_iter.user_key().size_); + PrepareNewOutputForFirstKey(c_iter); } if (UNLIKELY(is_range_del)) { + // c_iter may emit range deletion keys, so update + // `last_key_for_partitioner_` before returning. Only needed in the fallback + // mode; the prefix fast path does not use it. + if (partitioner_ && !partitioner_prefix_valid_) { + last_key_for_partitioner_.assign(c_iter.user_key().data_, + c_iter.user_key().size_); + } return s; } assert(builder_ != nullptr); const Slice& value = c_iter.value(); - s = current_output().validator.Add(key, value); - if (!s.ok()) { - return s; + bool retried_after_new_table = false; + while (true) { + s = current_output().validator.CanAdd(key); + if (!s.ok()) { + return s; + } + + TableBuilderAddContext add_context; + TableBuilderAddResult add_result = + builder_->TryAdd(key, value, &add_context); + + if (add_result == TableBuilderAddResult::kError) { + s = builder_->status(); + if (s.ok()) { + return Status::Corruption( + "TableBuilder returned kError without a non-OK status"); + } + return s; + } + + if (add_result == TableBuilderAddResult::kAdded) { + current_output().validator.AddValidated(key, value); + break; + } + + assert(add_result == TableBuilderAddResult::kRequiresNewTable); + s = builder_->status(); + if (!s.ok()) { + return s; + } + const Compaction* compaction = compaction_; + if (compaction == nullptr) { + return Status::Corruption("Compaction output is missing compaction"); + } + ColumnFamilyData* cfd = compaction->column_family_data(); + if (cfd == nullptr) { + return Status::Corruption( + "Compaction output is missing column family data"); + } + const bool allow_data_in_errors = + compaction->immutable_options().allow_data_in_errors; + const std::string key_for_diagnostics = + allow_data_in_errors ? key.ToString(/* hex */ true) : "hidden"; + if (compaction->output_level() == 0) { + const std::string detail = + "cf: " + cfd->GetName() + ", file_number: " + + std::to_string(current_output().meta.fd.GetNumber()) + ", reason: " + + TableBuilderNewTableReasonToString(add_context.new_table_reason) + + ", key: " + key_for_diagnostics; + return Status::Corruption( + "TableBuilder requested a new table while producing level 0 output", + detail); + } + if (retried_after_new_table || builder_->IsEmpty()) { + const std::string detail = + "cf: " + cfd->GetName() + ", file_number: " + + std::to_string(current_output().meta.fd.GetNumber()) + ", reason: " + + TableBuilderNewTableReasonToString(add_context.new_table_reason) + + ", key: " + key_for_diagnostics; + return Status::Corruption( + "TableBuilder requested a new table without accepting any rows", + detail); + } + + ROCKS_LOG_INFO( + compaction->immutable_options().info_log, + "[%s] Compaction output table builder requested a new table before " + "key %s for file #%" PRIu64 ". reason: %s", + cfd->GetName().c_str(), key_for_diagnostics.c_str(), + current_output().meta.fd.GetNumber(), + TableBuilderNewTableReasonToString(add_context.new_table_reason)); + + s = close_file_func(c_iter.InputStatus(), prev_iter_output_internal_key, + key, &c_iter, *this); + if (!s.ok()) { + return s; + } + + grandparent_boundary_switched_num_ = 0; + grandparent_overlapped_bytes_ = + GetCurrentKeyGrandparentOverlappedBytes(key); + range_tombstone_lower_bound_.Clear(); + + s = open_file_func(*this); + if (!s.ok()) { + return s; + } + PrepareNewOutputForFirstKey(c_iter); + retried_after_new_table = true; + } + + if (partitioner_ && !partitioner_prefix_valid_) { + last_key_for_partitioner_.assign(c_iter.user_key().data_, + c_iter.user_key().size_); } - builder_->Add(key, value); stats_.num_output_records++; current_output_file_size_ = builder_->EstimatedFileSize(); diff --git a/db/compaction/compaction_outputs.h b/db/compaction/compaction_outputs.h index 62b0063761..29fe094fba 100644 --- a/db/compaction/compaction_outputs.h +++ b/db/compaction/compaction_outputs.h @@ -245,6 +245,11 @@ class CompactionOutputs { // before processing the current key in compaction iterator. bool ShouldStopBefore(const CompactionIterator& c_iter); + // Re-establishes per-file state that depends on the first key accepted by a + // freshly opened builder. This keeps caller-requested file cuts equivalent to + // RocksDB's built-in file-boundary decisions for partition tracking. + void PrepareNewOutputForFirstKey(const CompactionIterator& c_iter); + void Cleanup() { if (builder_ != nullptr) { // May happen if we get a shutdown call in the middle of compaction diff --git a/db/db_compaction_test.cc b/db/db_compaction_test.cc index b6b8529e1b..304fde00d9 100644 --- a/db/db_compaction_test.cc +++ b/db/db_compaction_test.cc @@ -7,6 +7,9 @@ // Use of this source code is governed by a BSD-style license that can be // found in the LICENSE file. See the AUTHORS file for names of contributors. +#include +#include +#include #include #include @@ -26,6 +29,8 @@ #include "rocksdb/file_checksum.h" #include "rocksdb/iostats_context.h" #include "rocksdb/sst_file_writer.h" +#include "table/block_based/block_based_table_factory.h" +#include "table/table_builder.h" #include "test_util/mock_time_env.h" #include "test_util/sync_point.h" #include "test_util/testutil.h" @@ -173,6 +178,264 @@ class DBCompactionTest : public DBTestBase { } }; +struct FileCuttingTableFactoryState { + std::atomic cut_requests{0}; + int max_cut_requests = 1; + bool allow_empty_builder_cut = false; +}; + +class FileCuttingTableBuilder : public TableBuilder { + public: + FileCuttingTableBuilder(std::unique_ptr base_builder, + std::string cut_before_user_key, + std::shared_ptr state) + : base_builder_(std::move(base_builder)), + cut_before_user_key_(std::move(cut_before_user_key)), + state_(std::move(state)) {} + + void Add(const Slice& key, const Slice& value) override { + base_builder_->Add(key, value); + } + + TableBuilderAddResult TryAdd(const Slice& key, const Slice& value, + TableBuilderAddContext* context) override { + const Slice user_key = ExtractUserKey(key); + const bool cut_request_limit_not_reached = + state_->cut_requests.load() < state_->max_cut_requests; + const bool should_cut = + cut_request_limit_not_reached && + (state_->allow_empty_builder_cut || !base_builder_->IsEmpty()) && + user_key.compare(Slice(cut_before_user_key_)) == 0; + if (should_cut) { + state_->cut_requests++; + context->new_table_reason = TableBuilderNewTableReason::kBuilderPolicy; + return TableBuilderAddResult::kRequiresNewTable; + } + return base_builder_->TryAdd(key, value, context); + } + + Status status() const override { return base_builder_->status(); } + + IOStatus io_status() const override { return base_builder_->io_status(); } + + Status Finish() override { return base_builder_->Finish(); } + + void Abandon() override { base_builder_->Abandon(); } + + uint64_t NumEntries() const override { return base_builder_->NumEntries(); } + + bool IsEmpty() const override { return base_builder_->IsEmpty(); } + + uint64_t PreCompressionSize() const override { + return base_builder_->PreCompressionSize(); + } + + uint64_t FileSize() const override { return base_builder_->FileSize(); } + + uint64_t EstimatedFileSize() const override { + return base_builder_->EstimatedFileSize(); + } + + uint64_t EstimatedTailSize() const override { + return base_builder_->EstimatedTailSize(); + } + + uint64_t GetTailSize() const override { return base_builder_->GetTailSize(); } + + bool NeedCompact() const override { return base_builder_->NeedCompact(); } + + TableProperties GetTableProperties() const override { + return base_builder_->GetTableProperties(); + } + + std::string GetFileChecksum() const override { + return base_builder_->GetFileChecksum(); + } + + const char* GetFileChecksumFuncName() const override { + return base_builder_->GetFileChecksumFuncName(); + } + + void SetSeqnoTimeTableProperties(const SeqnoToTimeMapping& relevant_mapping, + uint64_t oldest_ancestor_time) override { + base_builder_->SetSeqnoTimeTableProperties(relevant_mapping, + oldest_ancestor_time); + } + + uint64_t GetWorkerCPUMicros() const override { + return base_builder_->GetWorkerCPUMicros(); + } + + private: + std::unique_ptr base_builder_; + std::string cut_before_user_key_; + std::shared_ptr state_; +}; + +class FileCuttingTableFactory : public TableFactory { + public: + FileCuttingTableFactory(std::string cut_before_user_key, + std::shared_ptr state) + : cut_before_user_key_(std::move(cut_before_user_key)), + state_(std::move(state)) {} + + const char* Name() const override { return "FileCuttingTableFactory"; } + + using TableFactory::NewTableReader; + + Status NewTableReader( + const ReadOptions& read_options, + const TableReaderOptions& table_reader_options, + std::unique_ptr&& file, uint64_t file_size, + std::unique_ptr* table_reader, + bool prefetch_index_and_filter_in_cache) const override { + return base_factory_.NewTableReader( + read_options, table_reader_options, std::move(file), file_size, + table_reader, prefetch_index_and_filter_in_cache); + } + + TableBuilder* NewTableBuilder( + const TableBuilderOptions& table_builder_options, + WritableFileWriter* file) const override { + return new FileCuttingTableBuilder( + std::unique_ptr( + base_factory_.NewTableBuilder(table_builder_options, file)), + cut_before_user_key_, state_); + } + + std::unique_ptr Clone() const override { + return std::make_unique(cut_before_user_key_, + state_); + } + + bool IsDeleteRangeSupported() const override { + return base_factory_.IsDeleteRangeSupported(); + } + + std::string GetPrintableOptions() const override { + return base_factory_.GetPrintableOptions(); + } + + private: + std::string cut_before_user_key_; + std::shared_ptr state_; + BlockBasedTableFactory base_factory_; +}; + +TEST_F(DBCompactionTest, TableBuilderCanRequestNewTableDuringCompaction) { + auto factory_state = std::make_shared(); + + Options options = CurrentOptions(); + options.disable_auto_compactions = true; + options.level0_file_num_compaction_trigger = 10; + options.target_file_size_base = 1 << 30; + options.table_factory = + std::make_shared(Key(5), factory_state); + Reopen(options); + + for (int key : {0, 2, 4, 6, 8}) { + ASSERT_OK(Put(Key(key), "value" + std::to_string(key))); + } + ASSERT_OK(Flush()); + for (int key : {1, 3, 5, 7, 9}) { + ASSERT_OK(Put(Key(key), "value" + std::to_string(key))); + } + ASSERT_OK(Flush()); + ASSERT_EQ(2, NumTableFilesAtLevel(0)); + + ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr)); + + ASSERT_EQ(1, factory_state->cut_requests.load()); + ASSERT_EQ(0, NumTableFilesAtLevel(0)); + ASSERT_EQ(2, NumTableFilesAtLevel(1)); + for (int key = 0; key < 10; ++key) { + ASSERT_EQ("value" + std::to_string(key), Get(Key(key))); + } +} + +TEST_F(DBCompactionTest, TableBuilderCutRequestForLevelZeroFailsCompaction) { + auto factory_state = std::make_shared(); + + Options options = CurrentOptions(); + options.compaction_style = kCompactionStyleUniversal; + options.disable_auto_compactions = true; + options.num_levels = 1; + options.table_factory = + std::make_shared(Key(5), factory_state); + Reopen(options); + + for (int key : {0, 2, 4, 6, 8}) { + ASSERT_OK(Put(Key(key), "value" + std::to_string(key))); + } + ASSERT_OK(Flush()); + for (int key : {1, 3, 5, 7, 9}) { + ASSERT_OK(Put(Key(key), "value" + std::to_string(key))); + } + ASSERT_OK(Flush()); + + Status status = db_->CompactRange(CompactRangeOptions(), nullptr, nullptr); + ASSERT_TRUE(status.IsCorruption()) << status.ToString(); + ASSERT_NE(std::string::npos, status.ToString().find("level 0 output")) + << status.ToString(); + ASSERT_EQ(1, factory_state->cut_requests.load()); +} + +TEST_F(DBCompactionTest, TableBuilderCutOnEmptyBuilderFailsCompaction) { + auto factory_state = std::make_shared(); + factory_state->allow_empty_builder_cut = true; + + Options options = CurrentOptions(); + options.disable_auto_compactions = true; + options.allow_data_in_errors = false; + options.level0_file_num_compaction_trigger = 10; + options.target_file_size_base = 1 << 30; + options.table_factory = + std::make_shared(Key(0), factory_state); + Reopen(options); + + for (int key : {0, 2, 4, 6, 8}) { + ASSERT_OK(Put(Key(key), "value" + std::to_string(key))); + } + ASSERT_OK(Flush()); + for (int key : {1, 3, 5, 7, 9}) { + ASSERT_OK(Put(Key(key), "value" + std::to_string(key))); + } + ASSERT_OK(Flush()); + + Status status = db_->CompactRange(CompactRangeOptions(), nullptr, nullptr); + ASSERT_TRUE(status.IsCorruption()) << status.ToString(); + ASSERT_NE(std::string::npos, status.ToString().find("key: hidden")) + << status.ToString(); + ASSERT_EQ(1, factory_state->cut_requests.load()); +} + +TEST_F(DBCompactionTest, TableBuilderRepeatedCutRequestFailsCompaction) { + auto factory_state = std::make_shared(); + factory_state->max_cut_requests = 2; + factory_state->allow_empty_builder_cut = true; + + Options options = CurrentOptions(); + options.disable_auto_compactions = true; + options.level0_file_num_compaction_trigger = 10; + options.target_file_size_base = 1 << 30; + options.table_factory = + std::make_shared(Key(5), factory_state); + Reopen(options); + + for (int key : {0, 2, 4, 6, 8}) { + ASSERT_OK(Put(Key(key), "value" + std::to_string(key))); + } + ASSERT_OK(Flush()); + for (int key : {1, 3, 5, 7, 9}) { + ASSERT_OK(Put(Key(key), "value" + std::to_string(key))); + } + ASSERT_OK(Flush()); + + Status status = db_->CompactRange(CompactRangeOptions(), nullptr, nullptr); + ASSERT_TRUE(status.IsCorruption()) << status.ToString(); + ASSERT_EQ(2, factory_state->cut_requests.load()); +} + class DBCompactionTestWithParam : public DBTestBase, public testing::WithParamInterface> { diff --git a/db/output_validator.cc b/db/output_validator.cc index 0c7109f3c2..6462925eb1 100644 --- a/db/output_validator.cc +++ b/db/output_validator.cc @@ -9,12 +9,7 @@ #include "util/hash.h" namespace ROCKSDB_NAMESPACE { -Status OutputValidator::Add(const Slice& key, const Slice& value) { - if (enable_hash_) { - // Generate a rolling 64-bit hash of the key and values - paranoid_hash_ = NPHash64(key.data(), key.size(), paranoid_hash_); - paranoid_hash_ = NPHash64(value.data(), value.size(), paranoid_hash_); - } +Status OutputValidator::CanAdd(const Slice& key) const { if (key.size() < kNumInternalBytes) { return Status::Corruption( "Compaction tries to write a key without internal bytes."); @@ -23,7 +18,24 @@ Status OutputValidator::Add(const Slice& key, const Slice& value) { if (!prev_key_.empty() && icmp_.Compare(key, prev_key_) < 0) { return Status::Corruption("Compaction sees out-of-order keys."); } - prev_key_.assign(key.data(), key.size()); return Status::OK(); } + +Status OutputValidator::Add(const Slice& key, const Slice& value) { + Status s = CanAdd(key); + if (!s.ok()) { + return s; + } + AddValidated(key, value); + return Status::OK(); +} + +void OutputValidator::AddValidated(const Slice& key, const Slice& value) { + if (enable_hash_) { + // Generate a rolling 64-bit hash of the key and values. + paranoid_hash_ = NPHash64(key.data(), key.size(), paranoid_hash_); + paranoid_hash_ = NPHash64(value.data(), value.size(), paranoid_hash_); + } + prev_key_.assign(key.data(), key.size()); +} } // namespace ROCKSDB_NAMESPACE diff --git a/db/output_validator.h b/db/output_validator.h index 1e7a8988ed..72dc632b55 100644 --- a/db/output_validator.h +++ b/db/output_validator.h @@ -25,6 +25,13 @@ class OutputValidator { // criteria, e.g. key is ordered. Status Add(const Slice& key, const Slice& value); + // Add a key that has already been checked with CanAdd(). + void AddValidated(const Slice& key, const Slice& value); + + // Validate whether Add() would accept this key without updating the + // validator state. + Status CanAdd(const Slice& key) const; + // Compare result of two key orders are the same. It can be used // to compare the keys inserted into a file, and what is read back. // Return true if the validation passes. diff --git a/table/block_based/block_based_table_builder.cc b/table/block_based/block_based_table_builder.cc index bc8a96af6e..2e05abcff4 100644 --- a/table/block_based/block_based_table_builder.cc +++ b/table/block_based/block_based_table_builder.cc @@ -2102,6 +2102,14 @@ void BlockBasedTableBuilder::Add(const Slice& ikey, const Slice& value) { } } +TableBuilderAddResult BlockBasedTableBuilder::TryAdd( + const Slice& ikey, const Slice& value, TableBuilderAddContext* context) { + assert(context != nullptr); + context->new_table_reason = TableBuilderNewTableReason::kNone; + Add(ikey, value); + return TableBuilderAddResult::kAdded; +} + void BlockBasedTableBuilder::Flush(const Slice* first_key_in_next_block) { Rep* r = rep_.get(); assert(rep_->state != Rep::State::kClosed); diff --git a/table/block_based/block_based_table_builder.h b/table/block_based/block_based_table_builder.h index eb50244213..6ee4f460fe 100644 --- a/table/block_based/block_based_table_builder.h +++ b/table/block_based/block_based_table_builder.h @@ -68,6 +68,13 @@ class BlockBasedTableBuilder : public TableBuilder { // REQUIRES: Finish(), Abandon() have not been called void Add(const Slice& key, const Slice& value) override; + // Keeps block-based builders behind the retry-capable TableBuilder contract + // even though the current format accepts every row. Future format-specific + // boundary policies can then ask callers to cut the file without changing + // compaction's Add() assumptions. + TableBuilderAddResult TryAdd(const Slice& key, const Slice& value, + TableBuilderAddContext* context) override; + // Return non-ok iff some error has been detected. Status status() const override; diff --git a/table/external_table.cc b/table/external_table.cc index 8e6984a73a..6cdd5617b8 100644 --- a/table/external_table.cc +++ b/table/external_table.cc @@ -35,6 +35,30 @@ namespace ROCKSDB_NAMESPACE { +static_assert(static_cast(ExternalTableBuilderNewTableReason::kNone) == + static_cast(TableBuilderNewTableReason::kNone)); +static_assert( + static_cast( + ExternalTableBuilderNewTableReason::kRowClassificationChanged) == + static_cast( + TableBuilderNewTableReason::kRowClassificationChanged)); +static_assert( + static_cast( + ExternalTableBuilderNewTableReason::kSchemaIdChanged) == + static_cast(TableBuilderNewTableReason::kSchemaIdChanged)); +static_assert( + static_cast( + ExternalTableBuilderNewTableReason::kSchemaVersionIncompatible) == + static_cast( + TableBuilderNewTableReason::kSchemaVersionIncompatible)); +static_assert( + static_cast( + ExternalTableBuilderNewTableReason::kUnsupportedEntryType) == + static_cast(TableBuilderNewTableReason::kUnsupportedEntryType)); +static_assert( + static_cast(ExternalTableBuilderNewTableReason::kBuilderPolicy) == + static_cast(TableBuilderNewTableReason::kBuilderPolicy)); + namespace { Status GetValueType(EntryType entry_type, ValueType* value_type) { @@ -774,78 +798,78 @@ class ExternalTableBuilderAdapter : public TableBuilder { } void Add(const Slice& key, const Slice& value) override { + TableBuilderAddContext context; + const TableBuilderAddResult result = TryAdd(key, value, &context); + if (result == TableBuilderAddResult::kRequiresNewTable) { + status_ = Status::InvalidArgument( + "External table builder requested a new table through Add"); + } else if (result == TableBuilderAddResult::kError && status_.ok()) { + status_ = Status::InvalidArgument( + "External table builder returned an " + "error without a non-OK status"); + } + } + + TableBuilderAddResult TryAdd(const Slice& key, const Slice& value, + TableBuilderAddContext* context) override { + assert(context != nullptr); + context->new_table_reason = TableBuilderNewTableReason::kNone; if (!status_.ok()) { - return; + return TableBuilderAddResult::kError; } ParsedEntryInfo entry; - status_ = ParseEntry(key, ioptions_.user_comparator, &entry); + ValueType value_type = kTypeMaxValid; + bool is_range_deletion = false; + status_ = PrepareEntry(key, entry, value_type, is_range_deletion); if (!status_.ok()) { - return; + return TableBuilderAddResult::kError; } - if (!entry.timestamp.empty()) { - status_ = Status::NotSupported( - "External table does not support user-defined timestamps"); - return; - } - const bool is_range_deletion = entry.type == kEntryRangeDeletion; - if constexpr (Mode == ExternalTableMode::kOnlyZeroSeqnoAndPuts) { - if (entry.sequence != 0 || - (entry.type != kEntryPut && - (!is_range_deletion || !support_range_deletions_))) { - if (support_range_deletions_) { - status_ = Status::NotSupported( - "Basic external table factory only supports sequence-zero Put " - "and range-deletion entries"); - } else { - status_ = Status::NotSupported( - "Basic external table factory only supports sequence-zero Put " - "entries"); - } - return; - } + if (value_type == kTypeMaxValid) { + status_ = Status::Corruption("External table entry value type is unset"); + return TableBuilderAddResult::kError; } - ValueType value_type; if (is_range_deletion) { - value_type = kTypeRangeDeletion; range_del_block_builder_.Add(key, value); } else { - status_ = GetValueType(entry.type, &value_type); - if (!status_.ok()) { - return; - } + ExternalTableBuilderAddContext external_context; + ExternalTableBuilderAddResult result; if constexpr (Mode == ExternalTableMode::kFull) { - builder_->Add(key, value); + result = builder_->TryAdd(key, value, &external_context); } else { - builder_->Add(entry.user_key, value); + result = builder_->TryAdd(entry.user_key, value, &external_context); } + + if (result == ExternalTableBuilderAddResult::kRequiresNewTable) { + status_ = builder_->status(); + if (!status_.ok()) { + return TableBuilderAddResult::kError; + } + context->new_table_reason = + ToTableBuilderNewTableReason(external_context.new_table_reason); + return TableBuilderAddResult::kRequiresNewTable; + } + + if (result == ExternalTableBuilderAddResult::kError) { + status_ = builder_->status(); + if (status_.ok()) { + status_ = Status::InvalidArgument( + "External table builder returned an error without a non-OK " + "status"); + } + return TableBuilderAddResult::kError; + } + + assert(result == ExternalTableBuilderAddResult::kAdded); + status_ = builder_->status(); } - status_ = builder_->status(); if (!status_.ok()) { - return; + return TableBuilderAddResult::kError; } - // Only entries accepted by the external builder contribute to properties - // and property collectors. - properties_.key_largest_seqno = - std::max(properties_.key_largest_seqno, entry.sequence); - properties_.key_smallest_seqno = - std::min(properties_.key_smallest_seqno, entry.sequence); - properties_.num_entries++; - properties_.raw_key_size += key.size(); - properties_.raw_value_size += value.size(); - if (value_type == kTypeDeletion || value_type == kTypeSingleDeletion) { - properties_.num_deletions++; - } else if (value_type == kTypeRangeDeletion) { - properties_.num_deletions++; - properties_.num_range_deletions++; - } else if (value_type == kTypeMerge) { - properties_.num_merge_operands++; - } - NotifyCollectTableCollectorsOnAdd(key, value, /*file_size=*/0, - table_properties_collectors_, - ioptions_.logger); + RecordAcceptedEntry(key, value, entry, value_type); + return TableBuilderAddResult::kAdded; } Status status() const override { @@ -909,6 +933,10 @@ class ExternalTableBuilderAdapter : public TableBuilder { uint64_t NumEntries() const override { return properties_.num_entries; } + bool IsEmpty() const override { + return properties_.num_entries == 0 && properties_.num_range_deletions == 0; + } + bool NeedCompact() const override { for (const auto& collector : table_properties_collectors_) { if (collector->NeedCompact()) { @@ -949,6 +977,88 @@ class ExternalTableBuilderAdapter : public TableBuilder { } private: + TableBuilderNewTableReason ToTableBuilderNewTableReason( + ExternalTableBuilderNewTableReason reason) const { + switch (reason) { + case ExternalTableBuilderNewTableReason::kNone: + return TableBuilderNewTableReason::kNone; + case ExternalTableBuilderNewTableReason::kRowClassificationChanged: + return TableBuilderNewTableReason::kRowClassificationChanged; + case ExternalTableBuilderNewTableReason::kSchemaIdChanged: + return TableBuilderNewTableReason::kSchemaIdChanged; + case ExternalTableBuilderNewTableReason::kSchemaVersionIncompatible: + return TableBuilderNewTableReason::kSchemaVersionIncompatible; + case ExternalTableBuilderNewTableReason::kUnsupportedEntryType: + return TableBuilderNewTableReason::kUnsupportedEntryType; + case ExternalTableBuilderNewTableReason::kBuilderPolicy: + return TableBuilderNewTableReason::kBuilderPolicy; + } + assert(false); + return TableBuilderNewTableReason::kBuilderPolicy; + } + + Status PrepareEntry(const Slice& key, ParsedEntryInfo& entry, + ValueType& value_type, bool& is_range_deletion) { + is_range_deletion = false; + + Status status = ParseEntry(key, ioptions_.user_comparator, &entry); + if (!status.ok()) { + return status; + } + if (!entry.timestamp.empty()) { + return Status::NotSupported( + "External table does not support user-defined timestamps"); + } + is_range_deletion = entry.type == kEntryRangeDeletion; + if constexpr (Mode == ExternalTableMode::kOnlyZeroSeqnoAndPuts) { + if (entry.sequence != 0 || + (entry.type != kEntryPut && + (!is_range_deletion || !support_range_deletions_))) { + if (support_range_deletions_) { + return Status::NotSupported( + "Basic external table factory only supports sequence-zero Put " + "and range-deletion entries"); + } else { + return Status::NotSupported( + "Basic external table factory only supports sequence-zero Put " + "entries"); + } + } + if (!is_range_deletion) { + value_type = kTypeValue; + return Status::OK(); + } + } + + if (is_range_deletion) { + value_type = kTypeRangeDeletion; + return Status::OK(); + } + + return GetValueType(entry.type, &value_type); + } + + void RecordAcceptedEntry(const Slice& key, const Slice& value, + const ParsedEntryInfo& entry, ValueType value_type) { + properties_.key_largest_seqno = + std::max(properties_.key_largest_seqno, entry.sequence); + properties_.key_smallest_seqno = + std::min(properties_.key_smallest_seqno, entry.sequence); + properties_.num_entries++; + properties_.raw_key_size += key.size(); + properties_.raw_value_size += value.size(); + if (value_type == kTypeDeletion || value_type == kTypeSingleDeletion) { + properties_.num_deletions++; + } else if (value_type == kTypeRangeDeletion) { + properties_.num_deletions++; + properties_.num_range_deletions++; + } else if (value_type == kTypeMerge) { + properties_.num_merge_operands++; + } + NotifyCollectTableCollectorsOnAdd(key, value, /*file_size=*/0, + table_properties_collectors_, + ioptions_.logger); + } Status status_; std::unique_ptr builder_; const ImmutableOptions& ioptions_; diff --git a/table/table_builder.h b/table/table_builder.h index 5e314e4fad..be9baf6692 100644 --- a/table/table_builder.h +++ b/table/table_builder.h @@ -11,6 +11,7 @@ #include +#include #include #include #include @@ -22,6 +23,7 @@ #include "options/cf_options.h" #include "rocksdb/options.h" #include "rocksdb/slice.h" +#include "rocksdb/status.h" #include "rocksdb/table_properties.h" #include "table/embedded_blob_sst.h" #include "table/unique_id_impl.h" @@ -31,7 +33,6 @@ namespace ROCKSDB_NAMESPACE { class Slice; -class Status; class BlobSource; struct TableReaderOptions { @@ -201,6 +202,50 @@ struct TableBuilderOptions : public TablePropertiesCollectorFactory::Context { const EmbeddedBlobSstBuilderOptions* embedded_blob_options; }; +enum class TableBuilderAddResult : uint8_t { + // The row was consumed by this builder. The caller may advance input. + kAdded, + // The row was not consumed. The caller should finish the current table, + // create another table through the same TableFactory, and retry the same row. + kRequiresNewTable, + // A real error occurred. The caller should abandon the current output. + kError, +}; + +enum class TableBuilderNewTableReason : uint8_t { + kNone = 0, + kRowClassificationChanged, + kSchemaIdChanged, + kSchemaVersionIncompatible, + kUnsupportedEntryType, + kBuilderPolicy, +}; + +inline const char* TableBuilderNewTableReasonToString( + TableBuilderNewTableReason reason) { + switch (reason) { + case TableBuilderNewTableReason::kNone: + return "None"; + case TableBuilderNewTableReason::kRowClassificationChanged: + return "RowClassificationChanged"; + case TableBuilderNewTableReason::kSchemaIdChanged: + return "SchemaIdChanged"; + case TableBuilderNewTableReason::kSchemaVersionIncompatible: + return "SchemaVersionIncompatible"; + case TableBuilderNewTableReason::kUnsupportedEntryType: + return "UnsupportedEntryType"; + case TableBuilderNewTableReason::kBuilderPolicy: + return "BuilderPolicy"; + } + assert(false); + return "Invalid"; +} + +struct TableBuilderAddContext { + TableBuilderNewTableReason new_table_reason = + TableBuilderNewTableReason::kNone; +}; + // TableBuilder provides the interface used to build a Table // (an immutable and sorted map from keys to values). // @@ -218,6 +263,26 @@ class TableBuilder { // REQUIRES: Finish(), Abandon() have not been called virtual void Add(const Slice& key, const Slice& value) = 0; + // Add key,value to the table being constructed and report whether the row was + // accepted or whether the caller should cut the current table and retry this + // same row with a new table builder. A kRequiresNewTable result is not an + // error: the row must not have been consumed and status() must remain OK. + // A kError result must be reported through status(). + // + // Existing table builders do not request table cuts; they inherit this + // default adapter around Add(). + // + // REQUIRES: key is after any previously added key according to comparator. + // REQUIRES: Finish(), Abandon() have not been called + // REQUIRES: context != nullptr + virtual TableBuilderAddResult TryAdd(const Slice& key, const Slice& value, + TableBuilderAddContext* context) { + assert(context != nullptr); + context->new_table_reason = TableBuilderNewTableReason::kNone; + Add(key, value); + return TableBuilderAddResult::kAdded; + } + // Return non-ok iff some error has been detected. virtual Status status() const = 0;