roachtest: retry fetching cluster and binary version

Author: DarrylWong

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
 	}
 
@@ -477,7 +502,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
 	}
@@ -490,7 +515,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

@@ -174,7 +174,7 @@ 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)
 		}
@@ -270,9 +270,13 @@ func (h *Helper) ExecWithGateway(
 // 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{},
+	rng *rand.Rand,
+	nodes option.NodeListOption,
+	retryOpts retry.Options,
+	query string,
+	args ...interface{},
 ) error {
-	return h.DefaultService().ExecWithGateway(rng, nodes, query, args...)
+	return h.DefaultService().ExecWithRetry(rng, nodes, retryOpts, query, args...)
 }
 
 // defaultTaskOptions returns the default options that are passed to all tasks

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 {
@@ -527,7 +527,7 @@ func (tr *testRunner) refreshBinaryVersions(ctx context.Context, service *servic
 	group := ctxgroup.WithContext(connectionCtx)
 	for j, node := range tr.getAvailableNodes(service.descriptor) {
 		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 {
 				return fmt.Errorf(
 					"failed to get binary version for node %d (%s): %w",
@@ -558,7 +558,7 @@ func (tr *testRunner) refreshClusterVersions(ctx context.Context, service *servi
 	group := ctxgroup.WithContext(connectionCtx)
 	for j, node := range tr.getAvailableNodes(service.descriptor) {
 		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 {
 				return fmt.Errorf(
 					"failed to get cluster version for node %d (%s): %w",

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

@@ -306,7 +306,7 @@ 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
 	}
@@ -558,7 +558,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)
 		}

pkg/cmd/roachtest/tests/mixed_version_backup.go

@@ -216,8 +216,10 @@ func sanitizeVersionForBackup(v *clusterupgrade.Version) string {
 // hasInternalSystemJobs returns true if the cluster is expected to
 // have the `crdb_internal.system_jobs` vtable. If so, it should be
 // used instead of `system.jobs` when querying job status.
-func hasInternalSystemJobs(ctx context.Context, rng *rand.Rand, db *gosql.DB) (bool, error) {
-	cv, err := clusterupgrade.ClusterVersion(ctx, db)
+func hasInternalSystemJobs(
+	ctx context.Context, l *logger.Logger, rng *rand.Rand, db *gosql.DB,
+) (bool, error) {
+	cv, err := clusterupgrade.ClusterVersion(ctx, l, db)
 	if err != nil {
 		return false, fmt.Errorf("failed to query cluster version: %w", err)
 	}
@@ -2211,7 +2213,7 @@ func (mvb *mixedVersionBackup) createBackupCollection(
 	backupNamePrefix := mvb.backupNamePrefix(h, label)
 	n, db := h.System.RandomDB(rng)
 	l.Printf("checking existence of crdb_internal.system_jobs via node %d", n)
-	internalSystemJobs, err := hasInternalSystemJobs(ctx, rng, db)
+	internalSystemJobs, err := hasInternalSystemJobs(ctx, l, rng, db)
 	if err != nil {
 		return err
 	}
@@ -2742,7 +2744,7 @@ func (mvb *mixedVersionBackup) verifySomeBackups(
 
 	n, db := h.System.RandomDB(rng)
 	l.Printf("checking existence of crdb_internal.system_jobs via node %d", n)
-	internalSystemJobs, err := hasInternalSystemJobs(ctx, rng, db)
+	internalSystemJobs, err := hasInternalSystemJobs(ctx, l, rng, db)
 	if err != nil {
 		return err
 	}
@@ -2823,7 +2825,7 @@ func (mvb *mixedVersionBackup) verifyAllBackups(
 
 			n, db := h.System.RandomDB(rng)
 			l.Printf("checking existence of crdb_internal.system_jobs via node %d", n)
-			internalSystemJobs, err := hasInternalSystemJobs(ctx, rng, db)
+			internalSystemJobs, err := hasInternalSystemJobs(ctx, l, rng, db)
 			if err != nil {
 				restoreErrors = append(restoreErrors, fmt.Errorf("error checking for internal system jobs: %w", err))
 				return

pkg/cmd/roachtest/tests/multitenant_upgrade.go

@@ -284,7 +284,7 @@ func runMultitenantUpgrade(ctx context.Context, t test.Test, c cluster.Cluster)
 					}
 					defer tenantDB.Close()
 
-					binaryVersion, err := clusterupgrade.BinaryVersion(ctx, tenantDB)
+					binaryVersion, err := clusterupgrade.BinaryVersion(ctx, l, tenantDB)
 					if err != nil {
 						return errors.Wrapf(err, "failed to query binary version for %s on n%d", tenant.name, n)
 					}

pkg/cmd/roachtest/tests/versionupgrade.go

@@ -263,7 +263,7 @@ func makeVersionFixtureAndFatal(
 	// cluster version might be 2.0, so we can only use the 2.0 or
 	// 2.1 binary, but not the 19.1 binary (as 19.1 and 2.0 are not
 	// compatible).
-	binaryVersion, err := clusterupgrade.BinaryVersion(ctx, dbFunc(1))
+	binaryVersion, err := clusterupgrade.BinaryVersion(ctx, t.L(), dbFunc(1))
 	if err != nil {
 		t.Fatalf("fetching binary version on n1: %v", err)
 	}