backup: add checkpointing to backup compactions

This patch adds periodic checkpointing to backup compaction jobs.

Epic: None

Release note: None

Author: kev-cao

Diff for: pkg/backup/compaction_job.go

@@ -40,6 +40,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
 	"github.com/gogo/protobuf/types"
@@ -218,8 +219,8 @@ func (b *backupResumer) ResumeCompaction(
 
 	var backupManifest *backuppb.BackupManifest
 	updatedDetails := initialDetails
+	testingKnobs := execCtx.ExecCfg().BackupRestoreTestingKnobs
 	if initialDetails.URI == "" {
-		testingKnobs := execCtx.ExecCfg().BackupRestoreTestingKnobs
 		if testingKnobs != nil && testingKnobs.RunBeforeResolvingCompactionDest != nil {
 			if err := testingKnobs.RunBeforeResolvingCompactionDest(); err != nil {
 				return err
@@ -334,6 +335,10 @@ func (b *backupResumer) ResumeCompaction(
 		}
 	}
 
+	if testingKnobs != nil && testingKnobs.AfterLoadingManifestOnResume != nil {
+		testingKnobs.AfterLoadingManifestOnResume(backupManifest)
+	}
+
 	// We retry on pretty generic failures -- any rpc error. If a worker node were
 	// to restart, it would produce this kind of error, but there may be other
 	// errors that are also rpc errors. Don't retry too aggressively.
@@ -342,9 +347,8 @@ func (b *backupResumer) ResumeCompaction(
 		MaxRetries: 5,
 	}
 
-	if execCtx.ExecCfg().BackupRestoreTestingKnobs != nil &&
-		execCtx.ExecCfg().BackupRestoreTestingKnobs.BackupDistSQLRetryPolicy != nil {
-		retryOpts = *execCtx.ExecCfg().BackupRestoreTestingKnobs.BackupDistSQLRetryPolicy
+	if testingKnobs != nil && testingKnobs.BackupDistSQLRetryPolicy != nil {
+		retryOpts = *testingKnobs.BackupDistSQLRetryPolicy
 	}
 
 	if err := execCtx.ExecCfg().JobRegistry.CheckPausepoint("backup.before.flow"); err != nil {
@@ -376,8 +380,6 @@ func (b *backupResumer) ResumeCompaction(
 
 		// Reload the backup manifest to pick up any spans we may have completed on
 		// previous attempts.
-		// TODO (kev-cao): Compactions currently do not create checkpoints, but this
-		// can be used to reload the manifest once we add checkpointing.
 		var reloadBackupErr error
 		mem.Shrink(ctx, memSize)
 		backupManifest, memSize, reloadBackupErr = b.readManifestOnResume(ctx, &mem, execCtx.ExecCfg(),
@@ -752,9 +754,13 @@ func concludeBackupCompaction(
 // the associated manifest.
 func processProgress(
 	ctx context.Context,
+	execCtx sql.JobExecContext,
+	details jobspb.BackupDetails,
 	manifest *backuppb.BackupManifest,
 	progCh <-chan *execinfrapb.RemoteProducerMetadata_BulkProcessorProgress,
+	kmsEnv cloud.KMSEnv,
 ) error {
+	var lastCheckpointTime time.Time
 	// When a processor is done exporting a span, it will send a progress update
 	// to progCh.
 	for progress := range progCh {
@@ -763,15 +769,22 @@ func processProgress(
 			log.Errorf(ctx, "unable to unmarshal backup progress details: %+v", err)
 			return err
 		}
-		for _, file := range progDetails.Files {
-			manifest.Files = append(manifest.Files, file)
-			manifest.EntryCounts.Add(file.EntryCounts)
-		}
+		updateManifestWithProgress(progDetails, manifest)
+
 		// TODO (kev-cao): Add per node progress updates.
+
+		if wroteCheckpoint, err := maybeWriteBackupCheckpoint(
+			ctx, execCtx, details, manifest, lastCheckpointTime, kmsEnv,
+		); err != nil {
+			log.Errorf(ctx, "unable to checkpoint compaction: %+v", err)
+		} else if wroteCheckpoint {
+			lastCheckpointTime = timeutil.Now()
+		}
 	}
 	return nil
 }
 
+// compactionJobDescription generates a redacted description of the job.
 func compactionJobDescription(details jobspb.BackupDetails) (string, error) {
 	fmtCtx := tree.NewFmtCtx(tree.FmtSimple)
 	redactedURIs, err := sanitizeURIList(details.Destination.To)
@@ -822,8 +835,7 @@ func doCompaction(
 		)
 	}
 	checkpointLoop := func(ctx context.Context) error {
-		// TODO (kev-cao): Add logic for checkpointing during loop.
-		return processProgress(ctx, manifest, progCh)
+		return processProgress(ctx, execCtx, details, manifest, progCh, kmsEnv)
 	}
 	// TODO (kev-cao): Add trace aggregator loop.
 
@@ -838,6 +850,49 @@ func doCompaction(
 	)
 }
 
+// updateManifestWithProgress takes a progress update from the processors and
+// updates the backup manifest accordingly.
+func updateManifestWithProgress(
+	progDetails backuppb.BackupManifest_Progress, manifest *backuppb.BackupManifest,
+) {
+	for _, file := range progDetails.Files {
+		manifest.Files = append(manifest.Files, file)
+		manifest.EntryCounts.Add(file.EntryCounts)
+	}
+}
+
+// maybeWriteBackupCheckpoint writes a checkpoint for the backup if
+// the time since the last checkpoint exceeds the configured interval. If a
+// checkpoint is written, the function returns true.
+func maybeWriteBackupCheckpoint(
+	ctx context.Context,
+	execCtx sql.JobExecContext,
+	details jobspb.BackupDetails,
+	manifest *backuppb.BackupManifest,
+	lastCheckpointTime time.Time,
+	kmsEnv cloud.KMSEnv,
+) (bool, error) {
+	if details.URI == "" {
+		return false, errors.New("backup details does not contain a default URI")
+	}
+	execCfg := execCtx.ExecCfg()
+	interval := BackupCheckpointInterval.Get(&execCfg.Settings.SV)
+	if timeutil.Since(lastCheckpointTime) < interval {
+		return false, nil
+	}
+	if err := backupinfo.WriteBackupManifestCheckpoint(
+		ctx, details.URI, details.EncryptionOptions, kmsEnv,
+		manifest, execCfg, execCtx.User(),
+	); err != nil {
+		return false, err
+	}
+	backupRestoreKnobs := execCfg.BackupRestoreTestingKnobs
+	if backupRestoreKnobs != nil && backupRestoreKnobs.AfterCompactBackupsCheckpoint != nil {
+		backupRestoreKnobs.AfterCompactBackupsCheckpoint()
+	}
+	return true, nil
+}
+
 func init() {
 	builtins.StartCompactionJob = StartCompactionJob
 }

Diff for: pkg/backup/compaction_test.go

@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"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/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
@@ -592,6 +593,110 @@ func TestScheduledBackupCompaction(t *testing.T) {
 	require.Equal(t, 4, numBackups)
 }
 
+func TestCompactionCheckpointing(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	// These two channels are for having the test wait until the compaction job
+	// writes a checkpoint and to pause the compaction job until the test is
+	// ready to resume, respectively.
+	// We block the job from continuing to prevent the job from completing before
+	// we can pause it.
+	hitCheckpoint := make(chan struct{})
+	blockCompaction := make(chan struct{})
+	doBlockCheckpoint := true
+
+	// These two channels are for having the test wait until the compaction job
+	// has loaded a manifest after resuming a job and to pause the job until the
+	// test is ready to resume, respectively.
+	// We block resuming because otherwise we have a race condition where the
+	// manifest is updated by the job before we can check its files.
+	loadedManifest := make(chan struct{})
+	blockResume := make(chan struct{})
+	doBlockResume := true
+	var manifest *backuppb.BackupManifest
+
+	defer func() {
+		close(hitCheckpoint)
+		close(blockCompaction)
+		close(loadedManifest)
+		close(blockResume)
+	}()
+	// Need a large enough backup so that it doesn't finish in one iteration
+	const numAccounts = 10000
+	tc, db, _, cleanup := backupRestoreTestSetupWithParams(
+		t, singleNode, numAccounts, InitManualReplication, base.TestClusterArgs{
+			ServerArgs: base.TestServerArgs{
+				Knobs: base.TestingKnobs{
+					BackupRestore: &sql.BackupRestoreTestingKnobs{
+						AfterLoadingManifestOnResume: func(m *backuppb.BackupManifest) {
+							if doBlockResume {
+								manifest = m
+								loadedManifest <- struct{}{}
+								<-blockResume
+							}
+						},
+						AfterCompactBackupsCheckpoint: func() {
+							if doBlockCheckpoint {
+								hitCheckpoint <- struct{}{}
+								<-blockCompaction
+							}
+						},
+					},
+				},
+			},
+		},
+	)
+	defer cleanup()
+	writeQueries := func() {
+		db.Exec(t, "UPDATE data.bank SET balance = balance + 1")
+	}
+	db.Exec(t, "SET CLUSTER SETTING bulkio.backup.checkpoint_interval = '10ms'")
+	start := getTime(t)
+	db.Exec(t, fmt.Sprintf("BACKUP INTO 'nodelocal://1/backup' AS OF SYSTEM TIME %d", start))
+	writeQueries()
+	db.Exec(t, "BACKUP INTO LATEST IN 'nodelocal://1/backup'")
+	writeQueries()
+	end := getTime(t)
+	db.Exec(t, fmt.Sprintf("BACKUP INTO LATEST IN 'nodelocal://1/backup' AS OF SYSTEM TIME %d", end))
+
+	var backupPath string
+	db.QueryRow(t, "SHOW BACKUPS IN 'nodelocal://1/backup'").Scan(&backupPath)
+
+	var jobID jobspb.JobID
+	db.QueryRow(
+		t,
+		"SELECT crdb_internal.backup_compaction(ARRAY['nodelocal://1/backup'], $1, ''::BYTES, $2, $3)",
+		backupPath, start, end,
+	).Scan(&jobID)
+	// Ensure that the very first manifest when the job is initially started has
+	// no files.
+	<-loadedManifest
+	require.Equal(t, 0, len(manifest.Files))
+	blockResume <- struct{}{}
+
+	// Pause after first checkpoint and then resume the job from that checkpoint.
+	<-hitCheckpoint
+	db.Exec(t, "PAUSE JOB $1", jobID)
+	// Wait for the job to pause.
+	jobutils.WaitForJobToPause(t, db, jobID)
+	doBlockCheckpoint = false // Don't bother pausing on other checkpoints.
+	blockCompaction <- struct{}{}
+	db.Exec(t, "RESUME JOB $1", jobID)
+	jobutils.WaitForJobToRun(t, db, jobID)
+
+	// Now that the job has been paused and resumed after previously hitting a
+	// checkpoint, the initial manifest at the start of the job should have
+	// some files from before the pause.
+	<-loadedManifest
+	require.NotEmpty(t, manifest.Files)
+	doBlockResume = false
+	blockResume <- struct{}{}
+
+	waitForSuccessfulJob(t, tc, jobID)
+	validateCompactedBackupForTables(t, db, []string{"bank"}, "'nodelocal://1/backup'", start, end)
+}
+
 // Start and end are unix epoch in nanoseconds.
 func validateCompactedBackupForTables(
 	t *testing.T, db *sqlutils.SQLRunner, tables []string, collectionURIs string, start, end int64,

Diff for: pkg/sql/BUILD.bazel

@@ -301,6 +301,7 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/sql",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/backup/backuppb",
         "//pkg/base",
         "//pkg/build",
         "//pkg/cloud",

Diff for: pkg/sql/exec_util.go

@@ -21,6 +21,7 @@ import (
 	"time"
 
 	apd "github.com/cockroachdb/apd/v3"
+	"github.com/cockroachdb/cockroach/pkg/backup/backuppb"
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/cloud/externalconn"
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
@@ -1873,6 +1874,14 @@ type BackupRestoreTestingKnobs struct {
 	// is written.
 	AfterBackupCheckpoint func()
 
+	// AfterCompactBackupsCheckpoint if set will be called after a BACKUP-CHECKPOINT
+	// for compaction is written.
+	AfterCompactBackupsCheckpoint func()
+
+	// AfterLoadingManifestOnResume is run once the backup manifest has been
+	// loaded/created on the resumption of a compaction job.
+	AfterLoadingManifestOnResume func(manifest *backuppb.BackupManifest)
+
 	// CaptureResolvedTableDescSpans allows for intercepting the spans which are
 	// resolved during backup planning, and will eventually be backed up during
 	// execution.