From e633dd8724a6911019b10d95c1914a6602fb2006 Mon Sep 17 00:00:00 2001
From: Steven Danna <danna@cockroachlabs.com>
Date: Mon, 4 Mar 2024 11:33:44 +0000
Subject: [PATCH] roachtest: restore: restart nodes with same options

Fixes #119802
Fixes #119803

Release note: None
---
 pkg/cmd/roachtest/tests/jobs_util.go | 19 ++++++++-------
 pkg/cmd/roachtest/tests/restore.go   | 35 +++++++++++++++++++++-------
 2 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/pkg/cmd/roachtest/tests/jobs_util.go b/pkg/cmd/roachtest/tests/jobs_util.go
index 6aeab961b26f..5c9bab1b4949 100644
--- a/pkg/cmd/roachtest/tests/jobs_util.go
+++ b/pkg/cmd/roachtest/tests/jobs_util.go
@@ -50,14 +50,7 @@ type jobStarter func(c cluster.Cluster, l *logger.Logger) (jobspb.JobID, error)
 func jobSurvivesNodeShutdown(
 	ctx context.Context, t test.Test, c cluster.Cluster, nodeToShutdown int, startJob jobStarter,
 ) {
-	cfg := nodeShutdownConfig{
-		shutdownNode:         nodeToShutdown,
-		watcherNode:          1 + (nodeToShutdown)%c.Spec().NodeCount,
-		crdbNodes:            c.All(),
-		waitFor3XReplication: true,
-		sleepBeforeShutdown:  30 * time.Second,
-	}
-	require.NoError(t, executeNodeShutdown(ctx, t, c, cfg, startJob))
+	require.NoError(t, executeNodeShutdown(ctx, t, c, defaultNodeShutdownConfig(c, nodeToShutdown), startJob))
 }
 
 type nodeShutdownConfig struct {
@@ -70,6 +63,16 @@ type nodeShutdownConfig struct {
 	rng                  *rand.Rand
 }
 
+func defaultNodeShutdownConfig(c cluster.Cluster, nodeToShutdown int) nodeShutdownConfig {
+	return nodeShutdownConfig{
+		shutdownNode:         nodeToShutdown,
+		watcherNode:          1 + (nodeToShutdown)%c.Spec().NodeCount,
+		crdbNodes:            c.All(),
+		waitFor3XReplication: true,
+		sleepBeforeShutdown:  30 * time.Second,
+	}
+}
+
 // executeNodeShutdown executes a node shutdown and returns all errors back to the caller.
 //
 // TODO(msbutler): ideally, t.L() is only passed to this function instead of t,
diff --git a/pkg/cmd/roachtest/tests/restore.go b/pkg/cmd/roachtest/tests/restore.go
index c8331cf6fbd8..28690102bcaa 100644
--- a/pkg/cmd/roachtest/tests/restore.go
+++ b/pkg/cmd/roachtest/tests/restore.go
@@ -81,8 +81,11 @@ func registerRestoreNodeShutdown(r registry.Registry) {
 
 			rd := makeRestoreDriver(t, c, sp)
 			rd.prepareCluster(ctx)
-			jobSurvivesNodeShutdown(ctx, t, c, nodeToShutdown, makeRestoreStarter(ctx, t, c,
-				gatewayNode, rd))
+			cfg := defaultNodeShutdownConfig(c, nodeToShutdown)
+			cfg.restartSettings = rd.defaultClusterSettings()
+			require.NoError(t,
+				executeNodeShutdown(ctx, t, c, cfg,
+					makeRestoreStarter(ctx, t, c, gatewayNode, rd)))
 			rd.checkFingerprint(ctx)
 		},
 	})
@@ -96,15 +99,16 @@ func registerRestoreNodeShutdown(r registry.Registry) {
 		Leases:           registry.MetamorphicLeases,
 		Timeout:          sp.timeout,
 		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
-
 			gatewayNode := 2
 			nodeToShutdown := 2
 
 			rd := makeRestoreDriver(t, c, sp)
 			rd.prepareCluster(ctx)
-
-			jobSurvivesNodeShutdown(ctx, t, c, nodeToShutdown, makeRestoreStarter(ctx, t, c,
-				gatewayNode, rd))
+			cfg := defaultNodeShutdownConfig(c, nodeToShutdown)
+			cfg.restartSettings = rd.defaultClusterSettings()
+			require.NoError(t,
+				executeNodeShutdown(ctx, t, c, cfg,
+					makeRestoreStarter(ctx, t, c, gatewayNode, rd)))
 			rd.checkFingerprint(ctx)
 		},
 	})
@@ -881,10 +885,23 @@ func makeRestoreDriver(t test.Test, c cluster.Cluster, sp restoreSpecs) restoreD
 	}
 }
 
-func (rd *restoreDriver) prepareCluster(ctx context.Context) {
+func (rd *restoreDriver) defaultClusterSettings() []install.ClusterSettingOption {
+	return []install.ClusterSettingOption{
+		install.SecureOption(false),
+	}
+}
+
+func (rd *restoreDriver) roachprodOpts() option.StartOpts {
 	opts := option.DefaultStartOptsNoBackups()
-	opts.RoachprodOpts.ExtraArgs = rd.sp.extraArgs
-	rd.c.Start(ctx, rd.t.L(), opts, install.MakeClusterSettings(install.SecureOption(false)), rd.sp.hardware.getCRDBNodes())
+	opts.RoachprodOpts.ExtraArgs = append(opts.RoachprodOpts.ExtraArgs, rd.sp.extraArgs...)
+	return opts
+}
+
+func (rd *restoreDriver) prepareCluster(ctx context.Context) {
+	rd.c.Start(ctx, rd.t.L(),
+		rd.roachprodOpts(),
+		install.MakeClusterSettings(rd.defaultClusterSettings()...),
+		rd.sp.hardware.getCRDBNodes())
 	rd.getAOST(ctx)
 }