backup: implement online restore in distsql flow

This implements online restore using the normal restore distsql flow,
rather than the bespoke online restore paths used in its initial prototype.
This is behind a setting for now, but eventually could be the basis for
extending online restore -- now unified with regular restore -- to support
ingesting rather than linking some data (such as revision-history layers).

Release note: none.
Epic: none.

Author: dt

pkg/backup/bench_covering_test.go

@@ -109,7 +109,9 @@ func BenchmarkRestoreEntryCover(b *testing.B) {
 													nil,
 													filter,
 													&inclusiveEndKeyComparator{},
-													spanCh)
+													spanCh,
+													false, /* useLink */
+												)
 											})
 
 											var cov []execinfrapb.RestoreSpanEntry

pkg/backup/compaction_dist.go

@@ -94,6 +94,7 @@ func runCompactionPlan(
 			filter,
 			fsc,
 			spanCh,
+			false, /* useLink */
 		), "generateAndSendImportSpans")
 	}
 	dsp := execCtx.DistSQLPlanner()

pkg/backup/compaction_processor.go

@@ -213,6 +213,7 @@ func (p *compactBackupsProcessor) runCompactBackups(ctx context.Context) error {
 			filter,
 			fsc,
 			entryCh,
+			false, /* useLink */
 		), "generate and send import spans")
 	}
 

pkg/backup/generative_split_and_scatter_processor.go

@@ -461,6 +461,7 @@ func runGenerativeSplitAndScatter(
 ) error {
 	log.Dev.Infof(ctx, "Running generative split and scatter with %d total spans, %d chunk size, %d nodes",
 		spec.NumEntries, spec.ChunkSize, spec.NumNodes)
+
 	g := ctxgroup.WithContext(ctx)
 
 	chunkSplitAndScatterWorkers := len(chunkSplitAndScatterers)
@@ -512,6 +513,7 @@ func runGenerativeSplitAndScatter(
 			filter,
 			fsc,
 			restoreSpanEntriesCh,
+			spec.UseLink,
 		), "generating and sending import spans")
 	})
 

pkg/backup/restore_data_processor.go

@@ -31,6 +31,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	bulkutil "github.com/cockroachdb/cockroach/pkg/util/bulk"
 	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -420,6 +421,27 @@ func (rd *restoreDataProcessor) runRestoreWorkers(
 				ctx, sp := tracing.ChildSpan(ctx, "restore.processRestoreSpanEntry")
 				defer sp.Finish()
 
+				// Check if the first file has UseLink set. If so, link all files
+				// in this entry via LinkExternalSSTable instead of ingesting.
+				// TODO(dt): support entries where some layers are link and some are not.
+				if entry.Files[0].UseLink {
+					for _, file := range entry.Files {
+						if !file.UseLink {
+							return done, errors.AssertionFailedf("invalid non-linked file in linking restore")
+						}
+					}
+					summary, err := rd.linkFiles(ctx, kr, entry)
+					if err != nil {
+						return done, errors.Wrap(err, "linking files")
+					}
+					select {
+					case rd.progCh <- makeProgressUpdate(summary, entry, rd.spec.PKIDs, rd.spec.RestoreTime):
+					case <-ctx.Done():
+						return done, errors.Wrap(ctx.Err(), "sending progress update")
+					}
+					return done, nil
+				}
+
 				var res *resumeEntry
 				for {
 					sstIter, res, err = rd.openSSTs(ctx, entry, res)
@@ -454,6 +476,134 @@ func (rd *restoreDataProcessor) runRestoreWorkers(
 	})
 }
 
+// linkFiles links all files in the entry via LinkExternalSSTable instead of
+// downloading and ingesting them. This is used for online restore when the
+// UseLink flag is set on files.
+func (rd *restoreDataProcessor) linkFiles(
+	ctx context.Context, kr *KeyRewriter, entry execinfrapb.RestoreSpanEntry,
+) (kvpb.BulkOpSummary, error) {
+	ctx, sp := tracing.ChildSpan(ctx, "restore.linkFiles")
+	defer sp.Finish()
+
+	var summary kvpb.BulkOpSummary
+	kvDB := rd.FlowCtx.Cfg.DB.KV()
+
+	if err := assertCommonPrefix(entry.Span, entry.ElidedPrefix); err != nil {
+		return summary, err
+	}
+
+	var currentLayer int32
+	for i := range entry.Files {
+		file := &entry.Files[i]
+		if file.Layer < currentLayer {
+			return summary, errors.AssertionFailedf("files not sorted by layer")
+		}
+		currentLayer = file.Layer
+		if file.HasRangeKeys {
+			return summary, errors.Wrapf(permanentRestoreError, "online restore of range keys not supported")
+		}
+		if err := assertCommonPrefix(file.BackupFileEntrySpan, entry.ElidedPrefix); err != nil {
+			return summary, err
+		}
+
+		restoringSubspan := file.BackupFileEntrySpan.Intersect(entry.Span)
+		if !restoringSubspan.Valid() {
+			return summary, errors.AssertionFailedf("file %s with span %s has no overlap with restore span %s",
+				file.Path,
+				file.BackupFileEntrySpan,
+				entry.Span,
+			)
+		}
+
+		var rewriteErr error
+		restoringSubspan, rewriteErr = rewriteSpan(kr, restoringSubspan.Clone(), entry.ElidedPrefix)
+		if rewriteErr != nil {
+			return summary, rewriteErr
+		}
+
+		log.VEventf(ctx, 2, "linking file %s (file span: %s) to span %s", file.Path, file.BackupFileEntrySpan, restoringSubspan)
+
+		counts := file.BackupFileEntryCounts
+		fileSize := file.ApproximatePhysicalSize
+		if fileSize == 0 {
+			fileSize = uint64(counts.DataSize)
+		}
+		if fileSize == 0 {
+			fileSize = 16 << 20 // guess
+		}
+
+		// The synthetic prefix is computed from the REWRITTEN span key, not the
+		// original file span key. This matches the old online restore path which
+		// overwrites file.BackupFileEntrySpan before computing the prefix.
+		syntheticPrefix, err := backupsink.ElidedPrefix(restoringSubspan.Key, entry.ElidedPrefix)
+		if err != nil {
+			return summary, err
+		}
+
+		fileStats := &enginepb.MVCCStats{
+			ContainsEstimates: 1,
+			KeyBytes:          counts.DataSize / 2,
+			ValBytes:          counts.DataSize / 2,
+			LiveBytes:         counts.DataSize,
+			KeyCount:          counts.Rows + counts.IndexEntries,
+			LiveCount:         counts.Rows + counts.IndexEntries,
+		}
+
+		var batchTimestamp hlc.Timestamp
+		if writeAtBatchTS(ctx, restoringSubspan, kr.fromSystemTenant) {
+			batchTimestamp = kvDB.Clock().Now()
+		}
+
+		loc := kvpb.LinkExternalSSTableRequest_ExternalFile{
+			Locator:                 file.Dir.URI,
+			Path:                    file.Path,
+			ApproximatePhysicalSize: fileSize,
+			BackingFileSize:         file.BackingFileSize,
+			SyntheticPrefix:         syntheticPrefix,
+			UseSyntheticSuffix:      batchTimestamp.IsSet(),
+			MVCCStats:               fileStats,
+		}
+
+		if err := kvDB.LinkExternalSSTable(ctx, restoringSubspan, loc, batchTimestamp); err != nil {
+			return summary, errors.Wrap(err, "linking external SSTable")
+		}
+
+		// Call testing knob after each file link, matching the old online restore path.
+		if restoreKnobs, ok := rd.FlowCtx.TestingKnobs().BackupRestoreTestingKnobs.(*sql.BackupRestoreTestingKnobs); ok {
+			if restoreKnobs.AfterAddRemoteSST != nil {
+				if err := restoreKnobs.AfterAddRemoteSST(); err != nil {
+					return summary, err
+				}
+			}
+		}
+
+		// Use ApproximatePhysicalSize for DataSize to match the coordinator path.
+		summary.DataSize += int64(fileSize)
+
+		// Populate EntryCounts so that countRows() can compute row counts for
+		// progress tracking. We need to use a key that's in pkIDs for rows to be
+		// counted correctly. Pick any key from pkIDs (the specific key doesn't
+		// matter since we're just summing totals).
+		if counts.Rows > 0 || counts.IndexEntries > 0 {
+			if summary.EntryCounts == nil {
+				summary.EntryCounts = make(map[uint64]int64)
+			}
+			// Use a pkID key for rows so countRows counts them as rows.
+			for pkID := range rd.spec.PKIDs {
+				summary.EntryCounts[pkID] += counts.Rows
+				break // only need one key
+			}
+			// Use key 0 for index entries (won't be in pkIDs, so counted as index entries).
+			if counts.IndexEntries > 0 {
+				summary.EntryCounts[0] += counts.IndexEntries
+			}
+		}
+	}
+
+	rd.progressMade = true
+	return summary, nil
+}
+
 var backupFileReadError = errors.New("error reading backup file")
 
 func (rd *restoreDataProcessor) processRestoreSpanEntry(

pkg/backup/restore_job.go

@@ -404,6 +404,10 @@ func restore(
 	job := resumer.job
 	details := job.Details().(jobspb.RestoreDetails)
 
+	// resolvedURIs holds the main URIs after resolving any external:// aliases.
+	// This is set in the OnlineImpl block below and used by runRestore.
+	resolvedURIs := details.URIs
+
 	if details.OnlineImpl() {
 		var linkPhaseComplete bool
 		if err := execCtx.ExecCfg().InternalDB.Txn(restoreCtx, func(ctx context.Context, txn isql.Txn) error {
@@ -417,6 +421,18 @@ func restore(
 		if linkPhaseComplete {
 			return emptyRowCount, nil
 		}
+
+		// If any URIs utilize external:// aliases, we need to resolve the alias
+		// to its underlying URI before feeding it into online restore. This
+		// applies to the main URIs, the locality info URIs, and the backup
+		// manifest Dir fields.
+		var err error
+		resolvedURIs, backupLocalityInfo, err = resolveExternalStorageURIs(
+			restoreCtx, execCtx, details.URIs, backupLocalityInfo, backupManifests,
+		)
+		if err != nil {
+			return emptyRowCount, errors.Wrap(err, "resolving external storage URIs for online restore")
+		}
 	}
 
 	// If we've already migrated some of the system tables we're about to
@@ -510,6 +526,7 @@ func restore(
 			filter,
 			fsc,
 			spanCh,
+			false, /* useLink */
 		), "generate and send import spans")
 	}
 
@@ -553,9 +570,15 @@ func restore(
 		tasks = append(tasks, jobProgressLoop)
 	}
 
+	// Check if online restore should use the distributed flow with file linking
+	// instead of the simpler sendAddRemoteSSTs loop.
+	useDistFlow := onlineRestoreUseDistFlow.Get(&execCtx.ExecCfg().Settings.SV)
+
 	progCh := make(chan *execinfrapb.RemoteProducerMetadata_BulkProcessorProgress)
-	if !details.OnlineImpl() {
-		// Online restore tracks progress by pinging requestFinishedCh instead
+	// Start the progress checkpoint loop if this is a traditional restore OR
+	// an online restore using the distributed flow (which also sends progress
+	// updates via progCh).
+	if !details.OnlineImpl() || useDistFlow {
 		generativeCheckpointLoop := func(ctx context.Context) error {
 			defer close(requestFinishedCh)
 			for progress := range progCh {
@@ -599,7 +622,10 @@ func restore(
 	}
 
 	resumeClusterVersion := execCtx.ExecCfg().Settings.Version.ActiveVersion(restoreCtx).Version
-	if clusterversion.V24_3.Version().LessEq(resumeClusterVersion) && !details.OnlineImpl() {
+	// Start the countCompletedProcLoop if this is a traditional restore OR
+	// an online restore using the distributed flow (which also sends processor
+	// completion signals via procCompleteCh).
+	if clusterversion.V24_3.Version().LessEq(resumeClusterVersion) && (!details.OnlineImpl() || useDistFlow) {
 		tasks = append(tasks, countCompletedProcLoop)
 	}
 
@@ -628,39 +654,78 @@ func restore(
 	runRestore := func(ctx context.Context) error {
 		if details.OnlineImpl() {
 			log.Dev.Warningf(ctx, "EXPERIMENTAL ONLINE RESTORE being used")
-			approxRows, approxDataSize, err := sendAddRemoteSSTs(
-				ctx,
-				execCtx,
-				job,
-				dataToRestore,
-				encryption,
-				details.URIs,
-				backupLocalityInfo,
-				requestFinishedCh,
-				tracingAggCh,
-				genSpan,
-			)
-			progressTracker.mu.Lock()
-			defer progressTracker.mu.Unlock()
-			//  During the link phase of online restore, we do not update stats
-			// progress as job occurs.  We merely reuse the `progressTracker.mu.res`
-			// var to reduce the number of local vars floating around in `restore`.
-			progressTracker.mu.res = roachpb.RowCount{Rows: approxRows, DataSize: approxDataSize}
-			return errors.Wrap(err, "sending remote AddSSTable requests")
+
+			// Use the bespoke online restore path that directly splits and links from
+			// the coordinator without a distSQL restore flow.
+			if !useDistFlow {
+				approxRows, approxDataSize, err := sendAddRemoteSSTs(
+					ctx,
+					execCtx,
+					job,
+					dataToRestore,
+					encryption,
+					resolvedURIs,
+					backupLocalityInfo,
+					requestFinishedCh,
+					tracingAggCh,
+					genSpan,
+				)
+				progressTracker.mu.Lock()
+				defer progressTracker.mu.Unlock()
+				//  During the link phase of online restore, we do not update stats
+				// progress as job occurs.  We merely reuse the `progressTracker.mu.res`
+				// var to reduce the number of local vars floating around in `restore`.
+				progressTracker.mu.res = roachpb.RowCount{Rows: approxRows, DataSize: approxDataSize}
+				return errors.Wrap(err, "sending remote AddSSTable requests")
+			}
+
+			// If we did not switch to the bespoke online restore path, we'll proceed
+			// to the normal restore distSQL flow and let its split and scattter
+			// processor direct the restore data processors to link files rather than
+			// ingest them. But we do need to pre-split at the top-level logical spans
+			// that are being restored, as stored in DownloadSpan, as these are what
+			// we will clear via kv/pebble excises if we fail at any point after we
+			// start linking files.
+			//
+			// TODO(dt): we should record in persisted progress when we are ready to
+			// enter the linking/ingesting phase, i.e. *after* we make these splits,
+			// so any failures prior to it can skip cleanup, as that cleanup could
+			// fail if we failed prior to making these splits.
+			var prevSplit roachpb.Key
+			for _, span := range details.DownloadSpans {
+				if !span.Key.Equal(prevSplit) {
+					if err := execCtx.ExecCfg().DB.AdminSplit(ctx, span.Key, hlc.MaxTimestamp); err != nil {
+						return errors.Wrapf(err, "pre-splitting at key %s", span.Key)
+					}
+				}
+				if err := execCtx.ExecCfg().DB.AdminSplit(ctx, span.EndKey, hlc.MaxTimestamp); err != nil {
+					return errors.Wrapf(err, "pre-splitting at key %s", span.EndKey)
+				}
+				prevSplit = span.EndKey
+			}
+		}
+
+		// Use the distributed restore flow. If this is an online restore (with
+		// useDistFlow, since we returned above otherwise), link files instead of
+		// ingesting them.
+		useLink := details.OnlineImpl()
+		if useLink {
+			log.Dev.Infof(ctx, "online restore using distributed flow with file linking")
 		}
 		md := restoreJobMetadata{
 			jobID:                job.ID(),
 			dataToRestore:        dataToRestore,
 			restoreTime:          endTime,
 			encryption:           encryption,
 			kmsEnv:               kmsEnv,
-			uris:                 details.URIs,
+			uris:                 resolvedURIs,
 			backupLocalityInfo:   backupLocalityInfo,
 			spanFilter:           filter,
 			numImportSpans:       numImportSpans,
 			execLocality:         details.ExecutionLocality,
 			exclusiveEndKeys:     fsc.isExclusive(),
 			resumeClusterVersion: resumeClusterVersion,
+			useLink:              useLink,
 		}
 		return errors.Wrap(distRestore(
 			ctx,