restore: fix retry revert stuck in reverting

Previously, if a reverting restore job was paused after dropping its
descriptors and then resumed, the restore job would get stuck in a retry
loop due to missing descriptors.

This commit adds a field in the restore details that can be set once
descriptors have been dropped during restore cleanup. This allows us to
avoid attempting to drop descriptors if they have already been dropped.

Fixes: #156019

Release note: Restore no longer gets stuck in a retry loop when reverts
are attempted twice.

Author: kev-cao

pkg/backup/restore_job.go

@@ -2820,6 +2820,13 @@ func (r *restoreResumer) OnFailOrCancel(
 		return err
 	}
 
+	testingKnobs := execCfg.BackupRestoreTestingKnobs
+	if testingKnobs != nil && testingKnobs.AfterRevertRestoreDropDescriptors != nil {
+		if err := testingKnobs.AfterRevertRestoreDropDescriptors(); err != nil {
+			return err
+		}
+	}
+
 	if details.DescriptorCoverage == tree.AllDescriptors {
 		// The temporary system table descriptors should already have been dropped
 		// in `dropDescriptors` but we still need to drop the temporary system db.
@@ -2865,6 +2872,12 @@ func (r *restoreResumer) dropDescriptors(
 		return nil
 	}
 
+	// Descriptors have already been dropped once before, this is a retry of the
+	// cleanup.
+	if details.DroppedDescsOnFail {
+		return nil
+	}
+
 	b := txn.KV().NewBatch()
 	const kvTrace = false
 	// Collect the tables into mutable versions.
@@ -3186,7 +3199,8 @@ func (r *restoreResumer) dropDescriptors(
 		return errors.Wrap(err, "dropping tables created at the start of restore caused by fail/cancel")
 	}
 
-	return nil
+	details.DroppedDescsOnFail = true
+	return errors.Wrap(r.job.WithTxn(txn).SetDetails(ctx, details), "checkpointing dropped descs on fail")
 }
 
 // removeExistingTypeBackReferences removes back references from types that

pkg/backup/restore_test.go

@@ -318,3 +318,53 @@ func TestRestoreDuplicateTempTables(t *testing.T) {
 	result := sqlDB.QueryStr(t, `SELECT table_name FROM [SHOW TABLES] ORDER BY table_name`)
 	require.Equal(t, [][]string{{"permanent_table"}}, result)
 }
+
+func TestRestoreRetryRevert(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	droppedDescs := make(chan struct{})
+	testKnobs := &sql.BackupRestoreTestingKnobs{
+		AfterRevertRestoreDropDescriptors: func() error {
+			close(droppedDescs)
+			return errors.New("injected error to test pause during revert")
+		},
+	}
+	var params base.TestClusterArgs
+	params.ServerArgs.Knobs.BackupRestore = testKnobs
+
+	const numAccounts = 1
+	_, sqlDB, _, cleanupFn := backuptestutils.StartBackupRestoreTestCluster(
+		t, singleNode, backuptestutils.WithParams(params),
+	)
+	defer cleanupFn()
+
+	// We create a variety of descriptors to ensure that missing descriptors
+	// during drop do not break revert.
+	sqlDB.Exec(t, "CREATE DATABASE foo")
+	sqlDB.Exec(t, "USE foo")
+	sqlDB.Exec(t, "CREATE OR REPLACE FUNCTION bar(a INT) RETURNS INT AS 'SELECT a*a' LANGUAGE SQL;")
+	sqlDB.Exec(t, "CREATE TYPE baz AS ENUM ('a', 'b', 'c')")
+	sqlDB.Exec(t, "CREATE TABLE qux (x INT)")
+	sqlDB.Exec(t, "BACKUP DATABASE foo INTO 'nodelocal://1/backup'")
+
+	// We need restore to publish descriptors so that they will be cleaned up
+	// during restore.
+	sqlDB.Exec(t, "SET CLUSTER SETTING jobs.debug.pausepoints = 'restore.after_publishing_descriptors'")
+
+	var restoreID jobspb.JobID
+	sqlDB.QueryRow(
+		t, "RESTORE DATABASE foo FROM LATEST IN 'nodelocal://1/backup' WITH detached, new_db_name='foo_restored'",
+	).Scan(&restoreID)
+
+	jobutils.WaitForJobToPause(t, sqlDB, restoreID)
+
+	sqlDB.Exec(t, "CANCEL JOB $1", restoreID)
+	<-droppedDescs
+	sqlDB.Exec(t, "PAUSE JOB $1", restoreID)
+	jobutils.WaitForJobToPause(t, sqlDB, restoreID)
+	testKnobs.AfterRevertRestoreDropDescriptors = nil
+
+	sqlDB.Exec(t, "RESUME JOB $1", restoreID)
+	jobutils.WaitForJobToCancel(t, sqlDB, restoreID)
+}

pkg/jobs/jobspb/jobs.proto

@@ -656,7 +656,11 @@ message RestoreDetails {
 
   bool experimental_copy = 37;
 
-  // NEXT ID: 38.
+	// DroppedDescsOnFail is set once the restore job has cleaned up descriptors
+	// on cancel/failure.
+	bool dropped_descs_on_fail = 38;
+
+  // NEXT ID: 39.
 }
 
 

pkg/sql/exec_util_backup.go

@@ -86,6 +86,10 @@ type BackupRestoreTestingKnobs struct {
 	// RunBeforeDownloadCleanup is called before we cleanup after all external
 	// files have been download.
 	RunBeforeDownloadCleanup func() error
+
+	// AfterRevertRestoreDropDescriptors is called after a reverting restore
+	// drops its descriptors.
+	AfterRevertRestoreDropDescriptors func() error
 }
 
 var _ base.ModuleTestingKnobs = &BackupRestoreTestingKnobs{}