diff --git a/db/compaction/compaction_iterator.cc b/db/compaction/compaction_iterator.cc index 909be149a1..9b060e2894 100644 --- a/db/compaction/compaction_iterator.cc +++ b/db/compaction/compaction_iterator.cc @@ -644,6 +644,14 @@ bool CompactionIterator::InvokeFilterIfNeeded(bool* need_skip, entity_deserialized_ = false; } + if (UNLIKELY(compaction_filter_value_.size() > + std::numeric_limits::max())) { + status_ = Status::Corruption( + "CompactionFilter result value size exceeds 4GB limit"); + validity_info_.Invalidate(); + return false; + } + return true; } diff --git a/db/db_compaction_test.cc b/db/db_compaction_test.cc index 58439a3384..048931e9cd 100644 --- a/db/db_compaction_test.cc +++ b/db/db_compaction_test.cc @@ -12121,6 +12121,56 @@ TEST_F(DBCompactionTest, VerifyAllOutputFlagsWithoutParanoidFileChecks) { } } +// Requires ~8GB+ RAM because the compaction filter internally creates a ~4GB +// value via std::string::assign(). +TEST_F(DBCompactionTest, CompactionFilterLargeValueRejected) { + if (!test::HasBigMem()) { + ROCKSDB_GTEST_BYPASS("insufficient memory for reliable continuous testing"); + return; + } + + // A compaction filter that inflates every value to a configurable size + class InflatingFilter : public CompactionFilter { + public: + std::atomic target_size{0}; + Decision FilterV2(int /*level*/, const Slice& /*key*/, + ValueType /*value_type*/, const Slice& /*existing_value*/, + std::string* new_value, + std::string* /*skip_until*/) const override { + new_value->assign(target_size.load(), 'X'); + return Decision::kChangeValue; + } + const char* Name() const override { return "InflatingFilter"; } + }; + + InflatingFilter inflating_filter; + Options options = CurrentOptions(); + options.compaction_filter = &inflating_filter; + options.disable_auto_compactions = true; + DestroyAndReopen(options); + + // Control: value at exactly 4GB - 1 should be accepted + inflating_filter.target_size = size_t{std::numeric_limits::max()}; + ASSERT_OK(Put("key", "small_value")); + ASSERT_OK(Flush()); + ASSERT_OK(Put("key", "small_value2")); + ASSERT_OK(Flush()); + ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr)); + + // Now test rejection: value at 4GB should be rejected + inflating_filter.target_size = + size_t{std::numeric_limits::max()} + 1; + DestroyAndReopen(options); + ASSERT_OK(Put("key", "small_value")); + ASSERT_OK(Flush()); + ASSERT_OK(Put("key", "small_value2")); + ASSERT_OK(Flush()); + + Status s = db_->CompactRange(CompactRangeOptions(), nullptr, nullptr); + ASSERT_TRUE(s.IsCorruption()) << s.ToString(); + ASSERT_TRUE(s.ToString().find("4GB") != std::string::npos) << s.ToString(); +} + } // namespace ROCKSDB_NAMESPACE int main(int argc, char** argv) { diff --git a/db/external_sst_file_basic_test.cc b/db/external_sst_file_basic_test.cc index 326b3d567a..c8390a4f5f 100644 --- a/db/external_sst_file_basic_test.cc +++ b/db/external_sst_file_basic_test.cc @@ -3272,6 +3272,56 @@ INSTANTIATE_TEST_CASE_P(ExternalSSTFileBasicTest, ExternalSSTFileBasicTest, std::make_tuple(false, true), std::make_tuple(false, false))); +// Uses anonymous mmap (lazy-zeroed) so the large data itself doesn't consume +// physical memory -- only the SstFileWriter's internal copy does (~4GB peak +// during the acceptance case). +TEST_F(ExternalSSTFileBasicTest, LargeSizeSstFileWriter) { + if (!test::HasBigMem()) { + ROCKSDB_GTEST_BYPASS("insufficient memory for reliable continuous testing"); + return; + } + + constexpr size_t kMaxKeySize = + size_t{std::numeric_limits::max()} - 8; + constexpr size_t kMaxValueSize = size_t{std::numeric_limits::max()}; + + // --- Large key --- + { + Options options = CurrentOptions(); + SstFileWriter sst_file_writer(EnvOptions(), options); + ASSERT_OK(sst_file_writer.Open(dbname_ + "/large_key.sst")); + + MemMapping mm = MemMapping::AllocateLazyZeroed(kMaxKeySize + 1); + ASSERT_NE(nullptr, mm.Get()); + + // A key one byte over the limit should be rejected + ASSERT_TRUE(sst_file_writer.Put(mm.AsSlice(), "val").IsInvalidArgument()); + + // A key at the limit should be accepted + ASSERT_OK( + sst_file_writer.Put(Slice(mm.AsSlice().data(), kMaxKeySize), "val")); + ASSERT_OK(sst_file_writer.Finish()); + } + + // --- Large value --- + { + Options options = CurrentOptions(); + SstFileWriter sst_file_writer(EnvOptions(), options); + ASSERT_OK(sst_file_writer.Open(dbname_ + "/large_value.sst")); + + MemMapping mm = MemMapping::AllocateLazyZeroed(kMaxValueSize + 1); + ASSERT_NE(nullptr, mm.Get()); + + // A value one byte over the limit should be rejected + ASSERT_TRUE(sst_file_writer.Put("key", mm.AsSlice()).IsInvalidArgument()); + + // A value at the limit should be accepted + ASSERT_OK( + sst_file_writer.Put("key", Slice(mm.AsSlice().data(), kMaxValueSize))); + ASSERT_OK(sst_file_writer.Finish()); + } +} + } // namespace ROCKSDB_NAMESPACE int main(int argc, char** argv) { diff --git a/db/merge_helper.cc b/db/merge_helper.cc index 3958a84d01..61fdbc0b09 100644 --- a/db/merge_helper.cc +++ b/db/merge_helper.cc @@ -163,10 +163,20 @@ Status MergeHelper::TimedFullMergeImpl( return Status::OK(); }}; - return TimedFullMergeCommonImpl(merge_operator, key, - std::move(existing_value), operands, logger, - statistics, clock, update_num_ops_stats, - op_failure_scope, std::move(visitor)); + Status s = TimedFullMergeCommonImpl( + merge_operator, key, std::move(existing_value), operands, logger, + statistics, clock, update_num_ops_stats, op_failure_scope, + std::move(visitor)); + if (s.ok()) { + // Check merge result fits in uint32_t (a BlockBuilder assumption) + size_t result_size = (result_operand && !result_operand->empty()) + ? result_operand->size() + : result->size(); + if (UNLIKELY(result_size > std::numeric_limits::max())) { + return Status::Corruption("Merge result exceeds 4GB limit"); + } + } + return s; } Status MergeHelper::TimedFullMergeImpl( @@ -239,10 +249,19 @@ Status MergeHelper::TimedFullMergeImpl( return Status::OK(); }}; - return TimedFullMergeCommonImpl(merge_operator, key, - std::move(existing_value), operands, logger, - statistics, clock, update_num_ops_stats, - op_failure_scope, std::move(visitor)); + Status s = TimedFullMergeCommonImpl( + merge_operator, key, std::move(existing_value), operands, logger, + statistics, clock, update_num_ops_stats, op_failure_scope, + std::move(visitor)); + if (s.ok()) { + // Check merge result fits in uint32_t (a BlockBuilder assumption) + size_t result_size = + result_value ? result_value->size() : result_entity->serialized_size(); + if (UNLIKELY(result_size > std::numeric_limits::max())) { + return Status::Corruption("Merge result exceeds 4GB"); + } + } + return s; } // PRE: iter points to the first merge type entry @@ -649,6 +668,10 @@ Status MergeHelper::MergeUntil(InternalIterator* iter, if (merge_success) { // Merging of operands (associative merge) was successful. // Replace operands with the merge result + if (UNLIKELY(merge_result.size() > + std::numeric_limits::max())) { + return Status::Corruption("PartialMerge result exceeds 4GB limit"); + } merge_context_.Clear(); merge_context_.PushOperand(merge_result); keys_.erase(keys_.begin(), keys_.end() - 1); diff --git a/db/merge_helper_test.cc b/db/merge_helper_test.cc index 05408d5b96..e4fd40ba5a 100644 --- a/db/merge_helper_test.cc +++ b/db/merge_helper_test.cc @@ -289,6 +289,72 @@ TEST_F(MergeHelperTest, DontFilterMergeOperandsBeforeSnapshotTest) { ASSERT_FALSE(merge_output_iter.Valid()); } +// Requires ~8GB+ RAM due to large std::string operands that are copied into +// the MergeHelper's internal iterator. +TEST_F(MergeHelperTest, LargePartialMergeResultRejected) { + if (!test::HasBigMem()) { + ROCKSDB_GTEST_BYPASS("insufficient memory for reliable continuous testing"); + return; + } + + // A merge operator whose PartialMergeMulti concatenates operands + class ConcatMergeOp : public MergeOperator { + public: + bool FullMergeV2(const MergeOperationInput& merge_in, + MergeOperationOutput* merge_out) const override { + merge_out->new_value.clear(); + if (merge_in.existing_value) { + merge_out->new_value.append(merge_in.existing_value->data(), + merge_in.existing_value->size()); + } + for (const auto& op : merge_in.operand_list) { + merge_out->new_value.append(op.data(), op.size()); + } + return true; + } + bool PartialMergeMulti(const Slice& /*key*/, + const std::deque& operand_list, + std::string* new_value, + Logger* /*logger*/) const override { + new_value->clear(); + for (const auto& op : operand_list) { + new_value->append(op.data(), op.size()); + } + return true; + } + const char* Name() const override { return "ConcatMergeOp"; } + }; + + merge_op_ = std::make_shared(); + + // Two merge operands that sum to > 4GB + constexpr size_t kHalfSize = + (size_t{std::numeric_limits::max()} + 1) / 2; + std::string op1(kHalfSize, 'A'); + std::string op2(kHalfSize, 'B'); + + AddKeyVal("a", 20, kTypeMerge, op1); + AddKeyVal("a", 10, kTypeMerge, op2); + AddKeyVal("b", 10, kTypeMerge, "x"); // different key to stop iteration + + // at_bottom=false forces PartialMergeMulti path (not FullMerge) + ASSERT_TRUE(Run(0, false).IsCorruption()); + + // Control: just under 4GB should be accepted + ks_.clear(); + vs_.clear(); + op2.resize(kHalfSize - 1); + AddKeyVal("a", 20, kTypeMerge, op1); + AddKeyVal("a", 10, kTypeMerge, op2); + AddKeyVal("b", 10, kTypeMerge, "x"); + + Status s = Run(0, false); + ASSERT_TRUE(s.IsMergeInProgress()) << s.ToString(); + ASSERT_EQ(1U, merge_helper_->keys().size()); + ASSERT_EQ(size_t{std::numeric_limits::max()}, + merge_helper_->values()[0].size()); +} + } // namespace ROCKSDB_NAMESPACE int main(int argc, char** argv) { diff --git a/db/merge_test.cc b/db/merge_test.cc index 5f3546d6ce..e00db966e8 100644 --- a/db/merge_test.cc +++ b/db/merge_test.cc @@ -9,6 +9,7 @@ #include "db/db_impl/db_impl.h" #include "db/dbformat.h" +#include "db/merge_helper.h" #include "db/write_batch_internal.h" #include "port/stack_trace.h" #include "rocksdb/cache.h" @@ -889,6 +890,64 @@ TEST_F(MergeTest, FullMergeV3FallbackFailure) { } } +// Uses anonymous mmap (lazy-zeroed) so the large operands don't consume +// physical memory -- only the merge result does (~4GB peak). +TEST_F(MergeTest, LargeMergeResultRejected) { + if (!test::HasBigMem()) { + ROCKSDB_GTEST_BYPASS("insufficient memory for reliable continuous testing"); + return; + } + + // A merge operator that produces a result > 4GB by concatenating operands + class ConcatOperator : public MergeOperator { + public: + bool FullMergeV2(const MergeOperationInput& merge_in, + MergeOperationOutput* merge_out) const override { + merge_out->new_value.clear(); + if (merge_in.existing_value) { + merge_out->new_value.append(merge_in.existing_value->data(), + merge_in.existing_value->size()); + } + for (const auto& op : merge_in.operand_list) { + merge_out->new_value.append(op.data(), op.size()); + } + return true; + } + const char* Name() const override { return "ConcatOperator"; } + }; + + // Use TimedFullMerge directly - two operands that sum to > 4GB + constexpr size_t kHalfSize = + (size_t{std::numeric_limits::max()} + 1) / 2; + MemMapping mm1 = MemMapping::AllocateLazyZeroed(kHalfSize); + ASSERT_NE(nullptr, mm1.Get()); + MemMapping mm2 = MemMapping::AllocateLazyZeroed(kHalfSize); + ASSERT_NE(nullptr, mm2.Get()); + std::vector operands{mm1.AsSlice(), mm2.AsSlice()}; + std::string result; + Slice result_operand; + ValueType result_type; + + ConcatOperator concat_op; + Status s = MergeHelper::TimedFullMerge( + &concat_op, "key", MergeHelper::kNoBaseValue, operands, nullptr, nullptr, + SystemClock::Default().get(), false, nullptr, &result, &result_operand, + &result_type); + ASSERT_TRUE(s.IsCorruption()); + ASSERT_TRUE(s.ToString().find("4GB") != std::string::npos); + + // Control: just under 4GB should be accepted + result.clear(); + result.shrink_to_fit(); + operands = {mm1.AsSlice(), Slice(mm2.AsSlice().data(), kHalfSize - 1)}; + s = MergeHelper::TimedFullMerge(&concat_op, "key", MergeHelper::kNoBaseValue, + operands, nullptr, nullptr, + SystemClock::Default().get(), false, nullptr, + &result, &result_operand, &result_type); + ASSERT_OK(s); + ASSERT_EQ(result.size(), size_t{std::numeric_limits::max()}); +} + } // namespace ROCKSDB_NAMESPACE int main(int argc, char** argv) { diff --git a/db/write_batch.cc b/db/write_batch.cc index 9bf1dad32e..806bd1f23d 100644 --- a/db/write_batch.cc +++ b/db/write_batch.cc @@ -94,6 +94,11 @@ enum ContentFlags : uint32_t { HAS_TIMED_PUT = 1 << 13, }; +// The user key will become an internal key by appending kNumInternalBytes, +// so ensure the result still fits in uint32_t (a BlockBuilder assumption). +constexpr size_t kMaxWriteBatchKeySize = + size_t{std::numeric_limits::max()} - kNumInternalBytes; + struct BatchContentClassifier : public WriteBatch::Handler { uint32_t content_flags = 0; @@ -852,7 +857,7 @@ Status CheckColumnFamilyTimestampSize(ColumnFamilyHandle* column_family, Status WriteBatchInternal::Put(WriteBatch* b, uint32_t column_family_id, const Slice& key, const Slice& value) { - if (key.size() > size_t{std::numeric_limits::max()}) { + if (key.size() > kMaxWriteBatchKeySize) { return Status::InvalidArgument("key is too large"); } if (value.size() > size_t{std::numeric_limits::max()}) { @@ -889,7 +894,7 @@ Status WriteBatchInternal::Put(WriteBatch* b, uint32_t column_family_id, Status WriteBatchInternal::TimedPut(WriteBatch* b, uint32_t column_family_id, const Slice& key, const Slice& value, uint64_t write_unix_time) { - if (key.size() > size_t{std::numeric_limits::max()}) { + if (key.size() > kMaxWriteBatchKeySize) { return Status::InvalidArgument("key is too large"); } if (value.size() > size_t{std::numeric_limits::max()}) { @@ -1002,7 +1007,7 @@ Status WriteBatchInternal::CheckSlicePartsLength(const SliceParts& key, for (int i = 0; i < key.num_parts; ++i) { total_key_bytes += key.parts[i].size(); } - if (total_key_bytes >= size_t{std::numeric_limits::max()}) { + if (total_key_bytes > kMaxWriteBatchKeySize) { return Status::InvalidArgument("key is too large"); } @@ -1010,7 +1015,7 @@ Status WriteBatchInternal::CheckSlicePartsLength(const SliceParts& key, for (int i = 0; i < value.num_parts; ++i) { total_value_bytes += value.parts[i].size(); } - if (total_value_bytes >= size_t{std::numeric_limits::max()}) { + if (total_value_bytes > size_t{std::numeric_limits::max()}) { return Status::InvalidArgument("value is too large"); } return Status::OK(); @@ -1077,7 +1082,7 @@ Status WriteBatchInternal::PutEntity(WriteBatch* b, uint32_t column_family_id, const WideColumns& columns) { assert(b); - if (key.size() > size_t{std::numeric_limits::max()}) { + if (key.size() > kMaxWriteBatchKeySize) { return Status::InvalidArgument("key is too large"); } @@ -1275,6 +1280,9 @@ Status WriteBatchInternal::MarkRollback(WriteBatch* b, const Slice& xid) { Status WriteBatchInternal::Delete(WriteBatch* b, uint32_t column_family_id, const Slice& key) { + if (key.size() > kMaxWriteBatchKeySize) { + return Status::InvalidArgument("key is too large"); + } LocalSavePoint save(b); WriteBatchInternal::SetCount(b, WriteBatchInternal::Count(b) + 1); if (column_family_id == 0) { @@ -1347,6 +1355,10 @@ Status WriteBatch::Delete(ColumnFamilyHandle* column_family, const Slice& key, Status WriteBatchInternal::Delete(WriteBatch* b, uint32_t column_family_id, const SliceParts& key) { + Status s = CheckSlicePartsLength(key, SliceParts(nullptr, 0)); + if (!s.ok()) { + return s; + } LocalSavePoint save(b); WriteBatchInternal::SetCount(b, WriteBatchInternal::Count(b) + 1); if (column_family_id == 0) { @@ -1401,6 +1413,9 @@ Status WriteBatch::Delete(ColumnFamilyHandle* column_family, Status WriteBatchInternal::SingleDelete(WriteBatch* b, uint32_t column_family_id, const Slice& key) { + if (key.size() > kMaxWriteBatchKeySize) { + return Status::InvalidArgument("key is too large"); + } LocalSavePoint save(b); WriteBatchInternal::SetCount(b, WriteBatchInternal::Count(b) + 1); if (column_family_id == 0) { @@ -1475,6 +1490,10 @@ Status WriteBatch::SingleDelete(ColumnFamilyHandle* column_family, Status WriteBatchInternal::SingleDelete(WriteBatch* b, uint32_t column_family_id, const SliceParts& key) { + Status s = CheckSlicePartsLength(key, SliceParts(nullptr, 0)); + if (!s.ok()) { + return s; + } LocalSavePoint save(b); WriteBatchInternal::SetCount(b, WriteBatchInternal::Count(b) + 1); if (column_family_id == 0) { @@ -1530,6 +1549,12 @@ Status WriteBatch::SingleDelete(ColumnFamilyHandle* column_family, Status WriteBatchInternal::DeleteRange(WriteBatch* b, uint32_t column_family_id, const Slice& begin_key, const Slice& end_key) { + if (begin_key.size() > kMaxWriteBatchKeySize) { + return Status::InvalidArgument("key is too large"); + } + if (end_key.size() > kMaxWriteBatchKeySize) { + return Status::InvalidArgument("end key is too large"); + } LocalSavePoint save(b); WriteBatchInternal::SetCount(b, WriteBatchInternal::Count(b) + 1); if (column_family_id == 0) { @@ -1611,6 +1636,10 @@ Status WriteBatch::DeleteRange(ColumnFamilyHandle* column_family, Status WriteBatchInternal::DeleteRange(WriteBatch* b, uint32_t column_family_id, const SliceParts& begin_key, const SliceParts& end_key) { + Status s = CheckSlicePartsLength(begin_key, end_key); + if (!s.ok()) { + return s; + } LocalSavePoint save(b); WriteBatchInternal::SetCount(b, WriteBatchInternal::Count(b) + 1); if (column_family_id == 0) { @@ -1665,7 +1694,7 @@ Status WriteBatch::DeleteRange(ColumnFamilyHandle* column_family, Status WriteBatchInternal::Merge(WriteBatch* b, uint32_t column_family_id, const Slice& key, const Slice& value) { - if (key.size() > size_t{std::numeric_limits::max()}) { + if (key.size() > kMaxWriteBatchKeySize) { return Status::InvalidArgument("key is too large"); } if (value.size() > size_t{std::numeric_limits::max()}) { diff --git a/db/write_batch_test.cc b/db/write_batch_test.cc index e55b626d4e..1d819eab1a 100644 --- a/db/write_batch_test.cc +++ b/db/write_batch_test.cc @@ -18,6 +18,7 @@ #include "db/wide/wide_columns_helper.h" #include "db/write_batch_internal.h" #include "dbformat.h" +#include "port/mmap.h" #include "rocksdb/comparator.h" #include "rocksdb/db.h" #include "rocksdb/env.h" @@ -732,6 +733,54 @@ TEST_F(WriteBatchTest, DISABLED_LargeKeyValue) { ASSERT_EQ(2, handler.num_seen); } +// Uses anonymous mmap (lazy-zeroed) so the large data itself doesn't consume +// physical memory -- only the destination copy does (~4GB peak during the +// acceptance case where batch.Put copies into batch.rep_). +TEST_F(WriteBatchTest, LargeKeyValueSizeLimit) { + if (!test::HasBigMem()) { + ROCKSDB_GTEST_BYPASS("insufficient memory for reliable continuous testing"); + return; + } + + constexpr size_t kMaxKeySize = + size_t{std::numeric_limits::max()} - 8; + constexpr size_t kMaxValueSize = size_t{std::numeric_limits::max()}; + + WriteBatch batch; + + // --- Large key --- + { + MemMapping mm = MemMapping::AllocateLazyZeroed(kMaxKeySize + 1); + ASSERT_NE(nullptr, mm.Get()); + + // A key at the limit should be accepted + ASSERT_OK(batch.Put(Slice(mm.AsSlice().data(), kMaxKeySize), "val")); + batch.Clear(); + + // A key one byte over the limit should be rejected + ASSERT_TRUE(batch.Put(mm.AsSlice(), "val").IsInvalidArgument()); + ASSERT_TRUE(batch.Merge(mm.AsSlice(), "val").IsInvalidArgument()); + ASSERT_TRUE(batch.Delete(mm.AsSlice()).IsInvalidArgument()); + ASSERT_TRUE(batch.SingleDelete(mm.AsSlice()).IsInvalidArgument()); + ASSERT_TRUE( + batch.DeleteRange(mm.AsSlice(), mm.AsSlice()).IsInvalidArgument()); + } + + // --- Large value --- + { + MemMapping mm = MemMapping::AllocateLazyZeroed(kMaxValueSize + 1); + ASSERT_NE(nullptr, mm.Get()); + + // A value at the limit should be accepted + ASSERT_OK(batch.Put("key", Slice(mm.AsSlice().data(), kMaxValueSize))); + batch.Clear(); + + // A value one byte over the limit should be rejected + ASSERT_TRUE(batch.Put("key", mm.AsSlice()).IsInvalidArgument()); + ASSERT_TRUE(batch.Merge("key", mm.AsSlice()).IsInvalidArgument()); + } +} + TEST_F(WriteBatchTest, Continue) { WriteBatch batch; diff --git a/port/mmap.h b/port/mmap.h index 0f385522fc..48b115f03c 100644 --- a/port/mmap.h +++ b/port/mmap.h @@ -17,6 +17,7 @@ #include #include "rocksdb/rocksdb_namespace.h" +#include "rocksdb/slice.h" namespace ROCKSDB_NAMESPACE { @@ -51,8 +52,15 @@ class MemMapping { ~MemMapping(); inline void* Get() const { return addr_; } + // The requested length of the mapping, which may be smaller than the + // actual usable length inline size_t Length() const { return length_; } + // Return the mapping as a Slice (zero-copy view of the mapped memory) + inline Slice AsSlice() const { + return Slice(static_cast(addr_), length_); + } + private: MemMapping() {} diff --git a/table/block_based/block_based_table_builder.cc b/table/block_based/block_based_table_builder.cc index dd51a33a5f..03d87e5520 100644 --- a/table/block_based/block_based_table_builder.cc +++ b/table/block_based/block_based_table_builder.cc @@ -1551,6 +1551,18 @@ void BlockBasedTableBuilder::Add(const Slice& ikey, const Slice& value) { if (UNLIKELY(!ok())) { return; } + if (UNLIKELY(value.size() > std::numeric_limits::max())) { + // NOTE: + // * WriteBatch and SstFileWriter check input value sizes + // * MergeHelper checks merge result size + // * AddWithLastKey() below requires < 4GB and we have seen random + // corruptions here + r->SetStatus(Status::Corruption( + "BlockBasedBuilder::Add() received a value sized >= 4GB. This could be " + "a random corruption. size=" + + std::to_string(value.size()))); + return; + } ValueType value_type; SequenceNumber seq; UnPackSequenceAndType(ExtractInternalKeyFooter(ikey), &seq, &value_type); @@ -1588,6 +1600,7 @@ void BlockBasedTableBuilder::Add(const Slice& ikey, const Slice& value) { } } + // NOTE: WriteBatch guarantees keys < 4GB; value size checked above r->data_block.AddWithLastKey(ikey, value, r->last_ikey); r->last_ikey.assign(ikey.data(), ikey.size()); assert(!r->last_ikey.empty()); @@ -1612,6 +1625,7 @@ void BlockBasedTableBuilder::Add(const Slice& ikey, const Slice& value) { if (r->ts_sz > 0 && !r->persist_user_defined_timestamps) { persisted_end = StripTimestampFromUserKey(value, r->ts_sz); } + // NOTE: WriteBatch guarantees keys < 4GB; here 'value' is also a key r->range_del_block.Add(ikey, persisted_end); // TODO offset passed in is not accurate for parallel compression case NotifyCollectTableCollectorsOnAdd(ikey, value, r->get_offset(), diff --git a/table/block_based/block_builder.cc b/table/block_based/block_builder.cc index dcfbcd815e..5d0ba93c48 100644 --- a/table/block_based/block_builder.cc +++ b/table/block_based/block_builder.cc @@ -269,6 +269,11 @@ inline void BlockBuilder::AddWithLastKeyImpl(const Slice& key, size_t buffer_size) { assert(!finished_); assert(counter_ <= block_restart_interval_); + // Verify < 4GB assumption (see API comments on Add()) + assert(key.size() < uint64_t{1} << 32); + assert(value.size() < uint64_t{1} << 32); + assert(last_key.size() < uint64_t{1} << 32); + assert(delta_value == nullptr || delta_value->size() < uint64_t{1} << 32); std::string key_buf; std::string last_key_buf; const Slice key_to_persist = MaybeStripTimestampFromKey(&key_buf, key); @@ -278,46 +283,56 @@ inline void BlockBuilder::AddWithLastKeyImpl(const Slice& key, last_key.size() == 0 ? last_key : MaybeStripTimestampFromKey(&last_key_buf, last_key); - size_t shared = 0; // number of bytes shared with prev key + + // FIXME: check/enforce that buffer_ hasn't exceeded 4GB. The concern + // with adding that check and propagating the result is inner-loop + // performance. This case is HIGH concern because blocks like range deletions + // and non-partitioned indexes could pile large keys together into one block. + const uint32_t buffer_size32 = static_cast(buffer_size); + // NOTE: assuming all slice sizes < 4GB (see API comments on Add()) + uint32_t shared = 0; // number of bytes shared with prev key if (counter_ >= block_restart_interval_) { // Restart compression - restarts_.push_back(static_cast(buffer_size)); + restarts_.push_back(buffer_size32); estimate_ += sizeof(uint32_t); counter_ = 0; } else if (use_delta_encoding_ && !skip_delta_encoding) { // See how much sharing to do with previous string - shared = key_to_persist.difference_offset(last_key_persisted); + shared = static_cast( + key_to_persist.difference_offset(last_key_persisted)); } - const size_t non_shared = key_to_persist.size() - shared; - const size_t previous_value_offset = values_buffer_.size(); + const uint32_t non_shared = + static_cast(key_to_persist.size()) - shared; + const size_t prev_values_size = values_buffer_.size(); + + // FIXME: check/enforce that values_buffer_ hasn't exceeded 4GB. The concern + // with adding that check and propagating the result is inner-loop + // performance. This case is low concern because it (at time of writing) only + // applies to data blocks and those are flushed as soon as the size exceeds + // the block size. + const uint32_t prev_values_size32 = static_cast(prev_values_size); + const uint32_t value_size = static_cast(value.size()); if (use_value_delta_encoding_) { if (use_separated_kv_storage_ && counter_ == 0) { // Add "" to buffer_ - PutVarint32(&buffer_, static_cast(shared), - static_cast(non_shared), - static_cast(values_buffer_.size())); + PutVarint32(&buffer_, shared, non_shared, prev_values_size32); } else { // Add "" to buffer_ - PutVarint32(&buffer_, static_cast(shared), - static_cast(non_shared)); + PutVarint32(&buffer_, shared, non_shared); } } else { if (use_separated_kv_storage_ && counter_ == 0) { // Add "" to buffer_ - PutVarint32(&buffer_, static_cast(shared), - static_cast(non_shared), - static_cast(value.size()), - static_cast(values_buffer_.size())); + PutVarint32(&buffer_, shared, non_shared, value_size, prev_values_size32); } else { // Add "" to buffer_ - PutVarint32(&buffer_, static_cast(shared), - static_cast(non_shared), - static_cast(value.size())); + PutVarint32(&buffer_, shared, non_shared, value_size); } } - // Add string delta to buffer_ + // Add string delta to buffer_ (using only bottom 32 bits of size for + // consistent treatment in case of corruption) buffer_.append(key_to_persist.data() + shared, non_shared); auto& values_buffer = use_separated_kv_storage_ ? values_buffer_ : buffer_; @@ -325,10 +340,15 @@ inline void BlockBuilder::AddWithLastKeyImpl(const Slice& key, // simplify the decoding, where it can figure which decoding to use simply by // looking at the shared bytes size. if (shared != 0 && use_value_delta_encoding_) { + // Using only bottom 32 bits of size for consistent treatment in case of + // corruption assert(delta_value != nullptr); - values_buffer.append(delta_value->data(), delta_value->size()); + values_buffer.append(delta_value->data(), + static_cast(delta_value->size())); } else { - values_buffer.append(value.data(), value.size()); + // Using only bottom 32 bits of size for consistent treatment in case of + // corruption + values_buffer.append(value.data(), value_size); } // TODO(yuzhangyu): make user defined timestamp work with block hash index. @@ -342,18 +362,20 @@ inline void BlockBuilder::AddWithLastKeyImpl(const Slice& key, } counter_++; - estimate_ += buffer_.size() - buffer_size + values_buffer_.size() - - previous_value_offset; + estimate_ += + buffer_.size() - buffer_size + values_buffer_.size() - prev_values_size; } const Slice BlockBuilder::MaybeStripTimestampFromKey(std::string* key_buf, const Slice& key) { - Slice stripped_key = key; + // Only use bottom 32 bits of size for internal consistency (see API + // comments on Add()) + Slice stripped_key(key.data(), static_cast(key.size())); if (strip_ts_sz_ > 0) { if (is_user_key_) { stripped_key.remove_suffix(strip_ts_sz_); } else { - StripTimestampFromInternalKey(key_buf, key, strip_ts_sz_); + StripTimestampFromInternalKey(key_buf, stripped_key, strip_ts_sz_); stripped_key = *key_buf; } } diff --git a/table/block_based/block_builder.h b/table/block_based/block_builder.h index b0b65283f3..02d2746bdd 100644 --- a/table/block_based/block_builder.h +++ b/table/block_based/block_builder.h @@ -47,6 +47,10 @@ class BlockBuilder { // DO NOT mix with AddWithLastKey() between Resets. For efficiency, use // AddWithLastKey() in contexts where previous added key is already known // and delta encoding might be used. + // For efficiency, the implementation assumes the sizes of the input slices + // are each < 4GB, and only uses the bottom 32 bits of each size. (Using a + // dedicated Slice32 type would likely incur data movement overheads for this + // inner-loop code.) void Add(const Slice& key, const Slice& value, const Slice* const delta_value = nullptr, bool skip_delta_encoding = false); @@ -60,6 +64,8 @@ class BlockBuilder { // is the key from most recent AddWithLastKey. (For convenience, last_key // is ignored on first call after creation or Reset().) // DO NOT mix with Add() between Resets. + // For efficiency, the implementation assumes the sizes of the input slices + // are each < 4GB, and only uses the bottom 32 bits of each size. void AddWithLastKey(const Slice& key, const Slice& value, const Slice& last_key, const Slice* const delta_value = nullptr, @@ -100,6 +106,9 @@ class BlockBuilder { Slice GetRestartKey(uint32_t index, const char* limit) const; + // Returns key with timestamp stripped if applicable. + // For efficiency and internal consistency, only uses the bottom 32 bits of + // the key size (see API comments on Add()). inline const Slice MaybeStripTimestampFromKey(std::string* key_buf, const Slice& key); @@ -130,8 +139,8 @@ class BlockBuilder { std::string buffer_; // Destination buffer std::vector restarts_; // Restart points size_t estimate_; - int counter_; // Number of entries emitted since restart - bool finished_; // Has Finish() been called? + int counter_; // Number of entries emitted since restart + bool finished_; // Has Finish() been called? bool is_uniform_; // Was the last Finish()'d block uniform? std::string last_key_; DataBlockHashIndexBuilder data_block_hash_index_builder_; diff --git a/table/block_based/index_builder.cc b/table/block_based/index_builder.cc index 5771bfb86d..f089060a3f 100644 --- a/table/block_based/index_builder.cc +++ b/table/block_based/index_builder.cc @@ -333,6 +333,7 @@ Status PartitionedIndexBuilder::Finish( last_partition_block_handle.size() - last_encoded_handle_.size()); last_encoded_handle_ = last_partition_block_handle; const Slice handle_delta_encoding_slice(handle_delta_encoding); + // NOTE: WriteBatch guarantees keys < 4GB; handle values are also small index_block_builder_.Add(last_entry.key, handle_encoding.AsSlice(), &handle_delta_encoding_slice); if (!must_use_separator_with_seq_.LoadRelaxed()) { diff --git a/table/block_based/index_builder.h b/table/block_based/index_builder.h index 92f51e6f0c..4d7df7607f 100644 --- a/table/block_based/index_builder.h +++ b/table/block_based/index_builder.h @@ -344,6 +344,7 @@ class ShortenedIndexBuilder : public IndexBuilder { // away the UDT from key in index block as data block does the same thing. // What are the implications if a "FindShortInternalKeySuccessor" // optimization is provided. + // NOTE: WriteBatch guarantees keys < 4GB; handle values are also small index_block_builder_.Add(separator_with_seq, encoded_entry, &delta_encoded_entry_slice, skip_delta_encoding); if (!must_use_separator_with_seq) { diff --git a/table/block_based/partitioned_filter_block.cc b/table/block_based/partitioned_filter_block.cc index e84f0597c5..a553fd62e0 100644 --- a/table/block_based/partitioned_filter_block.cc +++ b/table/block_based/partitioned_filter_block.cc @@ -291,6 +291,7 @@ Status PartitionedFilterBlockBuilder::Finish( last_encoded_handle_ = last_partition_block_handle; const Slice handle_delta_encoding_slice(handle_delta_encoding); + // NOTE: WriteBatch guarantees keys < 4GB; handle values are also small index_on_filter_block_builder_.Add(e.ikey, handle_encoding, &handle_delta_encoding_slice); if (!p_index_builder_->separator_is_key_plus_seq()) { diff --git a/table/meta_blocks.cc b/table/meta_blocks.cc index d32f9bb94b..6b4b117977 100644 --- a/table/meta_blocks.cc +++ b/table/meta_blocks.cc @@ -43,6 +43,7 @@ void MetaIndexBuilder::Add(const std::string& key, const BlockHandle& handle) { Slice MetaIndexBuilder::Finish() { for (const auto& metablock : meta_block_handles_) { + // NOTE: meta index keys and block handles are guaranteed < 4GB meta_index_block_->Add(metablock.first, metablock.second); } return meta_index_block_->Finish(); @@ -198,7 +199,11 @@ Slice PropertyBlockBuilder::Finish() { for (const auto& prop : props_) { assert(last_prop_added_to_block_.empty() || comparator_->Compare(prop.first, last_prop_added_to_block_) > 0); - properties_block_->Add(prop.first, prop.second); + // Use first 4GB of key and value strings to avoid overflow, e.g. from + // user property collector (see BlockBuilder::Add API comments) + Slice key(prop.first.data(), static_cast(prop.first.size())); + Slice value(prop.second.data(), static_cast(prop.second.size())); + properties_block_->Add(key, value); #ifndef NDEBUG last_prop_added_to_block_ = prop.first; #endif /* !NDEBUG */ diff --git a/table/sst_file_writer.cc b/table/sst_file_writer.cc index cf6c32cdf7..2d958ade08 100644 --- a/table/sst_file_writer.cc +++ b/table/sst_file_writer.cc @@ -80,6 +80,16 @@ struct SstFileWriter::Rep { return builder->status(); } + // user_key + kNumInternalBytes must fit in uint32_t (BlockBuilder + // assumption). Also check value size. + if (user_key.size() > + size_t{std::numeric_limits::max()} - kNumInternalBytes) { + return Status::InvalidArgument("key is too large"); + } + if (value.size() > size_t{std::numeric_limits::max()}) { + return Status::InvalidArgument("value is too large"); + } + assert(user_key.size() >= ts_sz); if (strip_timestamp) { // In this mode, we expect users to always provide a min timestamp. @@ -173,6 +183,16 @@ struct SstFileWriter::Rep { if (!builder) { return Status::InvalidArgument("File is not opened"); } + // begin_key + kNumInternalBytes must fit in uint32_t (BlockBuilder + // assumption). end_key is stored as the value in the range deletion + // block, so it only needs to fit in uint32_t. + if (begin_key.size() > + size_t{std::numeric_limits::max()} - kNumInternalBytes) { + return Status::InvalidArgument("key is too large"); + } + if (end_key.size() > size_t{std::numeric_limits::max()}) { + return Status::InvalidArgument("end key is too large"); + } int cmp = internal_comparator.user_comparator()->CompareWithoutTimestamp( begin_key, end_key); if (cmp > 0) { diff --git a/test_util/testharness.cc b/test_util/testharness.cc index d3dfba1559..8b8a43bf1d 100644 --- a/test_util/testharness.cc +++ b/test_util/testharness.cc @@ -9,10 +9,15 @@ #include "test_util/testharness.h" +#include #include #include #include +#ifndef OS_WIN +#include +#endif + #ifndef NDEBUG #include "test_util/sync_point.h" #endif @@ -98,6 +103,25 @@ int RandomSeed() { return result; } +bool HasBigMem() { + const char* env = getenv("ROCKSDB_BIGMEM_TESTS"); + if (env != nullptr && env[0] != '\0') { + return true; + } +#ifdef _SC_PHYS_PAGES + // Check whether the system has at least 128GB of physical RAM. + // _SC_PHYS_PAGES is available on Linux and macOS but is not standard POSIX. + long pages = sysconf(_SC_PHYS_PAGES); + long page_size = sysconf(_SC_PAGE_SIZE); + if (pages > 0 && page_size > 0) { + size_t total_bytes = + static_cast(pages) * static_cast(page_size); + return total_bytes >= size_t{128} << 30; + } +#endif + return false; +} + TestRegex::TestRegex(const std::string& pattern) : impl_(std::make_shared(pattern)), pattern_(pattern) {} TestRegex::TestRegex(const char* pattern) diff --git a/test_util/testharness.h b/test_util/testharness.h index d8b6c9679c..9379225372 100644 --- a/test_util/testharness.h +++ b/test_util/testharness.h @@ -77,6 +77,12 @@ std::string PerThreadDBPath(std::string dir, std::string name); // runs may be able to vary the seed. int RandomSeed(); +// Returns true if the system has enough memory to run expensive +// large-allocation tests (~128GB+ RAM or environment variable +// ROCKSDB_BIGMEM_TESTS is set). There should generally be only one "big mem" +// test per unit test file, to reduce OOM risks in parallel runs. +bool HasBigMem(); + ::testing::AssertionResult AssertStatus(const char* s_expr, const Status& s); #define ASSERT_OK(s) \