Upgrade versions of vendored libs

vendor/github.com/cockroachdb/pebble/checkpoint.go

@@ -5,10 +5,12 @@
 package pebble
 
 import (
+	"io"
 	"os"
 
 	"github.com/cockroachdb/errors/oserror"
 	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/record"
 	"github.com/cockroachdb/pebble/vfs"
 	"github.com/cockroachdb/pebble/vfs/atomicfs"
 )
@@ -19,6 +21,9 @@ type checkpointOptions struct {
 	// flushWAL set to true will force a flush and sync of the WAL prior to
 	// checkpointing.
 	flushWAL bool
+
+	// If set, any SSTs that don't overlap with these spans are excluded from a checkpoint.
+	restrictToSpans []CheckpointSpan
 }
 
 // CheckpointOption set optional parameters used by `DB.Checkpoint`.
@@ -39,6 +44,43 @@ func WithFlushedWAL() CheckpointOption {
 	}
 }
 
+// WithRestrictToSpans specifies spans of interest for the checkpoint. Any SSTs
+// that don't overlap with any of these spans are excluded from the checkpoint.
+//
+// Note that the checkpoint can still surface keys outside of these spans (from
+// the WAL and from SSTs that partially overlap with these spans). Moreover,
+// these surface keys aren't necessarily "valid" in that they could have been
+// modified but the SST containing the modification is excluded.
+func WithRestrictToSpans(spans []CheckpointSpan) CheckpointOption {
+	return func(opt *checkpointOptions) {
+		opt.restrictToSpans = spans
+	}
+}
+
+// CheckpointSpan is a key range [Start, End) (inclusive on Start, exclusive on
+// End) of interest for a checkpoint.
+type CheckpointSpan struct {
+	Start []byte
+	End   []byte
+}
+
+// excludeFromCheckpoint returns true if an SST file should be excluded from the
+// checkpoint because it does not overlap with the spans of interest
+// (opt.restrictToSpans).
+func excludeFromCheckpoint(f *fileMetadata, opt *checkpointOptions, cmp Compare) bool {
+	if len(opt.restrictToSpans) == 0 {
+		// Option not set; don't exclude anything.
+		return false
+	}
+	for _, s := range opt.restrictToSpans {
+		if f.Overlaps(cmp, s.Start, s.End, true /* exclusiveEnd */) {
+			return false
+		}
+	}
+	// None of the restrictToSpans overlapped; we can exclude this file.
+	return true
+}
+
 // mkdirAllAndSyncParents creates destDir and any of its missing parents.
 // Those missing parents, as well as the closest existing ancestor, are synced.
 // Returns a handle to the directory created at destDir.
@@ -93,6 +135,9 @@ func mkdirAllAndSyncParents(fs vfs.FS, destDir string) (vfs.File, error) {
 // space overhead for a checkpoint if hard links are disabled. Also beware that
 // even if hard links are used, the space overhead for the checkpoint will
 // increase over time as the DB performs compactions.
+//
+// TODO(bananabrick): Test checkpointing of virtual sstables once virtual
+// sstables is running e2e.
 func (d *DB) Checkpoint(
 	destDir string, opts ...CheckpointOption,
 ) (
@@ -146,7 +191,10 @@ func (d *DB) Checkpoint(
 	manifestFileNum := d.mu.versions.manifestFileNum
 	manifestSize := d.mu.versions.manifest.Size()
 	optionsFileNum := d.optionsFileNum
-
+	virtualBackingFiles := make(map[base.DiskFileNum]struct{})
+	for diskFileNum := range d.mu.versions.fileBackingMap {
+		virtualBackingFiles[diskFileNum] = struct{}{}
+	}
 	// Release the manifest and DB.mu so we don't block other operations on
 	// the database.
 	d.mu.versions.logUnlock()
@@ -154,13 +202,10 @@ func (d *DB) Checkpoint(
 
 	// Wrap the normal filesystem with one which wraps newly created files with
 	// vfs.NewSyncingFile.
-	fs := syncingFS{
-		FS: d.opts.FS,
-		syncOpts: vfs.SyncingFileOptions{
-			NoSyncOnClose: d.opts.NoSyncOnClose,
-			BytesPerSync:  d.opts.BytesPerSync,
-		},
-	}
+	fs := vfs.NewSyncingFS(d.opts.FS, vfs.SyncingFileOptions{
+		NoSyncOnClose: d.opts.NoSyncOnClose,
+		BytesPerSync:  d.opts.BytesPerSync,
+	})
 
 	// Create the dir and its parents (if necessary), and sync them.
 	var dir vfs.File
@@ -214,45 +259,34 @@ func (d *DB) Checkpoint(
 		}
 	}
 
-	{
-		// Copy the MANIFEST, and create a pointer to it. We copy rather
-		// than link because additional version edits added to the
-		// MANIFEST after we took our snapshot of the sstables will
-		// reference sstables that aren't in our checkpoint. For a
-		// similar reason, we need to limit how much of the MANIFEST we
-		// copy.
-		srcPath := base.MakeFilepath(fs, d.dirname, fileTypeManifest, manifestFileNum)
-		destPath := fs.PathJoin(destDir, fs.PathBase(srcPath))
-		ckErr = vfs.LimitedCopy(fs, srcPath, destPath, manifestSize)
-		if ckErr != nil {
-			return ckErr
-		}
-
-		// Recent format versions use an atomic marker for setting the
-		// active manifest. Older versions use the CURRENT file. The
-		// setCurrentFunc function will return a closure that will
-		// take the appropriate action for the database's format
-		// version.
-		var manifestMarker *atomicfs.Marker
-		manifestMarker, _, ckErr = atomicfs.LocateMarker(fs, destDir, manifestMarkerName)
-		if ckErr != nil {
-			return ckErr
-		}
-		ckErr = setCurrentFunc(formatVers, manifestMarker, fs, destDir, dir)(manifestFileNum)
-		if ckErr != nil {
-			return ckErr
-		}
-		ckErr = manifestMarker.Close()
-		if ckErr != nil {
-			return ckErr
-		}
-	}
-
+	var excludedFiles map[deletedFileEntry]*fileMetadata
+	// Set of FileBacking.DiskFileNum which will be required by virtual sstables
+	// in the checkpoint.
+	requiredVirtualBackingFiles := make(map[base.DiskFileNum]struct{})
 	// Link or copy the sstables.
 	for l := range current.Levels {
 		iter := current.Levels[l].Iter()
 		for f := iter.First(); f != nil; f = iter.Next() {
-			srcPath := base.MakeFilepath(fs, d.dirname, fileTypeTable, f.FileNum)
+			if excludeFromCheckpoint(f, opt, d.cmp) {
+				if excludedFiles == nil {
+					excludedFiles = make(map[deletedFileEntry]*fileMetadata)
+				}
+				excludedFiles[deletedFileEntry{
+					Level:   l,
+					FileNum: f.FileNum,
+				}] = f
+				continue
+			}
+
+			fileBacking := f.FileBacking
+			if f.Virtual {
+				if _, ok := requiredVirtualBackingFiles[fileBacking.DiskFileNum]; ok {
+					continue
+				}
+				requiredVirtualBackingFiles[fileBacking.DiskFileNum] = struct{}{}
+			}
+
+			srcPath := base.MakeFilepath(fs, d.dirname, fileTypeTable, fileBacking.DiskFileNum)
 			destPath := fs.PathJoin(destDir, fs.PathBase(srcPath))
 			ckErr = vfs.LinkOrCopy(fs, srcPath, destPath)
 			if ckErr != nil {
@@ -261,6 +295,23 @@ func (d *DB) Checkpoint(
 		}
 	}
 
+	var removeBackingTables []base.DiskFileNum
+	for diskFileNum := range virtualBackingFiles {
+		if _, ok := requiredVirtualBackingFiles[diskFileNum]; !ok {
+			// The backing sstable associated with fileNum is no longer
+			// required.
+			removeBackingTables = append(removeBackingTables, diskFileNum)
+		}
+	}
+
+	ckErr = d.writeCheckpointManifest(
+		fs, formatVers, destDir, dir, manifestFileNum.DiskFileNum(), manifestSize,
+		excludedFiles, removeBackingTables,
+	)
+	if ckErr != nil {
+		return ckErr
+	}
+
 	// Copy the WAL files. We copy rather than link because WAL file recycling
 	// will cause the WAL files to be reused which would invalidate the
 	// checkpoint.
@@ -269,7 +320,7 @@ func (d *DB) Checkpoint(
 		if logNum == 0 {
 			continue
 		}
-		srcPath := base.MakeFilepath(fs, d.walDirname, fileTypeLog, logNum)
+		srcPath := base.MakeFilepath(fs, d.walDirname, fileTypeLog, logNum.DiskFileNum())
 		destPath := fs.PathJoin(destDir, fs.PathBase(srcPath))
 		ckErr = vfs.Copy(fs, srcPath, destPath)
 		if ckErr != nil {
@@ -286,3 +337,99 @@ func (d *DB) Checkpoint(
 	dir = nil
 	return ckErr
 }
+
+func (d *DB) writeCheckpointManifest(
+	fs vfs.FS,
+	formatVers FormatMajorVersion,
+	destDirPath string,
+	destDir vfs.File,
+	manifestFileNum base.DiskFileNum,
+	manifestSize int64,
+	excludedFiles map[deletedFileEntry]*fileMetadata,
+	removeBackingTables []base.DiskFileNum,
+) error {
+	// Copy the MANIFEST, and create a pointer to it. We copy rather
+	// than link because additional version edits added to the
+	// MANIFEST after we took our snapshot of the sstables will
+	// reference sstables that aren't in our checkpoint. For a
+	// similar reason, we need to limit how much of the MANIFEST we
+	// copy.
+	// If some files are excluded from the checkpoint, also append a block that
+	// records those files as deleted.
+	if err := func() error {
+		srcPath := base.MakeFilepath(fs, d.dirname, fileTypeManifest, manifestFileNum)
+		destPath := fs.PathJoin(destDirPath, fs.PathBase(srcPath))
+		src, err := fs.Open(srcPath, vfs.SequentialReadsOption)
+		if err != nil {
+			return err
+		}
+		defer src.Close()
+
+		dst, err := fs.Create(destPath)
+		if err != nil {
+			return err
+		}
+		defer dst.Close()
+
+		// Copy all existing records. We need to copy at the record level in case we
+		// need to append another record with the excluded files (we cannot simply
+		// append a record after a raw data copy; see
+		// https://github.com/cockroachdb/cockroach/issues/100935).
+		r := record.NewReader(&io.LimitedReader{R: src, N: manifestSize}, manifestFileNum.FileNum())
+		w := record.NewWriter(dst)
+		for {
+			rr, err := r.Next()
+			if err != nil {
+				if err == io.EOF {
+					break
+				}
+				return err
+			}
+
+			rw, err := w.Next()
+			if err != nil {
+				return err
+			}
+			if _, err := io.Copy(rw, rr); err != nil {
+				return err
+			}
+		}
+
+		if len(excludedFiles) > 0 {
+			// Write out an additional VersionEdit that deletes the excluded SST files.
+			ve := versionEdit{
+				DeletedFiles:         excludedFiles,
+				RemovedBackingTables: removeBackingTables,
+			}
+
+			rw, err := w.Next()
+			if err != nil {
+				return err
+			}
+			if err := ve.Encode(rw); err != nil {
+				return err
+			}
+		}
+		if err := w.Close(); err != nil {
+			return err
+		}
+		return dst.Sync()
+	}(); err != nil {
+		return err
+	}
+
+	// Recent format versions use an atomic marker for setting the
+	// active manifest. Older versions use the CURRENT file. The
+	// setCurrentFunc function will return a closure that will
+	// take the appropriate action for the database's format
+	// version.
+	var manifestMarker *atomicfs.Marker
+	manifestMarker, _, err := atomicfs.LocateMarker(fs, destDirPath, manifestMarkerName)
+	if err != nil {
+		return err
+	}
+	if err := setCurrentFunc(formatVers, manifestMarker, fs, destDirPath, destDir)(manifestFileNum.FileNum()); err != nil {
+		return err
+	}
+	return manifestMarker.Close()
+}