backup: checkpoint after online restore link phase

This patch will prevent online restore from rerunning the link flow if it has
already completed. This will be useful for restore with experimental copy,
which runs the download phase in the same job as the link phase. At some point
in the future, we could add finer grain checkpointing to the link phase, but
that would require more work and thought.

Epic: none

Release note: none

Author: msbutler

pkg/backup/restore_job.go

@@ -297,11 +297,27 @@ func restore(
 		return emptyRowCount, nil
 	}
 
+	job := resumer.job
+	details := job.Details().(jobspb.RestoreDetails)
+
+	if details.OnlineImpl() {
+		var linkPhaseComplete bool
+		if err := execCtx.ExecCfg().InternalDB.Txn(restoreCtx, func(ctx context.Context, txn isql.Txn) error {
+			jobInfo := jobs.InfoStorageForJob(txn, resumer.job.ID())
+			_, ok, err := jobInfo.Get(ctx, "get-link-complete-key", linkCompleteKey)
+			linkPhaseComplete = ok
+			return err
+		}); err != nil {
+			log.Warningf(restoreCtx, "failed to get checkpoint for link phase %v", err)
+		}
+		if linkPhaseComplete {
+			return emptyRowCount, nil
+		}
+	}
+
 	// If we've already migrated some of the system tables we're about to
 	// restore, this implies that a previous attempt restored all of this data.
 	// We want to avoid restoring again since we'll be shadowing migrated keys.
-	job := resumer.job
-	details := job.Details().(jobspb.RestoreDetails)
 	if alreadyMigrated := checkForMigratedData(details, dataToRestore); alreadyMigrated {
 		return emptyRowCount, nil
 	}
@@ -1982,6 +1998,16 @@ func (r *restoreResumer) doResume(ctx context.Context, execCtx interface{}) erro
 	if err := insertStats(ctx, r.job, p.ExecCfg(), remappedStats); err != nil {
 		return errors.Wrap(err, "inserting table statistics")
 	}
+
+	if details.OnlineImpl() {
+		if err := p.ExecCfg().InternalDB.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+			jobInfo := jobs.InfoStorageForJob(txn, r.job.ID())
+			return jobInfo.Write(ctx, linkCompleteKey, []byte{})
+		}); err != nil {
+			log.Warningf(ctx, "failed to checkpoint link flow %v", err)
+		}
+	}
+
 	if details.ExperimentalCopy {
 		downloadSpans, err := getDownloadSpans(p.ExecCfg().Codec, preData, mainData)
 		if err != nil {

pkg/backup/restore_online.go

@@ -61,6 +61,8 @@ var onlineRestoreLayerLimit = settings.RegisterIntSetting(
 	settings.WithVisibility(settings.Reserved),
 )
 
+const linkCompleteKey = "link_complete"
+
 // splitAndScatter runs through all entries produced by genSpans splitting and
 // scattering the key-space designated by the passed rewriter such that if all
 // files in the entries in those spans were ingested the amount ingested between

pkg/backup/restore_online_test.go

@@ -8,6 +8,7 @@ package backup
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"reflect"
 	"strings"
 	"testing"
@@ -26,11 +27,27 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/kr/pretty"
 	"github.com/stretchr/testify/require"
 )
 
+var orParams = base.TestClusterArgs{
+	// Online restore is not supported in a secondary tenant yet.
+	ServerArgs: base.TestServerArgs{
+		DefaultTestTenant: base.TestIsSpecificToStorageLayerAndNeedsASystemTenant,
+	},
+}
+
+func onlineImpl(rng *rand.Rand) string {
+	opt := "EXPERIMENTAL DEFERRED COPY"
+	if rng.Intn(2) == 0 {
+		opt = "EXPERIMENTAL COPY"
+	}
+	return opt
+}
+
 func TestOnlineRestoreBasic(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
@@ -40,16 +57,11 @@ func TestOnlineRestoreBasic(t *testing.T) {
 	ctx := context.Background()
 
 	const numAccounts = 1000
-	params := base.TestClusterArgs{
-		// Online restore is not supported in a secondary tenant yet.
-		ServerArgs: base.TestServerArgs{
-			DefaultTestTenant: base.TestIsSpecificToStorageLayerAndNeedsASystemTenant,
-		},
-	}
-	tc, sqlDB, dir, cleanupFn := backupRestoreTestSetupWithParams(t, singleNode, numAccounts, InitManualReplication, params)
+
+	tc, sqlDB, dir, cleanupFn := backupRestoreTestSetupWithParams(t, singleNode, numAccounts, InitManualReplication, orParams)
 	defer cleanupFn()
 
-	rtc, rSQLDB, cleanupFnRestored := backupRestoreTestSetupEmpty(t, 1, dir, InitManualReplication, params)
+	rtc, rSQLDB, cleanupFnRestored := backupRestoreTestSetupEmpty(t, 1, dir, InitManualReplication, orParams)
 	defer cleanupFnRestored()
 
 	externalStorage := "nodelocal://1/backup"
@@ -118,6 +130,36 @@ func TestOnlineRestorePartitioned(t *testing.T) {
 	sqlDB.Exec(t, fmt.Sprintf(`SHOW JOB WHEN COMPLETE %s`, j[0][4]))
 }
 
+func TestOnlineRestoreLinkCheckpoint(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	defer nodelocal.ReplaceNodeLocalForTesting(t.TempDir())()
+
+	rng, _ := randutil.NewTestRand()
+
+	const numAccounts = 10
+	_, sqlDB, _, cleanupFn := backupRestoreTestSetupWithParams(
+		t,
+		singleNode,
+		numAccounts,
+		InitManualReplication,
+		orParams,
+	)
+	defer cleanupFn()
+	sqlDB.Exec(t, "BACKUP DATABASE data INTO $1", localFoo)
+	sqlDB.Exec(t, "SET CLUSTER SETTING jobs.debug.pausepoints  = 'restore.before_publishing_descriptors'")
+	var jobID jobspb.JobID
+	stmt := fmt.Sprintf("RESTORE DATABASE data FROM LATEST IN $1 WITH OPTIONS (new_db_name='data2', %s, detached)", onlineImpl(rng))
+	sqlDB.QueryRow(t, stmt, localFoo).Scan(&jobID)
+	jobutils.WaitForJobToPause(t, sqlDB, jobID)
+
+	// Set a pauspoint during the link phase which should not get hit because of
+	// checkpointing.
+	sqlDB.Exec(t, "SET CLUSTER SETTING jobs.debug.pausepoints  = 'restore.before_link'")
+	sqlDB.Exec(t, "RESUME JOB $1", jobID)
+	jobutils.WaitForJobToSucceed(t, sqlDB, jobID)
+}
+
 func TestOnlineRestoreStatementResult(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)