Repository: cockroachdb/pebble
Pull request #4508 (2025-04-08 22:26:13)

crl-release-23.2: rowblk: raise max block size to (1<<31)-1; error rather than panic
23.2 backport of #4503 and #4507.

----

Relax the maximum block size to reclaim 7-unused bits within the restart offsets. This provides additional leeway for massive blocks we're sometimes forced to construct due to very large key-value pairs. Additionally, this commit turns this check into an error that we propagate up and out of the sstable Writer rather than panicking with an assertion failure. In practice this means a DB that hits this limit may fail a compaction without bringing down the node. The DB may repeat the compaction in a busy loop, but it's possible for the DB to make some forward progress in the meantime.

cockroach-teamcity: 

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/cockroachdb/pebble/4508)
<!-- Reviewable:end -->

jbowens: TFTR!

Patch:
diff --git a/options.go b/options.go
index 7657ccba..93849071 100644
--- a/options.go
+++ b/options.go
@@ -447,8 +447,8 @@ func (o *LevelOptions) EnsureDefaults() *LevelOptions {
 	}
 	if o.BlockSize <= 0 {
 		o.BlockSize = base.DefaultBlockSize
-	} else if o.BlockSize > sstable.MaximumBlockSize {
-		panic(errors.Errorf("BlockSize %d exceeds MaximumBlockSize", o.BlockSize))
+	} else if o.BlockSize > sstable.MaximumRestartOffset {
+		panic(errors.Errorf("BlockSize %d exceeds MaximumRestartOffset", o.BlockSize))
 	}
 	if o.BlockSizeThreshold <= 0 {
 		o.BlockSizeThreshold = base.DefaultBlockSizeThreshold
diff --git a/sstable/block.go b/sstable/block.go
index 36a11d9a..e76dc570 100644
--- a/sstable/block.go
+++ b/sstable/block.go
@@ -73,14 +73,20 @@ func (w *blockWriter) clear() {
 	}
 }
 
-// MaximumBlockSize is an extremely generous maximum block size of 256MiB. We
-// explicitly place this limit to reserve a few bits in the restart for
-// internal use.
-const MaximumBlockSize = 1 << 28
+// MaximumRestartOffset indicates the maximum offset that we can encode
+// within a restart point of a row-oriented block. The last bit is reserved
+// for the setHasSameKeyPrefixSinceLastRestart flag within a restart point.
+// If a block exceeds this size and we attempt to add another KV pair, the
+// restart points table will be unable to express the position of the pair,
+// resulting in undefined behavior and arbitrary corruption.
+const MaximumRestartOffset = 1<<31 - 1
 const setHasSameKeyPrefixRestartMask uint32 = 1 << 31
 const restartMaskLittleEndianHighByteWithoutSetHasSamePrefix byte = 0b0111_1111
 const restartMaskLittleEndianHighByteOnlySetHasSamePrefix byte = 0b1000_0000
 
+// ErrBlockTooBig is surfaced when a block exceeds the maximum size.
+var ErrBlockTooBig = errors.New("block size exceeds maximum size")
+
 func (w *blockWriter) getCurKey() InternalKey {
 	k := base.DecodeInternalKey(w.curKey)
 	k.Trailer = k.Trailer & trailerObsoleteMask
@@ -103,10 +109,9 @@ func (w *blockWriter) storeWithOptionalValuePrefix(
 	addValuePrefix bool,
 	valuePrefix valuePrefix,
 	setHasSameKeyPrefix bool,
-) {
-	if len(w.buf) >= MaximumBlockSize {
-		panic(errors.AssertionFailedf("block: adding KV to %d-block; already exceeds %d-byte maximum",
-			len(w.buf), MaximumBlockSize))
+) error {
+	if len(w.buf) >= MaximumRestartOffset {
+		return errors.WithDetailf(ErrBlockTooBig, "block is %d bytes long", len(w.buf))
 	}
 
 	shared := 0
@@ -208,10 +213,11 @@ func (w *blockWriter) storeWithOptionalValuePrefix(
 	w.curValue = w.buf[n-len(value):]
 
 	w.nEntries++
+	return nil
 }
 
-func (w *blockWriter) add(key InternalKey, value []byte) {
-	w.addWithOptionalValuePrefix(
+func (w *blockWriter) add(key InternalKey, value []byte) error {
+	return w.addWithOptionalValuePrefix(
 		key, false, value, len(key.UserKey), false, 0, false)
 }
 
@@ -232,7 +238,7 @@ func (w *blockWriter) addWithOptionalValuePrefix(
 	addValuePrefix bool,
 	valuePrefix valuePrefix,
 	setHasSameKeyPrefix bool,
-) {
+) error {
 	w.curKey, w.prevKey = w.prevKey, w.curKey
 
 	size := key.Size()
@@ -245,7 +251,7 @@ func (w *blockWriter) addWithOptionalValuePrefix(
 	}
 	key.Encode(w.curKey)
 
-	w.storeWithOptionalValuePrefix(
+	return w.storeWithOptionalValuePrefix(
 		size, value, maxSharedKeyLen, addValuePrefix, valuePrefix, setHasSameKeyPrefix)
 }
 
diff --git a/sstable/reader.go b/sstable/reader.go
index 6a7027a0..d0e03cfb 100644
--- a/sstable/reader.go
+++ b/sstable/reader.go
@@ -685,13 +685,18 @@ func (r *Reader) transformRangeDelV1(b []byte) ([]byte, error) {
 	rangeDelBlock := blockWriter{
 		restartInterval: 1,
 	}
+
+	// NB: There's subtle error handling here. The Emit closure will set err to
+	// the first non-nil error encountered. After we call frag.Finish(), we check
+	// err to see if the closure set an error.
+	var err error
 	frag := keyspan.Fragmenter{
 		Cmp:    r.Compare,
 		Format: r.FormatKey,
 		Emit: func(s keyspan.Span) {
 			for _, k := range s.Keys {
 				startIK := InternalKey{UserKey: s.Start, Trailer: k.Trailer}
-				rangeDelBlock.add(startIK, s.End)
+				err = firstError(err, rangeDelBlock.add(startIK, s.End))
 			}
 		},
 	}
@@ -699,6 +704,9 @@ func (r *Reader) transformRangeDelV1(b []byte) ([]byte, error) {
 		frag.Add(tombstones[i])
 	}
 	frag.Finish()
+	if err != nil {
+		return nil, err
+	}
 
 	// Return the contents of the constructed v2 format range-del block.
 	return rangeDelBlock.finish(), nil
diff --git a/sstable/suffix_rewriter.go b/sstable/suffix_rewriter.go
index 25254cde..75224985 100644
--- a/sstable/suffix_rewriter.go
+++ b/sstable/suffix_rewriter.go
@@ -263,7 +263,9 @@ func rewriteBlocks(
 					return errors.Errorf("multiple keys with same key prefix")
 				}
 			}
-			bw.add(scratch, v)
+			if err := bw.add(scratch, v); err != nil {
+				return err
+			}
 			if output[i].start.UserKey == nil {
 				keyAlloc, output[i].start = cloneKeyWithBuf(scratch, keyAlloc)
 			}
diff --git a/sstable/writer.go b/sstable/writer.go
index 883ee7ce..3a8b544f 100644
--- a/sstable/writer.go
+++ b/sstable/writer.go
@@ -437,14 +437,17 @@ func (i *indexBlockBuf) shouldFlush(
 		int(nEntries), targetBlockSize, sizeThreshold)
 }
 
-func (i *indexBlockBuf) add(key InternalKey, value []byte, inflightSize int) {
-	i.block.add(key, value)
+func (i *indexBlockBuf) add(key InternalKey, value []byte, inflightSize int) error {
+	if err := i.block.add(key, value); err != nil {
+		return err
+	}
 	size := i.block.estimatedSize()
 	if i.size.useMutex {
 		i.size.mu.Lock()
 		defer i.size.mu.Unlock()
 	}
 	i.size.estimate.writtenWithTotal(uint64(size), inflightSize)
+	return nil
 }
 
 func (i *indexBlockBuf) finish() []byte {
@@ -1008,9 +1011,12 @@ func (w *Writer) addPoint(key InternalKey, value []byte, forceObsolete bool) err
 	}
 
 	w.maybeAddToFilter(key.UserKey)
-	w.dataBlockBuf.dataBlock.addWithOptionalValuePrefix(
+	err = w.dataBlockBuf.dataBlock.addWithOptionalValuePrefix(
 		key, isObsolete, valueStoredWithKey, maxSharedKeyLen, addPrefixToValueStoredWithKey, prefix,
 		setHasSameKeyPrefix)
+	if err != nil {
+		return err
+	}
 
 	w.meta.updateSeqNum(key.SeqNum())
 
@@ -1146,8 +1152,7 @@ func (w *Writer) addTombstone(key InternalKey, value []byte) error {
 	w.props.NumRangeDeletions++
 	w.props.RawKeySize += uint64(key.Size())
 	w.props.RawValueSize += uint64(len(value))
-	w.rangeDelBlock.add(key, value)
-	return nil
+	return w.rangeDelBlock.add(key, value)
 }
 
 // RangeKeySet sets a range between start (inclusive) and end (exclusive) with
@@ -1347,8 +1352,7 @@ func (w *Writer) addRangeKey(key InternalKey, value []byte) error {
 	}
 
 	// Add the key to the block.
-	w.rangeKeyBlock.add(key, value)
-	return nil
+	return w.rangeKeyBlock.add(key, value)
 }
 
 // tempRangeKeyBuf returns a slice of length n from the Writer's rkBuf byte
@@ -1572,9 +1576,7 @@ func (w *Writer) addIndexEntry(
 			return err
 		}
 	}
-
-	writeTo.add(sep, encoded, inflightSize)
-	return nil
+	return writeTo.add(sep, encoded, inflightSize)
 }
 
 func (w *Writer) addPrevDataBlockToIndexBlockProps() {
@@ -1738,7 +1740,9 @@ func (w *Writer) writeTwoLevelIndex() (BlockHandle, error) {
 			Props:       b.properties,
 		}
 		encoded := encodeBlockHandleWithProperties(w.blockBuf.tmp[:], bhp)
-		w.topLevelIndexBlock.add(b.sep, encoded)
+		if err := w.topLevelIndexBlock.add(b.sep, encoded); err != nil {
+			return BlockHandle{}, err
+		}
 	}
 
 	// NB: RocksDB includes the block trailer length in the index size


Test patch:
diff --git a/sstable/block_test.go b/sstable/block_test.go
index 7922aaf2..bed73ba6 100644
--- a/sstable/block_test.go
+++ b/sstable/block_test.go
@@ -17,6 +17,7 @@ import (
 	"unsafe"
 
 	"github.com/cockroachdb/datadriven"
+	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/rand"
@@ -58,8 +59,9 @@ func TestBlockWriterWithPrefix(t *testing.T) {
 		addValuePrefix bool,
 		valuePrefix valuePrefix,
 		setHasSameKeyPrefix bool) {
-		w.addWithOptionalValuePrefix(
+		err := w.addWithOptionalValuePrefix(
 			key, false, value, len(key.UserKey), addValuePrefix, valuePrefix, setHasSameKeyPrefix)
+		require.NoError(t, err)
 	}
 	addAdapter(
 		ikey("apple"), []byte("red"), false, 0, true)
@@ -126,10 +128,9 @@ func testBlockCleared(t *testing.T, w, b *blockWriter) {
 
 func TestBlockClear(t *testing.T) {
 	w := blockWriter{restartInterval: 16}
-	w.add(ikey("apple"), nil)
-	w.add(ikey("apricot"), nil)
-	w.add(ikey("banana"), nil)
-
+	require.NoError(t, w.add(ikey("apple"), nil))
+	require.NoError(t, w.add(ikey("apricot"), nil))
+	require.NoError(t, w.add(ikey("banana"), nil))
 	w.clear()
 
 	// Once a block is cleared, we expect its fields to be cleared, but we expect
@@ -241,7 +242,7 @@ func TestBlockIter2(t *testing.T) {
 				case "build":
 					w := &blockWriter{restartInterval: r}
 					for _, e := range strings.Split(strings.TrimSpace(d.Input), ",") {
-						w.add(makeIkey(e), nil)
+						require.NoError(t, w.add(makeIkey(e), nil))
 					}
 					block = w.finish()
 					return ""
@@ -310,7 +311,7 @@ func TestBlockIterKeyStability(t *testing.T) {
 		[]byte("banana"),
 	}
 	for i := range expected {
-		w.add(InternalKey{UserKey: expected[i]}, nil)
+		require.NoError(t, w.add(InternalKey{UserKey: expected[i]}, nil))
 	}
 	block := w.finish()
 
@@ -368,7 +369,7 @@ func TestBlockIterReverseDirections(t *testing.T) {
 		[]byte("carrot"),
 	}
 	for i := range keys {
-		w.add(InternalKey{UserKey: keys[i]}, nil)
+		require.NoError(t, w.add(InternalKey{UserKey: keys[i]}, nil))
 	}
 	block := w.finish()
 
@@ -394,8 +395,13 @@ func TestBlockIterReverseDirections(t *testing.T) {
 		})
 	}
 }
-func TestSingularKVBlockRestartsOverflow(t *testing.T) {
 
+// TestSingularKVBlockRestartsOverflow tests a scenario where a large key-value
+// pair is written to a block, such that the total block size exceeds 4GiB. This
+// works becasue the restart table never needs to encode a restart offset beyond
+// the 1st key-value pair. The offset of the restarts table itself may exceed
+// 2^32-1 but the iterator takes care to support this.
+func TestSingularKVBlockRestartsOverflow(t *testing.T) {
 	_, isCI := os.LookupEnv("CI")
 	if isCI {
 		t.Skip("Skipping test: requires too much memory for CI now.")
@@ -412,12 +418,11 @@ func TestSingularKVBlockRestartsOverflow(t *testing.T) {
 
 	var largeKeySize int64 = 2 << 30   // 2GB key size
 	var largeValueSize int64 = 2 << 30 // 2GB value size
-
 	largeKey := bytes.Repeat([]byte("k"), int(largeKeySize))
 	largeValue := bytes.Repeat([]byte("v"), int(largeValueSize))
 
 	writer := &blockWriter{restartInterval: 1}
-	writer.add(base.InternalKey{UserKey: largeKey}, largeValue)
+	require.NoError(t, writer.add(base.InternalKey{UserKey: largeKey}, largeValue))
 	blockData := writer.finish()
 	iter, err := newBlockIter(bytes.Compare, blockData)
 	require.NoError(t, err, "failed to create iterator for block")
@@ -441,25 +446,26 @@ func TestSingularKVBlockRestartsOverflow(t *testing.T) {
 	require.Equal(t, largeValue, value.ValueOrHandle, "unexpected value")
 }
 
-func TestBufferExceeding256MBShouldPanic(t *testing.T) {
-
+// TestExceedingMaximumRestartOffset tests that writing a block that exceeds the
+// maximum restart offset errors.
+func TestExceedingMaximumRestartOffset(t *testing.T) {
 	_, isCI := os.LookupEnv("CI")
 	if isCI {
 		t.Skip("Skipping test: requires too much memory for CI now.")
 	}
 
-	// Test that writing to a block that is already >= 256MiB
-	// causes a panic to occur.
-
+	// Test that writing to a block that is already >= 2GiB
+	// returns an error.
+	//
 	// Skip this test on 32-bit architectures because they may not
 	// have sufficient memory to reliably execute this test.
 	if runtime.GOARCH == "386" || runtime.GOARCH == "arm" || strconv.IntSize == 32 {
 		t.Skip("Skipping test: not supported on 32-bit architecture")
 	}
 
-	// Adding 64 KVs each with size 4MiB will create a block
-	// size of >= ~256MiB
-	const numKVs = 64
+	// Adding 512 KVs each with size 4MiB will create a block
+	// size of >= 2GiB.
+	const numKVs = 512
 	const valueSize = (1 << 20) * 4
 
 	type KVTestPair struct {
@@ -473,27 +479,31 @@ func TestBufferExceeding256MBShouldPanic(t *testing.T) {
 		key := fmt.Sprintf("key-%04d", i)
 		KVTestPairs[i] = KVTestPair{key: []byte(key), value: value4MB}
 	}
-
 	writer := &blockWriter{restartInterval: 1}
 	for _, KVPair := range KVTestPairs {
-		writer.add(base.InternalKey{UserKey: KVPair.key}, KVPair.value)
+		require.NoError(t, writer.add(base.InternalKey{UserKey: KVPair.key}, KVPair.value))
 	}
 
-	// Check that buffer is larger than 256MiB
-	require.Greater(t, len(writer.buf), MaximumBlockSize)
+	// Check that buffer is larger than 2GiB
+	require.Greater(t, len(writer.buf), MaximumRestartOffset)
 
-	// Check that a panic has occurred after the final write after the 256MiB
+	// Check that an error is returned after the final write after the 2GiB
 	// threshold has been crossed
-	defer func() {
-		if r := recover(); r == nil {
-			t.Fatalf("expected panic on the last write, but none occurred")
-		}
-	}()
-	writer.add(base.InternalKey{UserKey: []byte("arbitrary-last-key")}, []byte("arbitrary-last-value"))
+	err := writer.add(base.InternalKey{UserKey: []byte("arbitrary-last-key")}, []byte("arbitrary-last-value"))
+	require.NotNil(t, err)
+	require.True(t, errors.Is(err, ErrBlockTooBig))
 }
 
+// TestMultipleKVBlockRestartsOverflow tests that SeekGE() works when
+// iter.restarts is greater than math.MaxUint32 for multiple KVs. Test writes
+// <2GiB to the block and then 4GiB causing iter.restarts to be an int >
+// math.MaxUint32. Reaching just shy of 2GiB before adding 4GiB allows the
+// final write to succeed without surpassing 2GiB limit. Then verify that
+// SeekGE() returns valid output without integer overflow.
+//
+// Although the block exceeds math.MaxUint32 bytes, no individual KV pair has an
+// offset that exceeds MaximumRestartOffset.
 func TestMultipleKVBlockRestartsOverflow(t *testing.T) {
-
 	_, isCI := os.LookupEnv("CI")
 	if isCI {
 		t.Skip("Skipping test: requires too much memory for CI now.")
@@ -514,10 +524,10 @@ func TestMultipleKVBlockRestartsOverflow(t *testing.T) {
 		t.Skip("Skipping test: not supported on 32-bit architecture")
 	}
 
-	// Write just shy of 256MiB to the block 63 * 4MiB < 256MiB
+	// Write just shy of 2GiB to the block 63 * 4MiB < 2GiB
 	const numKVs = 63
 	const valueSize = 4 * (1 << 20)
-	var FourGB int64 = 4 * (1 << 30)
+	var fourGB int64 = 4 * (1 << 30)
 
 	type KVTestPair struct {
 		key   []byte
@@ -533,23 +543,23 @@ func TestMultipleKVBlockRestartsOverflow(t *testing.T) {
 
 	writer := &blockWriter{restartInterval: 1}
 	for _, KVPair := range KVTestPairs {
-		writer.add(base.InternalKey{UserKey: KVPair.key}, KVPair.value)
+		require.NoError(t, writer.add(base.InternalKey{UserKey: KVPair.key}, KVPair.value))
 	}
 
 	// Add the 4GiB KV, causing iter.restarts >= math.MaxUint32.
 	// Ensure that SeekGE() works thereafter without integer
 	// overflows.
-	writer.add(base.InternalKey{UserKey: []byte("large-kv")}, []byte(strings.Repeat("v", int(FourGB))))
+	require.NoError(t, writer.add(base.InternalKey{UserKey: []byte("large-kv")}, bytes.Repeat([]byte("v"), int(fourGB))))
 
 	blockData := writer.finish()
 	iter, err := newBlockIter(bytes.Compare, blockData)
 	require.NoError(t, err, "failed to create iterator for block")
-	require.Greater(t, int64(iter.restarts), int64(MaximumBlockSize), "check iter.restarts > 256MiB")
+	require.Greater(t, int64(iter.restarts), int64(MaximumRestartOffset), "check iter.restarts > 2GiB")
 	require.Greater(t, int64(iter.restarts), int64(math.MaxUint32), "check iter.restarts > 2^32-1")
 
 	for i := 0; i < numKVs; i++ {
 		expectedKey := []byte(fmt.Sprintf("key-%04d", i))
-		expectedValue := []byte(strings.Repeat("a", valueSize))
+		expectedValue := bytes.Repeat([]byte("a"), valueSize)
 		key, value := iter.SeekGE(expectedKey, base.SeekGEFlagsNone)
 		require.NotNil(t, key, "failed to find the large key")
 		require.Equal(t, expectedKey, key.UserKey, "unexpected key")
diff --git a/sstable/writer_test.go b/sstable/writer_test.go
index 20f9e90e..75ce5e18 100644
--- a/sstable/writer_test.go
+++ b/sstable/writer_test.go
@@ -423,8 +423,8 @@ func TestBlockBufClear(t *testing.T) {
 func TestClearDataBlockBuf(t *testing.T) {
 	d := newDataBlockBuf(1, ChecksumTypeCRC32c)
 	d.blockBuf.compressedBuf = make([]byte, 1)
-	d.dataBlock.add(ikey("apple"), nil)
-	d.dataBlock.add(ikey("banana"), nil)
+	require.NoError(t, d.dataBlock.add(ikey("apple"), nil))
+	require.NoError(t, d.dataBlock.add(ikey("banana"), nil))
 
 	d.clear()
 	testBlockCleared(t, &d.dataBlock, &blockWriter{})
@@ -435,8 +435,8 @@ func TestClearDataBlockBuf(t *testing.T) {
 
 func TestClearIndexBlockBuf(t *testing.T) {
 	i := newIndexBlockBuf(false)
-	i.block.add(ikey("apple"), nil)
-	i.block.add(ikey("banana"), nil)
+	require.NoError(t, i.block.add(ikey("apple"), nil))
+	require.NoError(t, i.block.add(ikey("banana"), nil))
 	i.clear()
 
 	testBlockCleared(t, &i.block, &blockWriter{})