Merge #144365

144365: drt: generalize pause job to support backup, restore and LDR r=shailendra-patel a=vidit-bhat

Previously, the `pause-job` operation was hardcoded to target only `LOGICAL REPLICATION` jobs.
This was inadequate because it prevented the reuse of logic for other long-running jobs like BACKUP and RESTORE. This patch generalizes the `pause-job` operation into a parameterized function supporting `LOGICAL REPLICATION`, `BACKUP`, and `RESTORE`.

Epic: none
Fixes: #138503 #138504 #138507 #138508
Release note: None

Co-authored-by: Vidit Bhat <[email protected]>

Author: craig[bot]

pkg/cmd/roachtest/operations/pause_job.go

@@ -18,66 +18,109 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 )
 
-type resumePausedJob struct {
+// resumeJob implements OperationCleanup and resumes a paused job by job ID.
+type resumeJob struct {
 	jobId string
 }
 
-func (r *resumePausedJob) Cleanup(ctx context.Context, o operation.Operation, c cluster.Cluster) {
+func (r *resumeJob) Cleanup(ctx context.Context, o operation.Operation, c cluster.Cluster) {
 	conn := c.Conn(ctx, o.L(), 1, option.VirtualClusterName(roachtestflags.VirtualCluster))
 	defer conn.Close()
 
+	o.Status(fmt.Sprintf("resuming job %s", r.jobId))
 	resumeJobStmt := fmt.Sprintf("RESUME JOB %s", r.jobId)
 	_, err := conn.ExecContext(ctx, resumeJobStmt)
 	if err != nil {
 		o.Fatal(err)
 	}
 }
 
-func pauseLDRJob(
-	ctx context.Context, o operation.Operation, c cluster.Cluster,
-) registry.OperationCleanup {
-	conn := c.Conn(ctx, o.L(), 1, option.VirtualClusterName(roachtestflags.VirtualCluster))
-	defer conn.Close()
+// runPauseJob returns a registry.OperationCleanup function that pauses a running job of the given type.
+func runPauseJob(
+	jobType string,
+) func(ctx context.Context, o operation.Operation, c cluster.Cluster) registry.OperationCleanup {
+	return func(ctx context.Context, o operation.Operation, c cluster.Cluster) registry.OperationCleanup {
+		conn := c.Conn(ctx, o.L(), 1, option.VirtualClusterName(roachtestflags.VirtualCluster))
+		defer conn.Close()
 
-	//fetch running ldr jobs
-	jobs, err := conn.QueryContext(ctx, "(WITH x AS (SHOW JOBS) SELECT job_id FROM x WHERE job_type = 'LOGICAL REPLICATION' AND status = 'running')")
-	if err != nil {
-		o.Fatal(err)
-	}
+		// Query for running jobs of the specified type
+		query := `
+			WITH jobs AS (SHOW JOBS)
+			SELECT job_id FROM jobs
+			WHERE job_type = $1 AND status = 'running'
+		`
+		rows, err := conn.QueryContext(ctx, query, jobType)
+		if err != nil {
+			o.Fatal(err)
+		}
+
+		var jobIds []string
+		for rows.Next() {
+			var jobId string
+			if err := rows.Scan(&jobId); err != nil {
+				o.Fatal(err)
+			}
+			jobIds = append(jobIds, jobId)
+		}
 
-	var jobIds []string
-	for jobs.Next() {
-		var jobId string
-		if err := jobs.Scan(&jobId); err != nil {
+		if len(jobIds) == 0 {
+			o.Fatal(fmt.Sprintf("no running %s jobs found", jobType))
+		}
+
+		// Randomly pick a job to pause
+		rng, _ := randutil.NewPseudoRand()
+		jobId := jobIds[rng.Intn(len(jobIds))]
+
+		o.Status(fmt.Sprintf("pausing %s job %s", jobType, jobId))
+		pauseStmt := fmt.Sprintf("PAUSE JOB %s WITH REASON = 'roachtest operation'", jobId)
+		if _, err := conn.ExecContext(ctx, pauseStmt); err != nil {
 			o.Fatal(err)
 		}
-		jobIds = append(jobIds, jobId)
-	}
 
-	//pick a random ldr job
-	rng, _ := randutil.NewPseudoRand()
-	jobId := jobIds[rng.Intn(len(jobIds))]
+		o.Status(fmt.Sprintf("paused %s job %s", jobType, jobId))
 
-	o.Status(fmt.Sprintf("pausing LDR job %s", jobId))
-	pauseJobStmt := fmt.Sprintf("PAUSE JOB %s WITH REASON = 'roachtest operation'", jobId)
-	_, err = conn.ExecContext(ctx, pauseJobStmt)
-	if err != nil {
-		o.Fatal(err)
+		return &resumeJob{jobId: jobId}
 	}
+}
 
-	o.Status(fmt.Sprintf("paused LDR job %s", jobId))
-	return &resumePausedJob{
-		jobId: jobId,
+// registerPauseJob registers pause-job operations for various job types (e.g., BACKUP, RESTORE, LOGICAL REPLICATION).
+func registerPauseJob(r registry.Registry) {
+	jobTypes := []struct {
+		JobType     string
+		OpName      string
+		Dependency  registry.OperationDependency
+		Owner       registry.Owner
+		Description string
+	}{
+		{
+			JobType:    "LOGICAL REPLICATION",
+			OpName:     "pause-job/logical-replication",
+			Dependency: registry.OperationRequiresLDRJobRunning,
+			Owner:      registry.OwnerDisasterRecovery,
+		},
+		{
+			JobType:    "BACKUP",
+			OpName:     "pause-job/backup",
+			Dependency: registry.OperationRequiresRunningBackupJob,
+			Owner:      registry.OwnerDisasterRecovery,
+		},
+		{
+			JobType:    "RESTORE",
+			OpName:     "pause-job/restore",
+			Dependency: registry.OperationRequiresRunningRestoreJob,
+			Owner:      registry.OwnerDisasterRecovery,
+		},
 	}
-}
 
-func registerPauseLDRJob(r registry.Registry) {
-	r.AddOperation(registry.OperationSpec{
-		Name:             "pause-ldr",
-		Owner:            registry.OwnerDisasterRecovery,
-		Timeout:          15 * time.Minute,
-		CompatibleClouds: registry.AllClouds,
-		Dependencies:     []registry.OperationDependency{registry.OperationRequiresLDRJobRunning},
-		Run:              pauseLDRJob,
-	})
+	for _, jt := range jobTypes {
+		r.AddOperation(registry.OperationSpec{
+			Name:               jt.OpName,
+			Owner:              jt.Owner,
+			Timeout:            15 * time.Minute,
+			CompatibleClouds:   registry.AllClouds,
+			Dependencies:       []registry.OperationDependency{jt.Dependency},
+			Run:                runPauseJob(jt.JobType),
+			CanRunConcurrently: registry.OperationCanRunConcurrently,
+		})
+	}
 }

pkg/cmd/roachtest/operations/register.go

@@ -27,7 +27,7 @@ func RegisterOperations(r registry.Registry) {
 	registerBackupRestore(r)
 	registerManualCompaction(r)
 	registerResize(r)
-	registerPauseLDRJob(r)
+	registerPauseJob(r)
 	registerLicenseThrottle(r)
 	registerSessionVariables(r)
 	registerDebugZip(r)

pkg/cmd/roachtest/registry/operation_spec.go

@@ -27,6 +27,8 @@ const (
 	OperationRequiresZeroUnavailableRanges
 	OperationRequiresZeroUnderreplicatedRanges
 	OperationRequiresLDRJobRunning
+	OperationRequiresRunningBackupJob
+	OperationRequiresRunningRestoreJob
 )
 
 // OperationIsolation specifies to what extent the operation runner will try

pkg/cmd/roachtest/roachtestutil/operations/dependency.go

@@ -52,6 +52,18 @@ func CheckDependencies(
 				return ok, err
 			}
 
+		case registry.OperationRequiresRunningBackupJob:
+			ok, err := checkBackupJobRunning(ctx, c, l)
+			if err != nil || !ok {
+				return ok, err
+			}
+
+		case registry.OperationRequiresRunningRestoreJob:
+			ok, err := checkRestoreJobRunning(ctx, c, l)
+			if err != nil || !ok {
+				return ok, err
+			}
+
 		default:
 			panic(fmt.Sprintf("unknown operation dependency %d", dep))
 		}
@@ -128,3 +140,37 @@ func checkLDRJobRunning(
 	_ = jobsCur.Scan(&jobId)
 	return jobId != "", nil
 }
+
+func checkBackupJobRunning(
+	ctx context.Context, c cluster.Cluster, l *logger.Logger,
+) (ok bool, _ error) {
+	conn := c.Conn(ctx, l, 1, option.VirtualClusterName("system"))
+	defer conn.Close()
+
+	jobsCur, err := conn.QueryContext(ctx,
+		"(WITH x AS (SHOW JOBS) SELECT job_id FROM x WHERE job_type = 'BACKUP' AND status = 'running' LIMIT 1)")
+	if err != nil {
+		return false, err
+	}
+	jobsCur.Next()
+	var jobId string
+	_ = jobsCur.Scan(&jobId)
+	return jobId != "", nil
+}
+
+func checkRestoreJobRunning(
+	ctx context.Context, c cluster.Cluster, l *logger.Logger,
+) (ok bool, _ error) {
+	conn := c.Conn(ctx, l, 1, option.VirtualClusterName("system"))
+	defer conn.Close()
+
+	jobsCur, err := conn.QueryContext(ctx,
+		"WITH x AS (SHOW JOBS) SELECT job_id FROM x WHERE job_type = 'RESTORE' AND status = 'running' LIMIT 1)")
+	if err != nil {
+		return false, err
+	}
+	jobsCur.Next()
+	var jobId string
+	_ = jobsCur.Scan(&jobId)
+	return jobId != "", nil
+}