From 9f474a1034edf2b865ad309fa03b0e0ff9ac4415 Mon Sep 17 00:00:00 2001 From: Peter Dillinger Date: Fri, 10 Apr 2026 17:04:47 -0700 Subject: [PATCH] Enforce < 4GB sizes for keys, values, and some BlockBuilder inputs (#14461) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: BlockBuilder was somewhat inconsistent in its treatment of Slices whose size exceeds 4GB, which in a random corruption case could lead to (for example) serializing only the bottom 32 bits of a value size (uncorrupted) but appending the full multi-GB (corrupted) size. Because this is inner loop code, we don't want to pay CPU for extra conversions, data movements, or Status plumbing (already ruled out by Xingbo). In this change we make BlockBuilder more internally consistent and lift the 32-bit size requirement to callers (of a few functions specifically, for now). To ensure that's satisfied, I've added additional checks near the perimeter of RocksDB to ensure keys and values do not exceed 4GB, plus an extra random corruption or backstop check in BlockBasedTableBuilder. In detail, BlockBuilder (block_builder.cc/h): * Add API comments documenting the < 4GB assumption on Add/AddWithLastKey * Add debug assertions verifying input slice sizes < 4GB * Simplify AddWithLastKeyImpl to use uint32_t locals, reducing static_cast * Use only bottom 32 bits when appending derivative slices for consistency * Update MaybeStripTimestampFromKey to also truncate to 32-bit size * Add FIXME comments where buffer_/values_buffer_ sizes are truncated BlockBasedTableBuilder (block_based_table_builder.cc): * Add value size check (> uint32_t max) as a safety net against random corruption, since we have seen such corruptions in production WriteBatch (write_batch.cc): * Tighten existing key size checks to account for kNumInternalBytes (8-byte internal key suffix), using new kMaxWriteBatchKeySize constant * Add missing size checks to Delete, SingleDelete, and DeleteRange (both Slice and SliceParts variants) SstFileWriter (sst_file_writer.cc): * Add key and value size checks in AddImpl and DeleteRangeImpl, which bypass WriteBatch and go directly to the table builder MergeHelper (merge_helper.cc): * Add merge result size check (> uint32_t max) in both overloads of TimedFullMergeImpl, returning Corruption if exceeded * Add PartialMergeMulti result size check in MergeUntil CompactionIterator (compaction_iterator.cc): * Add size check on compaction filter output values (kChangeValue and kChangeWideColumnEntity), returning Corruption if > 4GB meta_blocks.cc: * Use Slice with uint32_t-truncated sizes in PropertyBlockBuilder::Finish to handle potentially oversized user property collector output Needed follow-up: * Check for blocks that are mis-encoded due to overflowing 32 bits for restart point offsets (and similar). See FIXME comments in block_builder.cc for why this is tricky. Pull Request resolved: https://github.com/facebook/rocksdb/pull/14461 Test Plan: New extreme-size unit tests, along with some testing infrastructure improvements: Added test::HasBigMem() (in test_util/testharness.h) which returns true when the system has ≥128GB RAM (via sysconf(_SC_PHYS_PAGES) on Linux/macOS) or when ROCKSDB_BIGMEM_TESTS is set. All extreme-size tests use HasBigMem() to skip gracefully on smaller machines rather than being permanently disabled. Where possible, tests use MemMapping::AllocateLazyZeroed() to supply large keys/values as Slices backed by anonymous mmap (cleaner with new MemMapping::AsSlice()). On Linux, read-only access to these pages maps to the shared kernel zero page, so the source data consumes no physical RAM — only the destination copy (e.g., WriteBatch::rep_) materializes, cutting peak memory roughly in half vs. std::string. Tests (all enabled, skip via HasBigMem()): ./write_batch_test --gtest_filter='*LargeKeyValueSizeLimit*' ./external_sst_file_basic_test --gtest_filter='*LargeSizeSstFileWriter*' ./merge_test --gtest_filter='*LargeMergeResultRejected*' ./merge_helper_test --gtest_filter='*LargePartialMergeResultRejected*' ./db_compaction_test --gtest_filter='*CompactionFilterLargeValueRejected*' All 5 pass (verified on a 128GB+ machine). On smaller machines, all 5 bypass cleanly with "insufficient memory for reliable continuous testing". Reviewed By: xingbowang Differential Revision: D96521899 Pulled By: pdillinger fbshipit-source-id: 70f4b5e6a23ab074d60e653fbb7ddc5edbe162ab --- db/compaction/compaction_iterator.cc | 8 +++ db/db_compaction_test.cc | 50 +++++++++++++ db/external_sst_file_basic_test.cc | 50 +++++++++++++ db/merge_helper.cc | 39 ++++++++--- db/merge_helper_test.cc | 66 +++++++++++++++++ db/merge_test.cc | 59 ++++++++++++++++ db/write_batch.cc | 41 +++++++++-- db/write_batch_test.cc | 49 +++++++++++++ port/mmap.h | 8 +++ .../block_based/block_based_table_builder.cc | 14 ++++ table/block_based/block_builder.cc | 70 ++++++++++++------- table/block_based/block_builder.h | 13 +++- table/block_based/index_builder.cc | 1 + table/block_based/index_builder.h | 1 + table/block_based/partitioned_filter_block.cc | 1 + table/meta_blocks.cc | 7 +- table/sst_file_writer.cc | 20 ++++++ test_util/testharness.cc | 24 +++++++ test_util/testharness.h | 6 ++ 19 files changed, 486 insertions(+), 41 deletions(-) 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) \