backup: reset retry counter for download job when progress is detected

In #1448821, we added a retry loop to the download phase of online
restore, but the retry policy allows at most 5 retries. In the event of
download progress, we want to be more lenient with the retry policy, so
this commit teaches the download phase to reset the retry counter if
progress is detected.

Epic: CRDB-51394

Fixes: #149893

Author: kev-cao

pkg/backup/restore_job.go

@@ -672,10 +672,11 @@ func loadBackupSQLDescs(
 type restoreResumer struct {
 	job *jobs.Job
 
-	settings      *cluster.Settings
-	execCfg       *sql.ExecutorConfig
-	restoreStats  roachpb.RowCount
-	downloadJobID jobspb.JobID
+	settings        *cluster.Settings
+	execCfg         *sql.ExecutorConfig
+	restoreStats    roachpb.RowCount
+	downloadJobID   jobspb.JobID
+	downloadJobProg float32
 
 	mu struct {
 		syncutil.Mutex

pkg/backup/restore_online.go

@@ -705,6 +705,7 @@ func (r *restoreResumer) waitForDownloadToComplete(
 	// Download is already complete or there is nothing to be downloaded, in
 	// either case we can mark the job as done.
 	if total == 0 {
+		r.downloadJobProg = 1.0
 		return r.job.NoTxn().FractionProgressed(ctx, func(ctx context.Context, details jobspb.ProgressDetails) float32 {
 			return 1.0
 		})
@@ -731,6 +732,7 @@ func (r *restoreResumer) waitForDownloadToComplete(
 		log.VInfof(ctx, 1, "restore download phase, %s downloaded, %s remaining of %s total (%.2f complete)",
 			sz(total-remaining), sz(remaining), sz(total), fractionComplete,
 		)
+		r.downloadJobProg = fractionComplete
 
 		if remaining == 0 {
 			r.notifyStatsRefresherOfNewTables()
@@ -808,6 +810,7 @@ func (r *restoreResumer) doDownloadFilesWithRetry(
 	ctx context.Context, execCtx sql.JobExecContext,
 ) error {
 	var err error
+	var lastProgress float32
 	for rt := retry.StartWithCtx(ctx, retry.Options{
 		InitialBackoff: time.Millisecond * 100,
 		MaxBackoff:     time.Second,
@@ -817,7 +820,12 @@ func (r *restoreResumer) doDownloadFilesWithRetry(
 		if err == nil {
 			return nil
 		}
-		log.Warningf(ctx, "failed attempt #%d to download files: %v", rt.CurrentAttempt(), err)
+		log.Warningf(ctx, "failed attempt to download files: %v", err)
+		if lastProgress != r.downloadJobProg {
+			lastProgress = r.downloadJobProg
+			rt.Reset()
+			log.Infof(ctx, "download progress has advanced since last retry, resetting retry counter")
+		}
 	}
 	return errors.Wrapf(err, "retries exhausted for downloading files")
 }
@@ -843,6 +851,13 @@ func (r *restoreResumer) doDownloadFiles(ctx context.Context, execCtx sql.JobExe
 		}
 		return errors.Wrap(err, "failed to generate and send download spans")
 	}
+
+	testingKnobs := execCtx.ExecCfg().BackupRestoreTestingKnobs
+	if testingKnobs != nil && testingKnobs.RunBeforeDownloadCleanup != nil {
+		if err := testingKnobs.RunBeforeDownloadCleanup(); err != nil {
+			return errors.Wrap(err, "testing knob RunBeforeDownloadCleanup failed")
+		}
+	}
 	return r.cleanupAfterDownload(ctx, details)
 }
 

pkg/backup/restore_online_test.go

@@ -642,6 +642,62 @@ func TestOnlineRestoreRetryingDownloadRequests(t *testing.T) {
 	}
 }
 
+func TestOnlineRestoreDownloadRetryReset(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	defer nodelocal.ReplaceNodeLocalForTesting(t.TempDir())()
+
+	var attemptCount int
+	clusterArgs := base.TestClusterArgs{
+		ServerArgs: base.TestServerArgs{
+			DefaultTestTenant: base.TestIsSpecificToStorageLayerAndNeedsASystemTenant,
+			Knobs: base.TestingKnobs{
+				BackupRestore: &sql.BackupRestoreTestingKnobs{
+					// We want the retry loop to fail until its final attempt, and then
+					// succeed on the last attempt. This will allow the download job to
+					// make progress, in which case the retry loop _should_ reset. Then
+					// we continue allowing the retry loop to fail until its last
+					// attempt, in which case it will succeed again.
+					RunBeforeSendingDownloadSpan: func() error {
+						attemptCount++
+						if attemptCount < maxDownloadAttempts {
+							return errors.Newf("injected download request failure")
+						}
+						return nil
+					},
+					RunBeforeDownloadCleanup: func() error {
+						if attemptCount < maxDownloadAttempts*2 {
+							return errors.Newf("injected download cleanup failure")
+						}
+						return nil
+					},
+				},
+			},
+		},
+	}
+	const numAccounts = 1
+	_, sqlDB, _, cleanupFn := backupRestoreTestSetupWithParams(
+		t, singleNode, numAccounts, InitManualReplication, clusterArgs,
+	)
+	defer cleanupFn()
+
+	externalStorage := "nodelocal://1/backup"
+	sqlDB.Exec(t, fmt.Sprintf("BACKUP INTO '%s'", externalStorage))
+	sqlDB.Exec(
+		t,
+		fmt.Sprintf(`
+		RESTORE DATABASE data FROM LATEST IN '%s'
+		WITH EXPERIMENTAL DEFERRED COPY, new_db_name=data2
+		`, externalStorage),
+	)
+
+	var downloadJobID jobspb.JobID
+	sqlDB.QueryRow(t, latestDownloadJobIDQuery).Scan(&downloadJobID)
+	jobutils.WaitForJobToSucceed(t, sqlDB, downloadJobID)
+	require.Equal(t, maxDownloadAttempts*2, attemptCount)
+}
+
 func bankOnlineRestore(
 	t *testing.T,
 	sqlDB *sqlutils.SQLRunner,

pkg/sql/exec_util_backup.go

@@ -73,6 +73,10 @@ type BackupRestoreTestingKnobs struct {
 	// RunBeforeSendingDownloadSpan is called within the retry loop of the
 	// download span worker before sending the download span request.
 	RunBeforeSendingDownloadSpan func() error
+
+	// RunBeforeDownloadCleanup is called before we cleanup after all external
+	// files have been download.
+	RunBeforeDownloadCleanup func() error
 }
 
 var _ base.ModuleTestingKnobs = &BackupRestoreTestingKnobs{}