Merge #154180 #156795 #156975

154180: roachtest: add helper to retry cluster settings r=herkolategan,srosenberg a=DarrylWong

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.

Informs: #143791
Informs: #145092

156795: roachtest: fix allow_unsafe_internals edge cases r=herkolategan,srosenberg a=DarrylWong

In #156581, we started setting `allow_unsafe_internals` unconditionally for roachtests. However, there are a few edge cases where we can't set this. Namely when using `psql` to connect to a postgres database where the CRDB only parameter does not exist.

Fixes: #156702
Fixes: #156703
Fixes: #156706

Release note: none

156975: crosscluster/logical: test seq support r=jeffswenson a=msbutler

Fixes #155521

Release note (ops change): ensure the user can run CREATE LOGICAL REPLICATION STREAM where the source table has a column with a sequence expression.

Co-authored-by: DarrylWong <[email protected]>
Co-authored-by: Michael Butler <[email protected]>

Author: 3 people

pkg/cmd/roachtest/cluster.go

@@ -2721,9 +2721,6 @@ func (c *clusterImpl) InternalPGUrl(
 	return c.pgURLErr(ctx, l, nodes, opts)
 }
 
-// Silence unused warning.
-var _ = (&clusterImpl{}).InternalPGUrl
-
 // ExternalPGUrl returns the external Postgres endpoint for the specified nodes.
 func (c *clusterImpl) ExternalPGUrl(
 	ctx context.Context, l *logger.Logger, nodes option.NodeListOption, opts roachprod.PGURLOptions,

pkg/cmd/roachtest/roachtestutil/BUILD.bazel

@@ -39,6 +39,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/clusterupgrade/BUILD.bazel

@@ -9,6 +9,7 @@ go_library(
         "//pkg/build",
         "//pkg/cmd/roachtest/cluster",
         "//pkg/cmd/roachtest/option",
+        "//pkg/cmd/roachtest/roachtestutil",
         "//pkg/cmd/roachtest/test",
         "//pkg/roachpb",
         "//pkg/roachprod/install",

pkg/cmd/roachtest/roachtestutil/clusterupgrade/clusterupgrade.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/build"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestutil"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
@@ -156,10 +157,22 @@ func LatestPatchRelease(series string) (*Version, error) {
 // associated with the given database connection.
 // NB: version means major.minor[-internal]; the patch level isn't
 // returned. For example, a binary of version 19.2.4 will return 19.2.
-func BinaryVersion(ctx context.Context, db *gosql.DB) (roachpb.Version, error) {
+func BinaryVersion(ctx context.Context, l *logger.Logger, db *gosql.DB) (roachpb.Version, error) {
 	zero := roachpb.Version{}
 	var sv string
-	if err := db.QueryRowContext(ctx, `SELECT crdb_internal.node_executable_version();`).Scan(&sv); err != nil {
+	rows, err := roachtestutil.QueryWithRetry(
+		ctx, l, db, roachtestutil.ClusterSettingRetryOpts, `SELECT crdb_internal.node_executable_version();`,
+	)
+	if err != nil {
+		return zero, err
+	}
+	defer rows.Close()
+
+	if !rows.Next() {
+		return zero, fmt.Errorf("no rows returned")
+	}
+
+	if err := rows.Scan(&sv); err != nil {
 		return zero, err
 	}
 
@@ -176,10 +189,22 @@ func BinaryVersion(ctx context.Context, db *gosql.DB) (roachpb.Version, error) {
 // in the background plus gossip asynchronicity.
 // NB: cluster versions are always major.minor[-internal]; there isn't
 // a patch level.
-func ClusterVersion(ctx context.Context, db *gosql.DB) (roachpb.Version, error) {
+func ClusterVersion(ctx context.Context, l *logger.Logger, db *gosql.DB) (roachpb.Version, error) {
 	zero := roachpb.Version{}
 	var sv string
-	if err := db.QueryRowContext(ctx, `SHOW CLUSTER SETTING version`).Scan(&sv); err != nil {
+	rows, err := roachtestutil.QueryWithRetry(
+		ctx, l, db, roachtestutil.ClusterSettingRetryOpts, `SHOW CLUSTER SETTING version`,
+	)
+	if err != nil {
+		return zero, err
+	}
+	defer rows.Close()
+
+	if !rows.Next() {
+		return zero, fmt.Errorf("no rows returned")
+	}
+
+	if err := rows.Scan(&sv); err != nil {
 		return zero, err
 	}
 
@@ -483,7 +508,7 @@ func WaitForClusterUpgrade(
 	timeout time.Duration,
 ) error {
 	firstNode := nodes[0]
-	newVersion, err := BinaryVersion(ctx, dbFunc(firstNode))
+	newVersion, err := BinaryVersion(ctx, l, dbFunc(firstNode))
 	if err != nil {
 		return err
 	}
@@ -496,7 +521,7 @@ func WaitForClusterUpgrade(
 		retryCtx, cancel := context.WithTimeout(ctx, timeout)
 		defer cancel()
 		err := opts.Do(retryCtx, func(ctx context.Context) error {
-			currentVersion, err := ClusterVersion(ctx, dbFunc(node))
+			currentVersion, err := ClusterVersion(ctx, l, dbFunc(node))
 			if err != nil {
 				return err
 			}

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,11 +154,27 @@ 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)
 		s.stepLogger.Printf("querying cluster version through node %d", n)
-		cv, err := clusterupgrade.ClusterVersion(s.ctx, db)
+		cv, err := clusterupgrade.ClusterVersion(s.ctx, s.stepLogger, db)
 		if err != nil {
 			return roachpb.Version{}, fmt.Errorf("failed to query cluster version: %w", err)
 		}
@@ -249,6 +267,18 @@ 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,
+	retryOpts retry.Options,
+	query string,
+	args ...interface{},
+) error {
+	return h.DefaultService().ExecWithRetry(rng, nodes, retryOpts, 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/runner.go

@@ -79,7 +79,7 @@ var (
 	// for an internal query (i.e., performed by the framework) to
 	// complete. These queries are typically associated with gathering
 	// upgrade state data to be displayed during execution.
-	internalQueryTimeout = 30 * time.Second
+	internalQueryTimeout = 90 * time.Second
 )
 
 func newServiceRuntime(desc *ServiceDescriptor) *serviceRuntime {
@@ -534,7 +534,7 @@ func (tr *testRunner) refreshBinaryVersions(ctx context.Context, service *servic
 	availableNodes := tr.getAvailableNodes(service.descriptor)
 	for j, node := range service.descriptor.Nodes {
 		group.GoCtx(func(ctx context.Context) error {
-			bv, err := clusterupgrade.BinaryVersion(ctx, tr.conn(node, service.descriptor.Name))
+			bv, err := clusterupgrade.BinaryVersion(ctx, tr.logger, tr.conn(node, service.descriptor.Name))
 			if err != nil {
 				if !availableNodes.Contains(node) {
 					return nil
@@ -575,7 +575,7 @@ func (tr *testRunner) refreshClusterVersions(ctx context.Context, service *servi
 	availableNodes := tr.getAvailableNodes(service.descriptor)
 	for j, node := range service.descriptor.Nodes {
 		group.GoCtx(func(ctx context.Context) error {
-			cv, err := clusterupgrade.ClusterVersion(ctx, tr.conn(node, service.descriptor.Name))
+			cv, err := clusterupgrade.ClusterVersion(ctx, tr.logger, tr.conn(node, service.descriptor.Name))
 			if err != nil {
 				if !availableNodes.Contains(node) {
 					return nil

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

@@ -306,12 +306,14 @@ func (s preserveDowngradeOptionStep) Run(
 	service := serviceByName(h, s.virtualClusterName)
 	node, db := service.RandomDB(rng)
 	l.Printf("checking binary version (via node %d)", node)
-	bv, err := clusterupgrade.BinaryVersion(ctx, db)
+	bv, err := clusterupgrade.BinaryVersion(ctx, l, db)
 	if err != nil {
 		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(debug bool) 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",
 	)
 }
 
@@ -518,8 +522,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...,
 	)
 }
 
@@ -557,7 +561,7 @@ func (s setClusterVersionStep) Run(
 		node, db := service.RandomDB(rng)
 		l.Printf("fetching binary version via n%d", node)
 
-		bv, err := clusterupgrade.BinaryVersion(ctx, db)
+		bv, err := clusterupgrade.BinaryVersion(ctx, l, db)
 		if err != nil {
 			return errors.Wrapf(err, "getting binary version on n%d", node)
 		}
@@ -566,7 +570,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 {
@@ -590,8 +596,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,61 @@ 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
+}
+
+// QueryWithRetry queries the given SQL statement with the specified retry logic.
+func QueryWithRetry(
+	ctx context.Context,
+	l *logger.Logger,
+	db *gosql.DB,
+	retryOpts retry.Options,
+	query string,
+	args ...any,
+) (*gosql.Rows, error) {
+	var rows *gosql.Rows
+	err := retryOpts.Do(ctx, func(ctx context.Context) error {
+		var err error
+		rows, err = db.QueryContext(ctx, query, args...)
+		if err != nil {
+			l.Printf("%s failed (retrying): %v", query, err)
+			return err
+		}
+		return nil
+	})
+
+	return rows, 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

@@ -46,7 +46,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'")
 	}
@@ -107,14 +107,18 @@ FROM crdb_internal.check_consistency(false, '', '') as t;`)
 
 // CheckInvalidDescriptors returns an error if there exists any descriptors in
 // the crdb_internal.invalid_objects virtual table.
-func CheckInvalidDescriptors(ctx context.Context, db *gosql.DB) error {
+func CheckInvalidDescriptors(ctx context.Context, l *logger.Logger, db *gosql.DB) error {
 	var invalidIDs string
+	l.Printf("checking for invalid descriptors")
 	if err := timeutil.RunWithTimeout(ctx, "descriptor validation", time.Minute, func(ctx context.Context) error {
 		// Because crdb_internal.invalid_objects is a virtual table, by default, the
 		// query will take a lease on the database sqlDB is connected to and only run
 		// the query on the given database. The "" prefix prevents this lease
 		// acquisition and allows the query to fetch all descriptors in the cluster.
-		rows, err := db.QueryContext(ctx, `SELECT id, obj_name, error FROM "".crdb_internal.invalid_objects`)
+		//
+		// This often times out when the cluster is overloaded. Since this is just
+		// for validation, we retry with a relatively high backoff.
+		rows, err := QueryWithRetry(ctx, l, db, ClusterSettingRetryOpts, `SELECT id, obj_name, error FROM "".crdb_internal.invalid_objects`)
 		if err != nil {
 			return err
 		}

pkg/cmd/roachtest/test_runner.go

@@ -1763,7 +1763,7 @@ func (r *testRunner) postTestAssertions(
 				// NB: the invalid description checks should run at the system tenant level.
 				db := c.Conn(ctx, t.L(), validationNode, option.VirtualClusterName(install.SystemInterfaceName))
 				defer db.Close()
-				if err := roachtestutil.CheckInvalidDescriptors(ctx, db); err != nil {
+				if err := roachtestutil.CheckInvalidDescriptors(ctx, t.L(), db); err != nil {
 					postAssertionErr(errors.WithDetail(err, "invalid descriptors check failed"))
 				}
 			}()