manifest: validate BlobReferences and BlobFileSet agree

Sometimes in invariants builds, validate that the set of unique BlobFileIDs
contained in TableMetadata BlobReferences is exactly the set of BlobFileIDs in
the Version's BlobFileSet.

Author: jbowens

internal/manifest/blob_metadata.go

@@ -7,6 +7,7 @@ package manifest
 import (
 	stdcmp "cmp"
 	"fmt"
+	"iter"
 	"slices"
 	"sync/atomic"
 
@@ -330,6 +331,11 @@ func MakeBlobFileSet(entries []BlobFileMetadata) BlobFileSet {
 	return BlobFileSet{tree: makeBTree(btreeCmpBlobFileID, entries)}
 }
 
+// All returns an iterator over all the blob files in the set.
+func (s *BlobFileSet) All() iter.Seq[BlobFileMetadata] {
+	return s.tree.All()
+}
+
 // 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/version.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"fmt"
 	"iter"
+	"maps"
 	"slices"
 	"sort"
 	"strings"
@@ -600,6 +601,46 @@ func (v *Version) CheckOrdering() error {
 	return nil
 }
 
+// validateBlobFileInvariants validates invariants around blob files. Currently
+// it validates that the set of BlobFileIDs referenced by the Version's tables'
+// blob references is exactly the same as the set of BlobFileIDs present in the
+// Version's blob files B-Tree.
+func (v *Version) validateBlobFileInvariants() error {
+	// Collect all the blob file IDs that are referenced by the Version's
+	// tables' blob references.
+	var referencedFileIDs []base.BlobFileID
+	{
+		referencedFileIDsMap := make(map[base.BlobFileID]struct{}, v.BlobFiles.tree.Count())
+		for i := 0; i < len(v.Levels); i++ {
+			for table := range v.Levels[i].All() {
+				for _, br := range table.BlobReferences {
+					referencedFileIDsMap[br.FileID] = struct{}{}
+				}
+			}
+		}
+		referencedFileIDs = slices.Collect(maps.Keys(referencedFileIDsMap))
+		slices.Sort(referencedFileIDs)
+	}
+
+	// Collect all the blob file IDs that are present in the Version's blob
+	// files B-Tree.
+	var versionBlobFileIDs []base.BlobFileID
+	{
+		versionBlobFileIDsMap := make(map[base.BlobFileID]struct{}, v.BlobFiles.tree.Count())
+		for bf := range v.BlobFiles.All() {
+			versionBlobFileIDsMap[bf.FileID] = struct{}{}
+		}
+		versionBlobFileIDs = slices.Collect(maps.Keys(versionBlobFileIDsMap))
+		slices.Sort(versionBlobFileIDs)
+	}
+
+	if !slices.Equal(referencedFileIDs, versionBlobFileIDs) {
+		return base.AssertionFailedf("divergence between referenced BlobFileIDs and Version's BlobFiles B-Tree: %v vs %v",
+			referencedFileIDs, versionBlobFileIDs)
+	}
+	return nil
+}
+
 // VersionList holds a list of versions. The versions are ordered from oldest
 // to newest.
 type VersionList struct {

internal/manifest/version_edit.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/cockroachdb/pebble/internal/strparse"
 	"github.com/cockroachdb/pebble/sstable"
 )
@@ -1296,5 +1297,14 @@ func (b *BulkVersionEdit) Apply(curr *Version, readCompactionRate int64) (*Versi
 			}
 		}
 	}
+
+	// In invariants builds, sometimes check invariants across all blob files
+	// and their references.
+	if invariants.Sometimes(20) {
+		if err := v.validateBlobFileInvariants(); err != nil {
+			return nil, err
+		}
+	}
+
 	return v, nil
 }