roachtest: add helper to retry cluster settings

We occasionally see cluster setting statements time out
due to overloaded clusters (e.g. mutators in mixed version
tests, post test assertions).

This change adds a small helper to retry those cluster setting
statements.

Fixes: #143791

Author: DarrylWong

pkg/cmd/roachtest/roachtestutil/BUILD.bazel

@@ -37,6 +37,7 @@ go_library(
         "//pkg/util/httputil",
         "//pkg/util/humanizeutil",
         "//pkg/util/protoutil",
+        "//pkg/util/retry",
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "//pkg/workload/histogram/exporter",

pkg/cmd/roachtest/roachtestutil/mixedversion/helper.go

@@ -15,12 +15,14 @@ import (
 	"sync/atomic"
 
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil/clusterupgrade"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil/task"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
 	"github.com/cockroachdb/cockroach/pkg/testutils/release"
+	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/errors"
 )
 
@@ -152,6 +154,22 @@ func (s *Service) ExecWithGateway(
 	return err
 }
 
+func (s *Service) ExecWithRetry(
+	rng *rand.Rand,
+	nodes option.NodeListOption,
+	retryOpts retry.Options,
+	query string,
+	args ...interface{},
+) error {
+	db, err := s.prepareQuery(rng, nodes, query, args...)
+	if err != nil {
+		return err
+	}
+
+	_, err = roachtestutil.ExecWithRetry(s.ctx, s.stepLogger, db, retryOpts, query, args...)
+	return err
+}
+
 func (s *Service) ClusterVersion(rng *rand.Rand) (roachpb.Version, error) {
 	if s.Finalizing {
 		n, db := s.RandomDB(rng)
@@ -249,6 +267,14 @@ func (h *Helper) ExecWithGateway(
 	return h.DefaultService().ExecWithGateway(rng, nodes, query, args...)
 }
 
+// ExecWithRetry is like ExecWithGateway, but retries the execution of
+// the statement on errors, using the retry options provided.
+func (h *Helper) ExecWithRetry(
+	rng *rand.Rand, nodes option.NodeListOption, query string, args ...interface{},
+) error {
+	return h.DefaultService().ExecWithGateway(rng, nodes, query, args...)
+}
+
 // defaultTaskOptions returns the default options that are passed to all tasks
 // started by the helper.
 func (h *Helper) defaultTaskOptions() []task.Option {

pkg/cmd/roachtest/roachtestutil/mixedversion/steps.go

@@ -311,7 +311,9 @@ func (s preserveDowngradeOptionStep) Run(
 		return err
 	}
 
-	return service.Exec(rng, "SET CLUSTER SETTING cluster.preserve_downgrade_option = $1", bv.String())
+	return service.ExecWithRetry(rng, service.Descriptor.Nodes, roachtestutil.ClusterSettingRetryOpts,
+		"SET CLUSTER SETTING cluster.preserve_downgrade_option = $1", bv.String(),
+	)
 }
 
 func (s preserveDowngradeOptionStep) ConcurrencyDisabled() bool {
@@ -408,8 +410,10 @@ func (s allowUpgradeStep) Description() string {
 func (s allowUpgradeStep) Run(
 	ctx context.Context, l *logger.Logger, rng *rand.Rand, h *Helper,
 ) error {
-	return serviceByName(h, s.virtualClusterName).Exec(
-		rng, "RESET CLUSTER SETTING cluster.preserve_downgrade_option",
+	service := serviceByName(h, s.virtualClusterName)
+	return service.ExecWithRetry(
+		rng, service.Descriptor.Nodes, roachtestutil.ClusterSettingRetryOpts,
+		"RESET CLUSTER SETTING cluster.preserve_downgrade_option",
 	)
 }
 
@@ -515,8 +519,8 @@ func (s setClusterSettingStep) Run(
 		args = []interface{}{val}
 	}
 
-	return serviceByName(h, serviceName).ExecWithGateway(
-		rng, nodesRunningAtLeast(s.virtualClusterName, s.minVersion, h), stmt, args...,
+	return serviceByName(h, serviceName).ExecWithRetry(
+		rng, nodesRunningAtLeast(s.virtualClusterName, s.minVersion, h), roachtestutil.ClusterSettingRetryOpts, stmt, args...,
 	)
 }
 
@@ -563,7 +567,9 @@ func (s setClusterVersionStep) Run(
 	}
 
 	l.Printf("setting cluster version to '%s'", binaryVersion)
-	return service.Exec(rng, "SET CLUSTER SETTING version = $1", binaryVersion)
+	return service.ExecWithRetry(rng, service.Descriptor.Nodes, roachtestutil.ClusterSettingRetryOpts,
+		"SET CLUSTER SETTING version = $1", binaryVersion,
+	)
 }
 
 func (s setClusterVersionStep) ConcurrencyDisabled() bool {
@@ -587,8 +593,8 @@ func (s resetClusterSettingStep) Run(
 	ctx context.Context, l *logger.Logger, rng *rand.Rand, h *Helper,
 ) error {
 	stmt := fmt.Sprintf("RESET CLUSTER SETTING %s", s.name)
-	return serviceByName(h, s.virtualClusterName).ExecWithGateway(
-		rng, nodesRunningAtLeast(s.virtualClusterName, s.minVersion, h), stmt,
+	return serviceByName(h, s.virtualClusterName).ExecWithRetry(
+		rng, nodesRunningAtLeast(s.virtualClusterName, s.minVersion, h), roachtestutil.ClusterSettingRetryOpts, stmt,
 	)
 }
 

pkg/cmd/roachtest/roachtestutil/utils.go

@@ -27,6 +27,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/logger"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
@@ -336,3 +337,38 @@ func SimulateMultiRegionCluster(
 
 	return cleanupFunc, nil
 }
+
+// ExecWithRetry executes the given SQL statement with the specified retry logic.
+func ExecWithRetry(
+	ctx context.Context,
+	l *logger.Logger,
+	db *gosql.DB,
+	retryOpts retry.Options,
+	query string,
+	args ...any,
+) (gosql.Result, error) {
+	var result gosql.Result
+	err := retryOpts.Do(ctx, func(ctx context.Context) error {
+		var err error
+		result, err = db.ExecContext(ctx, query, args...)
+		if err != nil {
+			l.Printf("%s failed (retrying): %v", query, err)
+			return err
+		}
+		return nil
+	})
+
+	return result, err
+}
+
+// ClusterSettingRetryOpts are retry options intended for cluster setting operations.
+//
+// We use relatively high backoff parameters with the assumption that:
+//  1. If we fail, it's likely due to cluster overload, and we want to give
+//     the cluster adequate time to recover.
+//  2. Setting a cluster setting in roachtest is not latency sensitive.
+var ClusterSettingRetryOpts = retry.Options{
+	InitialBackoff: 3 * time.Second,
+	MaxBackoff:     5 * time.Second,
+	MaxRetries:     5,
+}

pkg/cmd/roachtest/roachtestutil/validation_check.go

@@ -42,7 +42,7 @@ func CheckReplicaDivergenceOnDB(ctx context.Context, l *logger.Logger, db *gosql
 	defer cancel()
 
 	// Speed up consistency checks. The test is done, so let's go full throttle.
-	_, err := db.ExecContext(ctx, "SET CLUSTER SETTING server.consistency_check.max_rate = '1GB'")
+	_, err := ExecWithRetry(ctx, l, db, ClusterSettingRetryOpts, "SET CLUSTER SETTING server.consistency_check.max_rate = '1GB'")
 	if err != nil {
 		return errors.Wrap(err, "unable to set 'server.consistency_check.max_rate'")
 	}