backup: hooked up compaction processor logic to job

Author: kev-cao

Diff for: pkg/backup/BUILD.bazel

@@ -14,7 +14,7 @@ go_library(
         "backup_processor_planning.go",
         "backup_span_coverage.go",
         "backup_telemetry.go",
-        "compaction.go",
+        "compaction_job.go",
         "compaction_policy.go",
         "compaction_processor.go",
         "compaction_processor_planning.go",
@@ -180,6 +180,7 @@ go_test(
         "bench_covering_test.go",
         "bench_test.go",
         "compaction_policy_test.go",
+        "compaction_processor_planning_test.go",
         "compaction_test.go",
         "create_scheduled_backup_test.go",
         "data_driven_generated_test.go",  # keep

Diff for: pkg/backup/compaction.go renamed to pkg/backup/compaction_job.go

@@ -36,6 +36,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
@@ -82,7 +83,10 @@ func maybeStartCompactionJob(
 		execCfg.InternalDB,
 		user,
 	)
-	chain, _, _, _, err := getBackupChain(ctx, execCfg, user, triggerJob, &kmsEnv)
+	chain, _, _, _, err := getBackupChain(
+		ctx, execCfg, user, triggerJob.Destination, triggerJob.EncryptionOptions,
+		triggerJob.EndTime, &kmsEnv,
+	)
 	if err != nil {
 		return 0, errors.Wrap(err, "failed to get backup chain")
 	}
@@ -351,7 +355,9 @@ func (b *backupResumer) ResumeCompaction(
 	// dying), so if we receive a retryable error, re-plan and retry the backup.
 	// TODO (kev-cao): Add progress tracking to compactions.
 	for r := retry.StartWithCtx(ctx, retryOpts); r.Next(); {
-		if err = compactChain.Compact(ctx, execCtx, updatedDetails, backupManifest, defaultStore, kmsEnv); err == nil {
+		if err = doCompaction(
+			ctx, execCtx, b.job.ID(), updatedDetails, compactChain, backupManifest, defaultStore, kmsEnv,
+		); err == nil {
 			break
 		}
 
@@ -674,11 +680,19 @@ func getBackupChain(
 			log.Warningf(ctx, "failed to cleanup incremental backup stores: %+v", err)
 		}
 	}()
-	encryption, err := backupencryption.GetEncryptionFromBaseStore(
-		ctx, baseStores[0], encryptionOpts, kmsEnv,
-	)
-	if err != nil {
-		return nil, nil, nil, nil, err
+	// If encryption keys have not already been computed, then we will compute
+	// it using the base store.
+	var encryption *jobspb.BackupEncryptionOptions
+	if encryptionOpts != nil && encryptionOpts.Mode != jobspb.EncryptionMode_None &&
+		(encryptionOpts.Key != nil || encryptionOpts.KMSInfo != nil) {
+		encryption = encryptionOpts
+	} else {
+		encryption, err = backupencryption.GetEncryptionFromBaseStore(
+			ctx, baseStores[0], encryptionOpts, kmsEnv,
+		)
+		if err != nil {
+			return nil, nil, nil, nil, err
+		}
 	}
 	mem := execCfg.RootMemoryMonitor.MakeBoundAccount()
 	defer mem.Close(ctx)
@@ -738,7 +752,7 @@ func concludeBackupCompaction(
 func processProgress(
 	ctx context.Context,
 	manifest *backuppb.BackupManifest,
-	progCh <-chan execinfrapb.RemoteProducerMetadata_BulkProcessorProgress,
+	progCh <-chan *execinfrapb.RemoteProducerMetadata_BulkProcessorProgress,
 ) error {
 	// When a processor is done exporting a span, it will send a progress update
 	// to progCh.
@@ -782,6 +796,37 @@ func compactionJobDescription(details jobspb.BackupDetails) (string, error) {
 	return fmtCtx.CloseAndGetString(), nil
 }
 
+func doCompaction(
+	ctx context.Context,
+	execCtx sql.JobExecContext,
+	jobID jobspb.JobID,
+	details jobspb.BackupDetails,
+	compactChain compactionChain,
+	manifest *backuppb.BackupManifest,
+	defaultStore cloud.ExternalStorage,
+	kmsEnv cloud.KMSEnv,
+) error {
+	progCh := make(chan *execinfrapb.RemoteProducerMetadata_BulkProcessorProgress)
+	runDistCompaction := func(ctx context.Context) error {
+		return runCompactionPlan(
+			ctx, execCtx, jobID, details, compactChain, manifest, defaultStore, kmsEnv, progCh,
+		)
+	}
+	checkpointLoop := func(ctx context.Context) error {
+		return processProgress(ctx, manifest, progCh)
+	}
+
+	if err := ctxgroup.GoAndWait(
+		ctx, runDistCompaction, checkpointLoop,
+	); err != nil {
+		return err
+	}
+
+	return concludeBackupCompaction(
+		ctx, execCtx, defaultStore, details.EncryptionOptions, kmsEnv, manifest,
+	)
+}
+
 func init() {
 	builtins.StartCompactionJob = StartCompactionJob
 }

Diff for: pkg/backup/compaction_processor.go

@@ -167,6 +167,9 @@ func runCompactBackups(
 	spec execinfrapb.CompactBackupsSpec,
 	progCh chan execinfrapb.RemoteProducerMetadata_BulkProcessorProgress,
 ) error {
+	if len(spec.AssignedSpans) == 0 {
+		return nil
+	}
 	user := spec.User()
 	execCfg, ok := flowCtx.Cfg.ExecutorConfig.(*sql.ExecutorConfig)
 	if !ok {
@@ -177,6 +180,9 @@ func runCompactBackups(
 		return errors.Wrapf(err, "export configuration")
 	}
 	defaultStore, err := execCfg.DistSQLSrv.ExternalStorage(ctx, defaultConf)
+	if err != nil {
+		return errors.Wrapf(err, "external storage")
+	}
 
 	compactChain, encryption, err := compactionChainFromSpec(ctx, execCfg, spec, user)
 	if err != nil {
@@ -186,6 +192,7 @@ func runCompactBackups(
 	intersectSpanCh := make(chan execinfrapb.RestoreSpanEntry, 1000)
 	tasks := []func(context.Context) error{
 		func(ctx context.Context) error {
+			defer close(intersectSpanCh)
 			return genIntersectingSpansForCompaction(ctx, execCfg, spec, compactChain, intersectSpanCh)
 		},
 		func(ctx context.Context) error {
@@ -284,9 +291,7 @@ func genIntersectingSpansForCompaction(
 	compactChain compactionChain,
 	intersectSpanCh chan execinfrapb.RestoreSpanEntry,
 ) error {
-	user := spec.User()
-
-	backupLocalityMap, err := makeBackupLocalityMap(compactChain.compactedLocalityInfo, user)
+	backupLocalityMap, err := makeBackupLocalityMap(compactChain.compactedLocalityInfo, spec.User())
 	if err != nil {
 		return err
 	}
@@ -299,7 +304,6 @@ func genIntersectingSpansForCompaction(
 	}
 	defer introducedSpanFrontier.Release()
 
-	entryCh := make(chan execinfrapb.RestoreSpanEntry, 1000)
 	targetSize := targetRestoreSpanSize.Get(&execCfg.Settings.SV)
 	maxFiles := maxFileCount.Get(&execCfg.Settings.SV)
 
@@ -311,12 +315,12 @@ func genIntersectingSpansForCompaction(
 		targetSize,
 		maxFiles,
 	)
+	if err != nil {
+		return err
+	}
 
-	genSpan := func(ctx context.Context) error {
+	genSpan := func(ctx context.Context, entryCh chan execinfrapb.RestoreSpanEntry) error {
 		defer close(entryCh)
-		if err != nil {
-			return err
-		}
 		return errors.Wrapf(generateAndSendImportSpans(
 			ctx,
 			spec.Spans,
@@ -329,7 +333,7 @@ func genIntersectingSpansForCompaction(
 		), "generate and send import spans")
 	}
 
-	filterIntersectingSpans := func(ctx context.Context) error {
+	filterIntersectingSpans := func(ctx context.Context, entryCh chan execinfrapb.RestoreSpanEntry) error {
 		for {
 			select {
 			case <-ctx.Done():
@@ -348,12 +352,13 @@ func genIntersectingSpansForCompaction(
 		}
 	}
 
+	entryCh := make(chan execinfrapb.RestoreSpanEntry, 1000)
 	tasks := []func(context.Context) error{
 		func(ctx context.Context) error {
-			return genSpan(ctx)
+			return genSpan(ctx, entryCh)
 		},
 		func(ctx context.Context) error {
-			return filterIntersectingSpans(ctx)
+			return filterIntersectingSpans(ctx, entryCh)
 		},
 	}
 	return ctxgroup.GoAndWait(ctx, tasks...)
@@ -368,7 +373,7 @@ func openSSTs(
 ) (mergedSST, error) {
 	var dirs []cloud.ExternalStorage
 	storeFiles := make([]storageccl.StoreFile, 0, len(entry.Files))
-	for idx := 0; idx < len(entry.Files); idx++ {
+	for idx := range entry.Files {
 		file := entry.Files[idx]
 		dir, err := execCfg.DistSQLSrv.ExternalStorage(ctx, file.Dir)
 		if err != nil {

Diff for: pkg/backup/compaction_processor_planning.go

@@ -32,27 +32,34 @@ import (
 // in the backup chain that need to be compacted. It sends updates from the
 // BulkProcessor to the provided progress channel. It is the caller's
 // responsibility to close the progress channel.
-func (c *compactionChain) runCompactionPlan(
+func runCompactionPlan(
 	ctx context.Context,
 	execCtx sql.JobExecContext,
 	jobID jobspb.JobID,
-	manifest *backuppb.BackupManifest,
 	details jobspb.BackupDetails,
+	compactChain compactionChain,
+	manifest *backuppb.BackupManifest,
 	defaultStore cloud.ExternalStorage,
 	kmsEnv cloud.KMSEnv,
 	progCh chan *execinfrapb.RemoteProducerMetadata_BulkProcessorProgress,
 ) error {
+	defer close(progCh)
 	log.Infof(
 		ctx, "planning compaction of %d backups: %s",
-		len(c.chainToCompact), util.Map(c.chainToCompact, func(m backuppb.BackupManifest) string {
+		len(compactChain.chainToCompact),
+		util.Map(compactChain.chainToCompact, func(m backuppb.BackupManifest) string {
 			return m.ID.String()
 		}),
 	)
-	backupLocalityMap, err := makeBackupLocalityMap(c.compactedLocalityInfo, execCtx.User())
+	backupLocalityMap, err := makeBackupLocalityMap(
+		compactChain.compactedLocalityInfo, execCtx.User(),
+	)
 	if err != nil {
 		return err
 	}
-	introducedSpanFrontier, err := createIntroducedSpanFrontier(c.backupChain, manifest.EndTime)
+	introducedSpanFrontier, err := createIntroducedSpanFrontier(
+		compactChain.backupChain, manifest.EndTime,
+	)
 	if err != nil {
 		return err
 	}
@@ -72,18 +79,17 @@ func (c *compactionChain) runCompactionPlan(
 	}
 
 	spansToCompact, err := getSpansToCompact(
-		ctx, execCtx, manifest, c.chainToCompact, details, defaultStore, kmsEnv,
+		ctx, execCtx, manifest, compactChain.chainToCompact, details, defaultStore, kmsEnv,
 	)
 	if err != nil {
 		return err
 	}
 	genSpan := func(ctx context.Context, spanCh chan execinfrapb.RestoreSpanEntry) error {
-		defer close(spanCh)
 		return errors.Wrap(generateAndSendImportSpans(
 			ctx,
 			spansToCompact,
-			c.chainToCompact,
-			c.compactedIterFactory,
+			compactChain.chainToCompact,
+			compactChain.compactedIterFactory,
 			backupLocalityMap,
 			filter,
 			fsc,
@@ -105,9 +111,10 @@ func (c *compactionChain) runCompactionPlan(
 		}
 		return nil
 	}
+	rowResultWriter := sql.NewRowResultWriter(nil)
 	recv := sql.MakeDistSQLReceiver(
 		ctx,
-		sql.NewMetadataCallbackWriter(nil, metaFn),
+		sql.NewMetadataCallbackWriter(rowResultWriter, metaFn),
 		tree.Rows,
 		nil, /* rangeCache */
 		nil, /* txn */
@@ -125,7 +132,7 @@ func (c *compactionChain) runCompactionPlan(
 	return nil
 }
 
-// createCompactionPlan creates an unfinalized physical plan that will
+// createCompactionPlan creates an un-finalized physical plan that will
 // distribute spans from a generator across the cluster for compaction.
 func createCompactionPlan(
 	ctx context.Context,
@@ -184,6 +191,7 @@ func countRestoreSpanEntries(
 			return nil
 		},
 		func(ctx context.Context) error {
+			defer close(countSpansCh)
 			return genSpan(ctx, countSpansCh)
 		},
 	}
@@ -194,7 +202,7 @@ func countRestoreSpanEntries(
 }
 
 // createCompactionCorePlacements takes spans from a generator and evenly
-// distributes them across nodes in the cluster, returning the core core placements
+// distributes them across nodes in the cluster, returning the core placements
 // reflecting that distribution.
 func createCompactionCorePlacements(
 	ctx context.Context,
@@ -209,17 +217,32 @@ func createCompactionCorePlacements(
 ) ([]physicalplan.ProcessorCorePlacement, error) {
 	numNodes := len(sqlInstanceIDs)
 	corePlacements := make([]physicalplan.ProcessorCorePlacement, numNodes)
+	for i := range corePlacements {
+		corePlacements[i].SQLInstanceID = sqlInstanceIDs[i]
+		corePlacements[i].Core.CompactBackups = &execinfrapb.CompactBackupsSpec{
+			JobID:       int64(jobID),
+			DefaultURI:  details.URI,
+			Destination: details.Destination,
+			Encryption:  details.EncryptionOptions,
+			StartTime:   details.StartTime,
+			EndTime:     details.EndTime,
+			ElideMode:   elideMode,
+			UserProto:   user.EncodeProto(),
+			Spans:       spansToCompact,
+		}
+	}
 
 	spanEntryCh := make(chan execinfrapb.RestoreSpanEntry, 1000)
 	var tasks []func(ctx context.Context) error
 	tasks = append(tasks, func(ctx context.Context) error {
+		defer close(spanEntryCh)
 		return genSpan(ctx, spanEntryCh)
 	})
 	tasks = append(tasks, func(ctx context.Context) error {
 		numEntriesPerNode := numEntries / numNodes
 		leftoverEntries := numEntries % numNodes
 		getTargetNumEntries := func(nodeIdx int) int {
-			if nodeIdx <= leftoverEntries {
+			if nodeIdx < leftoverEntries {
 				// This more evenly distributes the leftover entries across the nodes
 				// after doing integer division to assign the entries to the nodes.
 				return numEntriesPerNode + 1
@@ -231,29 +254,21 @@ func createCompactionCorePlacements(
 		targetNumEntries := getTargetNumEntries(currNode)
 
 		for entry := range spanEntryCh {
+			currEntries = append(currEntries, entry)
 			if len(currEntries) == targetNumEntries {
 				corePlacements[currNode].SQLInstanceID = sqlInstanceIDs[currNode]
-				corePlacements[currNode].Core.CompactBackups = &execinfrapb.CompactBackupsSpec{
-					JobID:       int64(jobID),
-					Destination: details.Destination,
-					Encryption:  details.EncryptionOptions,
-					StartTime:   details.StartTime,
-					EndTime:     details.EndTime,
-					ElideMode:   elideMode,
-					UserProto:   user.EncodeProto(),
-					Spans:       spansToCompact,
-					AssignedSpans: util.Map(currEntries, func(entry execinfrapb.RestoreSpanEntry) roachpb.Span {
+				corePlacements[currNode].Core.CompactBackups.AssignedSpans = util.Map(
+					currEntries,
+					func(entry execinfrapb.RestoreSpanEntry) roachpb.Span {
 						return entry.Span
-					}),
-				}
+					})
 				currNode++
 				targetNumEntries = getTargetNumEntries(currNode)
 				currEntries = currEntries[:0]
 			}
 			if currNode == numNodes {
 				return nil
 			}
-			currEntries = append(currEntries, entry)
 		}
 		return nil
 	})