storageccl: use the new PebbleIterator in ExternalSSTReader

This PR refactors all call sites of ExternalSSTReader(), to support using the
new PebbleIterator, which has baked in range key support. Most notably, this
PR replaces the multiIterator used in the restore data processor with the
PebbleSSTIterator.

This patch is apart of a larger effort to teach backup and restore about MVCC
bulk operations. Next, the readAsOfIterator will need to learn how to
deal with range keys.

Informs #71155

This PR addresses a bug created in #83984: loop variables in
ExternalSSTReader were captured by reference, leading to roachtest failures
(#84240, #84162).

Informs #71155i

Fixes: #84240, #84162, #84181

Release note: none

pkg/ccl/backupccl/backupinfo/backup_metadata.go

@@ -49,6 +49,12 @@ const (
 	sstTenantsPrefix = "tenant/"
 )
 
+var iterOpts = storage.IterOptions{
+KeyTypes:             storage.IterKeyTypePointsAndRanges,
+LowerBound:           keys.LocalMax,
+UpperBound:           keys.MaxKey,
+}
+
 // WriteBackupMetadataSST is responsible for constructing and writing the
 // `metadata.sst` to dest. This file contains the metadata corresponding to this
 // backup.
@@ -618,7 +624,7 @@ func debugDumpFileSST(
 		}
 		encOpts = &roachpb.FileEncryptionOptions{Key: key}
 	}
-	iter, err := storageccl.ExternalSSTReader(ctx, store, fileInfoPath, encOpts)
+	iter, err := storageccl.ExternalSSTReader(ctx, []cloud.ExternalStorage{store}, []string{fileInfoPath}, encOpts,iterOpts)
 	if err != nil {
 		return err
 	}
@@ -664,8 +670,8 @@ func DebugDumpMetadataSST(
 		}
 		encOpts = &roachpb.FileEncryptionOptions{Key: key}
 	}
-
-	iter, err := storageccl.ExternalSSTReader(ctx, store, path, encOpts)
+	iter, err := storageccl.ExternalSSTReader(ctx, []cloud.ExternalStorage{store}, []string{path},
+	encOpts,iterOpts)
 	if err != nil {
 		return err
 	}
@@ -804,8 +810,7 @@ func NewBackupMetadata(
 		}
 		encOpts = &roachpb.FileEncryptionOptions{Key: key}
 	}
-
-	iter, err := storageccl.ExternalSSTReader(ctx, exportStore, sstFileName, encOpts)
+	iter, err := storageccl.ExternalSSTReader(ctx, []cloud.ExternalStorage{exportStore}, []string{sstFileName}, encOpts,iterOpts)
 	if err != nil {
 		return nil, err
 	}
@@ -921,8 +926,7 @@ func (b *BackupMetadata) FileIter(ctx context.Context) FileIterator {
 		if err != nil {
 			break
 		}
-
-		iter, err := storageccl.ExternalSSTReader(ctx, b.store, path, encOpts)
+		iter, err := storageccl.ExternalSSTReader(ctx, []cloud.ExternalStorage{b.store}, []string{path}, encOpts,iterOpts)
 		if err != nil {
 			return FileIterator{err: err}
 		}
@@ -1232,7 +1236,7 @@ func makeBytesIter(
 		encOpts = &roachpb.FileEncryptionOptions{Key: key}
 	}
 
-	iter, err := storageccl.ExternalSSTReader(ctx, store, path, encOpts)
+	iter, err := storageccl.ExternalSSTReader(ctx, []cloud.ExternalStorage{store}, []string{path}, encOpts,iterOpts)
 	if err != nil {
 		return bytesIter{iterError: err}
 	}

pkg/ccl/backupccl/restore_data_processor.go

@@ -278,37 +278,27 @@ func (rd *restoreDataProcessor) openSSTs(
 ) error {
 	ctxDone := ctx.Done()
 
-	// The sstables only contain MVCC data and no intents, so using an MVCC
-	// iterator is sufficient.
-	var iters []storage.SimpleMVCCIterator
+	// TODO(msbutler): use a a map of external storage factories to avoid reopening the same dir
+	// in a given restore span entry
 	var dirs []cloud.ExternalStorage
 
 	// If we bail early and haven't handed off responsibility of the dirs/iters to
 	// the channel, close anything that we had open.
 	defer func() {
-		for _, iter := range iters {
-			iter.Close()
-		}
-
 		for _, dir := range dirs {
 			if err := dir.Close(); err != nil {
 				log.Warningf(ctx, "close export storage failed %v", err)
 			}
 		}
 	}()
 
-	// sendIters sends all of the currently accumulated iterators over the
+	// sendIter sends a multiplexed iterator covering the currently accumulated files over the
 	// channel.
-	sendIters := func(itersToSend []storage.SimpleMVCCIterator, dirsToSend []cloud.ExternalStorage) error {
-		multiIter := storage.MakeMultiIterator(itersToSend)
-		readAsOfIter := storage.NewReadAsOfIterator(multiIter, rd.spec.RestoreTime)
+	sendIter := func(iter storage.SimpleMVCCIterator, dirsToSend []cloud.ExternalStorage) error {
+		readAsOfIter := storage.NewReadAsOfIterator(iter, rd.spec.RestoreTime)
 
 		cleanup := func() {
 			readAsOfIter.Close()
-			multiIter.Close()
-			for _, iter := range itersToSend {
-				iter.Close()
-			}
 
 			for _, dir := range dirsToSend {
 				if err := dir.Close(); err != nil {
@@ -329,13 +319,13 @@ func (rd *restoreDataProcessor) openSSTs(
 			return ctx.Err()
 		}
 
-		iters = make([]storage.SimpleMVCCIterator, 0)
 		dirs = make([]cloud.ExternalStorage, 0)
 		return nil
 	}
 
 	log.VEventf(ctx, 1 /* level */, "ingesting span [%s-%s)", entry.Span.Key, entry.Span.EndKey)
 
+	filePaths := make([]string, 0, len(EntryFiles{}))
 	for _, file := range entry.Files {
 		log.VEventf(ctx, 2, "import file %s which starts at %s", file.Path, entry.Span.Key)
 
@@ -344,17 +334,23 @@ func (rd *restoreDataProcessor) openSSTs(
 			return err
 		}
 		dirs = append(dirs, dir)
+		filePaths = append(filePaths, file.Path)
 
 		// TODO(pbardea): When memory monitoring is added, send the currently
 		// accumulated iterators on the channel if we run into memory pressure.
-		iter, err := storageccl.ExternalSSTReader(ctx, dir, file.Path, rd.spec.Encryption)
-		if err != nil {
-			return err
-		}
-		iters = append(iters, iter)
 	}
-
-	return sendIters(iters, dirs)
+	iterOpts := storage.IterOptions{
+		RangeKeyMaskingBelow: rd.spec.RestoreTime,
+		KeyTypes:             storage.IterKeyTypePointsAndRanges,
+		LowerBound:           keys.LocalMax,
+		UpperBound:           keys.MaxKey,
+	}
+	iter, err := storageccl.ExternalSSTReader(ctx, dirs, filePaths, rd.spec.Encryption,
+		iterOpts)
+	if err != nil {
+		return err
+	}
+	return sendIter(iter, dirs)
 }
 
 func (rd *restoreDataProcessor) runRestoreWorkers(ctx context.Context, ssts chan mergedSST) error {

pkg/ccl/cliccl/debug_backup.go

@@ -591,7 +591,14 @@ func makeIters(
 			return nil, nil, errors.Wrapf(err, "making external storage")
 		}
 
-		iters[i], err = storageccl.ExternalSSTReader(ctx, dirStorage[i], file.Path, nil)
+		var iterOpts = storage.IterOptions{
+			KeyTypes:             storage.IterKeyTypePointsAndRanges,
+			LowerBound:           keys.LocalMax,
+			UpperBound:           keys.MaxKey,
+		}
+		iters[i], err = storageccl.ExternalSSTReader(ctx, []cloud.ExternalStorage{dirStorage[i]},
+		[]string{file.Path}, nil,iterOpts)
+
 		if err != nil {
 			return nil, nil, errors.Wrapf(err, "fetching sst reader")
 		}

pkg/ccl/storageccl/BUILD.bazel

@@ -35,14 +35,23 @@ go_test(
     ],
     embed = [":storageccl"],
     deps = [
+        "//pkg/base",
+        "//pkg/blobs",
+        "//pkg/cloud",
+        "//pkg/keys",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
+        "//pkg/security/username",
         "//pkg/server",
+        "//pkg/storage",
+        "//pkg/testutils",
         "//pkg/testutils/serverutils",
+        "//pkg/testutils/storageutils",
         "//pkg/testutils/testcluster",
         "//pkg/util/humanizeutil",
         "//pkg/util/ioctx",
         "//pkg/util/leaktest",
+        "//pkg/util/log",
         "//pkg/util/randutil",
         "@com_github_stretchr_testify//require",
     ],

pkg/ccl/storageccl/external_sst_reader.go

@@ -25,6 +25,8 @@ import (
 	"github.com/cockroachdb/pebble/sstable"
 )
 
+// RemoteSSTs lets external SSTables get iterated directly in some cases,
+// rather than being downloaded entirely first.
 var remoteSSTs = settings.RegisterBoolSetting(
 	settings.TenantWritable,
 	"kv.bulk_ingest.stream_external_ssts.enabled",
@@ -39,32 +41,44 @@ var remoteSSTSuffixCacheSize = settings.RegisterByteSizeSetting(
 	64<<10,
 )
 
-// ExternalSSTReader returns opens an SST in external storage, optionally
-// decrypting with the supplied parameters, and returns iterator over it.
-//
-// ctx is captured and used throughout the life of the returned iterator, until
-// the iterator's Close() method is called.
-func ExternalSSTReader(
-	ctx context.Context,
-	e cloud.ExternalStorage,
-	basename string,
-	encryption *roachpb.FileEncryptionOptions,
-) (storage.SimpleMVCCIterator, error) {
+func getFileWithRetry(
+	ctx context.Context, basename string, e cloud.ExternalStorage,
+) (ioctx.ReadCloserCtx, int64, error) {
 	// Do an initial read of the file, from the beginning, to get the file size as
 	// this is used e.g. to read the trailer.
 	var f ioctx.ReadCloserCtx
 	var sz int64
-
 	const maxAttempts = 3
 	if err := retry.WithMaxAttempts(ctx, base.DefaultRetryOptions(), maxAttempts, func() error {
 		var err error
 		f, sz, err = e.ReadFileAt(ctx, basename, 0)
 		return err
 	}); err != nil {
-		return nil, err
+		return nil, 0, err
 	}
+	return f, sz, nil
+}
 
-	if !remoteSSTs.Get(&e.Settings().SV) {
+// newMemPebbleSSTReader returns a PebbleSSTIterator for in-memory SSTs from
+// external storage, optionally decrypting with the supplied parameters.
+//
+// ctx is captured and used throughout the life of the returned iterator, until
+// the iterator's Close() method is called.
+func newMemPebbleSSTReader(
+	ctx context.Context,
+	e []cloud.ExternalStorage,
+	basenames []string,
+	encryption *roachpb.FileEncryptionOptions,
+	iterOps storage.IterOptions,
+) (storage.SimpleMVCCIterator, error) {
+
+	inMemorySSTs := make([][]byte, 0, len(basenames))
+
+	for i, basename := range basenames {
+		f, _, err := getFileWithRetry(ctx, basename, e[i])
+		if err != nil {
+			return nil, err
+		}
 		content, err := ioctx.ReadAll(ctx, f)
 		f.Close(ctx)
 		if err != nil {
@@ -76,44 +90,73 @@ func ExternalSSTReader(
 				return nil, err
 			}
 		}
-		return storage.NewMemSSTIterator(content, false)
+		inMemorySSTs = append(inMemorySSTs, content)
 	}
+	return storage.NewPebbleMultiMemSSTIterator(inMemorySSTs, false, iterOps)
+}
 
-	raw := &sstReader{
-		ctx:  ctx,
-		sz:   sizeStat(sz),
-		body: f,
-		openAt: func(offset int64) (ioctx.ReadCloserCtx, error) {
-			reader, _, err := e.ReadFileAt(ctx, basename, offset)
-			return reader, err
-		},
+// ExternalSSTReader returns a PebbleSSTIterator for the SSTs in external storage,
+// optionally decrypting with the supplied parameters.
+//
+// ctx is captured and used throughout the life of the returned iterator, until
+// the iterator's Close() method is called.
+func ExternalSSTReader(
+	ctx context.Context,
+	e []cloud.ExternalStorage,
+	filenames []string,
+	encryption *roachpb.FileEncryptionOptions,
+	iterOps storage.IterOptions,
+) (storage.SimpleMVCCIterator, error) {
+	if len(e) != len(filenames) {
+		return nil, errors.New("must supply an external store for each file")
+	}
+	if !remoteSSTs.Get(&e[0].Settings().SV) {
+		return newMemPebbleSSTReader(ctx, e, filenames, encryption, iterOps)
 	}
+	remoteCacheSize := remoteSSTSuffixCacheSize.Get(&e[0].Settings().SV)
+	readers := make([]sstable.ReadableFile, 0, len(filenames))
 
-	var reader sstable.ReadableFile = raw
+	for i, basename := range filenames {
+		// prevent capturing the loop variables by reference when defining openAt below.
+		basename := basename
+		i := i
 
-	if encryption != nil {
-		r, err := decryptingReader(raw, encryption.Key)
+		f, sz, err := getFileWithRetry(ctx, basename, e[i])
 		if err != nil {
-			f.Close(ctx)
 			return nil, err
 		}
-		reader = r
-	} else {
-		// We only explicitly buffer the suffix of the file when not decrypting as
-		// the decrypting reader has its own internal block buffer.
-		if err := raw.readAndCacheSuffix(remoteSSTSuffixCacheSize.Get(&e.Settings().SV)); err != nil {
-			f.Close(ctx)
-			return nil, err
+
+		raw := &sstReader{
+			ctx:  ctx,
+			sz:   sizeStat(sz),
+			body: f,
+			openAt: func(offset int64) (ioctx.ReadCloserCtx, error) {
+				reader, _, err := e[i].ReadFileAt(ctx, basename, offset)
+				return reader, err
+			},
 		}
-	}
 
-	iter, err := storage.NewSSTIterator(reader)
-	if err != nil {
-		reader.Close()
-		return nil, err
-	}
+		var reader sstable.ReadableFile
 
-	return iter, nil
+		if encryption != nil {
+			r, err := decryptingReader(raw, encryption.Key)
+			if err != nil {
+				f.Close(ctx)
+				return nil, err
+			}
+			reader = r
+		} else {
+			// We only explicitly buffer the suffix of the file when not decrypting as
+			// the decrypting reader has its own internal block buffer.
+			if err := raw.readAndCacheSuffix(remoteCacheSize); err != nil {
+				f.Close(ctx)
+				return nil, err
+			}
+			reader = raw
+		}
+		readers = append(readers, reader)
+	}
+	return storage.NewPebbleSSTIterator(readers, iterOps)
 }
 
 type sstReader struct {