blob,manifest: lookup physical blob file in Version

Configure blob.ValueFetcher to take a func to translate a base.BlobFileID to
the base.DiskFileNum of the backing blob file. Add a (*BlobFileSet).Lookup
function with the same signature, and configure the various users of
ValueFetcher to propagate their Version's BlobFileSet's Lookup func.

Informs #4802.

Author: jbowens

checkpoint.go

@@ -14,7 +14,6 @@ import (
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/manifest"
 	"github.com/cockroachdb/pebble/record"
-	"github.com/cockroachdb/pebble/sstable/blob"
 	"github.com/cockroachdb/pebble/vfs"
 	"github.com/cockroachdb/pebble/vfs/atomicfs"
 )
@@ -171,6 +170,9 @@ func (d *DB) Checkpoint(
 	}
 
 	// Disable file deletions.
+	// We acquire a reference on the version down below that will prevent any
+	// sstables or blob files from becoming "obsolete" and potentially deleted,
+	// but this doesn't protect the current WALs or manifests.
 	d.mu.Lock()
 	d.disableFileDeletions()
 	defer func() {
@@ -205,10 +207,17 @@ func (d *DB) Checkpoint(
 	// before our call to List.
 	allLogicalLogs := d.mu.log.manager.List()
 
-	// Release the manifest and DB.mu so we don't block other operations on
-	// the database.
+	// Release the manifest and DB.mu so we don't block other operations on the
+	// database.
+	//
+	// But first reference the version to ensure that the version's in-memory
+	// state and its physical files remain available for the checkpoint. In
+	// particular, the Version.BlobFileSet is only valid while a version is
+	// referenced.
+	current.Ref()
 	d.mu.versions.logUnlock()
 	d.mu.Unlock()
+	defer current.Unref()
 
 	// Wrap the normal filesystem with one which wraps newly created files with
 	// vfs.NewSyncingFile.
@@ -316,10 +325,11 @@ func (d *DB) Checkpoint(
 					if _, ok := includedBlobFiles[ref.FileID]; !ok {
 						includedBlobFiles[ref.FileID] = struct{}{}
 
-						// TODO(jackson): Perform a translation to the
-						// appropriate disk file number once we support blob
-						// file replacement.
-						diskFileNum := blob.DiskFileNumTODO(ref.FileID)
+						// Map the BlobFileID to a DiskFileNum in the current version.
+						diskFileNum, ok := current.BlobFiles.Lookup(ref.FileID)
+						if !ok {
+							return errors.Errorf("blob file %s not found", ref.FileID)
+						}
 						ckErr = copyFile(base.FileTypeBlob, diskFileNum)
 						if ckErr != nil {
 							return ckErr

compaction.go

@@ -3215,7 +3215,7 @@ func (d *DB) compactAndWrite(
 			categoryCompaction,
 		),
 	}
-	c.valueFetcher.Init(d.fileCache, blockReadEnv)
+	c.valueFetcher.Init(&c.version.BlobFiles, d.fileCache, blockReadEnv)
 	iiopts := internalIterOpts{
 		compaction:       true,
 		readEnv:          sstable.ReadEnv{Block: blockReadEnv},

compaction_test.go

@@ -592,7 +592,7 @@ func TestAutomaticFlush(t *testing.T) {
 					continue
 				}
 				err := func() error {
-					fetcher, blobContext := sstable.LoadValBlobContext(d.fileCache, &meta.BlobReferences)
+					fetcher, blobContext := sstable.LoadValBlobContext(&v.BlobFiles, d.fileCache, &meta.BlobReferences)
 					defer fetcher.Close()
 
 					f, err := provider.OpenForReading(context.Background(), base.FileTypeTable, meta.TableBacking.DiskFileNum, objstorage.OpenOptions{})

db.go

@@ -623,7 +623,7 @@ func (d *DB) getInternal(key []byte, b *Batch, s *Snapshot) ([]byte, io.Closer,
 		keyBuf:       buf.keyBuf,
 	}
 	// Set up a blob value fetcher to use for retrieving values from blob files.
-	i.blobValueFetcher.Init(d.fileCache, block.NoReadEnv)
+	i.blobValueFetcher.Init(&readState.current.BlobFiles, d.fileCache, block.NoReadEnv)
 	get.iiopts.blobValueFetcher = &i.blobValueFetcher
 
 	if !i.First() {
@@ -1317,15 +1317,18 @@ func (d *DB) newInternalIter(
 	// files in the associated version from being deleted if there is a current
 	// compaction. The readState is unref'd by Iterator.Close().
 	var readState *readState
+	var vers *manifest.Version
 	if sOpts.vers == nil {
 		if sOpts.readState != nil {
 			readState = sOpts.readState
 			readState.ref()
+			vers = readState.current
 		} else {
 			readState = d.loadReadState()
+			vers = readState.current
 		}
-	}
-	if sOpts.vers != nil {
+	} else {
+		vers = sOpts.vers
 		sOpts.vers.Ref()
 	}
 
@@ -1352,7 +1355,7 @@ func (d *DB) newInternalIter(
 		seqNum:          seqNum,
 		mergingIter:     &buf.merging,
 	}
-	dbi.blobValueFetcher.Init(d.fileCache, block.ReadEnv{})
+	dbi.blobValueFetcher.Init(&vers.BlobFiles, d.fileCache, block.ReadEnv{})
 
 	dbi.opts = *o
 	dbi.opts.logger = d.opts.Logger
@@ -1434,7 +1437,11 @@ func (i *Iterator) constructPointIter(
 			i.opts.Category,
 		),
 	}
-	i.blobValueFetcher.Init(i.fc, readEnv)
+	if i.readState != nil {
+		i.blobValueFetcher.Init(&i.readState.current.BlobFiles, i.fc, readEnv)
+	} else if i.version != nil {
+		i.blobValueFetcher.Init(&i.version.BlobFiles, i.fc, readEnv)
+	}
 	internalOpts := internalIterOpts{
 		readEnv:          sstable.ReadEnv{Block: readEnv},
 		blobValueFetcher: &i.blobValueFetcher,

internal/manifest/blob_metadata.go

@@ -343,6 +343,56 @@ func (s *BlobFileSet) Count() int {
 	return s.tree.Count()
 }
 
+// Lookup returns the file number of the physical blob file backing the given
+// file ID. It returns false for the second return value if the FileID is not
+// present in the set.
+func (s *BlobFileSet) Lookup(fileID base.BlobFileID) (base.DiskFileNum, bool) {
+	phys, ok := s.LookupPhysical(fileID)
+	if !ok {
+		return 0, false
+	}
+	return phys.FileNum, true
+}
+
+// LookupPhysical returns the *PhysicalBlobFile backing the given file ID. It
+// returns false for the second return value if the FileID is not present in the
+// set.
+func (s *BlobFileSet) LookupPhysical(fileID base.BlobFileID) (*PhysicalBlobFile, bool) {
+	// LookupPhysical is performed during value retrieval to determine the
+	// physical blob file that should be read, so it's considered to be
+	// performance sensitive. We manually inline the B-Tree traversal and binary
+	// search with this in mind.
+	n := s.tree.root
+	for n != nil {
+		var h int
+		// Logic copied from sort.Search.
+		i, j := 0, int(n.count)
+		for i < j {
+			h = int(uint(i+j) >> 1) // avoid overflow when computing h
+			// i ≤ h < j
+			v := stdcmp.Compare(fileID, n.items[h].FileID)
+			if v == 0 {
+				// Found the sought blob file.
+				return n.items[h].Physical, true
+			} else if v > 0 {
+				i = h + 1
+			} else {
+				j = h
+			}
+		}
+		// If we've reached a lead node without finding fileID, the file is not
+		// present.
+		if n.leaf {
+			return nil, false
+		}
+		n = n.children[i]
+	}
+	return nil, false
+}
+
+// Assert that (*BlobFileSet) implements blob.FileMapping.
+var _ blob.FileMapping = (*BlobFileSet)(nil)
+
 // clone returns a copy-on-write clone of the blob file set.
 func (s *BlobFileSet) clone() BlobFileSet {
 	return BlobFileSet{tree: s.tree.Clone()}

internal/manifest/blob_metadata_test.go

@@ -141,3 +141,46 @@ func TestCurrentBlobFileSet(t *testing.T) {
 		return ""
 	})
 }
+
+func TestBlobFileSet_Lookup(t *testing.T) {
+	const numBlobFiles = 10000
+	set, files := makeTestBlobFiles(numBlobFiles)
+	for i := 0; i < numBlobFiles; i++ {
+		fn, ok := set.Lookup(base.BlobFileID(i))
+		require.True(t, ok)
+		require.Equal(t, files[i].FileNum, fn)
+	}
+}
+
+func makeTestBlobFiles(numBlobFiles int) (BlobFileSet, []PhysicalBlobFile) {
+	files := make([]PhysicalBlobFile, numBlobFiles)
+	for i := 0; i < numBlobFiles; i++ {
+		fileNum := base.DiskFileNum(i)
+		if i%2 == 0 {
+			fileNum = base.DiskFileNum(2*numBlobFiles + i)
+		}
+		files[i] = PhysicalBlobFile{
+			FileNum:      fileNum,
+			Size:         uint64(i),
+			ValueSize:    uint64(i),
+			CreationTime: uint64(i),
+		}
+	}
+	set := MakeBlobFileSet(nil)
+	for i := 0; i < numBlobFiles; i++ {
+		set.insert(BlobFileMetadata{
+			FileID:   base.BlobFileID(i % numBlobFiles),
+			Physical: &files[i],
+		})
+	}
+	return set, files
+}
+
+func BenchmarkBlobFileSet_Lookup(b *testing.B) {
+	const numBlobFiles = 10000
+	set, _ := makeTestBlobFiles(numBlobFiles)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, _ = set.Lookup(base.BlobFileID(i % numBlobFiles))
+	}
+}

level_checker.go

@@ -549,7 +549,7 @@ func (d *DB) CheckLevels(stats *CheckLevelsStats) error {
 			// TODO(jackson): Add categorized stats.
 		},
 	}
-	checkConfig.blobValueFetcher.Init(d.fileCache, checkConfig.readEnv)
+	checkConfig.blobValueFetcher.Init(&readState.current.BlobFiles, d.fileCache, checkConfig.readEnv)
 	defer func() { _ = checkConfig.blobValueFetcher.Close() }()
 	return checkLevelsInternal(checkConfig)
 }

replay/replay.go

@@ -797,6 +797,7 @@ func (r *Runner) prepareWorkloadSteps(ctx context.Context) error {
 				}
 				var newFiles []base.DiskFileNum
 				blobRefMap := make(map[base.DiskFileNum]manifest.BlobReferences)
+				blobFileMap := make(map[base.BlobFileID]base.DiskFileNum)
 				for _, nf := range ve.NewTables {
 					newFiles = append(newFiles, nf.Meta.TableBacking.DiskFileNum)
 					if s.kind == ingestStepKind && (nf.Meta.SmallestSeqNum != nf.Meta.LargestSeqNum || nf.Level != 0) {
@@ -806,6 +807,9 @@ func (r *Runner) prepareWorkloadSteps(ctx context.Context) error {
 						blobRefMap[nf.Meta.TableBacking.DiskFileNum] = nf.Meta.BlobReferences
 					}
 				}
+				for _, bf := range ve.NewBlobFiles {
+					blobFileMap[bf.FileID] = bf.Physical.FileNum
+				}
 				if previousVersion != nil {
 					// previousVersion contains the current version, and so l0Organizer is
 					// consistent with it. The subsequent call to currentVersion will
@@ -861,7 +865,7 @@ func (r *Runner) prepareWorkloadSteps(ctx context.Context) error {
 					// Load all of the flushed sstables' keys into a batch.
 					s.flushBatch = r.d.NewBatch()
 					err := loadFlushedSSTableKeys(s.flushBatch, r.WorkloadFS, r.WorkloadPath,
-						newFiles, blobRefMap, provider, r.readerOpts, &flushBufs)
+						newFiles, blobRefMap, blobFileMap, provider, r.readerOpts, &flushBufs)
 					if err != nil {
 						return errors.Wrapf(err, "flush in %q at offset %d", manifestName, rr.Offset())
 					}
@@ -995,6 +999,13 @@ func findManifestStart(
 	return index, info.Size(), nil
 }
 
+type blobFileMap map[base.BlobFileID]base.DiskFileNum
+
+func (m blobFileMap) Lookup(fileID base.BlobFileID) (base.DiskFileNum, bool) {
+	diskFileNum, ok := m[fileID]
+	return diskFileNum, ok
+}
+
 // loadFlushedSSTableKeys copies keys from the sstables specified by `fileNums`
 // in the directory specified by `path` into the provided batch. Keys are
 // applied to the batch in the order dictated by their sequence numbers within
@@ -1011,6 +1022,7 @@ func loadFlushedSSTableKeys(
 	path string,
 	fileNums []base.DiskFileNum,
 	blobRefMap map[base.DiskFileNum]manifest.BlobReferences,
+	blobFileMap blobFileMap,
 	provider blob.ReaderProvider,
 	readOpts sstable.ReaderOptions,
 	bufs *flushBuffers,
@@ -1040,7 +1052,7 @@ func loadFlushedSSTableKeys(
 			if bf, ok := blobRefMap[fileNum]; ok {
 				blobRefs = &bf
 			}
-			vf, blobContext := sstable.LoadValBlobContext(provider, blobRefs)
+			vf, blobContext := sstable.LoadValBlobContext(blobFileMap, provider, blobRefs)
 			defer func() { _ = vf.Close() }()
 			iter, err := r.NewIter(sstable.NoTransforms, nil, nil, blobContext)
 			if err != nil {

replay/replay_test.go

@@ -211,7 +211,7 @@ func TestLoadFlushedSSTableKeys(t *testing.T) {
 				return err.Error()
 			}
 			defer closeFunc()
-			err = loadFlushedSSTableKeys(b, opts.FS, "", diskFileNums, nil /* blobRefMap */, provider,
+			err = loadFlushedSSTableKeys(b, opts.FS, "", diskFileNums, nil /* blobRefMap */, nil /* blobFileMap */, provider,
 				readerOpts, &flushBufs)
 			if err != nil {
 				b.Close()

sstable/blob/fetcher.go

@@ -39,20 +39,22 @@ type ValueReader interface {
 	ReadIndexBlock(context.Context, block.ReadEnv, objstorage.ReadHandle) (block.BufferHandle, error)
 }
 
+// A FileMapping defines the mapping between blob file IDs and disk file numbers.
+// It's implemented by *manifest.BlobFileSet.
+type FileMapping interface {
+	// Lookup returns the disk file number for the given blob file ID. It
+	// returns false for the second return value if the blob file ID is not
+	// present in the mapping.
+	Lookup(base.BlobFileID) (base.DiskFileNum, bool)
+}
+
 // A ReaderProvider is an interface that can be used to retrieve a ValueReader
 // for a given file number.
 type ReaderProvider interface {
 	// GetValueReader returns a ValueReader for the given file number.
 	GetValueReader(ctx context.Context, fileNum base.DiskFileNum) (r ValueReader, closeFunc func(), err error)
 }
 
-// DiskFileNumTODO is a temporary function to convert a BlobFileID to a
-// DiskFileNum. It should be removed once the manifest.Version contains a
-// mapping.
-func DiskFileNumTODO(blobFileID base.BlobFileID) base.DiskFileNum {
-	return base.DiskFileNum(blobFileID)
-}
-
 // A ValueFetcher retrieves values stored out-of-band in separate blob files.
 // The ValueFetcher caches accessed file readers to avoid redundant file cache
 // and block cache lookups when performing consecutive value retrievals.
@@ -63,6 +65,7 @@ func DiskFileNumTODO(blobFileID base.BlobFileID) base.DiskFileNum {
 // When finished with a ValueFetcher, one must call Close to release all cached
 // readers and block buffers.
 type ValueFetcher struct {
+	fileMapping    FileMapping
 	readerProvider ReaderProvider
 	env            block.ReadEnv
 	fetchCount     int
@@ -76,7 +79,8 @@ type ValueFetcher struct {
 var _ base.ValueFetcher = (*ValueFetcher)(nil)
 
 // Init initializes the ValueFetcher.
-func (r *ValueFetcher) Init(rp ReaderProvider, env block.ReadEnv) {
+func (r *ValueFetcher) Init(fm FileMapping, rp ReaderProvider, env block.ReadEnv) {
+	r.fileMapping = fm
 	r.readerProvider = rp
 	r.env = env
 	if r.readerProvider == nil {
@@ -129,7 +133,10 @@ func (r *ValueFetcher) retrieve(ctx context.Context, vh Handle) (val []byte, err
 				return nil, err
 			}
 		}
-		diskFileNum := DiskFileNumTODO(vh.BlobFileID)
+		diskFileNum, ok := r.fileMapping.Lookup(vh.BlobFileID)
+		if !ok {
+			return nil, errors.AssertionFailedf("blob file %s not found", vh.BlobFileID)
+		}
 		if cr.r, cr.closeFunc, err = r.readerProvider.GetValueReader(ctx, diskFileNum); err != nil {
 			return nil, err
 		}