roachtest: refactor online restore roachtest into helper functions

This commit cleans up the `runRestore` function for the online restore
roachtests by moving the restore driver logic into separate functions.
It also updates the monitor logic and moves away from the deprecated
monitor to the new `test.Go`.

Author: kev-cao

pkg/cmd/roachtest/tests/online_restore.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil/task"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
@@ -677,96 +678,145 @@ func runRestore(
 	statsCollector, err := createStatCollector(ctx, rd)
 	require.NoError(t, err)
 
-	m := c.NewDeprecatedMonitor(ctx, sp.hardware.getCRDBNodes())
-	var restoreStartTime, restoreEndTime time.Time
-	m.Go(func(ctx context.Context) error {
-		db, err := rd.c.ConnE(ctx, rd.t.L(), rd.c.Node(1)[0])
-		if err != nil {
-			return err
-		}
-		defer db.Close()
-		for _, setting := range clusterSettings {
-			if _, err := db.Exec(fmt.Sprintf("SET CLUSTER SETTING %s", setting)); err != nil {
-				return errors.Wrapf(err, "failed to set cluster setting %s", setting)
-			}
-		}
-		opts := "WITH UNSAFE_RESTORE_INCOMPATIBLE_VERSION"
-		if runOnline {
-			opts = "WITH EXPERIMENTAL DEFERRED COPY, UNSAFE_RESTORE_INCOMPATIBLE_VERSION"
-		}
-		if err := maybeAddSomeEmptyTables(ctx, rd); err != nil {
-			return errors.Wrapf(err, "failed to add some empty tables")
-		}
-		restoreStartTime = timeutil.Now()
-		restoreCmd := rd.restoreCmd(ctx, fmt.Sprintf("DATABASE %s", sp.backup.fixture.DatabaseName()), opts)
-		t.L().Printf("Running %s", restoreCmd)
-		if _, err = db.ExecContext(ctx, restoreCmd); err != nil {
-			return err
-		}
-		if runOnline && sp.linkPhaseTimeout != 0 && sp.linkPhaseTimeout < timeutil.Since(restoreStartTime) {
-			return errors.Newf("link phase took too long: %s greater than timeout %s", timeutil.Since(restoreStartTime), sp.linkPhaseTimeout)
+	restoreStartTime, restoreEndTime, err := executeTestRestorePhase(
+		ctx, t, c, sp, rd, runOnline, clusterSettings...,
+	)
+	require.NoError(t, err, "failed to execute restore phase")
+
+	downloadEndTimeLowerBound, workloadStartTime, workloadEndTime, err := executeTestDownloadPhase(
+		ctx, t, c, sp, rd, runOnline, runWorkload, testStartTime,
+	)
+	require.NoError(t, err, "failed to execute download phase")
+
+	return restoreStats{
+		collector:                 statsCollector,
+		restoreStartTime:          restoreStartTime,
+		restoreEndTime:            restoreEndTime,
+		workloadStartTime:         workloadStartTime,
+		workloadEndTime:           workloadEndTime,
+		downloadEndTimeLowerBound: downloadEndTimeLowerBound,
+	}
+}
+
+// executeTestRestorePhase executes the restore phase of the online restore
+// roachtests. If `runOnline` is not set, a conventional restore is run instead.
+// The start time and end time of the online restore link phase are returned (or
+// in the case of conventional restore, the start and end time of the entire
+// restore job).
+func executeTestRestorePhase(
+	ctx context.Context,
+	t test.Test,
+	c cluster.Cluster,
+	sp onlineRestoreSpecs,
+	rd restoreDriver,
+	runOnline bool,
+	clusterSettings ...string,
+) (time.Time, time.Time, error) {
+	db, err := rd.c.ConnE(ctx, t.L(), rd.c.Node(1)[0])
+	if err != nil {
+		return time.Time{}, time.Time{}, err
+	}
+	defer db.Close()
+	for _, setting := range clusterSettings {
+		if _, err := db.Exec(fmt.Sprintf("SET CLUSTER SETTING %s", setting)); err != nil {
+			return time.Time{}, time.Time{}, errors.Wrapf(err, "failed to set cluster setting %s", setting)
 		}
-		return nil
-	})
-	m.Wait()
-	restoreEndTime = timeutil.Now()
+	}
+	opts := "WITH UNSAFE_RESTORE_INCOMPATIBLE_VERSION"
+	if runOnline {
+		opts = "WITH EXPERIMENTAL DEFERRED COPY, UNSAFE_RESTORE_INCOMPATIBLE_VERSION"
+	}
+	if err := maybeAddSomeEmptyTables(ctx, rd); err != nil {
+		return time.Time{}, time.Time{}, errors.Wrapf(err, "failed to add some empty tables")
+	}
+	restoreStartTime := timeutil.Now()
+	restoreCmd := rd.restoreCmd(ctx, fmt.Sprintf("DATABASE %s", sp.backup.fixture.DatabaseName()), opts)
+	t.L().Printf("Running %s", restoreCmd)
+	if _, err = db.ExecContext(ctx, restoreCmd); err != nil {
+		return time.Time{}, time.Time{}, err
+	}
+	restoreEndTime := time.Now()
+	if runOnline && sp.linkPhaseTimeout != 0 && sp.linkPhaseTimeout < restoreEndTime.Sub(restoreStartTime) {
+		return restoreStartTime, restoreEndTime, errors.Newf(
+			"link phase took too long: %s greater than timeout %s",
+			timeutil.Since(restoreStartTime), sp.linkPhaseTimeout,
+		)
+	}
+	return restoreStartTime, restoreEndTime, err
+}
 
-	workloadCtx, workloadCancel := context.WithCancel(ctx)
-	mDownload := c.NewDeprecatedMonitor(workloadCtx, sp.hardware.getCRDBNodes())
+// executeTestDownloadPhase executes the download phase of the online restore
+// roachtest. `runWorkload` indicates whether a workload should be running
+// during the download phase. If `runOnline` is not set, no wait for the
+// download phase is performed, but the workload is still run for 5 minutes (or
+// the remaining time in the test, whichever is shorter).
+// The lower bound of the download job end time is returned, along with the
+// start and end time of the workload, if it was run.
+func executeTestDownloadPhase(
+	ctx context.Context,
+	t test.Test,
+	c cluster.Cluster,
+	sp onlineRestoreSpecs,
+	rd restoreDriver,
+	runOnline bool,
+	runWorkload bool,
+	testStartTime time.Time,
+) (time.Time, time.Time, time.Time, error) {
+	mon := t.NewErrorGroup(task.Logger(t.L()))
 
 	var workloadStartTime, workloadEndTime time.Time
-	mDownload.Go(func(ctx context.Context) error {
+	workloadCancel := mon.GoWithCancel(func(ctx context.Context, logger *logger.Logger) error {
 		if !runWorkload {
-			fmt.Printf("roachtest configured to skip running the foreground workload")
+			logger.Printf("roachtest configured to skip running the foreground workload")
 			return nil
 		}
 		workloadStartTime = timeutil.Now()
-		err := sp.workload.Run(ctx, t, c, sp.hardware)
 		// We expect the workload to return a context cancelled error because
 		// the roachtest driver cancels the monitor's context after the download job completes
-		if err != nil && ctx.Err() == nil {
+		if err := sp.workload.Run(ctx, t, c, sp.hardware); err != nil && ctx.Err() == nil {
 			// Implies the workload context was not cancelled and the workload cmd returned a
 			// different error.
 			return errors.Wrapf(err, `Workload context was not cancelled. Error returned by workload cmd`)
 		}
-		rd.t.L().Printf("workload successfully finished")
+		logger.Printf("workload successfully finished")
 		return nil
 	})
+
 	var downloadEndTimeLowerBound time.Time
-	mDownload.Go(func(ctx context.Context) error {
+	downloadStartTime := timeutil.Now()
+	mon.Go(func(ctx context.Context, logger *logger.Logger) error {
 		defer workloadCancel()
 		if runOnline {
-			downloadEndTimeLowerBound, err = waitForDownloadJob(ctx, c, t.L())
-			if err != nil {
+			var err error
+			if downloadEndTimeLowerBound, err = waitForDownloadJob(ctx, c, logger); err != nil {
 				return err
 			}
-			if sp.downloadPhaseTimeout != 0 && sp.downloadPhaseTimeout < timeutil.Since(restoreEndTime) {
-				return errors.Newf("download phase took too long: %s greater than timeout %s", timeutil.Since(restoreEndTime), sp.downloadPhaseTimeout)
+			downloadTime := downloadEndTimeLowerBound.Sub(downloadStartTime)
+			if sp.downloadPhaseTimeout != 0 && sp.downloadPhaseTimeout < downloadTime {
+				return errors.Newf(
+					"download phase took too long: %s greater than timeout %s",
+					downloadTime, sp.downloadPhaseTimeout,
+				)
 			}
 		}
+
 		if runWorkload {
-			// Run the workload for at most 5 minutes.
-			testRuntime := timeutil.Since(testStartTime)
-			workloadDuration := sp.timeout - (testRuntime + time.Minute)
+			// Remaining workload duration is capped by the test timeout
+			testRunTime := timeutil.Since(testStartTime)
+			testTimeoutRemaining := sp.timeout - (testRunTime + time.Minute)
+
+			// Run the workload for at most 5 more minutes.
 			maxWorkloadDuration := time.Minute * 5
-			if workloadDuration > maxWorkloadDuration {
-				workloadDuration = maxWorkloadDuration
-			}
-			t.L().Printf("let workload run for another %.2f minutes", workloadDuration.Minutes())
+
+			workloadDuration := min(testTimeoutRemaining, maxWorkloadDuration)
+			logger.Printf("let workload run for another %.2f minutes", workloadDuration.Minutes())
 			time.Sleep(workloadDuration)
 		}
 		return nil
 	})
-	mDownload.Wait()
-	if runWorkload {
-		workloadEndTime = timeutil.Now()
-	}
-	return restoreStats{
-		collector:                 statsCollector,
-		restoreStartTime:          restoreStartTime,
-		restoreEndTime:            restoreEndTime,
-		workloadStartTime:         workloadStartTime,
-		workloadEndTime:           workloadEndTime,
-		downloadEndTimeLowerBound: downloadEndTimeLowerBound,
+
+	if err := mon.WaitE(); err != nil {
+		return time.Time{}, time.Time{}, time.Time{}, err
 	}
+	return downloadEndTimeLowerBound, workloadStartTime, workloadEndTime, nil
 }