Merge cockroachdb#162100

162100: backup: implement online restore in distsql flow r=dt a=dt

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.

Co-authored-by: David Taylor <davidt@davidt.io>

Author: craig[bot]

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

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

pkg/backup/compaction_processor.go

@@ -205,6 +205,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,