Merge #159758

159758: backup: add support for execution locality to backup compactions r=andrew-r-thomas a=andrew-r-thomas

This patch teaches backup compactions about execution locality filters, allowing backup compactions to support backup schedules with execution locality set while respecting the locality filter.

Fixes: #149592

Release note: None

Co-authored-by: Andrew Thomas <[email protected]>

Author: craig[bot]

pkg/backup/compaction_dist.go

@@ -11,6 +11,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/backup/backupinfo"
 	"github.com/cockroachdb/cockroach/pkg/backup/backuppb"
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/ccl/kvccl/kvfollowerreadsccl"
 	"github.com/cockroachdb/cockroach/pkg/cloud"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobsprofiler"
@@ -152,10 +153,22 @@ func createCompactionPlan(
 		return nil, nil, errors.Wrap(err, "counting number of restore span entries")
 	}
 
-	// TODO (kev-cao): Add support for execution locality.
+	evalCtx := execCtx.ExtendedEvalContext()
+	locFilter := sql.SingleLocalityFilter(details.ExecutionLocality)
+	oracle := physicalplan.DefaultReplicaChooser
+	if useBulkOracle.Get(&evalCtx.Settings.SV) {
+		oracle, err = kvfollowerreadsccl.NewLocalityFilteringBulkOracle(
+			dsp.ReplicaOracleConfig(evalCtx.Locality),
+			locFilter,
+		)
+		if err != nil {
+			return nil, nil, errors.Wrap(err, "failed to create locality filtering bulk oracle")
+		}
+	}
+
 	planCtx, sqlInstanceIDs, err := dsp.SetupAllNodesPlanningWithOracle(
-		ctx, execCtx.ExtendedEvalContext(), execCtx.ExecCfg(),
-		physicalplan.DefaultReplicaChooser, []roachpb.Locality{}, sql.NoStrictLocalityFiltering,
+		ctx, evalCtx, execCtx.ExecCfg(),
+		oracle, locFilter, sql.NoStrictLocalityFiltering,
 	)
 	if err != nil {
 		return nil, nil, err

pkg/backup/compaction_job.go

@@ -73,8 +73,6 @@ func maybeStartCompactionJob(
 	}
 
 	switch {
-	case len(triggerJob.ExecutionLocality.Tiers) != 0:
-		return 0, errors.New("execution locality not supported for compaction")
 	case triggerJob.ScheduleID == 0:
 		return 0, errors.New("only scheduled backups can be compacted")
 	case len(triggerJob.SpecificTenantIds) != 0 || triggerJob.IncludeAllSecondaryTenants:
@@ -197,13 +195,47 @@ func StartCompactionJob(
 	scheduleID jobspb.ScheduleID,
 	collectionURI []string,
 	fullBackupPath string,
-	encryptionOpts jobspb.BackupEncryptionOptions,
+	options tree.BackupOptions,
 	start, end hlc.Timestamp,
 ) (jobspb.JobID, error) {
 	planHook, ok := planner.(sql.PlanHookState)
 	if !ok {
 		return 0, errors.New("missing job execution context")
 	}
+
+	var executionLocality roachpb.Locality
+	if options.ExecutionLocality != nil {
+		if strVal, ok := options.ExecutionLocality.(*tree.StrVal); ok {
+			s := strVal.RawString()
+			if s != "" {
+				if err := executionLocality.Set(s); err != nil {
+					return 0, errors.Wrap(err, "error setting execution locality")
+				}
+			}
+		} else {
+			return 0, errors.Newf(
+				"expected string value, got %+v", options.ExecutionLocality,
+			)
+		}
+	}
+
+	encryption := jobspb.BackupEncryptionOptions{
+		Mode: jobspb.EncryptionMode_None,
+	}
+	if options.EncryptionPassphrase != nil {
+		encryption.Mode = jobspb.EncryptionMode_Passphrase
+		encryption.RawPassphrase = tree.AsStringWithFlags(
+			options.EncryptionPassphrase,
+			tree.FmtBareStrings,
+		)
+	}
+	if options.EncryptionKMSURI != nil {
+		if encryption.Mode != jobspb.EncryptionMode_None {
+			return 0, errors.Newf("only one encryption mode can be specified")
+		}
+		encryption.RawKmsUris = builtins.ExprSliceToStrSlice(options.EncryptionKMSURI)
+	}
+
 	details := jobspb.BackupDetails{
 		ScheduleID: scheduleID,
 		StartTime:  start,
@@ -213,7 +245,8 @@ func StartCompactionJob(
 			Subdir: fullBackupPath,
 			Exists: true,
 		},
-		EncryptionOptions: &encryptionOpts,
+		EncryptionOptions: &encryption,
+		ExecutionLocality: executionLocality,
 		Compact:           true,
 	}
 	jobID := planHook.ExecCfg().JobRegistry.MakeJobID()
@@ -598,6 +631,7 @@ func updateCompactionBackupDetails(
 		ResolvedCompleteDbs: lastBackup.CompleteDbs,
 		FullCluster:         lastBackup.DescriptorCoverage == tree.AllDescriptors,
 		ScheduleID:          initialDetails.ScheduleID,
+		ExecutionLocality:   initialDetails.ExecutionLocality,
 		Compact:             true,
 	}
 	return compactedDetails, nil