server: use channel for DisableAutomaticVersionUpgrade

DisableAutomaticVersionUpgrade is an atomic integer which is rechecked
in a retry loop. This is not a very clean mechanism, and can lead to
issues where you're unknowingly dealing with a copy of the knobs and
setting the wrong atomic. The retry loop can also add unnecessary
delays in tests.

This commit changes DisableAutomaticVersionUpgrade from an atomic
integer to a channel. If the channel is set, auto-upgrade waits until
the channel is closed.

Release note: None

pkg/ccl/backupccl/insert_missing_public_schema_namespace_entry_restore_test.go

@@ -37,7 +37,7 @@ func TestInsertMissingPublicSchemaNamespaceEntry(t *testing.T) {
 			ExternalIODir: dir,
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride:          clusterversion.ByKey(clusterversion.InsertPublicSchemaNamespaceEntryOnRestore - 1),
 				},
 			},

pkg/ccl/backupccl/restore_old_versions_test.go

@@ -865,7 +865,7 @@ func restorePublicSchemaMixedVersion(exportDir string) func(t *testing.T) {
 					Knobs: base.TestingKnobs{
 						JobsTestingKnobs: jobs.NewTestingKnobsWithShortIntervals(),
 						Server: &server.TestingKnobs{
-							DisableAutomaticVersionUpgrade: 1,
+							DisableAutomaticVersionUpgrade: make(chan struct{}),
 							BinaryVersionOverride:          clusterversion.ByKey(clusterversion.PublicSchemasWithDescriptors - 1),
 						},
 					},
@@ -920,7 +920,7 @@ func restoreSyntheticPublicSchemaNamespaceEntry(exportDir string) func(t *testin
 					Knobs: base.TestingKnobs{
 						JobsTestingKnobs: jobs.NewTestingKnobsWithShortIntervals(),
 						Server: &server.TestingKnobs{
-							DisableAutomaticVersionUpgrade: 1,
+							DisableAutomaticVersionUpgrade: make(chan struct{}),
 							BinaryVersionOverride:          clusterversion.ByKey(clusterversion.PublicSchemasWithDescriptors - 1),
 						},
 					},
@@ -951,7 +951,7 @@ func restoreSyntheticPublicSchemaNamespaceEntryCleanupOnFail(exportDir string) f
 					Knobs: base.TestingKnobs{
 						JobsTestingKnobs: jobs.NewTestingKnobsWithShortIntervals(),
 						Server: &server.TestingKnobs{
-							DisableAutomaticVersionUpgrade: 1,
+							DisableAutomaticVersionUpgrade: make(chan struct{}),
 							BinaryVersionOverride:          clusterversion.ByKey(clusterversion.PublicSchemasWithDescriptors - 1),
 						},
 					},

pkg/ccl/importccl/import_stmt_test.go

@@ -7381,7 +7381,7 @@ func TestImportMixedVersion(t *testing.T) {
 			Knobs: base.TestingKnobs{
 				JobsTestingKnobs: jobs.NewTestingKnobsWithShortIntervals(),
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride:          clusterversion.ByKey(clusterversion.PublicSchemasWithDescriptors - 1),
 				},
 			},

pkg/ccl/kvccl/kvtenantccl/tenant_upgrade_test.go

@@ -58,7 +58,7 @@ func TestTenantUpgrade(t *testing.T) {
 			Settings: settings,
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride:          clusterversion.TestingBinaryMinSupportedVersion,
 				},
 			},
@@ -210,7 +210,7 @@ func TestTenantUpgradeFailure(t *testing.T) {
 			Settings: settings,
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride:          clusterversion.TestingBinaryMinSupportedVersion,
 				},
 			},

pkg/ccl/migrationccl/migrationsccl/seed_tenant_span_configs_external_test.go

@@ -38,7 +38,7 @@ func TestPreSeedSpanConfigsWrittenWhenActive(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.PreSeedTenantSpanConfigs,
 					),
@@ -90,7 +90,7 @@ func TestSeedTenantSpanConfigs(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.PreSeedTenantSpanConfigs - 1,
 					),
@@ -159,7 +159,7 @@ func TestSeedTenantSpanConfigsWithExistingEntry(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.PreSeedTenantSpanConfigs,
 					),

pkg/jobs/jobs_test.go

@@ -2471,7 +2471,7 @@ func TestStartableJobMixedVersion(t *testing.T) {
 		Knobs: base.TestingKnobs{
 			Server: &server.TestingKnobs{
 				BinaryVersionOverride:          clusterversion.TestingBinaryMinSupportedVersion,
-				DisableAutomaticVersionUpgrade: 1,
+				DisableAutomaticVersionUpgrade: make(chan struct{}),
 			},
 		},
 	})

pkg/kv/kvserver/client_migration_test.go

@@ -251,7 +251,7 @@ func TestMigrateWaitsForApplication(t *testing.T) {
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
 					BinaryVersionOverride:          startV,
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 				},
 				Store: &kvserver.StoreTestingKnobs{
 					TestingApplyFilter: func(args kvserverbase.ApplyFilterArgs) (int, *roachpb.Error) {

pkg/kv/kvserver/client_replica_test.go

@@ -4118,7 +4118,7 @@ func TestRangeMigration(t *testing.T) {
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
 					BinaryVersionOverride:          startV,
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 				},
 			},
 		},

pkg/migration/migrationmanager/manager_external_test.go

@@ -68,7 +68,7 @@ func TestAlreadyRunningJobsAreHandledProperly(t *testing.T) {
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
 					BinaryVersionOverride:          startCV.Version,
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 				},
 				MigrationManager: &migration.TestingKnobs{
 					ListBetweenOverride: func(from, to clusterversion.ClusterVersion) []clusterversion.ClusterVersion {
@@ -201,7 +201,7 @@ func TestMigrateUpdatesReplicaVersion(t *testing.T) {
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
 					BinaryVersionOverride:          startCV.Version,
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 				},
 				MigrationManager: &migration.TestingKnobs{
 					ListBetweenOverride: func(from, to clusterversion.ClusterVersion) []clusterversion.ClusterVersion {
@@ -319,7 +319,7 @@ func TestConcurrentMigrationAttempts(t *testing.T) {
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
 					BinaryVersionOverride:          versions[0].Version,
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 				},
 				MigrationManager: &migration.TestingKnobs{
 					ListBetweenOverride: func(from, to clusterversion.ClusterVersion) []clusterversion.ClusterVersion {
@@ -401,7 +401,7 @@ func TestPauseMigration(t *testing.T) {
 				JobsTestingKnobs: jobs.NewTestingKnobsWithShortIntervals(),
 				Server: &server.TestingKnobs{
 					BinaryVersionOverride:          startCV.Version,
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 				},
 				MigrationManager: &migration.TestingKnobs{
 					ListBetweenOverride: func(from, to clusterversion.ClusterVersion) []clusterversion.ClusterVersion {
@@ -521,7 +521,7 @@ func TestPrecondition(t *testing.T) {
 	}
 	knobs := base.TestingKnobs{
 		Server: &server.TestingKnobs{
-			DisableAutomaticVersionUpgrade: 1,
+			DisableAutomaticVersionUpgrade: make(chan struct{}),
 			BinaryVersionOverride:          v0.Version,
 		},
 		// Inject a migration which would run to upgrade the cluster.

pkg/migration/migrations/alter_statement_diagnostics_requests_test.go

@@ -39,7 +39,7 @@ func TestAlterSystemStmtDiagReqs(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride:          clusterversion.ByKey(clusterversion.AlterSystemStmtDiagReqs - 1),
 				},
 			},

pkg/migration/migrations/alter_table_protected_timestamp_records_test.go

@@ -38,7 +38,7 @@ func TestAlterSystemProtectedTimestampRecordsTable(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.AlterSystemProtectedTimestampAddColumn - 1),
 				},

pkg/migration/migrations/alter_table_statistics_avg_size_test.go

@@ -38,7 +38,7 @@ func TestAlterSystemTableStatisticsTable(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.AlterSystemTableStatisticsAddAvgSizeCol - 1),
 				},

pkg/migration/migrations/builtins_test.go

@@ -31,7 +31,7 @@ func TestIsAtLeastVersionBuiltin(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride:          clusterversion.ByKey(clusterversion.V21_2),
 				},
 			},

pkg/migration/migrations/comment_on_index_migration_external_test.go

@@ -37,7 +37,7 @@ func TestEnsureIndexesExistForComments(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.DeleteCommentsWithDroppedIndexes - 1),
 				},

pkg/migration/migrations/ensure_constraint_id_test.go

@@ -41,7 +41,7 @@ func TestEnsureConstraintIDs(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						tabledesc.ConstraintIDsAddedToTableDescsVersion - 1),
 				},

pkg/migration/migrations/ensure_no_draining_names_external_test.go

@@ -41,7 +41,7 @@ func TestEnsureNoDrainingNames(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.AvoidDrainingNames - 1),
 				},

pkg/migration/migrations/grant_option_migration_external_test.go

@@ -169,7 +169,7 @@ func TestGrantOptionMigration(t *testing.T) {
 				ServerArgs: base.TestServerArgs{
 					Knobs: base.TestingKnobs{
 						Server: &server.TestingKnobs{
-							DisableAutomaticVersionUpgrade: 1,
+							DisableAutomaticVersionUpgrade: make(chan struct{}),
 							BinaryVersionOverride:          clusterversion.ByKey(clusterversion.ValidateGrantOption - 1), // changed cluster version
 						},
 					},

pkg/migration/migrations/migrate_span_configs_test.go

@@ -44,7 +44,7 @@ func TestEnsureSpanConfigReconciliation(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.EnsureSpanConfigReconciliation - 1,
 					),
@@ -120,7 +120,7 @@ func TestEnsureSpanConfigReconciliationMultiNode(t *testing.T) {
 		serverArgs[i] = base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.EnsureSpanConfigReconciliation - 1,
 					),
@@ -133,7 +133,7 @@ func TestEnsureSpanConfigReconciliationMultiNode(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.EnsureSpanConfigReconciliation - 1,
 					),
@@ -196,7 +196,7 @@ func TestEnsureSpanConfigSubscription(t *testing.T) {
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.EnsureSpanConfigSubscription - 1,
 					),

pkg/migration/migrations/public_schema_migration_external_test.go

@@ -44,7 +44,7 @@ func publicSchemaMigrationTest(t *testing.T, ctx context.Context, numTables int)
 			Settings: settings,
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride:          clusterversion.ByKey(clusterversion.PublicSchemasWithDescriptors - 1),
 				},
 			},

pkg/migration/migrations/raft_applied_index_term_external_test.go

@@ -124,7 +124,7 @@ func TestRaftAppliedIndexTermMigration(t *testing.T) {
 			// Start at the version immediately preceding the migration.
 			BinaryVersionOverride: bootstrapVersion,
 			// We want to exercise manual control over the upgrade process.
-			DisableAutomaticVersionUpgrade: 1,
+			DisableAutomaticVersionUpgrade: make(chan struct{}),
 		}
 		return args
 	}
@@ -234,7 +234,7 @@ func TestLatestClusterDoesNotNeedRaftAppliedIndexTermMigration(t *testing.T) {
 		args.Knobs.Server = &server.TestingKnobs{
 			BinaryVersionOverride: binaryVersion,
 			// We want to exercise manual control over the upgrade process.
-			DisableAutomaticVersionUpgrade: 1,
+			DisableAutomaticVersionUpgrade: make(chan struct{}),
 		}
 		return args
 	}

pkg/migration/migrations/remove_invalid_database_privileges_external_test.go

@@ -38,7 +38,7 @@ func TestConvertIncompatibleDatabasePrivilegesToDefaultPrivileges(t *testing.T)
 		ServerArgs: base.TestServerArgs{
 			Knobs: base.TestingKnobs{
 				Server: &server.TestingKnobs{
-					DisableAutomaticVersionUpgrade: 1,
+					DisableAutomaticVersionUpgrade: make(chan struct{}),
 					BinaryVersionOverride: clusterversion.ByKey(
 						clusterversion.RemoveIncompatibleDatabasePrivileges - 1),
 				},

pkg/server/auto_upgrade.go

@@ -12,7 +12,6 @@ package server
 
 import (
 	"context"
-	"sync/atomic"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
@@ -36,16 +35,21 @@ func (s *Server) startAttemptUpgrade(ctx context.Context) {
 			Closer:         s.stopper.ShouldQuiesce(),
 		}
 
-		for r := retry.StartWithCtx(ctx, retryOpts); r.Next(); {
-			// Check if auto upgrade is disabled for test purposes.
-			if k := s.cfg.TestingKnobs.Server; k != nil {
-				upgradeTestingKnobs := k.(*TestingKnobs)
-				if disable := atomic.LoadInt32(&upgradeTestingKnobs.DisableAutomaticVersionUpgrade); disable == 1 {
-					log.Infof(ctx, "auto upgrade disabled by testing")
-					continue
+		// Check if auto upgrade is disabled for test purposes.
+		if k := s.cfg.TestingKnobs.Server; k != nil {
+			upgradeTestingKnobs := k.(*TestingKnobs)
+			if disableCh := upgradeTestingKnobs.DisableAutomaticVersionUpgrade; disableCh != nil {
+				log.Infof(ctx, "auto upgrade disabled by testing")
+				select {
+				case <-disableCh:
+					log.Infof(ctx, "auto upgrade no longer disabled by testing")
+				case <-s.stopper.ShouldQuiesce():
+					return
 				}
 			}
+		}
 
+		for r := retry.StartWithCtx(ctx, retryOpts); r.Next(); {
 			// Check if we should upgrade cluster version, keep checking upgrade
 			// status, or stop attempting upgrade.
 			if quit, err := s.upgradeStatus(ctx); err != nil {

pkg/server/migration_test.go

@@ -162,7 +162,7 @@ func TestBumpClusterVersion(t *testing.T) {
 						BinaryVersionOverride: test.activeClusterVersion.Version,
 						// We're bumping cluster versions manually ourselves. We
 						// want avoid racing with the auto-upgrade process.
-						DisableAutomaticVersionUpgrade: 1,
+						DisableAutomaticVersionUpgrade: make(chan struct{}),
 					},
 				},
 			})

pkg/server/testing_knobs.go

@@ -25,8 +25,8 @@ import (
 // TestingKnobs groups testing knobs for the Server.
 type TestingKnobs struct {
 	// DisableAutomaticVersionUpgrade, if set, temporarily disables the server's
-	// automatic version upgrade mechanism.
-	DisableAutomaticVersionUpgrade int32 // accessed atomically
+	// automatic version upgrade mechanism (until the channel is closed).
+	DisableAutomaticVersionUpgrade chan struct{}
 	// DefaultZoneConfigOverride, if set, overrides the default zone config
 	// defined in `pkg/config/zone.go`.
 	DefaultZoneConfigOverride *zonepb.ZoneConfig