blob: uncouple blob file rewriting logic

Uncouple the logic around actually writing out a rewritten blob file from the
logic around identifying which values remain live. This will help us test both
parts in isolation.

Additionally, this commit fixes a bug in which values with the same BlockID
could be split across different physical blocks if the blob.FileWriter's
FlushGovernor decided they should during a call to AddValue.

Author: jbowens

blob_rewrite.go

@@ -280,14 +280,20 @@ func (d *DB) runBlobFileRewriteLocked(
 		Path:    d.objProvider.Path(objMeta),
 		FileNum: objMeta.DiskFileNum,
 	})
-	// Initialize a blob file writer. We pass L6 to MakeBlobWriterOptions.
+	// Initialize a blob file rewriter. We pass L6 to MakeBlobWriterOptions.
 	// There's no single associated level with a blob file. A long-lived blob
 	// file that gets rewritten is likely to mostly be referenced from L6.
 	// TODO(jackson): Consider refactoring to remove the level association.
-	writer := blob.NewFileWriter(newDiskFileNum, writable, d.opts.MakeBlobWriterOptions(6))
-
+	rewriter := newBlobFileRewriter(
+		d.fileCache,
+		env,
+		newDiskFileNum,
+		writable,
+		d.opts.MakeBlobWriterOptions(6),
+		c.referencingTables,
+		c.input,
+	)
 	// Perform the rewrite.
-	rewriter := newBlobFileRewriter(d.fileCache, env, writer, c.referencingTables, c.input)
 	stats, err := rewriter.Rewrite(ctx)
 	if err != nil {
 		return objstorage.ObjectMetadata{}, nil, err
@@ -361,51 +367,35 @@ type blockValues struct {
 	liveValueIDs []int
 }
 
-// blobFileMapping implements blob.FileMapping to always map to the input blob
-// file.
-type blobFileMapping struct {
-	fileNum base.DiskFileNum
-}
-
-// Assert that (*blobFileMapping) implements blob.FileMapping.
-var _ blob.FileMapping = (*blobFileMapping)(nil)
-
-func (m *blobFileMapping) Lookup(fileID base.BlobFileID) (base.DiskFileNum, bool) {
-	return m.fileNum, true
-}
-
 // blobFileRewriter is responsible for rewriting blob files by combining and
 // processing blob reference liveness encodings from multiple SSTables. It
 // maintains state for writing to an output blob file.
 type blobFileRewriter struct {
-	fc       *fileCacheHandle
-	env      block.ReadEnv
-	sstables []*manifest.TableMetadata
-
-	inputBlob    manifest.BlobFileMetadata
-	valueFetcher blob.ValueFetcher
-	fileMapping  blobFileMapping
-	blkHeap      blockHeap
-
-	// Current blob writer state.
-	writer *blob.FileWriter
+	fc        *fileCacheHandle
+	readEnv   block.ReadEnv
+	sstables  []*manifest.TableMetadata
+	inputBlob manifest.BlobFileMetadata
+	rw        *blob.FileRewriter
+	blkHeap   blockHeap
 }
 
 func newBlobFileRewriter(
 	fc *fileCacheHandle,
-	env block.ReadEnv,
-	writer *blob.FileWriter,
+	readEnv block.ReadEnv,
+	outputFileNum base.DiskFileNum,
+	w objstorage.Writable,
+	opts blob.FileWriterOptions,
 	sstables []*manifest.TableMetadata,
 	inputBlob manifest.BlobFileMetadata,
 ) *blobFileRewriter {
+	rw := blob.NewFileRewriter(inputBlob.FileID, inputBlob.Physical.FileNum, fc, readEnv, outputFileNum, w, opts)
 	return &blobFileRewriter{
-		fc:          fc,
-		env:         env,
-		writer:      writer,
-		sstables:    sstables,
-		inputBlob:   inputBlob,
-		fileMapping: blobFileMapping{fileNum: inputBlob.Physical.FileNum},
-		blkHeap:     blockHeap{},
+		fc:        fc,
+		readEnv:   readEnv,
+		rw:        rw,
+		sstables:  sstables,
+		inputBlob: inputBlob,
+		blkHeap:   blockHeap{},
 	}
 }
 
@@ -426,7 +416,7 @@ func (rw *blobFileRewriter) generateHeap() error {
 			return errors.AssertionFailedf("table %s doesn't contain a reference to blob file %s",
 				sst.TableNum, rw.inputBlob.FileID)
 		}
-		err := rw.fc.withReader(ctx, rw.env, sst, func(r *sstable.Reader, readEnv sstable.ReadEnv) error {
+		err := rw.fc.withReader(ctx, rw.readEnv, sst, func(r *sstable.Reader, readEnv sstable.ReadEnv) error {
 			h, err := r.ReadBlobRefIndexBlock(ctx, readEnv.Block)
 			if err != nil {
 				return err
@@ -449,42 +439,8 @@ func (rw *blobFileRewriter) generateHeap() error {
 	return nil
 }
 
-// copyBlockValues copies the live values from the given block to the output
-// blob file, flushing the current block before if necessary.
-func (rw *blobFileRewriter) copyBlockValues(ctx context.Context, finishedBlock blockValues) error {
-	shouldFlush := rw.writer.ShouldFlushBefore(finishedBlock.valuesSize)
-	if shouldFlush {
-		rw.writer.ForceFlush()
-	}
-
-	// Record the mapping from the virtual block ID to the current physical
-	// block and value ID offset.
-	rw.writer.BeginNewVirtualBlock(finishedBlock.blockID)
-	slices.Sort(finishedBlock.liveValueIDs)
-
-	for i, valueID := range finishedBlock.liveValueIDs {
-		if i > 0 && finishedBlock.liveValueIDs[i-1]+1 != valueID {
-			// There's a gap in the referenced value IDs.
-			for missing := finishedBlock.liveValueIDs[i-1] + 1; missing < valueID; missing++ {
-				rw.writer.AddValue(nil)
-			}
-		}
-
-		value, _, err := rw.valueFetcher.Fetch(ctx, rw.inputBlob.FileID, finishedBlock.blockID, blob.BlockValueID(valueID))
-		if err != nil {
-			return err
-		}
-		rw.writer.AddValue(value)
-	}
-	return nil
-}
-
 func (rw *blobFileRewriter) Rewrite(ctx context.Context) (blob.FileWriterStats, error) {
-	rw.valueFetcher.Init(&rw.fileMapping, rw.fc, rw.env)
-	defer func() { _ = rw.valueFetcher.Close() }()
-
-	err := rw.generateHeap()
-	if err != nil {
+	if err := rw.generateHeap(); err != nil {
 		return blob.FileWriterStats{}, err
 	}
 	if rw.blkHeap.Len() == 0 {
@@ -494,7 +450,7 @@ func (rw *blobFileRewriter) Rewrite(ctx context.Context) (blob.FileWriterStats,
 	// Begin constructing our output blob file. We maintain a map of blockID
 	// to accumulated liveness data across all referencing sstables.
 	firstBlock := heap.Pop(&rw.blkHeap).(*sstable.BlobRefLivenessEncoding)
-	pendingBlock := blockValues{
+	pending := blockValues{
 		blockID:      firstBlock.BlockID,
 		valuesSize:   firstBlock.ValuesSize,
 		liveValueIDs: slices.Collect(sstable.IterSetBitsInRunLengthBitmap(firstBlock.Bitmap)),
@@ -504,21 +460,23 @@ func (rw *blobFileRewriter) Rewrite(ctx context.Context) (blob.FileWriterStats,
 
 		// If we are encountering a new block, write the last accumulated block
 		// to the blob file.
-		if pendingBlock.blockID != nextBlock.BlockID {
+		if pending.blockID != nextBlock.BlockID {
 			// Write the last accumulated block's values to the blob file.
-			if err := rw.copyBlockValues(ctx, pendingBlock); err != nil {
+			err := rw.rw.CopyBlock(ctx, pending.blockID, pending.valuesSize, pending.liveValueIDs)
+			if err != nil {
 				return blob.FileWriterStats{}, err
 			}
 		}
 		// Update the accumulated encoding for this block.
-		pendingBlock.valuesSize += nextBlock.ValuesSize
-		pendingBlock.liveValueIDs = slices.AppendSeq(pendingBlock.liveValueIDs,
+		pending.valuesSize += nextBlock.ValuesSize
+		pending.liveValueIDs = slices.AppendSeq(pending.liveValueIDs,
 			sstable.IterSetBitsInRunLengthBitmap(nextBlock.Bitmap))
 	}
 
 	// Copy the last accumulated block.
-	if err := rw.copyBlockValues(ctx, pendingBlock); err != nil {
+	err := rw.rw.CopyBlock(ctx, pending.blockID, pending.valuesSize, pending.liveValueIDs)
+	if err != nil {
 		return blob.FileWriterStats{}, err
 	}
-	return rw.writer.Close()
+	return rw.rw.Close()
 }

blob_rewrite_test.go

@@ -192,8 +192,8 @@ func TestBlobRewrite(t *testing.T) {
 				outputWritable, _, err := objStore.Create(ctx, base.FileTypeBlob, fn, objstorage.CreateOptions{})
 				require.NoError(t, err)
 
-				blobWriter := blob.NewFileWriter(fn, outputWritable, blob.FileWriterOptions{})
-				rewriter := newBlobFileRewriter(mockFC, block.ReadEnv{}, blobWriter, sstables, inputBlob)
+				rewriter := newBlobFileRewriter(mockFC, block.ReadEnv{}, fn, outputWritable,
+					blob.FileWriterOptions{}, sstables, inputBlob)
 				stats, err := rewriter.Rewrite(context.Background())
 				if err != nil {
 					fmt.Fprintf(&buf, "rewrite error: %v\n", err)

internal/blobtest/handles.go

@@ -243,7 +243,7 @@ func (bv *Values) WriteFiles(
 		prevID := -1
 		for i, handle := range handles {
 			if i > 0 && handles[i-1].BlockID != handle.BlockID {
-				writer.ForceFlush()
+				writer.FlushForTesting()
 				prevID = -1
 			}
 			// The user of a blobtest.Values may specify a value ID for a handle. If

sstable/blob/blob.go

@@ -162,19 +162,21 @@ func (w *FileWriter) AddValue(v []byte) Handle {
 	}
 }
 
-// BeginNewVirtualBlock adds a virtual block mapping to the current physical
+// beginNewVirtualBlock adds a virtual block mapping to the current physical
 // block and valueID offset within the block.
 //
-// When a blob file is rewritten, BeginNewVirtualBlock is called for each block
+// When a blob file is rewritten, beginNewVirtualBlock is called for each block
 // in the original blob file before adding any of the block's extant values.
-// BeginNewVirtualBlock records a mapping from the original block ID (referred
+// beginNewVirtualBlock records a mapping from the original block ID (referred
 // to as a virtual block) to a tuple of the physical block index and the offset
 // of the BlockValueIDs within the new physical block.
 //
 // This mapping is used by readers to determine which physical block contains a
 // given virtual block, and how to map BlockValueIDs from the given virtual
 // block to BlockValueIDs in the physical block.
-func (w *FileWriter) BeginNewVirtualBlock(vblockID BlockID) {
+func (w *FileWriter) beginNewVirtualBlock(vblockID BlockID) {
+	// TODO(jackson): Update tests to use the blob.FileRewriter type and move this
+	// into the FileRewriter.
 	w.indexEncoder.AddVirtualBlockMapping(vblockID, int(w.stats.BlockCount),
 		BlockValueID(w.valuesEncoder.Count()))
 }
@@ -195,29 +197,19 @@ func (w *FileWriter) EstimatedSize() uint64 {
 	return sz
 }
 
-// ForceFlush flushes the current block to the write queue. Writers should
-// generally not call ForceFlush, and instead let the heuristics configured
+// FlushForTesting flushes the current block to the write queue. Writers should
+// generally not call FlushForTesting, and instead let the heuristics configured
 // through FileWriterOptions handle flushing.
 //
-// For blob file rewriting, ForceFlush needs to be called to ensure that there
-// is a 1:1 remapping of virtual blocks to physical blocks.
-//
-// Otherwise, it's exposed so that tests can force flushes to construct blob
-// files with arbitrary structures.
-func (w *FileWriter) ForceFlush() {
+// It's exposed so that tests can force flushes to construct blob files with
+// arbitrary structures.
+func (w *FileWriter) FlushForTesting() {
 	if w.valuesEncoder.Count() == 0 {
 		return
 	}
 	w.flush()
 }
 
-// ShouldFlush returns true if the current block should be flushed before adding
-// newDataBytes of data.
-func (w *FileWriter) ShouldFlushBefore(newDataBytes int) bool {
-	size := w.valuesEncoder.size() + block.TrailerLen
-	return w.flushGov.ShouldFlush(size, size+newDataBytes)
-}
-
 // flush flushes the current block to the write queue.
 func (w *FileWriter) flush() {
 	if w.valuesEncoder.Count() == 0 {

sstable/blob/blob_test.go

@@ -52,7 +52,7 @@ func TestBlobWriter(t *testing.T) {
 				case l == "---flush---":
 					w.flush()
 				case l == "---add-vblock---":
-					w.BeginNewVirtualBlock(BlockID(vBlockID))
+					w.beginNewVirtualBlock(BlockID(vBlockID))
 					vBlockID++
 				default:
 					h := w.AddValue([]byte(l))

sstable/blob/rewrite.go

@@ -0,0 +1,109 @@
+// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package blob
+
+import (
+	"context"
+	"slices"
+
+	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/objstorage"
+	"github.com/cockroachdb/pebble/sstable/block"
+)
+
+// A FileRewriter copies values from an input blob file, outputting a new blob
+// file containing a subset of the original blob file's values. The original
+// Handles used to access values in the original blob file will continue to work
+// with the new blob file, as long as the value was copied during rewrite.
+type FileRewriter struct {
+	fileID base.BlobFileID
+	w      *FileWriter
+	f      ValueFetcher
+}
+
+// NewFileRewriter creates a new FileRewriter that will copy values from the
+// input blob file to the output blob file.
+func NewFileRewriter(
+	fileID base.BlobFileID,
+	inputFileNum base.DiskFileNum,
+	rp ReaderProvider,
+	readEnv block.ReadEnv,
+	outputFileNum base.DiskFileNum,
+	w objstorage.Writable,
+	opts FileWriterOptions,
+) *FileRewriter {
+	rw := &FileRewriter{
+		fileID: fileID,
+		w:      NewFileWriter(outputFileNum, w, opts),
+	}
+	rw.f.Init(inputFileMapping(inputFileNum), rp, readEnv)
+	return rw
+}
+
+// CopyBlock copies the values for the given blockID to the output blob file.
+// CopyBlock must be called with ascending blockIDs. The totalValueSize must be
+// the size of all the values indicated by valueIDs.
+func (rw *FileRewriter) CopyBlock(
+	ctx context.Context, blockID BlockID, totalValueSize int, valueIDs []int,
+) error {
+	slices.Sort(valueIDs)
+
+	// Consider whether we should flush the current physical block. We know
+	// we'll need to add totalValueSize worth of value data, and can make a
+	// decision up front. All values from the same original blockID must be
+	// located in the same physical block.
+	valuesInBlock := rw.w.valuesEncoder.Count()
+	if valuesInBlock > 0 {
+		currentBlockSize := rw.w.valuesEncoder.size() + block.TrailerLen
+		if rw.w.flushGov.ShouldFlush(currentBlockSize, currentBlockSize+totalValueSize) {
+			rw.w.flush()
+		}
+	}
+
+	// Record the mapping from the virtual block ID to the current physical
+	// block and offset within the block.
+	rw.w.beginNewVirtualBlock(blockID)
+
+	previousValueID := -1
+	for _, valueID := range valueIDs {
+		// If there is a gap in the referenced Value IDs within this block, we
+		// need to represent this sparseness as empty values within the block.
+		// We can represent sparseness at the tail of a block or between blocks
+		// more compactly, but not sparseless at the beginning of a virtual
+		// block. See the doc.go comment for more details on sparseness.
+		for missingValueID := previousValueID + 1; missingValueID < valueID; missingValueID++ {
+			rw.w.stats.ValueCount++
+			rw.w.valuesEncoder.AddValue(nil)
+		}
+
+		// Retrieve the value and copy it to the output blob file.
+		value, _, err := rw.f.Fetch(ctx, rw.fileID, blockID, BlockValueID(valueID))
+		if err != nil {
+			return err
+		}
+		rw.w.stats.ValueCount++
+		rw.w.stats.UncompressedValueBytes += uint64(len(value))
+		rw.w.valuesEncoder.AddValue(value)
+		previousValueID = valueID
+	}
+	return nil
+}
+
+// Close finishes writing the output blob file and releases resources.
+func (rw *FileRewriter) Close() (FileWriterStats, error) {
+	stats, err := rw.w.Close()
+	return stats, errors.CombineErrors(err, rw.f.Close())
+}
+
+// inputFileMapping implements blob.FileMapping and always maps to itself.
+type inputFileMapping base.DiskFileNum
+
+// Assert that (*inputFileMapping) implements blob.FileMapping.
+var _ FileMapping = inputFileMapping(0)
+
+func (m inputFileMapping) Lookup(fileID base.BlobFileID) (base.DiskFileNum, bool) {
+	return base.DiskFileNum(m), true
+}